import argparse
import os
import lasagne
import numpy as np

from madmom.features.notes import NotePeakPickingProcessor

from piano_transcription import BEST_MODEL_FILE_NAME, SETTINGS_FILE_NAME, LOSSES_FILE
from piano_transcription.utils import select_model
from piano_transcription.data import FEAT_SIZE, OUT_SIZE
from piano_transcription.data.annotations import write_txt_annotation
from piano_transcription.data.features import extract_features


def run(model_path, input_file):

    settings = np.load(os.path.join(model_path, SETTINGS_FILE_NAME))

    model = select_model(settings['model'])
    model_seq_len = model.MAX_PRED_SIZE

    features = extract_features(input_file)

    network = model.build_eval_model(model_seq_len, FEAT_SIZE, OUT_SIZE)

    with np.load(os.path.join(model_path, BEST_MODEL_FILE_NAME)) as f:
        param_values = [f['arr_%d' % i] for i in range(len(f.files))]
    lasagne.layers.set_all_param_values(network, param_values)

    # transcribe using model
    pred = model.predict(network, features, model_seq_len, OUT_SIZE)

    peak_picker = NotePeakPickingProcessor(pitch_offset=0)
    notes = peak_picker.process(pred)

    write_txt_annotation(open(input_file+'out.txt', 'w'), notes)


def main():
    # add argument parser
    parser = argparse.ArgumentParser(description='Transcribe piano tracks.')
    parser.add_argument('--file', help='file to be transcribed.')
    parser.add_argument('--model', help='path to trained model file.')
    args = parser.parse_args()

    model_path = args.model
    input_file = args.file

    assert os.path.exists(model_path)
    run(model_path, input_file)


if __name__ == '__main__':
    main()