#!/usr/bin/env python3
"""
Architecture Building Image Classifier — Gradio Space UI
Model : EfficientNetV2-S + Conv2D Head + GeM Pooling + Focal loss + SWA
Theme : Greyscope Labs - Green (#009925) inspired palette
Audit : v6 ipynb — 8 classes, parent class fix, flip-consistency, entropy certainty
"""

import json
import os
import io
import time
import base64

import numpy as np
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download

import tensorflow as tf

try:
    from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
except (ImportError, ModuleNotFoundError):
    from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras.layers import Layer

# Compatibility import — tf.keras.saving not available in all environments
try:
    from tensorflow.keras.saving import register_keras_serializable
except (ImportError, AttributeError):
    try:
        from keras.saving import register_keras_serializable
    except (ImportError, AttributeError):
        def register_keras_serializable(package=None):
            def decorator(cls):
                return cls
            return decorator

# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------
REPO_ID     = os.environ.get("REPO_ID",    "0xgr3y/Arch-Building-Image-Classification")
MODEL_FILE  = os.environ.get("MODEL_FILE", "tflite/model.tflite")
IMG_SIZE    = (320, 320)
HF_TOKEN    = os.environ.get("HF_TOKEN", "")

# 8 classes — alphabetical order matching flow_from_directory
LABELS = ["barn", "bridge", "castle", "mosque", "skyscraper", "stadium", "temple", "windmill"]

LOGO_LOCAL_FILE = "greyscope-labs-architecture-classification-transferlearning-efficientnetv2.jpg"

# --------------------------------------------------------------------------
# Dynamic label loading — reads label_mapping.json from the HF repo
# --------------------------------------------------------------------------
LABELS_FALLBACK = ["barn", "bridge", "castle", "mosque", "skyscraper", "stadium", "temple", "windmill"]

def _load_labels() -> list:
    """Try to load labels from label_mapping.json in the HF repo.
    Falls back to LABELS_FALLBACK (8-class training order)."""
    try:
        _hf_token = os.environ.get("HF_TOKEN", "") or None
        lm_path = hf_hub_download(
            repo_id=REPO_ID, filename="label_mapping.json",
            local_dir="/tmp/model", token=_hf_token,
        )
        import json as _json
        with open(lm_path) as f:
            mapping = _json.load(f)
        lbs = mapping.get("labels", [])
        if lbs:
            print(f"[INFO] Labels loaded from label_mapping.json: {lbs}")
            return lbs
    except Exception as e:
        print(f"[WARN] label_mapping.json not found or unreadable ({e}), using fallback labels")
    print(f"[INFO] Using fallback labels: {LABELS_FALLBACK}")
    return LABELS_FALLBACK


def _load_json_artifact(filename: str) -> dict:
    """Generic loader for evaluation artifacts exported by the
    training notebook (model_benchmark.json, confusion_pairs.json,
    class_confidence_stats.json, calibration_data.json).
    Returns empty dict if file isn't present in the HF repo —
    every downstream use must .get() with fallback."""
    try:
        _path = hf_hub_download(
            repo_id=REPO_ID, filename=filename,
            local_dir="/tmp/model",
            token=HF_TOKEN if HF_TOKEN else None,
        )
        with open(_path) as _f:
            _data = json.load(_f)
        print(f"[INFO] Loaded {filename}")
        return _data
    except Exception as _e:
        print(f"[WARN] {filename} not found ({_e}), related features will use fallbacks")
        return {}

GREEN_SCALE = ["#006619", "#007A1E", "#009925", "#1BA83A", "#4DBD5E", "#7DD18E", "#A8E0AE", "#C8EED0"]

# Short description shown for the predicted class
CLASS_INFO = {
    "barn":       "A traditional architectural style, a wooden building used for storage, or sheltering livestock.",
    "bridge":     "A structure built to span a gap — such as a valley, road, or body of water — connecting two points for passage.",
    "castle":     "A fortified residential structure, historically built for nobility and defense, often featuring towers and walls.",
    "mosque":     "An Islamic place of worship, often recognizable by domes, minarets, and ornamental geometric patterns.",
    "skyscraper": "A very tall, multi-storey building, typically a defining feature of dense modern urban skylines.",
    "stadium":    "A large, often open-air venue designed to host sporting events, concerts, and public gatherings.",
    "temple":     "A building dedicated to religious or spiritual practice, found across many cultures and architectural styles.",
    "windmill":   "A structure fitted with sails or vanes, historically used to convert wind energy into mechanical work for milling or pumping.",
}

# Hardcoded confusion-note fallbacks (used only when confusion_pairs.json hasn't been uploaded to the repo yet).
CONFUSION_NOTES_FALLBACK = {
    "barn":       None,
    "bridge":     None,
    "castle":     None,
    "mosque":     None,
    "skyscraper": None,
    "stadium":    None,
    "temple":     None,
    "windmill":   None,
}

CONFUSION_NOTES = CONFUSION_NOTES_FALLBACK
CONFUSION_DETAIL = {}


