import gradio as gr
from autogluon.tabular import TabularPredictor
import pandas as pd

# Load the model from the `model/` folder in this repo
predictor = TabularPredictor.load("model/")

key_center_mapping = {
    0: "A", 1: "Bb", 2: "B", 3: "C", 4: "Db", 5: "D",
    6: "Eb", 7: "E", 8: "F", 9: "Gb", 10: "G", 11: "Ab"
}
marking_mapping = {
    0: "Minuet", 1: "Allegro", 2: "Andante",
    3: "Moderato", 4: "Allegretto", 5: "Dance"
}

def predict_composer(rh, lh, measures, key_center, marking):
    df = pd.DataFrame({
        'right hand notes': [rh],
        'left hand notes': [lh],
        'measures': [measures],
        'Key Center': [key_center],
        'marking': [marking]
    })
    pred = predictor.predict(df)[0]
    probs = predictor.predict_proba(df).iloc[0].to_dict()
    return pred, probs

examples = [
    [108, 82, 16, 3, 1],
    [196, 136, 29, 2, 2],
    [96, 49, 13, 2, 4],
    [481, 561, 31, 5, 5],
    [174, 129, 31, 2, 1],
]

with gr.Blocks() as demo:
    gr.Markdown("# Classical Music Composer Classifier")
    gr.Markdown("Predict whether a piece was composed by **Mozart** or **Beethoven**.")
    with gr.Row():
        rh = gr.Number(150, label="Right Hand Notes")
        lh = gr.Number(100, label="Left Hand Notes")
        measures = gr.Number(20, label="Measures")
    with gr.Row():
        key_center = gr.Dropdown(list(key_center_mapping.keys()), value=3, label="Key Center")
        marking = gr.Dropdown(list(marking_mapping.keys()), value=1, label="Marking")
    out_label = gr.Textbox(label="Predicted Composer")
    out_probs = gr.Label(num_top_classes=2, label="Probabilities")
    for inp in [rh, lh, measures, key_center, marking]:
        inp.change(fn=predict_composer, inputs=[rh, lh, measures, key_center, marking], outputs=[out_label, out_probs])
    gr.Examples(examples, inputs=[rh, lh, measures, key_center, marking], outputs=[out_label, out_probs])
    demo.launch()