import torch
import gradio as gr
import torch.nn as nn
import torch.nn.functional as F

from PIL import Image
import torchvision.transforms as transforms
from torchvision.models import resnet50


# ============================================================
# CONFIG
# ============================================================
MODEL_PATH = r"C:\Users\LOQ\Desktop\Oral Diseases Image Classification\checkpoints\best_model.pth"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ============================================================
# LOAD MODEL
# ============================================================
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)

CLASS_NAMES = checkpoint["class_names"]
TEST_F1 = checkpoint["test_f1"]

model = resnet50(weights=None)
model.fc = nn.Sequential(
    nn.Dropout(0.3),
    nn.Linear(model.fc.in_features, len(CLASS_NAMES))
)
model.load_state_dict(checkpoint["state_dict"])
model.to(DEVICE)
model.eval()

eval_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])


# ============================================================
# INFERENCE
# ============================================================
def predict(image):
    if image is None:
        return (
            gr.update(value="—", visible=True),
            "—",
            {},
            gr.update(visible=False),
        )

    image = image.convert("RGB")
    tensor = eval_transform(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        output = model(tensor)
        probs = F.softmax(output, dim=1)[0]

    index = torch.argmax(probs).item()
    prediction = CLASS_NAMES[index]
    confidence = probs[index].item() * 100

    results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}

    # Build a small verdict badge depending on confidence level
    if confidence >= 85:
        badge = f'<div class="verdict verdict-high">✓ High Confidence — {confidence:.1f}%</div>'
    elif confidence >= 60:
        badge = f'<div class="verdict verdict-mid">! Moderate Confidence — {confidence:.1f}%</div>'
    else:
        badge = f'<div class="verdict verdict-low">? Low Confidence — {confidence:.1f}%</div>'

    return (
        prediction,
        f"{confidence:.2f}%",
        results,
        gr.update(value=badge, visible=True),
    )


def clear_all():
    return None, "—", "—", {}, gr.update(visible=False)


# ============================================================
# STYLING
# ============================================================
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&family=JetBrains+Mono:wght@500&display=swap');

:root {
    --bg-primary: #0a0e1a;
    --bg-secondary: #10162a;
    --bg-card: #131a30;
    --border-subtle: #232b45;
    --accent: #14b8a6;
    --accent-soft: #14b8a622;
    --accent-2: #6366f1;
    --text-primary: #e8ecf5;
    --text-secondary: #8892b0;
    --text-muted: #5b6485;
    --radius: 14px;
}

* { font-family: 'Inter', sans-serif !important; }

.gradio-container {
    background: radial-gradient(circle at 10% 0%, #101a33 0%, #080b14 55%, #05070d 100%) !important;
    max-width: 1180px !important;
    margin: 0 auto !important;
}

footer { display: none !important; }

/* ---------- HEADER ---------- */
.app-header {
    padding: 30px 8px 22px 8px;
    border-bottom: 1px solid var(--border-subtle);
    margin-bottom: 26px;
    display: flex;
    align-items: center;
    justify-content: space-between;
}
.app-header .brand {
    display: flex;
    align-items: center;
    gap: 14px;
}
.app-header .logo-badge {
    width: 46px;
    height: 46px;
    border-radius: 12px;
    background: linear-gradient(135deg, var(--accent), var(--accent-2));
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 22px;
    box-shadow: 0 8px 24px -6px #14b8a655;
    flex-shrink: 0;
}
.app-header h1 {
    font-size: 20px;
    font-weight: 700;
    color: var(--text-primary);
    margin: 0;
    letter-spacing: -0.02em;
}
.app-header p {
    font-size: 13px;
    color: var(--text-muted);
    margin: 2px 0 0 0;
}
.app-header .tag {
    font-size: 11px;
    font-weight: 600;
    color: var(--accent);
    background: var(--accent-soft);
    border: 1px solid #14b8a640;
    padding: 6px 14px;
    border-radius: 999px;
    letter-spacing: 0.03em;
    text-transform: uppercase;
}

/* ---------- CARDS ---------- */
.card {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: var(--radius) !important;
    padding: 18px !important;
}
.card-title {
    font-size: 13px;
    font-weight: 600;
    color: var(--text-secondary);
    text-transform: uppercase;
    letter-spacing: 0.04em;
    margin-bottom: 12px;
    display: flex;
    align-items: center;
    gap: 8px;
}
.card-title::before {
    content: "";
    width: 4px;
    height: 14px;
    background: var(--accent);
    border-radius: 2px;
    display: inline-block;
}

/* ---------- UPLOAD ZONE ---------- */
.upload-zone, .upload-zone > div {
    background: var(--bg-card) !important;
    border: 1.5px dashed #2b3454 !important;
    border-radius: var(--radius) !important;
}
.upload-zone:hover {
    border-color: var(--accent) !important;
}

/* ---------- BUTTONS ---------- */
#analyze-btn {
    background: linear-gradient(135deg, #14b8a6, #0d9488) !important;
    color: #05170f !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 10px !important;
    box-shadow: 0 10px 24px -8px #14b8a670 !important;
    letter-spacing: 0.01em;
    transition: transform .15s ease, box-shadow .15s ease;
}
#analyze-btn:hover {
    transform: translateY(-1px);
    box-shadow: 0 14px 28px -8px #14b8a690 !important;
}
#clear-btn {
    background: transparent !important;
    color: var(--text-secondary) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 10px !important;
}
#clear-btn:hover {
    border-color: #3a4468 !important;
    color: var(--text-primary) !important;
}

