#!/usr/bin/env python
"""
Context-CrackNet — Pavement Crack Analyzer
Gradio Web Application

Upload a pavement image, select a trained model, and get:
  • Segmentation mask overlay
  • Skeleton + width heatmap
  • Per-crack annotation
  • Pixel-space measurements: area, length, width, count, coverage
"""

import os
import sys
import warnings
import tempfile
import csv
import uuid
from pathlib import Path

import numpy as np
import cv2
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from PIL import Image
import gradio as gr

# ZeroGPU support — graceful fallback for local development
try:
    import spaces
except ImportError:
    import types as _types
    spaces = _types.ModuleType("spaces")
    def _noop_gpu(fn=None, *, duration=60):
        if fn is not None:
            return fn
        return lambda f: f
    spaces.GPU = _noop_gpu

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.models import Context_CrackNet, create_model

try:
    from skimage.morphology import skeletonize
    HAS_SKIMAGE = True
except ImportError:
    HAS_SKIMAGE = False
    warnings.warn("scikit-image not found. Using morphological fallback for skeleton.")

try:
    from scipy.ndimage import distance_transform_edt
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False

try:
    from huggingface_hub import hf_hub_download
    HAS_HF_HUB = True
except ImportError:
    HAS_HF_HUB = False

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# Set this to your HuggingFace model repo if hosting weights there.
# e.g.  "Blessing988/Context-CrackNet-weights"
# Checkpoint structure in the repo:
#   {DATASET}/{ARCHITECTURE}/best_attention.pth
HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "Blessing988/context-cracknet-weights")

ALL_ARCHITECTURES = [
    "Context_CrackNet",
    "Unet", "UnetPlusPlus", "PSPNet",
    "PAN", "MAnet", "Linknet",
    "FPN", "DeepLabV3Plus", "DeepLabV3",
]

ALL_DATASETS = [
    "CFD", "DeepCrack", "CRACK500", "cracktree200",
    "Eugen_Muller", "forest", "GAPS384",
    "Rissbilder", "Sylvie", "Volker",
]

DEFAULT_ARCHITECTURE = "Context_CrackNet"
DEFAULT_DATASET = os.environ.get("DEFAULT_DATASET", "cracktree200")
DEFAULT_CHECKPOINT_NAME = os.environ.get("DEFAULT_CHECKPOINT_NAME", "best_attention.pth")

SEVERITY_THRESHOLDS = {
    "Low":      (0,   5,   "#27ae60"),   # green
    "Moderate": (5,   15,  "#f39c12"),   # amber
    "High":     (15,  30,  "#e67e22"),   # orange
    "Severe":   (30,  100, "#c0392b"),   # red
}

# Global model cache  {(arch, dataset): (model, img_size)}
_MODEL_CACHE: dict = {}

# ---------------------------------------------------------------------------
# CSS
# ---------------------------------------------------------------------------

