import os
import shutil
import zipfile
import pathlib
import pandas
import gradio
import huggingface_hub
import autogluon.tabular

# Settings
MODEL_REPO_ID = "Iris314/classical-automl-model"
ZIP_FILENAME  = "lego_predictor_dir.zip"
CACHE_DIR = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"

# Feature column names and target column names
FEATURE_COLS = [
    "Max Length (cm)",
    "Max Height (cm)",
    "Width (cm)",
    "Studs",
]
TARGET_COL = "Type of LEGO Brick"


# Encoding for outcome questions
OUTCOME_LABELS = {
    0: "Standard",
    1: "Flat",
    2: "Sloped",
}

# Download & load the native predictor
def _prepare_predictor_dir() -> str:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    local_zip = huggingface_hub.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, exist_ok=True)
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(str(EXTRACT_DIR))
    contents = list(EXTRACT_DIR.iterdir())
    predictor_root = contents[0] if (len(contents) == 1 and contents[0].is_dir()) else EXTRACT_DIR
    return str(predictor_root)

PREDICTOR_DIR = _prepare_predictor_dir()
PREDICTOR = autogluon.tabular.TabularPredictor.load(PREDICTOR_DIR, require_py_version_match=False)

# A mapping utility to make it easier to encode the variables
def _human_label(c):
    try:
        ci = int(c)
        if ci in OUTCOME_LABELS:
            return OUTCOME_LABELS[ci]
    except Exception:
        pass
    if c in OUTCOME_LABELS:
        return OUTCOME_LABELS[c]
    return str(c)


def do_predict(Max_Length, Max_Height, Width, Studs):

    row = {
        FEATURE_COLS[0]: float(Max_Length),
        FEATURE_COLS[1]: float(Max_Height),
        FEATURE_COLS[2]: float(Width),
        FEATURE_COLS[3]: int(Studs),
    }
    X = pandas.DataFrame([row], columns=FEATURE_COLS)

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

    try:
        proba = PREDICTOR.predict_proba(X)
        if isinstance(proba, pandas.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 the predicted label and confidence
    confidence = round((proba_dict.get(pred_label, 1.0) if proba_dict else 1.0) * 100, 2)
    return pred_label, confidence

# Representative examples
EXAMPLES = [
    [2.4, 1, 1.5, 6, "Standard"],
    [1.5, 0.3, 1.5, 4, "Flat"],
    [1.5, 0.8, 1.5, 1, "Sloped"],
]

# Gradio UI
with gradio.Blocks() as demo:
    # Provide an introduction
    gradio.Markdown("# Predict the Type of LEGO Brick")
    gradio.Markdown("""
    This is a simple app that predicts the type of LEGO brick based on its dimensions and number of studs.
    Enter the values for the features below and click "Submit" to see the prediction.
    """)

    with gradio.Row():
        max_length = gradio.Slider(0, 10, step=0.1, value=2.4, label=FEATURE_COLS[0])
        max_height = gradio.Slider(0, 5, step=0.1, value=1.0, label=FEATURE_COLS[1])
        width = gradio.Slider(0, 5, step=0.1, value=1.5, label=FEATURE_COLS[2])
        studs = gradio.Number(value=6, precision=0, label=FEATURE_COLS[3])

    # Output labels
    predicted_label = gradio.Label(label="Predicted Type of LEGO Brick")
    confidence_output = gradio.Label(label="Confidence (%)")

    # Submit button
    submit_button = gradio.Button("Submit")

    # Link inputs and outputs to the predict function
    inputs = [max_length, max_height, width, studs]
    outputs = [predicted_label, confidence_output]

    submit_button.click(fn=do_predict, inputs=inputs, outputs=outputs)

    gradio.Examples(
        examples=EXAMPLES,
        inputs=inputs,
        outputs=outputs,
        label="Representative examples",
        examples_per_page=5,
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.launch() # Set debug=False for deployment