# --------------------------------------------------------------------------
# Logo helper
# --------------------------------------------------------------------------
def _resolve_logo_src() -> str:
    """Embed local logo as data-URI so it always renders."""
    if os.path.exists(LOGO_LOCAL_FILE):
        try:
            with open(LOGO_LOCAL_FILE, "rb") as f:
                encoded = base64.b64encode(f.read()).decode("utf-8")
            ext = os.path.splitext(LOGO_LOCAL_FILE)[1].lstrip(".").lower() or "png"
            return f"data:image/{ext};base64,{encoded}"
        except Exception:
            pass
    return ""


LOGO_SRC = _resolve_logo_src()


# --------------------------------------------------------------------------
# Custom objects — MUST match training code exactly for deserialization
# --------------------------------------------------------------------------
@register_keras_serializable(package='ArchClassifier')
class GeMPooling(Layer):
    """Generalized Mean Pooling (Radenovic et al., CVPR 2018)."""
    def __init__(self, p=3.0, eps=1e-6, **kwargs):
        super().__init__(**kwargs)
        self.p_init = p
        self.eps    = eps

    def build(self, input_shape):
        self.p = self.add_weight(
            name="gem_p", shape=(), dtype=tf.float32,
            initializer=tf.keras.initializers.Constant(self.p_init), trainable=True)
        super().build(input_shape)

    def call(self, x):
        x = tf.maximum(x, self.eps)
        x = tf.pow(x, self.p)
        x = tf.reduce_mean(x, axis=[1, 2], keepdims=False)
        x = tf.pow(x, 1.0 / self.p)
        return x

    def get_config(self):
        cfg = super().get_config()
        cfg.update({"p": self.p_init, "eps": self.eps})
        return cfg


@register_keras_serializable(package='ArchClassifier')
class FocalLoss(tf.keras.losses.Loss):
    """Focal Loss (Lin et al., ICCV 2017), gamma=2.0."""
    def __init__(self, gamma=2.0, alpha=None, label_smoothing=0.0, **kwargs):
        super().__init__(**kwargs)
        self.gamma           = gamma
        self.alpha           = alpha
        self.label_smoothing = label_smoothing

    def call(self, y_true, y_pred):
        y_pred = tf.clip_by_value(y_pred, 1e-7, 1.0 - 1e-7)
        if self.label_smoothing > 0:
            y_true = (y_true * (1.0 - self.label_smoothing)
                      + self.label_smoothing / tf.cast(tf.shape(y_true)[-1], tf.float32))
        ce     = -y_true * tf.math.log(y_pred)
        weight = tf.pow(1.0 - y_pred, self.gamma)
        fl     = weight * ce
        if self.alpha is not None:
            fl = y_true * self.alpha * fl
        return tf.reduce_mean(tf.reduce_sum(fl, axis=-1))

    def get_config(self):
        cfg = super().get_config()
        cfg.update({"gamma": self.gamma, "alpha": self.alpha,
                    "label_smoothing": self.label_smoothing})
        return cfg


@register_keras_serializable(package='ArchClassifier')
class DiscriminativeAdamW(tf.keras.optimizers.AdamW):
    """AdamW with per-variable LR scaling via update_step override.

    Inherits from AdamW (matching training code) — NOT from the base
    Optimizer class, which would be a class-hierarchy mismatch on deserialization.
    compile=False is used at inference so optimizer weights are not loaded, but
    the class signature must still match for correct custom_objects lookup.
    """
    def __init__(self, lr_multipliers=None, backbone_layer_idx=0, **kwargs):
        super().__init__(**kwargs)
        self.lr_multipliers    = lr_multipliers or {}
        self.backbone_layer_idx = backbone_layer_idx
        self._var_mult_cache   = {}

    def _build_var_cache(self, model):
        """Cache {id(var): multiplier} keyed on layer.name (stable post load_model)."""
        self._var_mult_cache = {}
        base_model = next((l for l in model.layers if isinstance(l, tf.keras.Model)), None)
        if base_model is None:
            base_model = model.layers[self.backbone_layer_idx]
        for layer in base_model.layers:
            mult = 1.0
            for pattern, m in self.lr_multipliers.items():
                if pattern in layer.name:
                    mult = m
                    break
            for var in layer.trainable_variables:
                self._var_mult_cache[id(var)] = mult

    def _get_multiplier(self, var):
        return self._var_mult_cache.get(id(var), 1.0)

    def update_step(self, gradient, variable, learning_rate):
        """Scale learning_rate per-variable — truly discriminative."""
        mult = self._get_multiplier(variable)
        effective_lr = learning_rate * mult
        return super().update_step(gradient, variable, effective_lr)

    def get_config(self):
        cfg = super().get_config()
        cfg.update({"lr_multipliers": self.lr_multipliers,
                    "backbone_layer_idx": self.backbone_layer_idx})
        return cfg


# --------------------------------------------------------------------------
# Model loading
# --------------------------------------------------------------------------

print("[INFO] Downloading TFLite model from HuggingFace Hub…")

