import gradio as gr
from gradio_client import Client, handle_file
from PIL import Image
import tempfile
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# ==============================
# Clients
# ==============================
cnn_client = Client("suyagi/cnn-isic2019")
resnet_client = Client("suyagi/Resnet-isic")
eff_client = Client("suyagi/EfficientNet-isic2019")
svmrf_client = Client("suyagi/SVM-RF-isic2019")

# ==============================
# Helpers
# ==============================
def parse_dl_result(result, model_name):
    rows = []
    for cls, conf in result.get("confidences", {}).items():
        rows.append([model_name, cls, conf * 100])

    return (
        result.get("predicted_class", "Unknown"),
        round(result.get("confidence", 0.0) * 100, 2),
        rows
    )


def parse_ml_result(result, model_name):
    rows = []
    label = result.get("label", "Unknown")
    conf = 0.0

    for item in result.get("confidences", []):
        rows.append([model_name, item["label"], item["confidence"] * 100])
        if item["label"] == label:
            conf = item["confidence"] * 100

    return label, round(conf, 2), rows


def soft_voting(df):
    pivot = df.pivot_table(
        index="Class",
        values="Confidence (%)",
        aggfunc="mean"
    ).reset_index()

    best = pivot.loc[pivot["Confidence (%)"].idxmax()]
    return best["Class"], round(best["Confidence (%)"], 2), pivot


def entropy(probs):
    probs = np.array(probs, dtype=np.float64)
    probs /= probs.sum()
    return -np.sum(probs * np.log2(probs + 1e-12))


def plot_heatmap(df):
    pivot = df.pivot_table(
        index="Model",
        columns="Class",
        values="Confidence (%)"
    )

    fig, ax = plt.subplots(figsize=(10, 4))
    im = ax.imshow(pivot.values, aspect="auto")

    ax.set_xticks(range(len(pivot.columns)))
    ax.set_yticks(range(len(pivot.index)))
    ax.set_xticklabels(pivot.columns, rotation=45, ha="right")
    ax.set_yticklabels(pivot.index)

    for i in range(len(pivot.index)):
        for j in range(len(pivot.columns)):
            v = pivot.values[i, j]
            if not np.isnan(v):
                ax.text(j, i, f"{v:.1f}", ha="center", va="center", fontsize=8)

    ax.set_title("Model vs Class Confidence Heatmap")
    fig.colorbar(im, ax=ax)
    plt.tight_layout()
    return fig


# ==============================
# MAIN
# ==============================
def predict_all(image):
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    image.save(tmp.name)
    img = handle_file(tmp.name)

    summary, probs, ent = [], [], []

    for name, client in [
        ("CNN", cnn_client),
        ("ResNet", resnet_client),
        ("EfficientNet", eff_client),
    ]:
        raw = client.predict(img, api_name="/predict_image")
        label, conf, rows = parse_dl_result(raw, name)
        summary.append([name, label, conf])
        probs.extend(rows)
        ent.append([name, entropy([r[2]/100 for r in rows])])

    svm, rf = svmrf_client.predict(img, api_name="/predict_both")

    for name, raw in [("SVM", svm), ("Random Forest", rf)]:
        label, conf, rows = parse_ml_result(raw, name)
        summary.append([name, label, conf])
        probs.extend(rows)
        ent.append([name, entropy([r[2]/100 for r in rows])])

    df_summary = pd.DataFrame(summary, columns=["Model", "Prediction", "Confidence (%)"])
    df_probs = pd.DataFrame(probs, columns=["Model", "Class", "Confidence (%)"])
    df_entropy = pd.DataFrame(ent, columns=["Model", "Entropy"]).sort_values("Entropy")

    best = df_summary.loc[df_summary["Confidence (%)"].idxmax()]
    ens_cls, ens_conf, ens_tbl = soft_voting(df_probs)

    warning = ""
    if df_summary["Prediction"].nunique() > 1:
        warning = "⚠️ **Model disagreement detected — manual review recommended**"

    md = f"""
### 🏆 Final Result
- **Best Single Model:** `{best['Model']}` → **{best['Prediction']} ({best['Confidence (%)']}%)**
- **Soft Voting:** **{ens_cls} ({ens_conf}%)**
- **Most Confident Model (Lowest Entropy):** `{df_entropy.iloc[0]['Model']}`

{warning}
"""

    return df_summary, df_probs, plot_heatmap(df_probs), ens_tbl, df_entropy, md


# ==============================
# UI (CLEAN & SIMPLE)
# ==============================
with gr.Blocks(title="Skin Disease Multi-Model Dashboard") as demo:

    gr.Markdown("## 🧬 Skin Disease Multi-Model Dashboard (ISIC 2019)")

    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(type="pil", label="Upload Skin Image")
            submit = gr.Button("🔍 Analyze")

        with gr.Column(scale=1):
            result_md = gr.Markdown()

    gr.Markdown("### 📊 Model Summary")
    out_summary = gr.Dataframe()

    with gr.Tabs():
        with gr.Tab("📈 Class Probabilities"):
            out_probs = gr.Dataframe()

        with gr.Tab("🔥 Heatmap"):
            out_heatmap = gr.Plot()

        with gr.Tab("🧠 Soft Voting"):
            out_vote = gr.Dataframe()

        with gr.Tab("📉 Uncertainty"):
            out_entropy = gr.Dataframe()

    submit.click(
        predict_all,
        inputs=img_input,
        outputs=[
            out_summary,
            out_probs,
            out_heatmap,
            out_vote,
            out_entropy,
            result_md
        ]
    )

demo.launch()