import os

import gradio as gr

from utils import (
    generate_song,
    remove_last_instrument,
    regenerate_last_instrument,
    change_tempo,
)


os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

DESCRIPTION = """

# 🎵 Multitrack Midi Generator 🎶
This interactive application uses an AI model to generate music sequences based on a chosen genre and various user inputs.

Features:
🎼 Select the genre for the music.
🌡️ Use the "Temperature" slider to adjust the randomness of the music generated (higher values will produce more random outputs).
⏱️ Adjust the "Tempo" slider to change the speed of the music.
🎹 Use the buttons to generate a new song from scratch, continue generation with the current settings, remove the last added instrument, regenerate the last added instrument with a new one, or change the tempo of the current song.
Outputs:
The app outputs the following:

🎧 The audio of the generated song.
📁 A MIDI file of the song.
📊 A plot of the song's sequence.
🎸 A list of the generated instruments.
📝 The text sequence of the song.
Enjoy creating your own AI-generated music! 🎵
"""

genres = ["ROCK", "POP", "OTHER", "R&B/SOUL", "JAZZ", "ELECTRONIC", "RANDOM"]

demo = gr.Blocks()

with demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            temp = gr.Slider(
                minimum=0, maximum=1, step=0.05, value=0.75, label="Temperature"
            )
            genre = gr.Dropdown(choices=genres, value="POP", label="Select the genre")
            with gr.Row():
                btn_from_scratch = gr.Button("Start from scratch")
                btn_continue = gr.Button("Continue Generation")
                btn_remove_last = gr.Button("Remove last instrument")
                btn_regenerate_last = gr.Button("Regenerate last instrument")
        with gr.Column():
            with gr.Box():
                audio_output = gr.Video()
                midi_file = gr.File()
                with gr.Row():
                    qpm = gr.Slider(
                        minimum=60, maximum=140, step=10, value=120, label="Tempo"
                    )
                    btn_qpm = gr.Button("Change Tempo")
    with gr.Row():
        with gr.Column():
            plot_output = gr.Plot()
        with gr.Column():
            instruments_output = gr.Markdown("# List of generated instruments")
    with gr.Row():
        text_sequence = gr.Text()
        empty_sequence = gr.Text(visible=False)
    with gr.Row():
        num_tokens = gr.Text()
    btn_from_scratch.click(
        fn=generate_song,
        inputs=[genre, temp, empty_sequence, qpm],
        outputs=[
            audio_output,
            midi_file,
            plot_output,
            instruments_output,
            text_sequence,
            num_tokens,
        ],
    )
    btn_continue.click(
        fn=generate_song,
        inputs=[genre, temp, text_sequence, qpm],
        outputs=[
            audio_output,
            midi_file,
            plot_output,
            instruments_output,
            text_sequence,
            num_tokens,
        ],
    )
    btn_remove_last.click(
        fn=remove_last_instrument,
        inputs=[text_sequence, qpm],
        outputs=[
            audio_output,
            midi_file,
            plot_output,
            instruments_output,
            text_sequence,
            num_tokens,
        ],
    )
    btn_regenerate_last.click(
        fn=regenerate_last_instrument,
        inputs=[text_sequence, qpm],
        outputs=[
            audio_output,
            midi_file,
            plot_output,
            instruments_output,
            text_sequence,
            num_tokens,
        ],
    )
    btn_qpm.click(
        fn=change_tempo,
        inputs=[text_sequence, qpm],
        outputs=[
            audio_output,
            midi_file,
            plot_output,
            instruments_output,
            text_sequence,
            num_tokens,
        ],
    )

demo.launch(debug=True)