import pathlib, shutil, zipfile, os, traceback
import pandas as pd
import gradio as gr

from huggingface_hub import hf_hub_download
from autogluon.tabular import TabularPredictor

# ---------------- UI copy ----------------
TITLE = "🧱 LEGO Brick Classifier"
DESC  = "Predicts whether a LEGO piece is Standard, Flat, or Sloped from basic dimensions."

# ---------------- Settings ----------------
MODEL_REPO_ID = "Iris314/classical-automl-model"
ZIP_FILENAME  = "lego_predictor_dir.zip"

# UI → model feature name mapping
COLUMN_ALIAS = {
    "Length": "Max Length (cm)",
    "Height": "Max Height (cm)",
    "Width":  "Width (cm)",
    "Studs":  "Studs",
}
FEATURE_COLS_UI = ["Length", "Height", "Width", "Studs"]

# ---------------- Load predictor ----------------
CACHE_DIR = pathlib.Path("hf_cache"); EXTRACT_DIR = CACHE_DIR / "predictor"
CACHE_DIR.mkdir(exist_ok=True, parents=True)

def load_predictor():
    local_zip = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
        local_dir=str(CACHE_DIR),
        local_dir_use_symlinks=False,
    )
    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True)
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(EXTRACT_DIR)
    kids = list(EXTRACT_DIR.iterdir())
    path = kids[0] if len(kids) == 1 and kids[0].is_dir() else EXTRACT_DIR
    return TabularPredictor.load(str(path), require_py_version_match=False)

try:
    PREDICTOR = load_predictor()
except Exception as e:
    PREDICTOR = None
    print("Failed to load predictor:", e)

# ---------------- Helpers ----------------
def _cast_and_rename(row_dict):
    row = dict(row_dict)
    row["Length"] = float(row["Length"])
    row["Height"] = float(row["Height"])
    row["Width"]  = float(row["Width"])
    # gr.Number returns float; round & cast for integer feature
    row["Studs"]  = int(round(float(row["Studs"])))
    X_ui = pd.DataFrame([row], columns=FEATURE_COLS_UI)
    X_model = X_ui.rename(columns=COLUMN_ALIAS)
    return X_model

def classify_brick(length, height, width, studs):
    try:
        if PREDICTOR is None:
            raise RuntimeError("Model failed to load on startup. Check model artifact path & runtime deps.")

        X = _cast_and_rename({
            "Length": length, "Height": height, "Width": width, "Studs": studs
        })

        # Try probabilities; fall back to label
        try:
            proba = PREDICTOR.predict_proba(X)
            s = proba.iloc[0] if hasattr(proba, "iloc") else proba
            s = s.sort_values(ascending=False)
            s.index = [str(k) for k in s.index]  # ensure JSON-serializable keys
            return {k: float(v) for k, v in s.items()}
        except Exception:
            pred = PREDICTOR.predict(X)
            pred_val = pred.iloc[0] if hasattr(pred, "iloc") else pred
            return {"prediction": str(pred_val)}
    except Exception as e:
        return {
            "error": f"{type(e).__name__}: {e}",
            "traceback": traceback.format_exc(limit=1)
        }

# ---------------- Gradio ----------------
demo = gr.Interface(
    fn=classify_brick,
    inputs=[
        gr.Slider(1, 10, step=0.1, value=4,   label="Length"),
        gr.Slider(0.2, 5, step=0.1, value=1.2, label="Height"),
        gr.Slider(1, 10, step=0.1, value=2,   label="Width"),
        gr.Number(value=4, precision=0, label="Studs"),
    ],
    outputs=gr.Label(num_top_classes=3, label="Predicted Class / Probabilities"),
    examples=[[4, 1.2, 2, 4], [2, 0.6, 2, 2], [3, 2.0, 2, 2]],
    title=TITLE,
    description=DESC
)

if __name__ == "__main__":
    # In Spaces, no share=True needed
    demo.launch()