#!python
# encoding: utf-8
"""
BarTracker (down-)beat detection algorithm.

"""
from __future__ import absolute_import, division, print_function

import argparse

from madmom.audio.signal import SignalProcessor
from madmom.features import ActivationsProcessor
from madmom.features.downbeats import (DBNBarTrackingProcessor,
                                       LoadBeatsProcessor, RNNBarProcessor)
from madmom.io import write_beats, write_downbeats
from madmom.processors import IOProcessor, ParallelProcessor, io_arguments


def main():
    """BarTracker"""

    # define parser
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description='''
    The BarTracker program detects all downbeats in an audio file given
    detected beats according to the method described in:

    "Downbeat Tracking Using Beat-Synchronous Features and Recurrent Networks"
    Florian Krebs, Sebastian Böck, Matthias Dorfer and Gerhard Widmer.
    Proceedings of the 17th International Society for Music Information
    Retrieval Conference (ISMIR), 2016.

    This program can be run in 'single' file mode to process a single audio
    file and read the beat positions from STDIN or file and write the detected
    beats to STDOUT or the given output file.

      $ BarTracker single INFILE [-o OUTFILE] < BEATS
      $ BarTracker --beats BEATS single INFILE [-o OUTFILE]

    If multiple audio files should be processed, the program can also be run
    in 'batch' mode to save the detected beats to files with the given suffix.
    The beat positions are read from files matching the given extension.

      $ BarTracker batch [-o OUTPUT_DIR] [-s OUTPUT_SUFFIX] FILES AND BEATS

    If no output directory is given, the program writes the files with the
    detected beats 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='BarTracker.2016')
    # add processor arguments
    io_arguments(p, output_suffix='.beats.txt')
    ActivationsProcessor.add_arguments(p)
    LoadBeatsProcessor.add_arguments(p)
    DBNBarTrackingProcessor.add_arguments(p, beats_per_bar=[3, 4],
                                          meter_change_prob=0.001)

    # parse arguments
    args = p.parse_args()

    # set immutable arguments
    args.fps = 100
    args.num_channels = 1

    # print arguments
    if args.verbose:
        print(args)

    # input processor
    if args.load:
        # load the activations from file
        in_processor = ActivationsProcessor(mode='r', **vars(args))
    else:
        # signal processor
        signal_processor = SignalProcessor(**vars(args))
        # beats processor
        if hasattr(args, 'infile'):
            # single mode: read the beats from STDIN
            load_beats_processor = LoadBeatsProcessor(**vars(args))
        elif hasattr(args, 'files'):
            # batch mode: load the beats from file (set default extension here)
            args.beats = args.files
            load_beats_processor = LoadBeatsProcessor(**vars(args))
            # remove the beats files from the list of files to be processed
            args.files = list(set(args.files) -
                              set(load_beats_processor.beats))
        # combine the signal and the beats
        beats_processor = ParallelProcessor([signal_processor,
                                             load_beats_processor])
        # sync the beat features and use a RNN to predict the beats
        rnn_processor = RNNBarProcessor(**vars(args))
        in_processor = [beats_processor, rnn_processor]

    # output processor
    if args.save:
        # save the RNN bar activations to file
        out_processor = ActivationsProcessor(mode='w', **vars(args))
    else:
        # track the (down-)beats in the bar activation function and output them
        out_processor = [DBNBarTrackingProcessor(**vars(args))]
        if args.downbeats:
            out_processor.append(write_downbeats)
        else:
            out_processor.append(write_beats)

    # create an IOProcessor
    processor = IOProcessor(in_processor, out_processor)

    # and call the processing function
    args.func(processor, **vars(args))


if __name__ == '__main__':
    main()