model_path = hf_hub_download(
    repo_id=REPO_ID, filename=MODEL_FILE, local_dir="/tmp/model",
    token=HF_TOKEN if HF_TOKEN else None,
)
print(f"[INFO] Model path: {model_path}")

# Labels resolved dynamically from label_mapping.json (uploaded alongside model)
LABELS = _load_labels()

print("[INFO] Loading TFLite interpreter …")
_t0 = time.time()
interpreter = tf.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()
_input_details  = interpreter.get_input_details()
_output_details = interpreter.get_output_details()
print(f"[INFO] TFLite interpreter ready in {time.time() - _t0:.1f}s")
print(f"[INFO]   Input:  shape={_input_details[0]['shape']}, dtype={_input_details[0]['dtype']}")
print(f"[INFO]   Output: shape={_output_details[0]['shape']}, dtype={_output_details[0]['dtype']}")


def _tflite_predict(arr: np.ndarray) -> np.ndarray:
    """Run a single forward pass via the TFLite interpreter.

    Applies Temperature Scaling (Guo et al., ICML 2017) post-hoc if
    CALIBRATION_TEMP != 1.0. TFLite outputs softmax probabilities; we
    recover logits via log(probs) and re-softmax with T scaling:
        softmax(log(probs) / T) == softmax(logits / T)
    (shift invariance of softmax — constant offset cancels).
    """
    interpreter.set_tensor(_input_details[0]["index"], arr)
    interpreter.invoke()
    probs = interpreter.get_tensor(_output_details[0]["index"])[0]
    if CALIBRATION_TEMP != 1.0:
        log_probs = np.log(np.clip(probs, 1e-7, 1.0))
        scaled = log_probs / CALIBRATION_TEMP
        e = np.exp(scaled - np.max(scaled))
        probs = e / e.sum()
    return probs

# ---------------------------------------------------------------------------
# Evaluation artifacts — exported by the training notebook alongside
# the model. All four are optional: Space still works (with fallbacks) if
# a given file isn't present in the repo yet.
# ---------------------------------------------------------------------------
MODEL_BENCHMARK         = _load_json_artifact("model_benchmark.json")
CONFUSION_PAIRS_DATA    = _load_json_artifact("confusion_pairs.json")
CLASS_CONFIDENCE_STATS  = _load_json_artifact("class_confidence_stats.json")
CALIBRATION_DATA        = _load_json_artifact("calibration_data.json")

# Temperature Scaling parameter (Guo et al., ICML 2017)
CALIBRATION_TEMP = float(CALIBRATION_DATA.get("temperature", 1.0))

# Prefer empirically-measured confusion notes from confusion_pairs.json
CONFUSION_NOTES = CONFUSION_PAIRS_DATA.get("confusion_notes") or CONFUSION_NOTES_FALLBACK
CONFUSION_DETAIL = CONFUSION_PAIRS_DATA.get("confusion_notes_detailed", {})


# --------------------------------------------------------------------------
# Inference helpers
# --------------------------------------------------------------------------
def _preprocess_arr(pil_img: Image.Image) -> np.ndarray:
    """Resize to 320×320 and apply EfficientNetV2-S preprocessing — matches training code exactly."""
    arr = np.array(pil_img.resize(IMG_SIZE), dtype=np.float32)
    return np.expand_dims(preprocess_input(arr), axis=0)


def _shannon_certainty(probs: np.ndarray) -> float:
    """Normalised certainty derived from Shannon entropy.

    certainty = 1 - H / H_max  ∈ [0, 1]
    H_max = log(n_classes) nats — uniform distribution.
    Returns a percentage 0–100 (100 = completely certain, 0 = uniform).
    """
    n   = len(probs)
    H   = -float(np.sum(probs * np.log(np.clip(probs, 1e-10, 1.0))))
    H_max = np.log(n)
    return max(0.0, (1.0 - H / H_max)) * 100


def _flip_predict(pil_img: Image.Image) -> tuple:
    """Single extra forward pass on the horizontally-flipped image.

    Directly mirrors TTA variant 2:
        tf.image.resize(tf.image.flip_left_right(image), (IMG_H, IMG_W))
    Returns (flip_label, flip_conf) for the flipped prediction.
    """
    flipped = pil_img.transpose(Image.FLIP_LEFT_RIGHT)
    arr     = _preprocess_arr(flipped)
    p       = _tflite_predict(arr)
    idx     = int(np.argmax(p))
    return LABELS[idx], float(p[idx]) * 100