/* ---------- RESULT FIELDS ---------- */
#pred-box textarea, #conf-box textarea {
    background: #0d1326 !important;
    border: 1px solid var(--border-subtle) !important;
    color: var(--text-primary) !important;
    font-weight: 700 !important;
    font-size: 17px !important;
    border-radius: 10px !important;
}
#conf-box textarea {
    color: var(--accent) !important;
    font-family: 'JetBrains Mono', monospace !important;
}
label span {
    color: var(--text-muted) !important;
    font-size: 11.5px !important;
    text-transform: uppercase;
    letter-spacing: 0.05em;
    font-weight: 600 !important;
}

/* ---------- VERDICT BADGE ---------- */
.verdict {
    padding: 10px 16px;
    border-radius: 10px;
    font-size: 13px;
    font-weight: 600;
    text-align: center;
    margin-bottom: 14px;
    border: 1px solid transparent;
}
.verdict-high { background: #14b8a61a; color: #2dd4bf; border-color: #14b8a640; }
.verdict-mid  { background: #f59e0b1a; color: #fbbf24; border-color: #f59e0b40; }
.verdict-low  { background: #ef44441a; color: #f87171; border-color: #ef444440; }

/* ---------- PROBABILITY BARS (gr.Label) ---------- */
.label-wrap {
    background: transparent !important;
    border: none !important;
}
#prob-label .container {
    background: transparent !important;
}

/* ---------- FOOTER ---------- */
.app-footer {
    margin-top: 30px;
    padding: 18px 4px 10px 4px;
    border-top: 1px solid var(--border-subtle);
    display: flex;
    justify-content: space-between;
    align-items: center;
    flex-wrap: wrap;
    gap: 10px;
}
.app-footer .meta {
    font-size: 12px;
    color: var(--text-muted);
    font-family: 'JetBrains Mono', monospace;
}
.app-footer .meta b { color: var(--text-secondary); }
.app-footer .credit {
    font-size: 12px;
    color: var(--text-muted);
}
.app-footer .credit b { color: var(--text-secondary); }

.disclaimer {
    font-size: 11.5px;
    color: var(--text-muted);
    background: #0d132666;
    border: 1px solid var(--border-subtle);
    border-radius: 10px;
    padding: 10px 14px;
    margin-top: 16px;
    line-height: 1.6;
}
"""


# ============================================================
# UI
# ============================================================
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate"),
    css=CUSTOM_CSS,
    title="Oral Disease Classifier"
) as demo:

    gr.HTML(
        f"""
        <div class="app-header">
            <div class="brand">
                <div class="logo-badge">🦷</div>
                <div>
                    <h1>Oral Disease Classification</h1>
                    <p>Computer-vision assisted screening · ResNet50 backbone</p>
                </div>
            </div>
            <div class="tag">Model F1 · {TEST_F1:.3f}</div>
        </div>
        """
    )

    with gr.Row(equal_height=True):

        with gr.Column(scale=5):
            gr.HTML('<div class="card-title">Input Image</div>')
            image_input = gr.Image(
                type="pil",
                label="",
                show_label=False,
                elem_classes="upload-zone",
                height=340,
            )

            with gr.Row():
                clear_btn = gr.Button("Clear", elem_id="clear-btn")
                button = gr.Button("Analyze Image", elem_id="analyze-btn")

            gr.HTML(
                """
                <div class="disclaimer">
                    ⚠ Decision-support tool only. Predictions are generated by an automated
                    model and are not a substitute for professional clinical diagnosis.
                </div>
                """
            )

        with gr.Column(scale=5):
            gr.HTML('<div class="card-title">Analysis Result</div>')

            verdict_html = gr.HTML(visible=False)

            with gr.Row():
                prediction = gr.Textbox(label="Predicted Class", elem_id="pred-box", interactive=False)
                confidence = gr.Textbox(label="Confidence", elem_id="conf-box", interactive=False)

            gr.HTML('<div class="card-title" style="margin-top:6px;">Class Probability Distribution</div>')
            probabilities = gr.Label(
                label="",
                show_label=False,
                elem_id="prob-label",
                num_top_classes=len(CLASS_NAMES),
            )

    gr.HTML(
        f"""
        <div class="app-footer">
            <div class="meta">ARCH · <b>ResNet50</b> &nbsp;|&nbsp; DEVICE · <b>{DEVICE.upper()}</b> &nbsp;|&nbsp; CLASSES · <b>{len(CLASS_NAMES)}</b></div>
            <div class="credit">Built by <b>Nasr Mohamed</b> — AI Engineer &nbsp;·&nbsp; © 2026</div>
        </div>
        """
    )

    button.click(
        predict,
        inputs=image_input,
        outputs=[prediction, confidence, probabilities, verdict_html],
    )

    clear_btn.click(
        clear_all,
        inputs=None,
        outputs=[image_input, prediction, confidence, probabilities, verdict_html],
    )


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