import gradio as gr
import autogluon.tabular
from autogluon.tabular import TabularPredictor
import pandas as pd
import os
import shutil
import zipfile
import pathlib
from huggingface_hub import HfApi, Repository, create_repo
import huggingface_hub as h

model_repo_id = "madhavkarthi/24679-HW2-tabular-autolguon-predictor"
zip_filename  = "autogluon_predictor_dir.zip"
cache_dir = pathlib.Path("hf_assests")
extract_dir = cache_dir / "predictor_native"

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

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"
}

feature_col = ["right hand notes", "left hand notes",
               "measures", "Key Center", "marking",
]
target_col = "Target (Composer)"

outcome_lab = {
    0: "Beethoven",
    1: "Mozart"
}

def human_label(c):
    try:
        ci = int(c)
        if ci in outcome_lab:
            return outcome_lab[ci]
    except Exception:
        pass
    if c in outcome_lab:
        return outcome_lab[c]
    return str(c)

def do_predict(right_hand_notes, left_hand_notes, measures, Key_Center, marking):

    row = {
        feature_col[0]: int(right_hand_notes),
        feature_col[1]: int(left_hand_notes),
        feature_col[2]: int(measures),
        feature_col[3]: int(Key_Center),
        feature_col[4]: int(marking),
    }
    X = pd.DataFrame([row], columns=feature_col)

    pred_series = predictor.predict(X)
    raw_pred = pred_series.iloc[0]

    try:
        proba = predictor.predict_proba(X)
        if isinstance(proba, pd.Series):
            proba = proba.to_frame().T
    except Exception:
        proba = None

    pred_label = human_label(raw_pred)

    proba_dict = None
    if proba is not None:
        row0 = proba.iloc[0]
        tmp = {}
        for cls, val in row0.items():
            key = human_label(cls)
            tmp[key] = float(val) + float(tmp.get(key, 0.0))
        proba_dict = dict(sorted(tmp.items(), key=lambda kv: kv[1], reverse=True))

    return pred_label, proba_dict

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(value=150,
            precision=0,
            label="Right Hand Notes",
            info="Number of notes played by right hand")
        lh = gr.Number(value=100,
            precision=0,
            label="Left Hand Notes",
            info="Number of notes played by left hand")
        measures = gr.Number(value=20,
            precision=0,
            label="Measures",
            info="Number of musical measures")
    with gr.Row():
        key_center = gr.Dropdown(choices=[(f"{v} ({k})", k) for k, v in key_center_mapping.items()],
            value=3,
            label="Key Center",
            info="Musical key of the piece")
        marking = gr.Dropdown(choices=[(f"{v} ({k})", k) for k, v in marking_mapping.items()],
            value=1,
            label="Musical Marking",
            info="Tempo/style marking of the piece")

    out_label = gr.Textbox(label="Predicted Composer")
    out_probs = gr.Label(num_top_classes=2, label="Probabilities")

    inputs_for_predict = [rh, lh, measures, key_center, marking]
    outputs_for_predict = [out_label, out_probs]

    for comp in inputs_for_predict:
        comp.change(
            fn=do_predict,
            inputs=inputs_for_predict,
            outputs=outputs_for_predict
        )
    gr.Examples(examples, inputs=inputs_for_predict, outputs=outputs_for_predict)


if __name__ == "__main__":
    demo.launch()