# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1pj05m47J2hNpEuUjfEoP7Jk3UaGpUjPE
"""

import gradio as gr
import joblib
import torch
import torch.nn as nn
import numpy as np
from PIL import Image

import torchvision.transforms as T
from torchvision.models import resnet18, ResNet18_Weights

device = "cuda" if torch.cuda.is_available() else "cpu"

cnn = resnet18(weights=ResNet18_Weights.DEFAULT)
cnn.fc = nn.Identity()
cnn.eval().to(device)

for p in cnn.parameters():
    p.requires_grad = False

transform = T.Compose([
    T.Resize((640, 640)),
    T.ToTensor()
])

def extract_feature(image: Image.Image):
    img = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        feat = cnn(img).cpu().numpy().flatten()
    return feat

MODELS = {
    "XGBoost": joblib.load("saved_models/xgb/hier_xgb_bundle.pkl"),
    "Logistic": joblib.load("saved_models/logistic/hier_logistic_bundle.pkl"),
    "MLP": joblib.load("saved_models/mlp/hier_mlp_bundle.pkl"),
}

def predict_brix(model_name, image, weight, width, height):

    if image is None:
        return "❌ Please upload an image", {"Low": 0, "Medium": 0, "High": 0}

    bundle = MODELS[model_name]

    model_A = bundle["model_A"]
    model_B = bundle["model_B"]
    scaler  = bundle["scaler"]
    th      = bundle["thresholds"]["BEST_TH"]

    # ---- Feature extraction ----
    feat_img = extract_feature(image)

    num = np.array([[weight, width, height]])
    num = scaler.transform(num).flatten()

    x = np.concatenate([feat_img, num]).reshape(1, -1)

    # ---- Stage A ----
    p_sweet = model_A.predict_proba(x)[0][1]

    probs = {"Low": 0.0, "Medium": 0.0, "High": 0.0}

    if p_sweet < (1 - th):
        label = "🍈 Low Brix"
        probs["Low"] = float(1 - p_sweet)

    elif p_sweet > th:
        cls = model_B.predict(x)[0] + 1
        p2 = model_B.predict_proba(x)[0]

        probs["Medium"] = float(p2[0])
        probs["High"]   = float(p2[1])

        label = "🍈 Medium Brix" if cls == 1 else "🍈 High Brix"

    else:
        label = "⚠️ Abstain (Uncertain)"

    return label, probs

with gr.Blocks() as demo:

    gr.Markdown(
        """
        # 🍈 Melon Brix Classifier
        **Hierarchical Classification with Abstention**
        CNN Feature Extraction + Tabular Data

        เลือกโมเดล → ใส่ข้อมูล → ระบบจะทำนายความหวาน
        ถ้าไม่มั่นใจ ระบบจะ *Abstain* แทนการเดา
        """
    )

    with gr.Tab("🍈 Prediction"):   # ✅ แท็บเดียว
        with gr.Row():
            with gr.Column(scale=1):
                model_dd = gr.Dropdown(
                    ["XGBoost", "Logistic", "MLP"],
                    value="MLP",
                    label="Select Model"
                )

                image_in = gr.Image(
                      type="pil",
                      image_mode="RGB",
                      label="Melon Image (.jpg / .png)",
                      sources=["upload"])

                weight_in = gr.Number(label="Weight (kg)", value=1.5)
                width_in  = gr.Number(label="Width (cm)", value=15.0)
                height_in = gr.Number(label="Height (cm)", value=15.0)

                btn = gr.Button("🔍 Predict")

            with gr.Column(scale=1):
                pred_out = gr.Label(label="Prediction")

                prob_out = gr.BarPlot(
                    x=["Low", "Medium", "High"],
                    y="value",
                    title="Class Probability",
                    height=300
                )

        btn.click(
            fn=predict_brix,
            inputs=[model_dd, image_in, weight_in, width_in, height_in],
            outputs=[pred_out, prob_out]
        )

demo.launch()