CSS = """
/* ── Global ── */
body { font-family: 'Segoe UI', Arial, sans-serif; }

/* ── Header ── */
#header {
    background: linear-gradient(135deg, #0d2137 0%, #1a5276 60%, #2a7da8 100%);
    border-radius: 12px;
    padding: 28px 32px 20px 32px;
    margin-bottom: 8px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.35);
}
#header h1 {
    color: #ffffff !important;
    font-size: 2rem !important;
    font-weight: 700 !important;
    margin: 0 0 6px 0 !important;
    letter-spacing: -0.5px;
}
#header p {
    color: #a8d4f0 !important;
    margin: 0 !important;
    font-size: 0.92rem !important;
    line-height: 1.5;
}
#header .badge {
    display: inline-block;
    background: rgba(255,255,255,0.12);
    border: 1px solid rgba(255,255,255,0.25);
    border-radius: 20px;
    padding: 3px 12px;
    font-size: 0.78rem;
    color: #d6eaf8;
    margin-top: 10px;
    margin-right: 6px;
}

/* ── Section labels ── */
.section-label {
    font-size: 0.78rem;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    color: #5d6d7e;
    margin-bottom: 4px;
}

/* ── Metric cards ── */
.metric-row { display: flex; gap: 10px; margin: 12px 0 4px 0; flex-wrap: wrap; }
.metric-card {
    flex: 1; min-width: 110px;
    background: #fff;
    border-radius: 10px;
    padding: 14px 16px 12px 16px;
    box-shadow: 0 2px 8px rgba(0,0,0,0.09);
    border-top: 4px solid #ccc;
    text-align: center;
}
.metric-card .m-label { font-size: 0.72rem; color: #7f8c8d; font-weight: 600; text-transform: uppercase; letter-spacing: 0.06em; margin-bottom: 4px; }
.metric-card .m-value { font-size: 1.6rem; font-weight: 700; color: #1a252f; line-height: 1; }
.metric-card .m-unit  { font-size: 0.7rem; color: #95a5a6; margin-top: 2px; }
.metric-card.blue   { border-top-color: #2980b9; }
.metric-card.purple { border-top-color: #8e44ad; }
.metric-card.teal   { border-top-color: #16a085; }
.metric-card.orange { border-top-color: #d35400; }
.metric-card.severity-low      { border-top-color: #27ae60; }
.metric-card.severity-moderate { border-top-color: #f39c12; }
.metric-card.severity-high     { border-top-color: #e67e22; }
.metric-card.severity-severe   { border-top-color: #c0392b; }
.m-value.severity-low      { color: #27ae60; }
.m-value.severity-moderate { color: #f39c12; }
.m-value.severity-high     { color: #e67e22; }
.m-value.severity-severe   { color: #c0392b; }

/* ── Status pill ── */
.status-pill {
    display: inline-flex; align-items: center; gap: 6px;
    border-radius: 20px; padding: 5px 14px;
    font-size: 0.82rem; font-weight: 600;
}
.status-pill.ready   { background: #eafaf1; color: #1e8449; border: 1px solid #a9dfbf; }
.status-pill.waiting { background: #fef9e7; color: #9a7d0a; border: 1px solid #f9e79f; }
.status-pill.error   { background: #fdedec; color: #922b21; border: 1px solid #f5b7b1; }

/* ── Analyze button ── */
#analyze-btn {
    background: linear-gradient(135deg, #1a5276 0%, #2e86c1 100%) !important;
    color: #fff !important;
    font-size: 1rem !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
    padding: 14px 0 !important;
    border: none !important;
    box-shadow: 0 3px 12px rgba(30,80,130,0.35) !important;
    transition: all 0.2s !important;
}
#analyze-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 5px 18px rgba(30,80,130,0.45) !important;
}

/* ── Load model button ── */
#load-btn {
    background: #eafaf1 !important;
    color: #1e8449 !important;
    font-weight: 600 !important;
    border: 1.5px solid #a9dfbf !important;
    border-radius: 8px !important;
}

/* ── Tab styling ── */
.tab-nav button {
    font-weight: 600 !important;
    font-size: 0.88rem !important;
}

/* ── Visualization panel ── */
#vis-panel img { border-radius: 8px; }

/* ── Table ── */
.crack-table table { font-size: 0.82rem; }
.crack-table th { background: #1a5276 !important; color: white !important; font-weight: 600 !important; }
.crack-table tr:nth-child(even) { background: #f2f7fb !important; }

/* ── Footer ── */
#footer {
    text-align: center;
    color: #7f8c8d;
    font-size: 0.78rem;
    margin-top: 16px;
    padding: 12px;
    border-top: 1px solid #eaecee;
}
"""

# ---------------------------------------------------------------------------
# Measurement helpers  (identical logic to real_world_eval.py)
# ---------------------------------------------------------------------------

def _morphological_skeleton(binary: np.ndarray) -> np.ndarray:
    img = binary.copy()
    skel = np.zeros_like(img)
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    while True:
        eroded = cv2.erode(img, kernel)
        temp = cv2.subtract(img, cv2.dilate(eroded, kernel))
        skel = cv2.bitwise_or(skel, temp)
        img = eroded.copy()
        if cv2.countNonZero(img) == 0:
            break
    return skel