# --------------------------------------------------------------------------
# CSS — Greyscope Labs - Color Green
# --------------------------------------------------------------------------
CUSTOM_CSS = """
:root {
  --bg:          #F4F8F5;
  --surface:     #FFFFFF;
  --ink:         #1A1F1C;
  --muted:       #5F6B63;
  --faint:       #A3B0A8;
  --accent:      #009925;
  --accent-dark: #007A1E;
  --accent-soft: rgba(0,153,37,.08);
  --border:      #D5DDD9;
  --sh: 0 1px 2px rgba(26,31,28,.06), 0 6px 18px rgba(26,31,28,.05);
  --font: 'Google Sans','Product Sans',-apple-system,'Segoe UI',Roboto,Helvetica,Arial,sans-serif;
}

.gradio-container {
  max-width: 960px !important; margin: 0 auto !important;
  background: var(--bg) !important; font-family: var(--font) !important;
  color: var(--ink) !important;
}

footer { display: none !important; }

/* ── Header ──────────────────────────────── */
.hdr-wrap { text-align: center; padding: 8px 8px 26px; }

.logo-banner {
  display: block; width: 100%; aspect-ratio: 2750/1530; object-fit: cover;
  margin: 0 auto 22px; border-radius: 14px; border: 1px solid var(--border);
  box-shadow: var(--sh);
}

.hdr-wrap h1 {
  font-family: var(--font); font-size: 28px; font-weight: 700;
  color: var(--ink); letter-spacing: -.01em; line-height: 1.25; margin: 0;
}

.hdr-sub { margin-top: 9px; font-size: 13px; color: var(--muted); font-weight: 400; }

.hdr-divider {
  width: 44px; height: 2px; background: var(--accent);
  margin: 18px auto; border-radius: 1px;
}

.hdr-badges { display: flex; justify-content: center; gap: 9px; flex-wrap: wrap; }

.hdr-badge {
  font-size: 11px; font-weight: 500; color: var(--accent-dark);
  background: var(--accent-soft); border: 1px solid rgba(0,153,37,.25);
  padding: 5px 13px; border-radius: 20px;
}

.hdr-note { margin-top: 12px; font-size: 11px; color: var(--faint); letter-spacing: .02em; }

/* ── Section label ───────────────────────── */
.section-label {
  font-size: 11px; font-weight: 700; text-transform: uppercase;
  letter-spacing: .14em; color: var(--muted); margin: 4px 0 12px 4px;
}

/* ── Cards ───────────────────────────────── */
.grn-card {
  background: var(--surface) !important; border: 1px solid var(--border) !important;
  border-radius: 14px !important; box-shadow: var(--sh) !important; padding: 18px !important;
}

/* ── Upload area ─────────────────────────── */
.grn-upload .image-frame,
.grn-upload [data-testid="image"],
.grn-upload .upload-container {
  border-radius: 10px !important; border: 1.5px dashed var(--border) !important;
  background: var(--bg) !important;
}
.upload-caption {
  margin-top: 10px; font-size: 11px; color: var(--faint);
  text-align: center; letter-spacing: .02em;
}

/* ── Button ──────────────────────────────── */
.grn-btn, .grn-btn button {
  background: var(--accent) !important; border: none !important;
  color: #FFFFFF !important; font-family: var(--font) !important;
  font-weight: 600 !important; font-size: 14px !important;
  border-radius: 10px !important; box-shadow: 0 4px 14px rgba(0,153,37,.28) !important;
}
.grn-btn:hover button, .grn-btn button:hover {
  background: var(--accent-dark) !important;
}

/* ── Result card ─────────────────────────── */
.result-box {
  display: flex; align-items: center; gap: 16px; padding: 14px;
  background: var(--bg); border: 1px solid var(--border);
  border-left: 4px solid var(--accent); border-radius: 10px; margin-bottom: 18px;
}
.result-thumb-wrap { display: flex; flex-direction: column; align-items: center; gap: 5px; flex-shrink: 0; }
.result-thumb {
  width: 104px; height: 104px; border-radius: 8px;
  object-fit: cover; border: 1px solid var(--border);
}
.thumb-label {
  font-size: 9px; font-weight: 500; color: var(--faint);
  text-align: center; line-height: 1.4; letter-spacing: .03em; text-transform: uppercase;
}
.result-info { flex: 1; min-width: 0; }
.res-name { font-size: 24px; font-weight: 700; color: var(--ink); line-height: 1.15; }
.res-meta {
  font-size: 12.5px; color: var(--muted); margin-top: 5px;
  line-height: 1.6; display: flex; flex-direction: column; gap: 1px;
}
.res-meta strong { color: var(--accent-dark); }

.confidence-badge {
  display: inline-block; margin-top: 8px; font-size: 10.5px; font-weight: 600;
  letter-spacing: .03em; padding: 3px 11px; border-radius: 20px;
  border: 1px solid rgba(0,153,37,.25); background: var(--accent-soft); color: var(--accent-dark);
}
.confidence-badge.medium { border-color: var(--border); background: #E6F0E9; color: #3D7A50; }
.confidence-badge.low    { border-color: var(--border); background: #EEF1EF; color: var(--muted); }

/* ── Class description ───────────────────── */
.class-desc {
  margin-top: 14px; padding-top: 12px; border-top: 1px solid var(--border);
  font-size: 12.5px; color: var(--muted); line-height: 1.65;
}
.class-desc strong { color: var(--ink); }

/* ── Probability bars ────────────────────── */
.plist { display: flex; flex-direction: column; gap: 10px; }
.prow  { display: flex; align-items: center; gap: 10px; }
.plabel { font-size: 12.5px; color: var(--ink); width: 92px; flex-shrink: 0; }
.pbg { flex: 1; height: 9px; background: #E2E9E4; border-radius: 5px; overflow: hidden; }
.pfill { height: 100%; border-radius: 5px; }
.ppct  { font-size: 12px; width: 46px; text-align: right; flex-shrink: 0; }

/* ── Live Prediction Analysis panel ─────── */
.lpa-panel {
  margin-top: 16px;
  padding: 18px 20px;
  background: var(--bg);
  border: 1px solid var(--border);
  border-radius: 10px;
}

.lpa-title {
  font-size: 10px; font-weight: 700; text-transform: uppercase;
  letter-spacing: .14em; color: var(--muted); margin-bottom: 14px;
  text-align: center;
}

.lpa-grid { display: grid; grid-template-columns: 1fr 1fr 1fr 1fr; gap: 10px; justify-items: center; }

.lpa-metric { display: flex; flex-direction: column; gap: 4px; align-items: center; text-align: center; }

.lpa-label {
  font-size: 9.5px; font-weight: 600; text-transform: uppercase;
  letter-spacing: .1em; color: var(--faint);
}

.lpa-value { font-size: 18px; font-weight: 700; color: var(--ink); line-height: 1; }

.lpa-value.good   { color: #007A1E; }
.lpa-value.warn   { color: #C47D05; }
.lpa-value.bad    { color: #C0392B; }

.lpa-bar-bg {
  height: 5px; background: #E2E9E4;
  border-radius: 3px; overflow: hidden; margin-top: 2px; width: 100%;
}
.lpa-bar-fill { height: 100%; border-radius: 3px; }

.lpa-sublabel { font-size: 10px; color: var(--faint); margin-top: 2px; }

.lpa-divider {
  border: none; border-top: 1px solid var(--border);
  margin: 12px 0 10px;
}

.lpa-insight {
  font-size: 12px; color: var(--muted); line-height: 1.6; text-align: center;
}
.lpa-insight strong { color: var(--ink); }
.lpa-insight.ok   { color: #007A1E; }
.lpa-insight.warn { color: #C47D05; }
.lpa-insight.bad  { color: #C0392B; }

/* ── Placeholder ─────────────────────────── */
.result-placeholder {
  min-height: 320px; display: flex; flex-direction: column;
  align-items: center; justify-content: center; gap: 14px;
  border: 1.5px dashed var(--border); border-radius: 10px;
  background: var(--bg); text-align: center; padding: 24px;
}
.result-placeholder p { font-size: 13px; color: var(--muted); line-height: 1.7; margin: 0; }
.result-placeholder strong { color: var(--accent-dark); }

/* ── Footer ──────────────────────────────── */
.grn-footer {
  margin-top: 28px; padding-top: 20px; border-top: 1px solid var(--border);
  text-align: center; font-size: 11.5px; color: var(--faint); line-height: 2.1;
}
.grn-footer a { color: var(--muted); text-decoration: none; }
.grn-footer a:hover { color: var(--accent); }

@media (max-width: 760px) {
  .result-box { flex-direction: column; text-align: center; }
  .lpa-grid   { grid-template-columns: 1fr 1fr; }
}
"""


