import gradio as gr, pandas as pd
from huggingface_hub import list_repo_files, hf_hub_download
import zipfile, cloudpickle, shutil
from pathlib import Path
import autogluon.tabular as ag

MODEL_REPO_ID = "rlogh/cheese-texture-autogluon-classifier"
ZIP_CANDIDATES = ['cheese_texture_predictor_dir.zip', 'autogluon_predictor_dir.zip', 'predictor_dir.zip']
PKL_CANDIDATES = ['cheese_texture_predictor.pkl', 'autogluon_predictor.pkl', 'predictor.pkl']
CACHE_DIR = Path("hf_assets"); CACHE_DIR.mkdir(parents=True, exist_ok=True)
EXTRACT_DIR = CACHE_DIR / "predictor_native"

def _pick_first(files, cands):
    for n in cands:
        if n in files: return n
    return None

def load_predictor(repo_id):
    files = list_repo_files(repo_id, repo_type="model")
    zn = _pick_first(files, ZIP_CANDIDATES)
    pn = _pick_first(files, PKL_CANDIDATES)
    if zn:
        zp = hf_hub_download(repo_id, filename=zn, 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, exist_ok=True)
        with zipfile.ZipFile(zp, "r") as zf: zf.extractall(str(EXTRACT_DIR))
        kids = list(EXTRACT_DIR.iterdir())
        root = kids[0] if (len(kids)==1 and kids[0].is_dir()) else EXTRACT_DIR
        return ag.TabularPredictor.load(str(root), require_py_version_match=False)
    elif pn:
        pp = hf_hub_download(repo_id, filename=pn, repo_type="model", local_dir=str(CACHE_DIR), local_dir_use_symlinks=False)
        with open(pp, "rb") as f: return cloudpickle.load(f)
    raise FileNotFoundError("No predictor artifact found.")

PREDICTOR = load_predictor(MODEL_REPO_ID)
FEATURES = ["fat","origin","holed","price","protein"]
HOLED_MAP = {"Yes (has holes)":1,"No (solid)":0}
COMMON_ORIGINS = ["Italy","France","USA","Switzerland","Netherlands","Spain","Greece","Mexico","Denmark","Norway","India","Cyprus"]

def predict(fat, origin_pick, origin_free, holed_label, price, protein, k):
    origin = origin_free.strip() if origin_free and origin_free.strip() else origin_pick
    row = {"fat":float(fat), "origin":origin, "holed":HOLED_MAP[holed_label], "price":float(price), "protein":float(protein)}
    X = pd.DataFrame([row], columns=FEATURES)
    pred = PREDICTOR.predict(X).iloc[0]
    try:
        proba = PREDICTOR.predict_proba(X)
        if hasattr(proba, "to_frame"): proba = proba.to_frame().T if proba.ndim==1 else proba
        row0 = proba.iloc[0].to_dict()
        probs = {str(k): float(v) for k,v in row0.items()}
        probs = dict(sorted(probs.items(), key=lambda kv: kv[1], reverse=True)[:int(k)])
    except Exception:
        probs = {}
    return f"**Prediction:** {pred}", probs

with gr.Blocks() as demo:
    gr.Markdown("# Cheese Texture Classifier")
    fat     = gr.Slider(0, 60, value=25.0, step=0.1, label="Fat (g/100g)")
    protein = gr.Slider(0, 40, value=22.0, step=0.1, label="Protein (g/100g)")
    price   = gr.Slider(0, 10, value=2.50, step=0.01, label="Price")
    holed   = gr.Radio(list(HOLED_MAP.keys()), value="No (solid)", label="Has holes?")
    origin_pick = gr.Dropdown(COMMON_ORIGINS, value="Italy", label="Origin (common)")
    origin_free = gr.Textbox(value="", placeholder="Or type a custom origin", label="Custom origin (optional)")
    topk = gr.Slider(1, 4, value=4, step=1, label="Top-K")

    summary = gr.Markdown()
    probs   = gr.Label(num_top_classes=4, label="Class probabilities")

    gr.Button("Predict").click(predict,
        [fat, origin_pick, origin_free, holed, price, protein, topk],
        [summary, probs])

demo.launch()