#!python
# encoding: utf-8
"""
Global Key Recognition using a Convolutional Neural Network.

"""

from __future__ import absolute_import, division, print_function

import argparse

from madmom.features import ActivationsProcessor
from madmom.features.key import (CNNKeyRecognitionProcessor,
                                 key_prediction_to_label)
from madmom.io import write_key
from madmom.processors import io_arguments, IOProcessor


def main():
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description='''
    The KeyRecognition program recognises the global key of a musical piece
    using a Convolutional Neural Network, as described in

    Filip Korzeniowski and Gerhard Widmer,
    "Genre-Agnostic Key Classification with Convolutional Neural Networks",
    In Proceedings of the 19th International Society for Music Information
    Retrieval Conference (ISMIR), Paris, France, 2018.

    This program can be run in 'single' file mode to process a single audio
    file and write the recognised chords to STDOUT or the given output file.

      $ KeyRecognition single INFILE [-o OUTFILE]

    The program can also be run in 'batch' mode to process multiple audio
    files and save the chords to files with the given suffix.

      $ KeyRecognition batch [-o OUTPUT_DIR] [-s OUTPUT_SUFFIX] FILES

    If no output directory is given, the program writes the files with the
    extracted key to the same location as the audio files.

    The 'pickle' mode can be used to store the used parameters to be able to
    exactly reproduce experiments.

    '''
    )
    # version
    p.add_argument('--version', action='version',
                   version='KeyRecognition.2018')
    io_arguments(p, output_suffix='.key.txt')
    ActivationsProcessor.add_arguments(p)

    args = p.parse_args()

    if args.verbose:
        print(args)

    if args.load:
        # load activations from a file
        in_processor = ActivationsProcessor(mode='r', fps=0, **vars(args))
    else:
        in_processor = CNNKeyRecognitionProcessor(**vars(args))

    if args.save:
        # save activations to a file
        out_processor = ActivationsProcessor(mode='w', fps=0, **vars(args))
    else:
        out_processor = [key_prediction_to_label, write_key]

    processor = IOProcessor(in_processor, out_processor)
    args.func(processor, **vars(args))


if __name__ == '__main__':
    main()