# --------------------------------------------------------------------------
# Dynamic HTML fragments (metrics badges reflect real test-set data
# when model_benchmark.json is present, otherwise generic badges)
# --------------------------------------------------------------------------
def _build_header_html() -> str:
    metrics = MODEL_BENCHMARK.get("metrics", {})
    if metrics and metrics.get("test_accuracy", 0) > 0:
        ta = metrics.get("test_accuracy", 0) * 100
        tta = metrics.get("tta_accuracy", 0) * 100
        t3 = metrics.get("top3_accuracy", 0) * 100
        auc = metrics.get("macro_auc", 0)
        auc_str = f" | AUC {auc:.4f}" if auc else ""
        badges = (f'<span class="hdr-badge">Test Accuracy {ta:.1f}%</span>'
                  f'<span class="hdr-badge">TTA Accuracy {tta:.1f}%</span>'
                  f'<span class="hdr-badge">Top-3 Accuracy {t3:.1f}%</span>'
                  f'<span class="hdr-badge">Macro AUC {auc:.4f}</span>')
        note = f"Evaluated on held-out test set &middot; {len(LABELS)} architectural classes"
    else:
        badges = ('<span class="hdr-badge">8 Classes</span>'
                  '<span class="hdr-badge">EfficientNetV2-S + GeM + SWA</span>'
                  '<span class="hdr-badge">13,440 Images</span>')
        note = ("Barn &middot; Bridge &middot; Castle &middot; Mosque "
                "&middot; Skyscraper &middot; Stadium &middot; Temple &middot; Windmill")
    return (
        f'<div class="hdr-wrap">\n'
        f'  <img class="logo-banner" src="{LOGO_SRC}" alt="Architecture Building Image Classifier">\n'
        f'  <h1>Architecture Building<br>Image Classifier</h1>\n'
        f'  <div class="hdr-sub">Fine-Grained World Architecture Image Classification</div>\n'
        f'  <div class="hdr-divider"></div>\n'
        f'  <div class="hdr-badges">{badges}</div>\n'
        f'  <div class="hdr-note">{note}</div>\n'
        f'</div>'
    )

