import os
import gradio as gr
from midi_ddsp.midi_ddsp_synthesize import load_pretrained_model
from midi_ddsp.utils.midi_synthesis_utils import synthesize_mono_midi
from midi_ddsp.data_handling.instrument_name_utils import INST_NAME_TO_ID_DICT, INST_NAME_LIST

os.system('midi_ddsp_download_model_weights')
synthesis_generator, expression_generator = load_pretrained_model()

def midi_ddsp_synth(instrument, midi_file):
    '''
    Paras:
        midi_file: the path of the midi file
        instrument: instrument type, choose from 'violin', 'viola', 'cello', 'double bass',
                                            'flute', 'oboe', 'clarinet', 'saxophone',
                                            'bassoon', 'trumpet', 'horn', 'trombone', 'tuba'                                            
    Returns:
        output_file_path: path of the output wav file
    '''
    if not (os.path.isfile(midi_file.name) and midi_file.name[-4:] == '.mid'):
        raise FileNotFoundError('Error: not a valid midi file')
    if instrument not in INST_NAME_LIST:
        raise ValueError('Error: not an available instrument type')
    instrument_id = INST_NAME_TO_ID_DICT[instrument]
    _ = synthesize_mono_midi(synthesis_generator, expression_generator, midi_file.name,
                             instrument_id, output_dir='/tmp', pitch_offset=0, speed_rate=1)
    return os.path.join('/tmp', os.path.basename(midi_file.name).replace('.mid', '.wav'))
  
title = "Midi-DDSP Demo"

article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.09312' target='_blank'>Paper</a>\
            | <a href='https://github.com/magenta/midi-ddsp' target='_blank'>Github Repo</a>\
            | <a href='https://huggingface.co/spaces/lukewys/midi-ddsp' target='_blank'>Official Hugging Face Space</a></p>"

gr.Interface(
    fn = midi_ddsp_synth,
    inputs = [
        gr.inputs.Textbox(
            lines=1, placeholder=None, default='', label='Input the instrument type (choose from: violin, viola, cello, double bass,\
                                                        flute, oboe, clarinet, saxophone, bassoon, trumpet, horn, trombone, tuba)'),
        gr.File(type="file", label="Upload .mid file here")
    ],
    outputs = "audio",
    title = title,
    article = article,
    ).launch(debug=True)