def measure_cracks(mask_255: np.ndarray) -> dict:
    binary = (mask_255 > 127).astype(np.uint8)
    H, W = binary.shape
    total_pixels = H * W
    area_px2 = int(binary.sum())
    coverage_pct = (area_px2 / total_pixels * 100) if total_pixels > 0 else 0.0

    num_labels, labels = cv2.connectedComponents(binary, connectivity=8)
    crack_count = max(0, num_labels - 1)

    if HAS_SCIPY and area_px2 > 0:
        dist = distance_transform_edt(binary).astype(np.float32)
    else:
        dist = cv2.distanceTransform(binary, cv2.DIST_L2, 5).astype(np.float32)

    skeleton_img = np.zeros_like(binary, dtype=np.uint8)
    length_px = 0
    mean_width_px = 0.0
    max_width_px = 0.0

    if area_px2 > 0:
        skel = (skeletonize(binary.astype(bool)).astype(np.uint8)
                if HAS_SKIMAGE else _morphological_skeleton(binary))
        skeleton_img = skel * 255
        skel_pixels = skel > 0
        length_px = int(skel_pixels.sum())
        if length_px > 0:
            widths = dist[skel_pixels] * 2.0
            mean_width_px = float(np.mean(widths))
            max_width_px = float(np.max(widths))
        else:
            max_width_px = float(dist.max()) * 2.0

    return dict(
        area_px2=area_px2,
        length_px=length_px,
        mean_width_px=round(mean_width_px, 2),
        max_width_px=round(max_width_px, 2),
        crack_count=crack_count,
        coverage_pct=round(coverage_pct, 3),
        skeleton=skeleton_img,
        dist_transform=dist,
        labels=labels,
    )


def measure_per_crack(labels: np.ndarray, dist: np.ndarray) -> list:
    results = []
    for lbl in range(1, labels.max() + 1):
        comp = (labels == lbl).astype(np.uint8)
        area = int(comp.sum())
        if area == 0:
            continue
        ys, xs = np.where(comp)
        x1, y1 = int(xs.min()), int(ys.min())
        x2, y2 = int(xs.max()), int(ys.max())
        skel = (skeletonize(comp.astype(bool)).astype(np.uint8)
                if HAS_SKIMAGE else _morphological_skeleton(comp))
        skel_mask = skel > 0
        length = int(skel_mask.sum())
        if length > 0:
            widths = dist[skel_mask] * 2.0
            mean_w = round(float(np.mean(widths)), 2)
            max_w  = round(float(np.max(widths)), 2)
        else:
            mean_w = 0.0
            max_w  = round(float(dist[comp > 0].max()) * 2.0, 2) if area > 0 else 0.0
        results.append(dict(
            crack_id=lbl,
            area_px2=area,
            length_px=length,
            mean_width_px=mean_w,
            max_width_px=max_w,
            bbox_x=x1, bbox_y=y1,
            bbox_w=x2 - x1, bbox_h=y2 - y1,
        ))
    results.sort(key=lambda r: r['area_px2'], reverse=True)
    return results


# ---------------------------------------------------------------------------
# Model helpers
# ---------------------------------------------------------------------------

def _infer_img_size(state_dict: dict) -> int:
    key = 'linformer.self_attn.proj_E'
    if key in state_dict:
        seq_len = state_dict[key].shape[0]
        return int(round(seq_len ** 0.5)) * 16
    return 448


def load_model_from_file(checkpoint_path: str, architecture: str):
    """Load any architecture from a .pth checkpoint file."""
    state_dict = torch.load(checkpoint_path, map_location='cpu')

    if architecture == 'Context_CrackNet':
        img_size = _infer_img_size(state_dict)
        model = Context_CrackNet(
            in_channels=3, out_channels=1, img_size=img_size,
            num_heads=8, ff_dim=2048, linformer_k=256, pretrained=False
        )
    else:
        img_size = None
        model = create_model(
            architecture=architecture, encoder_name='resnet50',
            in_channels=3, num_classes=1, encoder_weights=None
        )

    model.load_state_dict(state_dict)
    model.eval()
    return model, img_size


def load_model_from_hf(architecture: str, dataset: str):
    """Download checkpoint from HF Hub model repo and load it."""
    if not HAS_HF_HUB:
        raise RuntimeError("huggingface_hub not installed.")
    if not HF_MODEL_REPO:
        raise RuntimeError("HF_MODEL_REPO environment variable not set.")
    path = f"{dataset}/{architecture}/{DEFAULT_CHECKPOINT_NAME}"
    try:
        local = hf_hub_download(repo_id=HF_MODEL_REPO, filename=path)
    except Exception as e:
        raise RuntimeError(
            f"Hosted checkpoint not found at {HF_MODEL_REPO}/{path}. "
            "Upload it to the model repo or provide a manual checkpoint override."
        ) from e
    return load_model_from_file(local, architecture)


