#!/usr/bin/env python
# encoding: utf-8
"""
GMMPatternTracker for tracking (down-)beats based on rhythmic patterns.

"""

from __future__ import absolute_import, division, print_function

import argparse

from madmom.processors import IOProcessor, io_arguments
from madmom.audio.signal import SignalProcessor, FramedSignalProcessor
from madmom.audio.stft import ShortTimeFourierTransformProcessor
from madmom.audio.spectrogram import (FilteredSpectrogramProcessor,
                                      LogarithmicSpectrogramProcessor,
                                      SpectrogramDifferenceProcessor,
                                      MultiBandSpectrogramProcessor)
from madmom.features import ActivationsProcessor
from madmom.features.beats import PatternTrackingProcessor
from madmom.models import PATTERNS_BALLROOM


def main():
    """GMMPatternTracker"""

    # define parser
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description='''
    The GMMPatternTracker program detects rhythmic patterns in an audio file
    and reports the (down-)beats according to the method described in:

    "Rhythmic Pattern Modelling for Beat and Downbeat Tracking in Musical
     Audio"
    Florian Krebs, Sebastian Böck and Gerhard Widmer.
    Proceedings of the 14th International Society for Music Information
    Retrieval Conference (ISMIR), 2013.

    Instead of the originally proposed state space and transition model for the
    DBN, the following is used:

    "An Efficient State Space Model for Joint Tempo and Meter Tracking"
    Florian Krebs, Sebastian Böck and Gerhard Widmer.
    Proceedings of the 16th International Society for Music Information
    Retrieval Conference (ISMIR), 2015.

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

      $ GMMPatternTracker 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.

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

    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='GMMPatternTracker.2013')
    # add arguments
    io_arguments(p, output_suffix='.beats.txt')
    ActivationsProcessor.add_arguments(p)
    SignalProcessor.add_arguments(p, norm=False, gain=0)
    PatternTrackingProcessor.add_arguments(p)

    # parse arguments
    args = p.parse_args()

    # set immutable defaults
    args.num_channels = 1
    args.sample_rate = 44100
    args.fps = 50
    args.num_bands = 12
    args.fmin = 30
    args.fmax = 17000
    args.norm_filters = False
    args.mul = 1
    args.add = 1
    args.diff_ratio = 0.5
    args.positive_diffs = True
    args.crossover_frequencies = [270]
    args.pattern_files = PATTERNS_BALLROOM

    # 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:
        # define an input processor
        sig = SignalProcessor(**vars(args))
        frames = FramedSignalProcessor(**vars(args))
        stft = ShortTimeFourierTransformProcessor(**vars(args))
        filt = FilteredSpectrogramProcessor(**vars(args))
        log = LogarithmicSpectrogramProcessor(**vars(args))
        diff = SpectrogramDifferenceProcessor(**vars(args))
        mb = MultiBandSpectrogramProcessor(**vars(args))
        in_processor = [sig, frames, stft, filt, log, diff, mb]

    # output processor
    if args.save:
        # save the multiband features to file
        out_processor = ActivationsProcessor(mode='w', **vars(args))
    else:
        # downbeat processor
        downbeat_processor = PatternTrackingProcessor(**vars(args))
        if args.downbeats:
            # simply write the timestamps
            from madmom.utils import write_events as writer
        else:
            # borrow the note writer for outputting timestamps + beat numbers
            from madmom.features.notes import write_notes as writer
        # sequentially process the features
        out_processor = [downbeat_processor, writer]

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

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


if __name__ == '__main__':
    main()