HEADER_HTML = _build_header_html()

def _build_footer_html() -> str:
    params = MODEL_BENCHMARK.get("total_params", 0)
    param_str = f"| {params/1e6:.1f}M Parameters" if params else ""
    cal_temp = CALIBRATION_DATA.get("temperature", 1.0)
    cal_note = ""
    if cal_temp != 1.0:
        cal_note = f'\n  &middot; <strong>Temperature Scaling T={cal_temp:.4f} applied at inference</strong> (raw ECE 18.13% → 0.95%)'
    return (
        '<div class="grn-footer">\n'
        '  Model: <a href="https://huggingface.co/0xgr3y/Arch-Building-Image-Classification" target="_blank">0xgr3y/Arch-Building-Image-Classification</a>\n'
        '  &middot; Dataset: <a href="https://huggingface.co/datasets/0xgr3y/arch-building-dataset" target="_blank">0xgr3y/arch-building-dataset</a>\n'
        '  &middot; GitHub: <a href="https://github.com/arcxteam/building-architectural-image-classifier" target="_blank">arcxteam/building-architectural-image-classifier</a><br>\n'
        f'  Built by <a href="https://huggingface.co/0xgr3y" target="_blank">Saugani</a>\n'
        f'  &middot; EfficientNetV2-S | GeM Pooling | Focal Loss | SWA\n'
        f'{" &middot; " + param_str if param_str else ""}\n'
        '  &middot; Apache-2.0 License\n'
        f'{cal_note}\n'
        '</div>'
    )


FOOTER_HTML = _build_footer_html()

PLACEHOLDER_HTML = """
<div class="result-placeholder">
  <svg width="40" height="40" viewBox="0 0 40 40" fill="none">
    <rect x="3" y="3" width="34" height="34" rx="7"
          stroke="#009925" stroke-width="1.4" stroke-dasharray="3 3" opacity="0.5"/>
    <path d="M13 27L17 21L21 25L25 18L29 24"
          stroke="#009925" stroke-width="1.6" stroke-linecap="round" stroke-linejoin="round" opacity="0.6"/>
    <circle cx="15" cy="14" r="2" fill="#009925" opacity="0.5"/>
  </svg>
  <p>Upload a building photo on the left, then press<br><strong>Classify Building</strong> to see the result here.</p>
</div>
"""