def get_or_load_model(architecture: str, dataset: str, checkpoint_file=None):
    """
    Return (model, img_size).
    Priority: uploaded file → HF Hub → error.
    Results are cached by (arch, dataset) key.
    """
    cache_key = (architecture, dataset, checkpoint_file)

    if cache_key in _MODEL_CACHE:
        return _MODEL_CACHE[cache_key]

    if checkpoint_file is not None:
        model, img_size = load_model_from_file(checkpoint_file, architecture)
    elif HF_MODEL_REPO:
        model, img_size = load_model_from_hf(architecture, dataset)
    else:
        raise RuntimeError(
            "No checkpoint provided. Upload a .pth file or configure HF_MODEL_REPO."
        )

    _MODEL_CACHE[cache_key] = (model, img_size)
    return model, img_size


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def get_transforms():
    return A.Compose([
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])


TRANSFORMS = get_transforms()


@spaces.GPU(duration=120)
def _gpu_forward(model, tensor):
    """GPU-isolated forward pass — decorated for ZeroGPU allocation."""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    with torch.no_grad():
        out = model(tensor.to(device))
        if isinstance(out, tuple):
            out = out[0]
    model.cpu()
    return out.cpu()


def predict_mask(model, image_rgb: np.ndarray, img_size, threshold: float) -> np.ndarray:
    orig_h, orig_w = image_rgb.shape[:2]
    inp = cv2.resize(image_rgb, (img_size, img_size)) if img_size else image_rgb
    tensor = TRANSFORMS(image=inp)['image'].unsqueeze(0)

    out = _gpu_forward(model, tensor)

    prob = torch.sigmoid(out)
    binary = (prob > threshold).squeeze().cpu().numpy().astype(np.uint8)
    mask = binary * 255

    if img_size and (orig_h != img_size or orig_w != img_size):
        mask = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    return mask


# ---------------------------------------------------------------------------
# Visualization
# ---------------------------------------------------------------------------

OVERLAY_ALPHA = 0.45


def make_overlay(image_rgb, mask_255):
    overlay = image_rgb.copy().astype(np.float32)
    crack = mask_255 > 127
    overlay[crack] = overlay[crack] * (1 - OVERLAY_ALPHA) + np.array([255, 30, 30]) * OVERLAY_ALPHA
    return np.clip(overlay, 0, 255).astype(np.uint8)


def make_combined_panel(image_rgb, mask, skeleton, dist) -> np.ndarray:
    """
    Returns a single 1×2 panel figure as an RGB numpy array:
      [Crack Overlay] [Skel+Width]
    """
    fig = plt.figure(figsize=(14, 5.5), facecolor='#0d1b2a')
    gs = gridspec.GridSpec(1, 2, figure=fig, wspace=0.06,
                           left=0.02, right=0.98, top=0.9, bottom=0.04)

    panel_cfg = [
        (0, 0, make_overlay(image_rgb, mask), "Crack Segmentation Overlay", None),
        (0, 1, None, "Skeleton + Crack Width (px)", "skeleton"),
    ]

    for row, col, img, title, special in panel_cfg:
        ax = fig.add_subplot(gs[row, col])
        ax.set_facecolor('#0d1b2a')

        if special == "skeleton":
            ax.imshow(image_rgb)
            skel_pts = np.where(skeleton > 0)
            if len(skel_pts[0]) > 0:
                widths = dist[skel_pts] * 2.0
                sc = ax.scatter(skel_pts[1], skel_pts[0], c=widths,
                                cmap='plasma', s=1.5, linewidths=0,
                                vmin=0, vmax=max(widths.max(), 1))
                cb = fig.colorbar(sc, ax=ax, fraction=0.035, pad=0.01)
                cb.set_label('Width (px)', color='#a8d4f0', fontsize=7)
                cb.ax.yaxis.set_tick_params(color='#a8d4f0', labelsize=6)
                plt.setp(cb.ax.yaxis.get_ticklabels(), color='#a8d4f0')
        else:
            ax.imshow(img)

        ax.set_title(title, color='#d6eaf8', fontsize=9, fontweight='bold', pad=5)
        ax.axis('off')

    fig.canvas.draw()

    if hasattr(fig.canvas, "buffer_rgba"):
        # Matplotlib 3.10+ exposes the RGBA buffer API on Agg canvases.
        buf = np.asarray(fig.canvas.buffer_rgba(), dtype=np.uint8)[..., :3].copy()
    else:
        buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,))

    plt.close(fig)
    return buf


# ---------------------------------------------------------------------------
# Severity helper
# ---------------------------------------------------------------------------

