import replicate
import gradio as gr

def run(prompt, text_chords, bpm, time_sig, audio_chords, duration):
    if audio_chords and text_chords:
        raise gr.Error("You must choose either only one of Audio Chords or Text Chords.")

    input_dict = {
        prompt: prompt,
        "top_k": 250,
        "top_p": 0,
        "audio_start": 0,
        "temperature": 1,
        "continuation": False,
        "output_format": "wav",
        "chroma_coefficient": 1,
        "multi_band_diffusion": False,
        "normalization_strategy": "loudness",
        "classifier_free_guidance": 3,
        "bpm": bpm,
        "time_sig": time_sig,
        "duration": duration
    }
    if text_chords:
        try:
            for chord in text_chords.split(" "):
                if chord == "":
                    raise gr.Error("Invalid text chords format. Make sure you have no extra spaces before/after chords!")
        except Exception as e:
            raise gr.Error("Invalid text chords format. Please refer to the label for more information.")
        input_dict["text_chords"] = text_chords
    elif audio_chords:
        input_dict["audio_chords"] = open(audio_chords, "rb")




    output = replicate.run(
        "outputinc/cog-musicgen-chord-trap:3e95bf45179ff9f924ed7e2025a23910c074d2df292987a37b62b4757311b264",
        input=input_dict
    )
    return output



def ui():
    with gr.Blocks() as musicgen_ui:
        gr.Markdown(
        """
        # MusicGen Chord Fine-tuned on Rhodes
        Please be patient as it will take a couple minutes to boot up the machine before even running inference
        """)
        prompt = gr.Textbox(lines=2, label="Prompt", value="smooth rhodes no drums")
        gr.Markdown(
            """
            Chord Selection
            ### You must select only one of Audio Chords or Text Chords. Remove the other input if present.
            """
        )
        with gr.Row():
            with gr.Column():
                text_chords = gr.Textbox(lines=2, value="C D:min7 G:7 A:min", label="Text Chords ")
                label = gr.Textbox(label="Text chord info", value="Single uppercase alphabet character(eg. `C`) is considered as a major chord. Chord attributes like(`maj`, `min`, `dim`, `aug`, `min6`, `maj6`, `min7`, `minmaj7`, `maj7`, `7`, `dim7`, `hdim7`, `sus2` and `sus4`) can be added to the root alphabet character after `:`.(eg. `A:min7`) Each chord token splitted by `SPACE` is allocated to a single bar. If more than one chord must be allocated to a single bar, cluster the chords adding with `,` without any `SPACE`.(eg. `C,C:7 G, E:min A:min`).")
            audio_chords = gr.File(label="Audio Chords", value=None)
        with gr.Row():
            bpm = gr.Number(label="BPM", value=120)
            time_sig = gr.Textbox(label="Time Signature", value="4/4")
            duration = gr.Slider(minimum=2, maximum=20, value=8, label="Duration (seconds)")
        run_button = gr.Button()
        output_audio = gr.Audio(label="Output Audio")

        run_button.click(run, [prompt, text_chords, bpm, time_sig, audio_chords, duration], outputs=output_audio)



    return musicgen_ui


if __name__ == "__main__":
    ui().launch(auth=("output", "becreative"))