# --------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------
def classify(image):
    if image is None:
        return ('<div class="result-placeholder">'
                '<p>&#9888; Please upload a building photo first.</p>'
                '</div>')

    t_start = time.time()

    pil_img = Image.fromarray(image).convert("RGB")
    resized  = pil_img.resize(IMG_SIZE)
    arr      = _preprocess_arr(pil_img)
    probs    = _tflite_predict(arr)

    inference_ms = (time.time() - t_start) * 1000

    sorted_idx  = np.argsort(probs)[::-1]
    pred_idx    = int(sorted_idx[0])
    pred_label  = LABELS[pred_idx]
    pred_display = pred_label.title()
    pred_conf   = float(probs[pred_idx]) * 100

    buf = io.BytesIO()
    resized.save(buf, format="JPEG", quality=90)
    thumb_src = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("utf-8")

    # Confidence level badge
    if pred_conf >= 80:
        conf_class, conf_text = "high",   "High Confidence"
    elif pred_conf >= 50:
        conf_class, conf_text = "medium", "Medium Confidence"
    else:
        conf_class, conf_text = "low",    "Low Confidence"

    # Probability bars
    n_labels = len(LABELS)
    green    = GREEN_SCALE + ["#D8EDDB"] * max(0, n_labels - len(GREEN_SCALE))
    rows     = []
    for rank, idx in enumerate(sorted_idx):
        idx   = int(idx)
        pct   = float(probs[idx]) * 100
        color = green[rank]
        bold  = "700" if rank == 0 else "400"
        pcol  = "var(--accent-dark)" if rank == 0 else "var(--muted)"
        rows.append(
            f'<div class="prow">'
            f'<span class="plabel" style="font-weight:{bold}">{LABELS[idx].title()}</span>'
            f'<div class="pbg"><div class="pfill" style="width:{pct:.1f}%;background:{color}"></div></div>'
            f'<span class="ppct" style="color:{pcol};font-weight:{bold}">{pct:.1f}%</span>'
            f'</div>'
        )

    # Live Prediction AnalysiS

    # 1. Shannon-entropy certainty score
    certainty_pct = _shannon_certainty(probs)
    cert_color    = ("#007A1E" if certainty_pct >= 80
                     else "#C47D05" if certainty_pct >= 50
                     else "#C0392B")
    cert_class    = ("good" if certainty_pct >= 80
                     else "warn" if certainty_pct >= 50
                     else "bad")
    cert_bar_pct  = f"{certainty_pct:.1f}"

    # 2. Decision margin (top-1 vs runner-up probability)
    runner_up_idx   = int(sorted_idx[1])
    runner_up_label = LABELS[runner_up_idx]
    runner_up_display = runner_up_label.title()
    margin_pts      = (float(probs[pred_idx]) - float(probs[runner_up_idx])) * 100
    margin_class    = ("good" if margin_pts >= 20 else "warn" if margin_pts >= 8 else "bad")

    # 3. Flip consistency — extra forward pass on horizontally-flipped image
    t_flip = time.time()
    flip_label, flip_conf = _flip_predict(pil_img)
    flip_ms = (time.time() - t_flip) * 1000

    flip_match   = (flip_label == pred_label)
    flip_icon    = "&#10003;" if flip_match else "&#10007;"
    flip_display = flip_label.title()
    flip_cls     = "ok" if flip_match else "bad"
    flip_txt     = (f"Stable — flip also predicts <strong>{flip_display}</strong> ({flip_conf:.0f}%)"
                    if flip_match
                    else f"Disagrees — flip predicts <strong>{flip_display}</strong> ({flip_conf:.0f}%)")

    # 4. Context-aware insight note (margin + known confusion pairs)
    confused_with = CONFUSION_NOTES.get(pred_label)
    if margin_pts >= 20:
        insight_icon = "&#10003;"
        insight_txt  = (f"Clear margin over <strong>{runner_up_display}</strong> "
                        f"— model is confident about this image.")
        insight_cls  = "ok"
    elif confused_with and confused_with == runner_up_label:
        insight_icon = "&#9888;"
        insight_txt  = (f"Close call: <strong>{pred_display}</strong> vs "
                        f"<strong>{runner_up_display}</strong> — a pair known to be "
                        f"confused in evaluation.")
        insight_cls  = "warn"
    else:
        insight_icon = "&#9888;"
        insight_txt  = (f"Narrow margin over <strong>{runner_up_display}</strong> "
                        f"— image may have ambiguous features.")
        insight_cls  = "warn"

    class_stats = CLASS_CONFIDENCE_STATS.get(pred_label, {})
    if class_stats:
        _bl_mean = class_stats.get("mean_confidence", 0) * 100
        _bl_p5   = class_stats.get("p5", 0) * 100
        if pred_conf >= _bl_mean:
            baseline_label, baseline_cls = "Above Average", "good"
        elif pred_conf >= _bl_p5:
            baseline_label, baseline_cls = "Typical Range", "warn"
        else:
            baseline_label, baseline_cls = "Below Typical", "bad"
        baseline_sub = f"vs {_bl_mean:.0f}% class mean"
    else:
        baseline_label, baseline_cls, baseline_sub = "N/A", "warn", "no test-set data"

    cal_ece_raw = CALIBRATION_DATA.get("ece", None)
    cal_ece_after = CALIBRATION_DATA.get("ece_after_t_scaling", None)
    cal_temp = CALIBRATION_DATA.get("temperature", 1.0)
    # Use calibrated ECE (after T-scaling) as primary — predictions are already calibrated
    cal_ece = cal_ece_after if cal_ece_after is not None else cal_ece_raw
    if cal_ece is not None:
        ece_pct = cal_ece * 100
        # Determine direction from calibration bins (signed gap = accuracy - confidence)
        bin_accs = CALIBRATION_DATA.get("bin_accuracies", [])
        bin_confs = CALIBRATION_DATA.get("bin_confidences", [])
        bin_counts = CALIBRATION_DATA.get("bin_counts", [])
        total_n = sum(bin_counts)
        if total_n > 0 and len(bin_accs) == len(bin_confs) == len(bin_counts):
            signed_gap = sum((a - c) * n for a, c, n in zip(bin_accs, bin_confs, bin_counts)) / total_n
        else:
            signed_gap = 0
        underconfident = signed_gap > 0
        if ece_pct <= 5:
            ece_cls = "good"
            ece_note = "Well calibrated — confidence matches accuracy"
        elif ece_pct <= 10:
            ece_cls = "good"
            ece_note = "Confidence scores closely match actual accuracy"
        elif ece_pct <= 20:
            ece_cls = "warn"
            if underconfident:
                ece_note = "Model tends to be underconfident — actual accuracy may be higher than reported"
            else:
                ece_note = "Model tends to be overconfident — actual accuracy may be lower than reported"
        else:
            ece_cls = "bad"
            if underconfident:
                ece_note = "Model is significantly underconfident — actual accuracy is higher than confidence suggests"
            else:
                ece_note = "Model is significantly overconfident — trust confidence with caution"
        # Add T-scaling note if applied
        if cal_temp != 1.0 and cal_ece_raw is not None:
            raw_pct = cal_ece_raw * 100
            ece_note += f" (T-scaling T={cal_temp:.4f} applied, raw ECE: {raw_pct:.1f}%)"
            ece_note += " — raw model probabilities are overconfident; Temperature Scaling is REQUIRED for production inference"
    else:
        ece_pct, ece_cls, ece_note = 0, "warn", "Calibration data unavailable"

    class_desc = CLASS_INFO.get(pred_label, "")

    return f"""
    <div class="result-box">
      <div class="result-thumb-wrap">
        <img class="result-thumb" src="{thumb_src}" alt="Model input 320x320">
        <div class="thumb-label">320&times;320px<br>model input</div>
      </div>
      <div class="result-info">
        <div class="res-name">{pred_display}</div>
        <div class="res-meta">
          <span>Confidence: <strong>{pred_conf:.1f}%</strong></span>
          <span>Inference time: {inference_ms:.0f} ms</span>
        </div>
        <span class="confidence-badge {conf_class}">{conf_text}</span>
      </div>
    </div>

    <div class="class-desc"><strong>{pred_display}</strong> &mdash; {class_desc}</div>

    <div class="section-label" style="margin-top:16px">Probability Distribution</div>
    <div class="plist">{''.join(rows)}</div>

    <div class="lpa-panel">
      <div class="lpa-title">Live Prediction Analysis</div>
      <div class="lpa-grid">

        <div class="lpa-metric">
          <div class="lpa-label">Certainty Score</div>
          <div class="lpa-value {cert_class}">{certainty_pct:.0f}%</div>
          <div class="lpa-bar-bg">
            <div class="lpa-bar-fill" style="width:{cert_bar_pct}%;background:{cert_color}"></div>
          </div>
          <div class="lpa-sublabel">Shannon entropy</div>
        </div>

        <div class="lpa-metric">
          <div class="lpa-label">Decision Margin</div>
          <div class="lpa-value {margin_class}">+{margin_pts:.1f} pt</div>
          <div class="lpa-bar-bg">
            <div class="lpa-bar-fill"
                 style="width:{min(margin_pts*3,100):.0f}%;background:{cert_color}"></div>
          </div>
          <div class="lpa-sublabel">vs {runner_up_display}</div>
        </div>

        <div class="lpa-metric">
          <div class="lpa-label">Flip Check Classify</div>
          <div class="lpa-value {'good' if flip_match else 'bad'}">{flip_icon}</div>
          <div class="lpa-sublabel">{'Consistent' if flip_match else 'Inconsistent'} ({flip_ms:.0f} ms)</div>
        </div>

        <div class="lpa-metric">
          <div class="lpa-label">Test Baseline</div>
          <div class="lpa-value {baseline_cls}" style="font-size:14px">{baseline_label}</div>
          <div class="lpa-sublabel">{baseline_sub}</div>
        </div>

      </div>
      <div style="display:flex;align-items:center;justify-content:center;gap:10px;margin-top:10px;padding:8px 10px;background:var(--surface);border:1px solid var(--border);border-radius:8px">
        <div style="font-size:10px;font-weight:700;text-transform:uppercase;letter-spacing:.1em;color:var(--faint);flex-shrink:0">ECE</div>
        <div class="lpa-value {ece_cls}" style="font-size:15px">{ece_pct:.1f}%</div>
        <div style="flex:1;height:5px;background:#E2E9E4;border-radius:3px;overflow:hidden">
          <div style="height:100%;border-radius:3px;width:{min(ece_pct * 3, 100):.0f}%;background:{'#007A1E' if ece_cls == 'good' else '#C47D05' if ece_cls == 'warn' else '#C0392B'}"></div>
        </div>
        <div style="font-size:11px;color:var(--muted);line-height:1.4">{ece_note}</div>
      </div>
      <hr class="lpa-divider">
      <div class="lpa-insight {insight_cls}">{insight_icon} {insight_txt}</div>
      <div class="lpa-insight {flip_cls}" style="margin-top:4px">{flip_icon} Flip: {flip_txt}</div>
    </div>
    """


# --------------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------------
with gr.Blocks(title="Fine-Grained World Architecture Image Classification") as demo:

    gr.HTML(HEADER_HTML)

    with gr.Row():
        with gr.Column(scale=1, elem_classes=["grn-card"]):
            gr.HTML('<div class="section-label">Upload Photo</div>')
            inp_image = gr.Image(
                type="numpy", show_label=False, height=280,
                elem_classes=["grn-upload"],
            )
            gr.HTML(
                '<div class="upload-caption">'
                + ' &middot; '.join(l.title() for l in LABELS)
                + '</div>'
            )
            btn_run = gr.Button(
                "Classify Building", variant="primary", size="lg",
                elem_classes=["grn-btn"],
            )

        with gr.Column(scale=1, elem_classes=["grn-card"]):
            gr.HTML('<div class="section-label">Classification Result</div>')
            out_html = gr.HTML(PLACEHOLDER_HTML)

    btn_run.click(fn=classify, inputs=[inp_image], outputs=[out_html])

    gr.HTML(FOOTER_HTML)


if __name__ == "__main__":
    demo.launch(show_error=True, css=CUSTOM_CSS)