def get_severity(coverage_pct: float) -> tuple[str, str]:
    for label, (lo, hi, color) in SEVERITY_THRESHOLDS.items():
        if lo <= coverage_pct < hi:
            return label, color
    return "Severe", "#c0392b"


# ---------------------------------------------------------------------------
# HTML builders
# ---------------------------------------------------------------------------

def build_metric_cards(stats: dict) -> str:
    sev_label, sev_color = get_severity(stats['coverage_pct'])
    sev_class = f"severity-{sev_label.lower()}"

    cards = [
        ("blue",   "🔢 Crack Count",   stats['crack_count'],          "instances",  ""),
        ("purple", "📐 Crack Area",    f"{stats['area_px2']:,}",       "px²",        ""),
        ("teal",   "📏 Crack Length",  f"{stats['length_px']:,}",      "px",         ""),
        ("orange", "↔ Mean Width",    f"{stats['mean_width_px']:.1f}", "px",         ""),
        ("orange", "↕ Max Width",     f"{stats['max_width_px']:.1f}",  "px",         ""),
        (sev_class,"📊 Coverage",     f"{stats['coverage_pct']:.2f}",  "%",
         f'<div class="section-label" style="margin-top:2px;color:{sev_color}">'
         f'{sev_label}</div>'),
    ]

    html = '<div class="metric-row">'
    for cls, label, value, unit, extra in cards:
        html += f"""
        <div class="metric-card {cls}">
            <div class="m-label">{label}</div>
            <div class="m-value {cls}">{value}</div>
            <div class="m-unit">{unit}</div>
            {extra}
        </div>"""
    html += '</div>'
    return html


def build_table_html(per_crack: list) -> str:
    """Render per-crack list as a styled HTML table (avoids gr.DataFrame schema bugs)."""
    if not per_crack:
        return '<p style="color:#7f8c8d;padding:8px">No cracks detected.</p>'
    cols = ['Crack #', 'Area (px²)', 'Length (px)', 'Mean Width (px)',
            'Max Width (px)', 'BBox X', 'BBox Y', 'BBox W', 'BBox H']
    keys = ['crack_id', 'area_px2', 'length_px', 'mean_width_px',
            'max_width_px', 'bbox_x', 'bbox_y', 'bbox_w', 'bbox_h']
    hdr = "".join(f'<th style="background:#1a5276;color:#fff;padding:6px 10px;'
                  f'text-align:left;font-size:0.78rem">{c}</th>' for c in cols)
    rows = ""
    for i, c in enumerate(per_crack):
        bg = "#f2f7fb" if i % 2 == 0 else "#fff"
        cells = "".join(
            f'<td style="padding:5px 10px;font-size:0.8rem;border-bottom:1px solid #eaecee">'
            f'{c.get(k, "")}</td>' for k in keys
        )
        rows += f'<tr style="background:{bg}">{cells}</tr>'
    return (f'<div style="overflow-x:auto;border-radius:8px;border:1px solid #eaecee">'
            f'<table style="width:100%;border-collapse:collapse">'
            f'<thead><tr>{hdr}</tr></thead><tbody>{rows}</tbody></table></div>')


EXPORT_HEADERS = [
    "Crack #",
    "Area (px^2)",
    "Length (px)",
    "Mean Width (px)",
    "Max Width (px)",
    "BBox X",
    "BBox Y",
    "BBox W",
    "BBox H",
]

EXPORT_KEYS = [
    "crack_id",
    "area_px2",
    "length_px",
    "mean_width_px",
    "max_width_px",
    "bbox_x",
    "bbox_y",
    "bbox_w",
    "bbox_h",
]


def export_per_crack_details(per_crack: list) -> tuple[str, str]:
    from openpyxl import Workbook

    export_dir = Path(tempfile.gettempdir()) / "context_cracknet_exports"
    export_dir.mkdir(parents=True, exist_ok=True)

    export_id = uuid.uuid4().hex
    csv_path = export_dir / f"per_crack_details_{export_id}.csv"
    xlsx_path = export_dir / f"per_crack_details_{export_id}.xlsx"

    rows = [[item.get(key, "") for key in EXPORT_KEYS] for item in per_crack]

    with csv_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(EXPORT_HEADERS)
        writer.writerows(rows)

    workbook = Workbook()
    sheet = workbook.active
    sheet.title = "Per Crack Details"
    sheet.append(EXPORT_HEADERS)
    for row in rows:
        sheet.append(row)
    workbook.save(xlsx_path)

    return str(csv_path), str(xlsx_path)


