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



def inner_loop(nb_tokens, temperature, chord_progression, tempo, midi_file, bar_range):
    top_p = 1.0
    seed = 0

    print(midi_file)
    # Initialize the MusicLangPredictor
    ml = MusicLangPredictor('musiclang/musiclang-v2')

    if midi_file is not None:
        # 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))
    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:
        # 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,
            time_signature=(4, 4),
            temperature=temperature,
            topp=top_p,
            rng_seed=seed
        )

    chord_repr = generated_score.to_chord_repr()

    # Save the generated score as a MIDI file
    midi_path = 'test.mid'
    generated_score.to_midi(midi_path, tempo=tempo, time_signature=(4, 4))

    # Convert MIDI to WAV then WAV to MP3 for playback
    wav_path = 'result.wav'
    mp3_path = 'result.mp3'
    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}')
    return mp3_path, midi_path, chord_repr

def musiclang(nb_tokens, temperature, chord_progression, tempo, midi_file, bar_range):
    exception = None
    mp3_path, midi_path, chord_repr = None, None, None
    try:
        mp3_path, midi_path, chord_repr = inner_loop(nb_tokens, temperature, chord_progression, tempo, midi_file, 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, chord_repr, exception

# Update Gradio interface to include MIDI file upload and bar range selection
iface = gr.Interface(
    fn=musiclang,
    inputs=[
        gr.Number(label="Number of Tokens", value=1024, minimum=256, maximum=2048, step=256),
        gr.Slider(label="Temperature", value=0.9, minimum=0.1, maximum=1.0, step=0.1),
        gr.Textbox(label="Chord Progression", placeholder="Am CM Dm/F E7 Am", lines=2, value=""),
        gr.Slider(label="Tempo", value=120, minimum=60, maximum=240, step=1),
        gr.File(label="Upload MIDI File", type="filepath", file_types=[".mid", ".midi"]),
        gr.Textbox(label="Bar Range", placeholder="0-4", value="0-4")
    ],
    outputs=[
        gr.Audio(label="Generated Music"),
        gr.File(label="Download MIDI"),
        gr.Textbox(label="Inferred output Chord Progression", lines=2, value=""),
        gr.Textbox(label="Info Message")  # Initially hidden, shown only if there's an error
    ],
    title="Controllable Symbolic Music Generation with MusicLang Predict",
    description="""
    \n Simple music gen application that wraps <a href="https://github.com/musiclang/musiclang_predict">musiclang predict</a>.
    Beware that the model generates a score, not audio : rendered audio is only to provide a quick preview of the generated music.
    \n Customize the music generation by specifying the number of tokens, temperature, chord progression, tempo, and optionally uploading a MIDI file to use as a prompt. Specify the bar range for the MIDI prompt.
    \nChord qualities: M, m, 7, m7, m7b5, sus2, sus4, M7, dim, dim7. You can also specify the bass if it belongs to the chord (e.g., Bm/D).
    If no chord progression or MIDI file is given, it generates a free sample with the specified number of tokens."""
)

iface.launch()