import gradio as gr
from musiclang_predict import MusicLangPredictor
from musiclang import Score
from midi2audio import FluidSynth
import os
import tempfile


def inner_loop(midi_file, chord_progression, tempo, temperature, nb_tokens, bar_range):
    top_p = 0.98
    seed = 0
    # Initialize the MusicLangPredictor
    ml = MusicLangPredictor('musiclang/musiclang-v2')
    tempo_message = ""  # Default message if no MIDI file is uploaded
    time_signature = (4, 4)
    if midi_file is not None and midi_file != "":
        # Load the MIDI file and use it as the score prompt
        filepath = midi_file
        start_bar, end_bar = map(int, bar_range.split("-"))
        score = Score.from_midi(filepath, chord_range=(start_bar, end_bar))
        tempo = score.config['tempo'] # Use the tempo from the MIDI file and change input
        time_signature = score.config['time_signature']
        time_signature = (time_signature[1], time_signature[2])
        tempo_message = f"Warning : real tempo of file is : {int(tempo)} BPM."  # Update message based on MIDI file
    else:
        score = None  # Default score is None if no MIDI file is uploaded

    # Generate the score based on provided inputs and the uploaded MIDI file if available
    if chord_progression.strip() == "" and score is None:
        # Generate without specific chord progression or MIDI prompt
        generated_score = ml.predict(
            nb_tokens=int(nb_tokens),
            temperature=float(temperature),
            topp=top_p,
            rng_seed=seed
        )
    elif score is not None and chord_progression.strip() == "":
        # Generate using the uploaded MIDI file as a prompt
        generated_score = ml.predict(
            score=score,  # Use the uploaded MIDI as the score prompt
            nb_tokens=int(nb_tokens),
            temperature=float(temperature),
            topp=top_p,
            rng_seed=seed
        )
    else:
        # Generate with specific chord progression
        generated_score = ml.predict_chords(
            chord_progression,
            score=score,  # Use the uploaded MIDI as the score prompt
            time_signature=time_signature,
            temperature=temperature,
            topp=top_p,
            rng_seed=seed
        )

    chord_repr = generated_score.to_chord_repr()

    # Save the generated score as a MIDI file
    temp_midi_file = tempfile.NamedTemporaryFile(suffix=".mid", delete=False)
    midi_path = temp_midi_file.name
    generated_score.to_midi(midi_path, tempo=tempo, time_signature=time_signature)

    # Convert MIDI to WAV then WAV to MP3 for playback
    temp_wav_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_mp3_file = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
    wav_path = temp_wav_file.name
    mp3_path = temp_mp3_file.name
    FluidSynth("/usr/share/sounds/sf2/FluidR3_GM.sf2").midi_to_audio(midi_path, wav_path)
    os.system(f'ffmpeg -i {wav_path} -acodec libmp3lame -y -loglevel quiet -stats {mp3_path}')

    # Remove the temporary WAV file
    os.remove(wav_path)

    return mp3_path, midi_path, chord_repr, tempo_message

def musiclang(midi_file, chord_progression, tempo, temperature, nb_tokens, bar_range):
    exception = None
    mp3_path, midi_path, chord_repr, tempo_message = None, None, None, ""
    try:
        mp3_path, midi_path, chord_repr, tempo_message = inner_loop(midi_file, chord_progression, tempo, temperature, nb_tokens, bar_range)
    except Exception as e:
        exception = "Error : " + e.__class__.__name__ + " " + str(e)
    # Return the MP3 path for Gradio to display and the MIDI file path for download
    return mp3_path, midi_path, exception


with gr.Blocks() as demo:
    # Introductory text
    gr.Markdown("""
    # Controllable Symbolic Music Generation with MusicLang Predict
    [MusicLang Predict](https://github.com/musiclang/musiclang_predict) offers advanced controllability features and high-quality music generation by manipulating symbolic music.
    You can for example use it to continue your composition with a specific chord progression.
    """)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                midi_file = gr.File(label="Prompt MIDI File (Optional)", type="filepath", file_types=[".mid", ".midi"],
                                    elem_id='midi_file_input')
                with gr.Column():
                    bar_range = gr.Textbox(label="Bar Range of input file (eg: 0-4 for first four bars)", placeholder="0-4",
                                           value="0-4", elem_id='bar_range_input')
                    nb_tokens = gr.Number(label="Nb Tokens",
                                          value=512, minimum=256, maximum=2048, step=256, elem_id='nb_tokens_input')
                    temperature = gr.Slider(
                        label="Temperature",
                        value=0.95,
                        visible=False,
                        minimum=0.1, maximum=1.0, step=0.1, elem_id='temperature_input')
                    tempo = gr.Slider(label="Tempo", value=120, minimum=60, maximum=240, step=1, elem_id='tempo_input')

            with gr.Row():
                chord_progression = gr.Textbox(
                    label="Chord Progression (Optional)",
                    placeholder="Am CM Dm7/F E7 Asus4", lines=2, value="", elem_id='chord_progression_input')

            with gr.Row():
                generate_btn = gr.Button("Generate", elem_id='generate_button')

        with gr.Column():
            info_message = gr.Textbox(label="Info Message", elem_id='info_message_output')
            generated_music = gr.Audio(label="Preview generated Music", elem_id='generated_music_output')
            generated_midi = gr.File(label="Download MIDI", elem_id='generated_midi_output')
            generate_btn.click(
                fn=musiclang,
                inputs=[midi_file, chord_progression, tempo, temperature, nb_tokens, bar_range],
                outputs=[generated_music, generated_midi, info_message]
            )
            

    with gr.Row():
        with gr.Column():
            gr.Markdown("## Examples")
            gr.Examples(
                examples=[["/home/user/app/bach_847.mid", "", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/bach_847.mid", "Cm C7/E Fm F#dim G7", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/boney_m_ma_baker.mid", "", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/eminem_slim_shady.mid", "Cm AbM BbM G7 Cm", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/mozart_alla_turca.mid", "", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/mozart_alla_turca.mid", "Am Em CM G7 E7 Am Am E7 Am", 120, 0.95, 512, "0-4"],
                          ["/home/user/app/daft_punk_around_the_world.mid", "", 120, 0.95, 512, "0-4"],
                          ],
                inputs=[midi_file, chord_progression, tempo, temperature, nb_tokens, bar_range],
                outputs=[generated_music, generated_midi, info_message],
                fn=musiclang,
                cache_examples=True,
            )

demo.launch()