#!/usr/bin/env python
# encoding: utf-8
"""
meter and tempo estimator for Greek folk music
"""

import argparse
from os import listdir, makedirs
from os.path import isfile, join, basename, exists
from tempocnn.classifier import MeterClassifier, TempoClassifier
from tempocnn.feature import read_features
from tempocnn.version import __version__ as package_version


def main():
    """greekfolk"""

    # define parser
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description='''
    The program 'greekfolk' estimates the global tempo and the numerator of a
    global meter for Greek folk music tracks.
    
    The used models are based on an approach described in detail in:

    Hendrik Schreiber, Meinard Müller,
    "A single-step approach to musical meter estimation using a
    convolutional neural network"
    Proceedings of the 19th International Society for Music Information
    Retrieval Conference (ISMIR), Paris, France, Sept. 2018.
    
    For the purpose of estimating meter and tempo for Greek folk music,
    transfer learning on a small dataset was conducted.
    For details see:
    
    Hendrik Schreiber,
    "Technical Report: Tempo and Meter Estimation for Greek Folk Music Using
    Convolutional Neural Networks and Transfer Learning"
    8th International Workshop on Folk Music Analysis (FMA),
    Thessaloniki, Greece, June 2018.

    License: GNU Affero General Public License v3
    ''')

    parser.add_argument('-v', '--version', action='version', version=f'greekfolk {package_version}')
    parser.add_argument('input', help='input folder with wav files')
    parser.add_argument('output', help='output folder')

    # parse arguments
    args = parser.parse_args()

    if not exists(args.output):
        print('Creating output dir: ' + args.output)
        makedirs(args.output)

    # load models
    print('Loading models...')
    meter_classifier = MeterClassifier('fma2018-meter')
    tempo_classifier = TempoClassifier('fma2018')

    print('Processing file(s)...')

    wav_files = [join(args.input, f) for f in listdir(args.input) if f.endswith('.wav') and isfile(join(args.input, f))]
    if len(wav_files) == 0:
        print("No .wav files found in " + args.input)
    for input_file in wav_files:
        print('Analyzing: ' + input_file)
        meter_features = read_features(input_file, frames=512, hop_length=256)
        meter_result = str(meter_classifier.estimate_meter(meter_features))

        tempo_features = read_features(input_file)
        tempo_result = str(round(tempo_classifier.estimate_tempo(tempo_features), 1))

        output_file = join(args.output, basename(input_file).replace('.wav', '.txt'))
        with open(output_file, mode='w') as f:
            f.write(tempo_result + '\t' + meter_result + '\n')
    print('\nDone')


if __name__ == '__main__':
    main()
