"""
DeepNAPSI – automated NAPSI scoring from a hand photo.

Gradio 6.x Blocks UI.  Model: BEiT-base-384 quantised to INT8 ONNX.
"""

from __future__ import annotations

import numpy as np
import gradio as gr

from backend import Backend, FINGER_NAMES

# ---------------------------------------------------------------------------
# Load model once at startup
# ---------------------------------------------------------------------------

backend = Backend()

NAPSI_DESC = (
    "**NAPSI (Nail Psoriasis Severity Index)** scores each nail 0–4 based "
    "on nail-bed and nail-matrix changes. The total score across all 10 nails "
    "ranges from 0 (no disease) to 40 (maximum disease)."
)

DISCLAIMER = (
    "⚠️ **Not a medical product.** This tool is for research purposes only and "
    "must not be used for patient diagnosis or treatment decisions."
)

# ---------------------------------------------------------------------------
# Predict function
# ---------------------------------------------------------------------------


def predict(image: np.ndarray):
    """Called by Gradio on button click or example selection."""
    if image is None:
        empty = np.zeros((64, 64, 3), dtype=np.uint8)
        return (
            *([empty, "–"] * 5),  # 5× (nail crop, score label)
            "–",  # total
            "Please upload an image.",
        )

    result = backend.predict(image)

    outputs = []
    for nail_img, score in zip(result["nails"], result["napsi_scores"]):
        outputs.append(nail_img)
        outputs.append(str(score) if score >= 0 else "–")

    total = result["napsi_sum"]
    outputs.append(str(total) if total >= 0 else "–")
    outputs.append(result["error"] or "")
    return tuple(outputs)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(
    title="DeepNAPSI",
    analytics_enabled=False,
) as demo:
    gr.Markdown("# 🤚 DeepNAPSI")
    gr.Markdown(NAPSI_DESC)
    gr.Markdown(DISCLAIMER)

    # ── Top row: input only (annotated hand preview removed) ──────────────
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Hand photo",
                type="numpy",
                sources=["upload", "clipboard", "webcam"],
            )
            predict_btn = gr.Button("Predict NAPSI", variant="primary")

            gr.Examples(
                examples=[
                    ["assets/example_1.jpg"],
                ],
                inputs=image_input,
                label="Example images",
            )

    # ── Bottom rows: nail crops + scores ──────────────────────────────────
    with gr.Row():
        nail_images = [
            gr.Image(label=f, type="numpy", interactive=False, height=160)
            for f in FINGER_NAMES
        ]
    with gr.Row():
        nail_scores = [
            gr.Textbox(label=f"NAPSI {f}", interactive=False) for f in FINGER_NAMES
        ]

    with gr.Row():
        total_score = gr.Textbox(
            label="DeepNAPSI Total (one hand, 0–20)", interactive=False
        )

    error_box = gr.Textbox(label="Status", interactive=False, visible=True)

    # Wire outputs into a flat list matching predict() return order
    all_outputs = [x for pair in zip(nail_images, nail_scores) for x in pair] + [
        total_score,
        error_box,
    ]

    predict_btn.click(fn=predict, inputs=image_input, outputs=all_outputs)

# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        favicon_path="assets/favicon-32x32.png",
        theme=gr.themes.Soft(),
    )