def build_status(msg: str, kind: str = "ready") -> str:
    icons = {"ready": "✅", "waiting": "⏳", "error": "❌"}
    return (f'<span class="status-pill {kind}">'
            f'{icons.get(kind, "")} {msg}</span>')


# ---------------------------------------------------------------------------
# Main analysis function (called by Gradio)
# ---------------------------------------------------------------------------

def analyze(
    image_pil,
    architecture: str,
    dataset: str,
    threshold: float,
    checkpoint_file,
):
    if image_pil is None:
        return (
            None,
            build_status("No image uploaded.", "error"),
            build_metric_cards({
                'crack_count': 0, 'area_px2': 0, 'length_px': 0,
                'mean_width_px': 0, 'max_width_px': 0, 'coverage_pct': 0
            }),
            None,
            gr.Accordion(visible=False, open=False),
            gr.DownloadButton(visible=False),
            gr.DownloadButton(visible=False),
        )

    # Be defensive across Gradio payload formats.
    if isinstance(checkpoint_file, dict):
        checkpoint_file = checkpoint_file.get("path") or checkpoint_file.get("name")

    # Load model
    ckpt_path = checkpoint_file if checkpoint_file else None
    try:
        model, img_size = get_or_load_model(architecture, dataset, ckpt_path)
    except Exception as e:
        return (
            None,
            build_status(f"Model error: {e}", "error"),
            build_metric_cards({
                'crack_count': 0, 'area_px2': 0, 'length_px': 0,
                'mean_width_px': 0, 'max_width_px': 0, 'coverage_pct': 0
            }),
            None,
            gr.Accordion(visible=False, open=False),
            gr.DownloadButton(visible=False),
            gr.DownloadButton(visible=False),
        )

    # Prepare image
    image_rgb = np.array(image_pil.convert("RGB"))

    # Inference
    try:
        mask = predict_mask(model, image_rgb, img_size, threshold)
    except Exception as e:
        return (
            None,
            build_status(f"Inference failed: {e}", "error"),
            build_metric_cards({
                'crack_count': 0, 'area_px2': 0, 'length_px': 0,
                'mean_width_px': 0, 'max_width_px': 0, 'coverage_pct': 0
            }),
            None,
            gr.Accordion(visible=False, open=False),
            gr.DownloadButton(visible=False),
            gr.DownloadButton(visible=False),
        )

    # Measurements
    stats = measure_cracks(mask)
    per_crack = measure_per_crack(stats['labels'], stats['dist_transform'])

    # Visualization panel
    panel = make_combined_panel(image_rgb, mask, stats['skeleton'], stats['dist_transform'])
    panel_pil = Image.fromarray(panel)

    # Status
    sev_label, _ = get_severity(stats['coverage_pct'])
    status_html = build_status(
        f"Analysis complete — {stats['crack_count']} crack(s) detected "
        f"| Severity: {sev_label} | Model: {architecture} trained on {dataset}",
        "ready"
    )

    # Metric cards
    metrics_html = build_metric_cards(stats)

    # Keep the detail table in HTML so UI behavior is stable across Gradio releases.
    table_html = build_table_html(per_crack)
    csv_path, xlsx_path = export_per_crack_details(per_crack)

    return (
        panel_pil,
        status_html,
        metrics_html,
        table_html,
        gr.Accordion(visible=True, open=False),
        gr.DownloadButton(label="Download CSV", value=csv_path, visible=True),
        gr.DownloadButton(label="Download Excel", value=xlsx_path, visible=True),
    )


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

HEADER_HTML = """
<div id="header">
  <h1>🔬 Context-CrackNet &nbsp;—&nbsp; Pavement Crack Analyzer</h1>
  <p>
    A Context-Aware Framework for Precise Segmentation of Tiny Cracks in Pavement Images<br>
    <em>Construction and Building Materials · Volume 484 · 2025 · DOI: 10.1016/j.conbuildmat.2025.141583</em>
  </p>
  <span class="badge">🏛 North Dakota State University</span>
  <span class="badge">✍ Kyem et al.</span>
  <span class="badge">⚡ ResNet50 + CAGM + RFEM</span>
</div>
"""

