import gradio as gr
import pandas as pd
from autogluon.tabular import TabularPredictor
from huggingface_hub import hf_hub_download
import os, zipfile

REPO_ID = "Iris314/classical-automl-model"
ZIP_FILE = "lego_predictor_dir.zip"

local_zip = hf_hub_download(repo_id=REPO_ID, filename=ZIP_FILE)

extract_dir = "lego_predictor_dir"
os.makedirs(extract_dir, exist_ok=True)

with zipfile.ZipFile(local_zip, 'r') as zip_ref:
    zip_ref.extractall(extract_dir)

predictor = TabularPredictor.load(extract_dir, require_py_version_match=False)


def predict_brick(length, height, width, studs):
    record = pd.DataFrame([{
        "Max Length (cm)": length,
        "Max Height (cm)": height,
        "Width (cm)": width,
        "Studs": studs
    }])
    pred = predictor.predict(record)[0]
    proba = predictor.predict_proba(record).iloc[0].to_dict()
    return f"Prediction: {pred}", proba

with gr.Blocks(title="LEGO Brick Classifier") as demo:
    gr.Markdown("## LEGO Brick Classification\nPredict Standard / Flat / Sloped")
    with gr.Row():
        with gr.Column():
            length = gr.Slider(1, 10, step=0.5, value=4, label="Length")
            height = gr.Slider(0.5, 5, step=0.1, value=1.2, label="Height")
            width = gr.Slider(1, 10, step=0.5, value=2, label="Width")
            studs = gr.Slider(0, 12, step=1, value=4, label="Studs")
            btn = gr.Button("Predict")
        with gr.Column():
            out_label = gr.Textbox(label="Prediction")
            out_probs = gr.Label(label="Class Probabilities")

    btn.click(predict_brick, [length, height, width, studs], [out_label, out_probs])

    gr.Examples(
        examples=[
            [4, 1.2, 2, 4],
            [6, 0.5, 2, 6],
            [3, 2.0, 2, 2]
        ],
        inputs=[length, height, width, studs]
    )

demo.launch()