import pathlib, shutil, zipfile, tempfile
import pandas
import gradio
import huggingface_hub
import autogluon.multimodal
import PIL.Image

MODEL_REPO_ID = "george2cool36/hw2_image_automl_autogluon"
ZIP_FILENAME  = "ag_image_predictor_dir.zip"
CACHE_DIR = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"

CLASS_LABELS = {0: "🛑 Has Stop Sign", 1: "✅ No Stop Sign"}

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.multimodal.MultiModalPredictor.load(PREDICTOR_DIR)

def _human_label(c):
    try:
        ci = int(c)
        return CLASS_LABELS.get(ci, str(c))
    except Exception:
        return CLASS_LABELS.get(c, str(c))

def do_predict(pil_img: PIL.Image.Image):
    if pil_img is None:
        return {}, "No image provided."

    tmpdir = pathlib.Path(tempfile.mkdtemp())
    img_path = tmpdir / "input.png"
    pil_img.save(img_path)

    df = pandas.DataFrame({"image": [str(img_path)]})

    proba_df = PREDICTOR.predict_proba(df)
    proba_df = proba_df.rename(columns={0: "🛑 Has Stop Sign (0)", 1: "✅ No Stop Sign (1)"})
    row = proba_df.iloc[0]

    pretty_dict = {
        "🛑 Has Stop Sign": float(row.get("🛑 Has Stop Sign (0)", 0.0)),
        "✅ No Stop Sign": float(row.get("✅ No Stop Sign (1)", 0.0)),
    }

    predicted_class = PREDICTOR.predict(df).iloc[0]
    pred_label = _human_label(predicted_class)

    md = f"**Prediction:** {pred_label}"
    if pretty_dict:
        md += f"  \n**Confidence:** {round(pretty_dict.get(pred_label, 0.0) * 100, 2)}%"


    return pretty_dict, md

EXAMPLES = [
    ["https://www.kingsrivercasting.com/images/stories/virtuemart/product/STOP%20SIGN%20(5).jpg"],
    ["https://www.trafficsafetywarehouse.com/Resources/images/traffic-sign-shapes.jpeg"],
    ["https://di-uploads-pod16.dealerinspire.com/toyotaofnorthcharlotte/uploads/2020/08/yield-road-sign.jpg"]
]


with gradio.Blocks() as demo:
    gradio.Markdown("# Has Stop Sign or Not?")
    gradio.Markdown(
        "This is a simple app that demonstrates how to use an autogluon multimodal"
        "predictor in a gradio space to predict whether an image contains a stop sign. To use,"
        "just upload a photo. The result should be generated automatically."
    )

    image_in = gradio.Image(type="pil", label="Input image", sources=["upload", "webcam"])

    proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities")
    prediction_output = gradio.Markdown()


    inputs = [image_in]
    outputs = [proba_pretty, prediction_output]
    for comp in inputs:
        comp.change(fn=do_predict, inputs=inputs, outputs=outputs)

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

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