ABOUT_MD = """
## About Context-CrackNet

**Context-CrackNet** is a novel deep-learning architecture for crack segmentation in pavement images.
It combines:

| Module | Role |
|--------|------|
| **ResNet50 Encoder** | Hierarchical feature extraction (5 scales) |
| **CAGM** (Context-Aware Global Module) | Linformer self-attention — captures long-range crack context with O(n·k) complexity |
| **RFEM** (Region Focused Enhancement Module) | Attention-gated skip connections — focuses on fine crack details |

### Pixel-Space Measurements

All crack measurements are reported in **pixel units** — no camera calibration required.

| Metric | How Computed |
|--------|-------------|
| **Crack Area** (px²) | Total crack pixels in binary mask |
| **Crack Length** (px) | Skeleton pixel count via Lee's (1994) thinning algorithm |
| **Mean Width** (px) | 2 × mean Euclidean Distance Transform at skeleton pixels |
| **Max Width** (px) | 2 × max EDT value at skeleton pixels |
| **Crack Count** | 8-connected component labelling |
| **Coverage (%)** | Area ÷ (H × W) × 100 |

### Mathematical Foundation

The pipeline is grounded in **Blum's Medial Axis Transform (1967)**:
> Any 2D shape can be completely reconstructed from its skeleton + inscribed circle radii.
> Therefore, skeleton length = crack length and 2 × EDT at skeleton = crack width.

### Citation
```bibtex
@article{Context-CrackNet,
  title   = {Context-CrackNet: A context-aware framework for precise segmentation of tiny cracks in pavement images},
  journal = {Construction and Building Materials},
  volume  = {484}, pages = {141583}, year = {2025},
  doi     = {10.1016/j.conbuildmat.2025.141583},
  author  = {Blessing {Agyei Kyem} and Joshua Kofi Asamoah and Armstrong Aboah}
}
```
"""

HOW_TO_MD = """
## How to Use

### Step 1 — Choose a Hosted Model
- By default the app uses **Context-CrackNet** hosted from Hugging Face Hub.
- **Training dataset**: Choose the dataset-specific hosted checkpoint to use.
- **Advanced options**: Upload a custom checkpoint or switch to a baseline model only if you need to override the default hosted setup.

### Step 2 — Upload a Pavement Image
- Drag & drop or click to upload any road/pavement photograph.
- The app handles any resolution — the image is internally resized to the model's training resolution for inference and results are returned in **original pixel space**.

### Step 3 — Adjust Threshold (Optional)
- Default **0.5** = standard binary decision boundary.
- Lower → more cracks detected (higher recall, lower precision).
- Higher → fewer but more confident detections.

### Step 4 — Analyze
- Click **🚀 Analyze Cracks** and wait a few seconds.

### Reading the Results
| Panel | What it shows |
|-------|---------------|
| **Crack Segmentation Overlay** | Red = predicted crack region |
| **Skeleton + Width** | Crack centrelines coloured by local width (plasma scale) |

The **metric cards** show global statistics; the **Per-Crack Details** table shows each individual crack.
You can also download the per-crack details as **CSV** or **Excel** after each analysis.
"""


def build_ui():
    with gr.Blocks(css=CSS, title="Context-CrackNet — Crack Analyzer") as demo:

        # ── Header ──
        gr.HTML(HEADER_HTML)

        with gr.Tabs(elem_classes=["tab-nav"]):

            # ═══════════════════════════════════════════════
            # TAB 1 — ANALYZE
            # ═══════════════════════════════════════════════
            with gr.TabItem("🔍 Analyze"):

                # ── Top row: inputs ──
                with gr.Row(equal_height=False):

                    # Left column — image upload
                    with gr.Column(scale=5):
                        gr.HTML('<div class="section-label">📤 Pavement Image</div>')
                        image_input = gr.Image(
                            type="pil",
                            label="Upload pavement / road image",
                            height=310,
                        )
                        gr.HTML('<div class="section-label" style="margin-top:6px;color:#7f8c8d">'
                                'Supports JPG, PNG, TIFF — any resolution</div>')

                    # Right column — model settings
                    with gr.Column(scale=4):
                        gr.HTML('<div class="section-label">⚙️ Model Configuration</div>')

                        dataset = gr.Dropdown(
                            choices=ALL_DATASETS,
                            value=DEFAULT_DATASET if DEFAULT_DATASET in ALL_DATASETS else ALL_DATASETS[0],
                            label="Training Dataset",
                            info="Select the hosted Context-CrackNet checkpoint to use.",
                        )
                        threshold = gr.Slider(
                            minimum=0.1, maximum=0.9, value=0.5, step=0.05,
                            label="Segmentation Threshold",
                            info="0.5 = standard • lower = more cracks • higher = fewer cracks",
                        )

                        if HF_MODEL_REPO:
                            gr.HTML(
                                f'<div class="status-pill ready" style="margin-top:6px">'
                                f'✅ Default hosted model: <code>{DEFAULT_DATASET}/{DEFAULT_ARCHITECTURE}/{DEFAULT_CHECKPOINT_NAME}</code>'
                                f'</div>'
                            )
                            gr.HTML(
                                f'<div class="section-label" style="margin-top:6px;color:#7f8c8d">'
                                f'Loaded from <code>{HF_MODEL_REPO}</code> when no override is provided</div>'
                            )
                        else:
                            gr.HTML(
                                '<div class="status-pill waiting" style="margin-top:6px">'
                                '⏳ No hosted model repo configured yet. Use a manual checkpoint override below.'
                                '</div>'
                            )

                        with gr.Accordion("Advanced Model Options", open=False):
                            architecture = gr.Dropdown(
                                choices=ALL_ARCHITECTURES,
                                value=DEFAULT_ARCHITECTURE,
                                label="Architecture",
                                info="Leave this on Context-CrackNet unless you want to try a baseline.",
                            )
                            checkpoint_file = gr.File(
                                label="Optional checkpoint override (.pth)",
                                file_types=[".pth", ".pt"],
                                type="filepath",
                            )
                            gr.HTML(
                                '<div class="section-label" style="margin-top:4px;color:#7f8c8d">'
                                'Use this section for baseline models or for testing a checkpoint that is not in the hosted model repo.</div>'
                            )

                        analyze_btn = gr.Button(
                            "🚀 Analyze Cracks",
                            elem_id="analyze-btn",
                            variant="primary",
                        )

                # ── Status bar ──
                status_out = gr.HTML(
                    value=build_status(
                        (
                            f"Upload an image and click Analyze. "
                            f"Default hosted model: {DEFAULT_ARCHITECTURE} trained on {DEFAULT_DATASET}."
                            if HF_MODEL_REPO else
                            "Upload an image and click Analyze. Add a checkpoint override in Advanced Model Options if needed."
                        ),
                        "waiting",
                    )
                )

                # ── Metric cards ──
                gr.HTML('<div class="section-label" style="margin-top:6px">📊 Crack Measurements</div>')
                metrics_out = gr.HTML(
                    value=build_metric_cards({
                        'crack_count': 0, 'area_px2': 0, 'length_px': 0,
                        'mean_width_px': 0, 'max_width_px': 0, 'coverage_pct': 0
                    })
                )

                # ── Visualization panel ──
                gr.HTML('<div class="section-label" style="margin-top:10px">🖼️ Visual Analysis</div>')
                vis_out = gr.Image(
                    label="",
                    show_label=False,
                    elem_id="vis-panel",
                    height=500,
                )

                # ── Per-crack table ──
                with gr.Accordion("📋 Per-Crack Details", open=False, visible=False) as table_acc:
                    table_out = gr.HTML(value="")

                with gr.Row():
                    csv_download = gr.DownloadButton(
                        "Download CSV",
                        visible=False,
                    )
                    excel_download = gr.DownloadButton(
                        "Download Excel",
                        visible=False,
                    )

                # ── Wire up ──
                analyze_btn.click(
                    fn=analyze,
                    inputs=[image_input, architecture, dataset, threshold, checkpoint_file],
                    outputs=[vis_out, status_out, metrics_out, table_out, table_acc, csv_download, excel_download],
                )

            # ═══════════════════════════════════════════════
            # TAB 2 — HOW TO USE
            # ═══════════════════════════════════════════════
            with gr.TabItem("📖 How to Use"):
                gr.Markdown(HOW_TO_MD)

            # ═══════════════════════════════════════════════
            # TAB 3 — ABOUT
            # ═══════════════════════════════════════════════
            with gr.TabItem("ℹ️ About"):
                gr.Markdown(ABOUT_MD)

        # ── Footer ──
        gr.HTML("""
        <div id="footer">
            Context-CrackNet · Construction and Building Materials 2025 ·
            North Dakota State University ·
            <a href="https://doi.org/10.1016/j.conbuildmat.2025.141583" target="_blank">
                DOI: 10.1016/j.conbuildmat.2025.141583
            </a>
        </div>
        """)

    return demo


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

if __name__ == "__main__":
    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )