app.py

"""Gradio app for WoundNetB7 DFU Analysis — Hugging Face Spaces deployment.

Pipeline visualization:
  1. Binary ulcer segmentation (WoundNetB7 + ASPP + CBAM + CoordAttention + TAM)
  2. Multiclass segmentation (background / foot / perilesion / ulcer)
  3. Fitzpatrick/ITA skin type estimation
  4. PWAT scores (raw) + PWAT adjusted by Fitzpatrick debiasing
  5. Downloadable clinical report (PDF + JSON)
  6. Guided camera capture with foot silhouette overlay

Launch locally: python app.py
Deploy to HF: push this repo to a Hugging Face Space (GPU recommended).
"""
import gradio as gr
import numpy as np
import cv2
import json
import tempfile
import os
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
from fpdf import FPDF
from pipeline import WoundNetB7Pipeline
from src.pwat_estimator import ITEM_NAMES

pipe = WoundNetB7Pipeline(models_dir="models", use_tta=True)

FITZ_COLORS = {
    "I": "#fef3c7", "II": "#fde68a", "III": "#fbbf24",
    "IV": "#b45309", "V": "#78350f", "VI": "#451a03",
}

FITZ_TEXT_COLORS = {
    "I": "#1f2937", "II": "#1f2937", "III": "#1f2937",
    "IV": "#ffffff", "V": "#ffffff", "VI": "#ffffff",
}

FITZ_RGB = {
    "I": (254, 243, 199), "II": (253, 230, 138), "III": (251, 191, 36),
    "IV": (180, 83, 9), "V": (120, 53, 15), "VI": (69, 26, 3),
}

FITZ_TEXT_RGB = {
    "I": (31, 41, 55), "II": (31, 41, 55), "III": (31, 41, 55),
    "IV": (255, 255, 255), "V": (255, 255, 255), "VI": (255, 255, 255),
}


# ── Foot Guide Overlay ───────────────────────────────────────────────────────

def generate_foot_guide(width=640, height=480):
    """Generate a semi-transparent foot silhouette guide overlay for camera capture."""
    guide = np.zeros((height, width, 4), dtype=np.uint8)

    cx, cy = width // 2, height // 2
    scale_x = width / 640
    scale_y = height / 480

    # Foot outline points (plantar view, normalized for 640x480)
    foot_points = np.array([
        # Right side (lateral)
        (370, 420), (380, 380), (385, 340), (385, 300), (382, 260),
        (378, 220), (370, 180), (360, 150), (348, 120), (335, 100),
        (325, 85), (318, 72),
        # Toes (right to left)
        (320, 60), (325, 48), (318, 38), (305, 42),  # 5th toe
        (305, 35), (310, 22), (300, 18), (290, 28),   # 4th toe
        (288, 20), (292, 8), (280, 5), (272, 18),     # 3rd toe
        (268, 12), (270, -2), (258, -5), (250, 10),   # 2nd toe
        (245, 5), (242, -10), (228, -8), (230, 12),   # Big toe
        # Left side (medial)
        (225, 30), (218, 55), (215, 80), (218, 110),
        (222, 140), (228, 180), (235, 220), (240, 260),
        (245, 300), (248, 340), (250, 380), (255, 420),
        # Heel
        (270, 445), (300, 455), (330, 450), (355, 435),
    ], dtype=np.float32)

    # Center and scale
    foot_center = foot_points.mean(axis=0)
    foot_points -= foot_center
    foot_points[:, 0] *= scale_x * 0.85
    foot_points[:, 1] *= scale_y * 0.85
    foot_points += [cx, cy]
    foot_pts = foot_points.astype(np.int32)

    # Draw filled semi-transparent foot area
    foot_mask = np.zeros((height, width), dtype=np.uint8)
    cv2.fillPoly(foot_mask, [foot_pts], 255)

    # Semi-transparent green fill
    guide[foot_mask > 0] = [0, 200, 100, 35]

    # Foot outline (bright green, dashed effect via thick line)
    cv2.polylines(guide, [foot_pts], True, (0, 220, 120, 200), 3, cv2.LINE_AA)

    # Center crosshair
    cross_len = 20
    cv2.line(guide, (cx - cross_len, cy), (cx + cross_len, cy), (255, 255, 255, 150), 1)
    cv2.line(guide, (cx, cy - cross_len), (cx, cy + cross_len), (255, 255, 255, 150), 1)

    # Corner brackets for framing
    bracket_len = 40
    bracket_color = (0, 220, 120, 200)
    bw = 2
    margin = 30
    corners = [
        (margin, margin),
        (width - margin, margin),
        (margin, height - margin),
        (width - margin, height - margin),
    ]
    for (x, y) in corners:
        dx = bracket_len if x < width // 2 else -bracket_len
        dy = bracket_len if y < height // 2 else -bracket_len
        cv2.line(guide, (x, y), (x + dx, y), bracket_color, bw)
        cv2.line(guide, (x, y), (x, y + dy), bracket_color, bw)

    return guide


def apply_foot_guide(frame):
    """Apply the foot guide overlay to a camera frame."""
    if frame is None:
        return None

    h, w = frame.shape[:2]
    guide = generate_foot_guide(w, h)

    # Composite RGBA guide over RGB frame
    frame_rgba = cv2.cvtColor(frame, cv2.COLOR_RGB2RGBA)
    alpha = guide[:, :, 3:4].astype(np.float32) / 255.0
    blended = frame_rgba.astype(np.float32)
    overlay = guide.astype(np.float32)
    blended[:, :, :3] = blended[:, :, :3] * (1 - alpha) + overlay[:, :, :3] * alpha
    blended[:, :, 3] = 255

    result = blended[:, :, :3].astype(np.uint8)

    # Add instruction text at top
    cv2.putText(result, "Position the foot inside the guide",
                (w // 2 - 200, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 220, 120), 2, cv2.LINE_AA)
    cv2.putText(result, "Distance: 30-40 cm | Uniform lighting",
                (w // 2 - 230, h - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 1, cv2.LINE_AA)

    return result


def generate_static_guide():
    """Generate a static reference guide image with instructions."""
    W, H = 500, 700
    img = Image.new("RGB", (W, H), (248, 250, 252))
    draw = ImageDraw.Draw(img)

    font_title = _get_font(24)
    font_body = _get_font_regular(16)
    font_small = _get_font_regular(14)

    # Title
    draw.text((W // 2 - 130, 15), "DFU Capture Guide", fill=(31, 41, 55), font=font_title)

    # Draw foot silhouette (simplified)
    foot_guide_rgba = generate_foot_guide(400, 350)
    foot_rgb = foot_guide_rgba[:, :, :3]
    # Make non-zero areas visible on white background
    mask = foot_guide_rgba[:, :, 3] > 0
    bg_section = np.full((350, 400, 3), 245, dtype=np.uint8)
    bg_section[mask] = foot_rgb[mask]
    # Draw the outline more visibly
    foot_pil = Image.fromarray(bg_section)
    img.paste(foot_pil, (50, 55))

    # Border around foot area
    draw.rectangle([(48, 53), (452, 407)], outline=(209, 213, 219), width=2)

    # Instructions
    y = 425
    instructions = [
        ("1.", "Plantar view of the foot facing the camera"),
        ("2.", "Distance: 30-40 cm from the lens"),
        ("3.", "Uniform lighting, no harsh shadows"),
        ("4.", "Neutral background (white or blue sheet)"),
        ("5.", "Include the full ulcer + 3-5 cm of healthy skin"),
        ("6.", "Avoid direct flash (causes glare)"),
        ("7.", "Center the foot inside the green silhouette"),
    ]
    for num, text in instructions:
        draw.text((30, y), num, fill=(5, 150, 105), font=font_title)
        draw.text((60, y + 2), text, fill=(55, 65, 81), font=font_body)
        y += 30

    # Bottom note
    draw.line([(30, y + 5), (W - 30, y + 5)], fill=(229, 231, 235), width=1)
    draw.text((30, y + 12),
              "Tip: For best results capture with diffuse natural",
              fill=(107, 114, 128), font=font_small)
    draw.text((30, y + 32),
              "light. Avoid overhead lights that create shadows.",
              fill=(107, 114, 128), font=font_small)

    return np.array(img)


# ── PDF Report Generation ────────────────────────────────────────────────────

def _get_font(size):
    for name in ["DejaVuSans-Bold.ttf", "DejaVuSans.ttf", "arial.ttf", "LiberationSans-Bold.ttf"]:
        try:
            return ImageFont.truetype(name, size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


def _get_font_regular(size):
    for name in ["DejaVuSans.ttf", "arial.ttf", "LiberationSans-Regular.ttf"]:
        try:
            return ImageFont.truetype(name, size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


class DFUReport(FPDF):
    """Custom PDF report for DFU analysis results."""

    def __init__(self):
        super().__init__(orientation="P", unit="mm", format="A4")
        self.set_auto_page_break(auto=True, margin=15)
        self._setup_fonts()

    def _setup_fonts(self):
        """Register Unicode font if available, otherwise use built-in."""
        font_paths = [
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
            "/usr/share/fonts/TTF/DejaVuSans.ttf",
            "C:/Windows/Fonts/arial.ttf",
        ]
        self._has_unicode = False
        for fp in font_paths:
            if os.path.exists(fp):
                try:
                    self.add_font("CustomFont", "", fp, uni=True)
                    bold_fp = fp.replace("DejaVuSans.ttf", "DejaVuSans-Bold.ttf").replace("arial.ttf", "arialbd.ttf")
                    if os.path.exists(bold_fp):
                        self.add_font("CustomFont", "B", bold_fp, uni=True)
                    else:
                        self.add_font("CustomFont", "B", fp, uni=True)
                    self._has_unicode = True
                    break
                except Exception:
                    continue

    def _font(self, style="", size=10):
        if self._has_unicode:
            self.set_font("CustomFont", style, size)
        else:
            self.set_font("Helvetica", style, size)

    def header(self):
        self.set_fill_color(31, 41, 55)
        self.rect(0, 0, 210, 22, "F")
        self._font("B", 14)
        self.set_text_color(255, 255, 255)
        self.set_xy(10, 4)
        self.cell(0, 8, "WoundNetB7 - Integrated DFU Assessment Report", 0, 0, "L")
        self._font("", 8)
        self.set_text_color(156, 163, 175)
        self.set_xy(10, 13)
        self.cell(0, 6, "EfficientNet-B7 + ASPP + CBAM + CoordAttention + TAM  |  Ulcer Dice: 0.927", 0, 0, "L")
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
        self.set_xy(160, 4)
        self.cell(40, 8, timestamp, 0, 0, "R")
        self.ln(20)

    def footer(self):
        self.set_y(-12)
        self._font("", 7)
        self.set_text_color(156, 163, 175)
        self.cell(0, 5, "WoundNetB7 | Doctoral Thesis | Marcelo Marquez-Murillo | "
                  "Dice 0.927 (95% CI: [0.917, 0.936]) | Debiasing: 46.6% gap reduction (p < 1e-55)", 0, 0, "C")
        self.cell(0, 5, f"Page {self.page_no()}/{{nb}}", 0, 0, "R")

    def section_title(self, number, title):
        self._font("B", 11)
        self.set_text_color(31, 41, 55)
        self.set_fill_color(243, 244, 246)
        self.cell(8, 7, str(number), 0, 0, "C", fill=False)
        self.cell(0, 7, f"  {title}", 0, 1, "L")
        self.ln(2)

    def add_image_pair(self, img1_path, label1, img2_path, label2):
        """Add two images side by side with labels."""
        self._font("", 8)
        self.set_text_color(107, 114, 128)
        x = self.get_x()
        y = self.get_y()
        img_w = 90
        img_h = 60

        self.cell(img_w, 4, label1, 0, 0, "C")
        self.cell(5, 4, "", 0, 0)
        self.cell(img_w, 4, label2, 0, 1, "C")

        self.image(img1_path, x=x, y=self.get_y(), w=img_w, h=img_h)
        self.image(img2_path, x=x + img_w + 5, y=self.get_y(), w=img_w, h=img_h)
        self.ln(img_h + 3)


def generate_pdf_report(image_rgb, binary_overlay, multiclass_overlay, result):
    """Generate a clinical PDF report with all analysis results."""
    tmpdir = tempfile.mkdtemp(prefix="woundnetb7_report_")

    # Save temp images for embedding in PDF
    orig_path = os.path.join(tmpdir, "_orig.png")
    binary_path = os.path.join(tmpdir, "_binary.png")
    multi_path = os.path.join(tmpdir, "_multi.png")

    Image.fromarray(image_rgb).save(orig_path)
    Image.fromarray(binary_overlay).save(binary_path)
    Image.fromarray(multiclass_overlay).save(multi_path)

    pdf = DFUReport()
    pdf.alias_nb_pages()
    pdf.add_page()

    # ── Section 1: Images ──
    pdf.section_title(1, "Segmentation")
    pdf.add_image_pair(orig_path, "Original Image", binary_path, "Binary Ulcer Segmentation")
    pdf.ln(2)

    # Multiclass + legend
    pdf._font("", 8)
    pdf.set_text_color(107, 114, 128)
    x_start = pdf.get_x()
    y_start = pdf.get_y()
    pdf.cell(90, 4, "Multi-Class Segmentation", 0, 0, "C")
    pdf.cell(5, 4, "", 0, 0)
    pdf.cell(90, 4, "Class Area Distribution", 0, 1, "C")

    pdf.image(multi_path, x=x_start, y=pdf.get_y(), w=90, h=60)

    # Class distribution on the right
    legend_x = x_start + 95 + 5
    legend_y = pdf.get_y() + 5

    class_info = [
        ("Foot", result.class_distribution.get("foot", 0), (34, 197, 94)),
        ("Perilesional", result.class_distribution.get("perilesion", 0), (249, 115, 22)),
        ("Ulcer", result.class_distribution.get("ulcer", 0), (239, 68, 68)),
        ("Background", result.class_distribution.get("background", 0), (107, 114, 128)),
    ]

    for cls_name, pct, (r, g, b) in class_info:
        pdf.set_xy(legend_x, legend_y)
        pdf.set_fill_color(r, g, b)
        pdf.rect(legend_x, legend_y + 1, 4, 4, "F")
        pdf._font("B", 9)
        pdf.set_text_color(r, g, b)
        pdf.set_xy(legend_x + 6, legend_y)
        pdf.cell(30, 5, cls_name, 0, 0)
        pdf._font("", 9)
        pdf.set_text_color(50, 50, 50)
        pdf.cell(20, 5, f"{pct:.1f}%", 0, 0)
        # Mini bar
        bar_x = legend_x + 56
        bar_w = 30
        pdf.set_fill_color(229, 231, 235)
        pdf.rect(bar_x, legend_y + 1, bar_w, 4, "F")
        pdf.set_fill_color(r, g, b)
        pdf.rect(bar_x, legend_y + 1, max(0.5, bar_w * pct / 100), 4, "F")
        legend_y += 10

    pdf.ln(62)

    # Image metadata
    h_img, w_img = result.image_size
    pdf._font("", 8)
    pdf.set_text_color(107, 114, 128)
    pdf.cell(0, 4, f"Resolution: {w_img}x{h_img} px  |  Device: {result.device}  |  "
             f"Ulcer area: {result.class_distribution.get('ulcer', 0):.1f}%", 0, 1)
    pdf.ln(4)

    # ── Section 2: Fitzpatrick ──
    pdf.section_title(2, "Fitzpatrick / ITA Skin Type Estimation")
    fitz = result.fitzpatrick
    if fitz and fitz.confidence > 0:
        ftype = fitz.fitzpatrick_type
        bg = FITZ_RGB.get(ftype, (229, 231, 235))
        fg = FITZ_TEXT_RGB.get(ftype, (50, 50, 50))

        # Lighting warning in PDF
        lighting_quality = getattr(fitz, "lighting_quality", "good")
        lighting_warning = getattr(fitz, "lighting_warning", "")
        if lighting_quality == "insufficient":
            pdf.set_fill_color(254, 242, 242)
            pdf.set_draw_color(252, 165, 165)
            pdf.set_text_color(220, 38, 38)
            pdf._font("B", 8)
            y_warn = pdf.get_y()
            pdf.rect(pdf.get_x(), y_warn, 185, 10, "DF")
            pdf.set_xy(pdf.get_x() + 2, y_warn + 1)
            pdf.cell(0, 4, "WARNING: Insufficient lighting — Fitzpatrick type may be overestimated", 0, 1)
            pdf._font("", 7)
            pdf.set_text_color(153, 27, 27)
            pdf.cell(0, 3, lighting_warning, 0, 1)
            pdf.ln(3)
        elif lighting_quality == "low":
            pdf.set_fill_color(255, 251, 235)
            pdf.set_draw_color(252, 211, 77)
            pdf.set_text_color(217, 119, 6)
            pdf._font("B", 8)
            y_warn = pdf.get_y()
            pdf.rect(pdf.get_x(), y_warn, 185, 10, "DF")
            pdf.set_xy(pdf.get_x() + 2, y_warn + 1)
            pdf.cell(0, 4, "CAUTION: Suboptimal lighting — result may be off by 1-2 levels", 0, 1)
            pdf._font("", 7)
            pdf.set_text_color(146, 64, 14)
            pdf.cell(0, 3, lighting_warning, 0, 1)
            pdf.ln(3)

        # Badge
        x_badge = pdf.get_x()
        y_badge = pdf.get_y()
        pdf.set_fill_color(*bg)
        pdf.set_draw_color(180, 180, 180)
        pdf.rect(x_badge, y_badge, 35, 20, "DF")
        pdf._font("B", 16)
        pdf.set_text_color(*fg)
        pdf.set_xy(x_badge, y_badge + 2)
        pdf.cell(35, 9, f"Type {ftype}", 0, 0, "C")
        pdf._font("", 8)
        pdf.set_xy(x_badge, y_badge + 12)
        pdf.cell(35, 6, fitz.fitzpatrick_label, 0, 0, "C")

        # Details table
        pdf.set_text_color(50, 50, 50)
        pdf._font("", 9)
        det_x = x_badge + 40
        det_y = y_badge
        l_scene = getattr(fitz, "l_scene_mean", 0)
        details = [
            ("ITA", f"{fitz.ita_angle:.1f} +/- {fitz.ita_std:.1f} deg"),
            ("L* mean (healthy skin)", f"{fitz.l_skin_mean:.1f}"),
            ("L* scene (global)", f"{l_scene:.1f}"),
            ("b* mean (healthy skin)", f"{fitz.b_skin_mean:.1f}"),
            ("Healthy pixels", f"{fitz.healthy_pixels:,}"),
            ("Confidence", f"{fitz.confidence:.0%}"),
        ]
        for label, value in details:
            pdf.set_xy(det_x, det_y)
            pdf._font("B", 8)
            pdf.cell(42, 4, f"{label}:", 0, 0)
            pdf._font("", 8)
            pdf.cell(50, 4, value, 0, 0)
            det_y += 4.5

        pdf.set_y(y_badge + 22)
    else:
        pdf._font("", 9)
        pdf.set_text_color(107, 114, 128)
        pdf.cell(0, 5, "Not estimable (insufficient healthy-skin pixels).", 0, 1)
    pdf.ln(4)

    # ── Section 3: PWAT ──
    pdf.section_title(3, "PWAT — Raw vs Fitzpatrick-Adjusted Scores")
    pwat = result.pwat
    if pwat and pwat.scores_raw:
        ftype_str = pwat.fitzpatrick_type or "III"

        # Table header
        pdf.set_fill_color(243, 244, 246)
        pdf._font("B", 9)
        pdf.set_text_color(55, 65, 81)
        col_widths = [55, 25, 25, 25, 20, 35]
        headers = ["PWAT Item", "Raw", "Adj.", "Delta", "Scale", ""]
        for w, h in zip(col_widths, headers):
            pdf.cell(w, 6, h, 1, 0, "C", fill=True)
        pdf.ln()

        # Table rows
        for item in [3, 4, 5, 6, 7, 8]:
            name = ITEM_NAMES.get(item, f"Item {item}")
            raw = pwat.scores_raw.get(item, 0)
            adj = pwat.scores_adjusted.get(item, 0.0)
            diff = adj - raw
            diff_str = f"{diff:+.1f}" if abs(diff) > 0.01 else "0.0"

            pdf._font("", 9)
            pdf.set_text_color(50, 50, 50)
            pdf.cell(col_widths[0], 6, name, "LB", 0, "L")
            pdf.cell(col_widths[1], 6, str(raw), "B", 0, "C")
            pdf.cell(col_widths[2], 6, f"{adj:.1f}", "B", 0, "C")

            if diff < -0.05:
                pdf.set_text_color(5, 150, 105)
            else:
                pdf.set_text_color(107, 114, 128)
            pdf._font("B", 9)
            pdf.cell(col_widths[3], 6, diff_str, "B", 0, "C")

            # Visual bar
            pdf.set_text_color(50, 50, 50)
            pdf._font("", 7)
            bar_x = pdf.get_x() + 2
            bar_y = pdf.get_y() + 1.5
            pdf.set_fill_color(229, 231, 235)
            pdf.rect(bar_x, bar_y, col_widths[4] - 4, 3, "F")
            pdf.set_fill_color(239, 68, 68)
            pdf.rect(bar_x, bar_y, max(0.3, (col_widths[4] - 4) * raw / 4), 3, "F")
            pdf.cell(col_widths[4], 6, "", "B", 0)

            # Severity label
            pdf._font("", 7)
            sev_labels = {0: "Normal", 1: "Mild", 2: "Moderate", 3: "Severe", 4: "Extreme"}
            pdf.set_text_color(107, 114, 128)
            pdf.cell(col_widths[5], 6, sev_labels.get(raw, ""), "RB", 0, "L")
            pdf.ln()

        # Total row
        pdf.set_fill_color(31, 41, 55)
        pdf._font("B", 10)
        pdf.set_text_color(255, 255, 255)
        pdf.cell(col_widths[0], 7, "TOTAL", 1, 0, "L", fill=True)
        pdf.cell(col_widths[1], 7, str(pwat.total_raw), 1, 0, "C", fill=True)
        pdf.cell(col_widths[2], 7, f"{pwat.total_adjusted:.1f}", 1, 0, "C", fill=True)
        total_diff = pwat.total_adjusted - pwat.total_raw
        total_diff_str = f"{total_diff:+.1f}" if abs(total_diff) > 0.01 else "0.0"
        pdf.cell(col_widths[3], 7, total_diff_str, 1, 0, "C", fill=True)
        pdf.cell(col_widths[4] + col_widths[5], 7, f"Fitzpatrick {ftype_str}", 1, 0, "C", fill=True)
        pdf.ln(10)

        # Score interpretation
        pdf._font("", 8)
        pdf.set_text_color(107, 114, 128)
        pdf.cell(0, 4, "Scale: 0 (normal) — 4 (extreme) per item. Total: 0-24.", 0, 1)
        pdf.cell(0, 4, f"Bias correction applied for Fitzpatrick type {ftype_str} "
                 "(calibrated on 61 images, r=0.975).", 0, 1)

        # Interpretation ranges
        pdf.ln(2)
        pdf._font("B", 8)
        pdf.set_text_color(55, 65, 81)
        pdf.cell(0, 4, "Interpretation of total score:", 0, 1)
        pdf._font("", 8)
        ranges = [
            ("0-6:",   "Wound healing well",                                     (34, 197, 94)),
            ("7-12:",  "Moderate compromise — clinical follow-up required",      (249, 115, 22)),
            ("13-18:", "Severe compromise — adjust treatment",                   (239, 68, 68)),
            ("19-24:", "Critical wound — urgent reassessment",                   (180, 30, 30)),
        ]
        for label, desc, (r, g, b) in ranges:
            pdf.set_fill_color(r, g, b)
            pdf.rect(pdf.get_x(), pdf.get_y() + 0.5, 3, 3, "F")
            pdf._font("B", 8)
            pdf.set_text_color(r, g, b)
            pdf.set_x(pdf.get_x() + 5)
            pdf.cell(15, 4, label, 0, 0)
            pdf._font("", 8)
            pdf.set_text_color(80, 80, 80)
            pdf.cell(0, 4, desc, 0, 1)

    else:
        pdf._font("", 9)
        pdf.set_text_color(107, 114, 128)
        pdf.cell(0, 5, "Not estimable (ulcer not detected or area too small).", 0, 1)

    # Save PDF
    pdf_path = os.path.join(tmpdir, "WoundNetB7_DFU_Report.pdf")
    pdf.output(pdf_path)

    # Cleanup temp images
    for p in [orig_path, binary_path, multi_path]:
        try:
            os.remove(p)
        except OSError:
            pass

    return pdf_path


def generate_report_files(image_rgb, binary_overlay, multiclass_overlay, result):
    """Generate downloadable report files (PDF + JSON)."""
    tmpdir = tempfile.mkdtemp(prefix="woundnetb7_report_")

    # PDF report
    pdf_path = generate_pdf_report(image_rgb, binary_overlay, multiclass_overlay, result)

    # JSON report
    report_data = result.to_dict()
    report_data["report_metadata"] = {
        "generated_at": datetime.now().isoformat(),
        "model": "WoundNetB7 (EfficientNet-B7 + ASPP + CBAM + CoordAttention + TAM)",
        "ulcer_dice": 0.927,
        "dice_ci_95": [0.917, 0.936],
        "tta_folds": 6,
        "debiasing": "Fitzpatrick-calibrated ITA (86.9% accuracy, r=0.975)",
    }
    json_path = os.path.join(tmpdir, "WoundNetB7_DFU_Report.json")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(report_data, f, indent=2, ensure_ascii=False)

    return [pdf_path, json_path]


# ── Gradio callbacks ─────────────────────────────────────────────────────────

_last_analysis = {}


def analyze_image(image):
    """Main analysis function called by Gradio."""
    if image is None:
        empty = np.zeros((100, 100, 3), dtype=np.uint8)
        _last_analysis.clear()
        return empty, empty, empty, "", "", "", "{}"

    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    result = pipe.analyze(img_bgr, use_tta=True)

    binary_overlay = pipe.visualize_binary(img_bgr, result)
    multiclass_overlay = pipe.visualize_multiclass(img_bgr, result)
    dashboard = pipe.render_integrated_report(img_bgr, result)

    _last_analysis["image_rgb"] = image
    _last_analysis["binary"] = binary_overlay
    _last_analysis["multiclass"] = multiclass_overlay
    _last_analysis["dashboard"] = dashboard
    _last_analysis["result"] = result

    seg_stats = build_seg_stats_html(result)
    fitz_html = build_fitz_html(result.fitzpatrick)
    pwat_html = build_pwat_html(result.pwat)
    json_out = json.dumps(result.to_dict(), indent=2, ensure_ascii=False)

    return dashboard, binary_overlay, multiclass_overlay, seg_stats, fitz_html, pwat_html, json_out


def analyze_from_camera(image):
    """Same analysis but from camera capture (routes to same pipeline)."""
    return analyze_image(image)


def download_report():
    if not _last_analysis:
        return None
    return generate_report_files(
        _last_analysis["image_rgb"],
        _last_analysis["binary"],
        _last_analysis["multiclass"],
        _last_analysis["result"],
    )


# ── HTML builders ─────────────────────────────────────────────────────────────

def build_fitz_html(fitz):
    if fitz is None or fitz.confidence == 0:
        return "<p style='color:#6b7280;'>Not estimable (insufficient healthy-skin pixels).</p>"
    bg = FITZ_COLORS.get(fitz.fitzpatrick_type, "#e5e7eb")
    fg = FITZ_TEXT_COLORS.get(fitz.fitzpatrick_type, "#1f2937")

    # Lighting warning banner
    warning_html = ""
    lighting_warning = getattr(fitz, "lighting_warning", "")
    lighting_quality = getattr(fitz, "lighting_quality", "good")
    l_scene = getattr(fitz, "l_scene_mean", 0)

    if lighting_quality == "insufficient":
        warning_html = f"""
        <div style="background:#fef2f2; border:1px solid #fca5a5; border-radius:8px;
                    padding:12px 16px; margin-bottom:12px; font-size:0.9em;">
          <span style="color:#dc2626; font-weight:700;">&#9888; Insufficient lighting</span><br>
          <span style="color:#991b1b;">{lighting_warning}</span><br>
          <span style="color:#6b7280; font-size:0.85em;">L* scene: {l_scene:.0f} (recommended minimum: 35)</span>
        </div>"""
    elif lighting_quality == "low":
        warning_html = f"""
        <div style="background:#fffbeb; border:1px solid #fcd34d; border-radius:8px;
                    padding:12px 16px; margin-bottom:12px; font-size:0.9em;">
          <span style="color:#d97706; font-weight:700;">&#9888; Suboptimal lighting</span><br>
          <span style="color:#92400e;">{lighting_warning}</span><br>
          <span style="color:#6b7280; font-size:0.85em;">L* scene: {l_scene:.0f} (recommended: >50)</span>
        </div>"""

    return f"""
    {warning_html}
    <div style="display:flex; gap:16px; align-items:center; flex-wrap:wrap;">
      <div style="background:{bg}; color:{fg}; border-radius:12px; padding:18px 28px;
                  font-size:1.5em; font-weight:700; min-width:120px; text-align:center;
                  border:2px solid rgba(0,0,0,0.1);">
        Type {fitz.fitzpatrick_type}<br>
        <span style="font-size:0.55em; font-weight:400;">{fitz.fitzpatrick_label}</span>
      </div>
      <div style="font-size:0.95em; line-height:1.8;">
        <b>ITA:</b> {fitz.ita_angle:.1f}&deg; &plusmn; {fitz.ita_std:.1f}&deg;<br>
        <b>L* healthy skin:</b> {fitz.l_skin_mean:.1f}<br>
        <b>L* scene:</b> {l_scene:.1f}<br>
        <b>Healthy pixels:</b> {fitz.healthy_pixels:,}<br>
        <b>Confidence:</b> {fitz.confidence:.0%}
      </div>
    </div>"""


def build_pwat_html(pwat):
    if pwat is None or not pwat.scores_raw:
        return "<p style='color:#6b7280;'>PWAT not estimable (ulcer not detected or area too small).</p>"
    rows = ""
    for item in [3, 4, 5, 6, 7, 8]:
        name = ITEM_NAMES.get(item, f"Item {item}")
        raw = pwat.scores_raw.get(item, 0)
        adj = pwat.scores_adjusted.get(item, 0.0)
        diff = adj - raw
        diff_color = "#059669" if diff < -0.05 else "#6b7280"
        diff_str = f"{diff:+.1f}" if abs(diff) > 0.01 else "0.0"
        raw_pct = raw / 4 * 100
        adj_pct = adj / 4 * 100
        rows += f"""
        <tr>
          <td style="padding:8px 12px; font-weight:500;">{name}</td>
          <td style="padding:8px 12px; text-align:center;">
            <div style="display:flex; align-items:center; gap:8px;">
              <div style="background:#e5e7eb; border-radius:4px; height:14px; width:80px; overflow:hidden;">
                <div style="background:#ef4444; height:100%; width:{raw_pct}%; border-radius:4px;"></div>
              </div>
              <span style="font-weight:600; min-width:20px;">{raw}</span>
            </div>
          </td>
          <td style="padding:8px 12px; text-align:center;">
            <div style="display:flex; align-items:center; gap:8px;">
              <div style="background:#e5e7eb; border-radius:4px; height:14px; width:80px; overflow:hidden;">
                <div style="background:#3b82f6; height:100%; width:{adj_pct}%; border-radius:4px;"></div>
              </div>
              <span style="font-weight:600; min-width:30px;">{adj:.1f}</span>
            </div>
          </td>
          <td style="padding:8px 12px; text-align:center; color:{diff_color}; font-weight:600;">{diff_str}</td>
        </tr>"""
    total_diff = pwat.total_adjusted - pwat.total_raw
    total_color = "#059669" if total_diff < -0.05 else "#6b7280"
    total_diff_str = f"{total_diff:+.1f}" if abs(total_diff) > 0.01 else "0.0"
    return f"""
    <table style="width:100%; border-collapse:collapse; font-size:0.92em;">
      <thead>
        <tr style="border-bottom:2px solid #d1d5db;">
          <th style="padding:10px 12px; text-align:left;">PWAT Item</th>
          <th style="padding:10px 12px; text-align:center;">Raw Score</th>
          <th style="padding:10px 12px; text-align:center;">Adjusted Score</th>
          <th style="padding:10px 12px; text-align:center;">&Delta;</th>
        </tr>
      </thead>
      <tbody>{rows}
        <tr style="border-top:2px solid #374151; font-weight:700; font-size:1.05em;">
          <td style="padding:10px 12px;">TOTAL</td>
          <td style="padding:10px 12px; text-align:center;">{pwat.total_raw}</td>
          <td style="padding:10px 12px; text-align:center;">{pwat.total_adjusted:.1f}</td>
          <td style="padding:10px 12px; text-align:center; color:{total_color};">{total_diff_str}</td>
        </tr>
      </tbody>
    </table>
    <p style="font-size:0.82em; color:#6b7280; margin-top:8px;">
      Scale: 0 (best) — 4 (worst) per item |
      Fitzpatrick type {pwat.fitzpatrick_type} correction applied |
      Items: 3=Necrotic Type, 4=Necrotic Amount, 5=Granulation Type,
      6=Granulation Amount, 7=Edges, 8=Periulcer Skin
    </p>"""


def build_seg_stats_html(result):
    dist = result.class_distribution
    colors = {"background": "#374151", "foot": "#22c55e", "perilesion": "#f97316", "ulcer": "#ef4444"}
    bars = ""
    for cls_name in ["foot", "perilesion", "ulcer"]:
        pct = dist.get(cls_name, 0)
        color = colors.get(cls_name, "#6b7280")
        label = {"foot": "Foot", "perilesion": "Perilesional", "ulcer": "Ulcer"}.get(cls_name, cls_name)
        bars += f"""
        <div style="margin-bottom:6px;">
          <div style="display:flex; justify-content:space-between; font-size:0.9em; margin-bottom:2px;">
            <span style="color:{color}; font-weight:600;">{label}</span>
            <span>{pct:.1f}%</span>
          </div>
          <div style="background:#e5e7eb; border-radius:4px; height:12px; overflow:hidden;">
            <div style="background:{color}; height:100%; width:{pct}%; border-radius:4px;"></div>
          </div>
        </div>"""
    return f"""
    <div style="padding:4px 0;">
      <p style="font-size:0.85em; color:#6b7280; margin-bottom:10px;">
        Image: {result.image_size[1]}x{result.image_size[0]} | Device: {result.device}
      </p>
      {bars}
    </div>"""


# ── Gradio UI ────────────────────────────────────────────────────────────────

css = """
.step-header {
    display: flex; align-items: center; gap: 10px; margin-bottom: 12px;
}
.step-number {
    background: #1f2937; color: white; border-radius: 50%;
    width: 30px; height: 30px; display: flex; align-items: center;
    justify-content: center; font-weight: 700; font-size: 0.9em; flex-shrink: 0;
}
.step-title { font-weight: 600; font-size: 1.1em; }
"""

with gr.Blocks(
    title="WoundNetB7 DFU Analysis Pipeline",
    theme=gr.themes.Soft(),
    css=css,
) as demo:

    gr.HTML("""
    <div style="text-align:center; padding:20px 0 10px;">
      <h1 style="font-size:1.8em; margin:0;">WoundNetB7 — DFU Analysis Pipeline</h1>
      <p style="color:#6b7280; font-size:1em; margin-top:6px;">
        EfficientNet-B7 + ASPP + CBAM + CoordAttention + TAM &bull; Ulcer Dice: 0.927
      </p>
    </div>
    """)

    with gr.Tabs():
        # ══════════════════════════════════════════════════════════════════════
        # TAB 1: DFU Analysis (upload or gallery)
        # ══════════════════════════════════════════════════════════════════════
        with gr.Tab("DFU Analysis"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(label="DFU Image", type="numpy",
                                           sources=["upload", "clipboard"])
                    analyze_btn = gr.Button("Analyze", variant="primary", size="lg")
                    gr.HTML("""
                    <div style="font-size:0.82em; color:#6b7280; margin-top:8px; line-height:1.6;">
                      <b>Pipeline:</b> the image goes through 4 sequential stages.<br>
                      <b>Model:</b> WoundNetB7 with Combo Loss + Small Object Loss.
                      Attention modules: CBAM, CoordAttention, TAM (fractal + Euler).<br>
                      <b>TTA:</b> 6-fold test-time augmentation.
                    </div>
                    """)

            # Integrated clinical dashboard (primary nurse-facing output)
            gr.HTML("""<div class="step-header" style="margin-top:8px;">
              <div class="step-number" style="background:#0e7490;">&#9733;</div>
              <div class="step-title">Integrated Clinical Assessment Report</div></div>
            <p style="font-size:0.88em; color:#6b7280; margin-bottom:8px;">
              Single-page summary combining segmentation, Fitzpatrick / ITA estimation
              and Fitzpatrick-adjusted PWAT scoring. Designed for clinical staff.
            </p>""")
            output_dashboard = gr.Image(label="Integrated DFU Assessment Report",
                                        show_download_button=True, height=720)

            # Step 1
            gr.HTML("""<div class="step-header"><div class="step-number">1</div>
              <div class="step-title">Binary Ulcer Segmentation</div></div>""")
            with gr.Row():
                with gr.Column(scale=1):
                    output_binary = gr.Image(label="Binary Ulcer Mask (WoundNetB7)")
                with gr.Column(scale=1):
                    output_seg_stats = gr.HTML(label="Segmentation Statistics")

            # Step 2
            gr.HTML("""<div class="step-header" style="margin-top:12px;"><div class="step-number">2</div>
              <div class="step-title">Multi-Class Segmentation (4 classes)</div></div>""")
            with gr.Row():
                with gr.Column(scale=1):
                    output_multiclass = gr.Image(label="Multi-Class Overlay")
                with gr.Column(scale=1):
                    gr.HTML("""
                    <div style="padding:12px;">
                      <p style="font-weight:600; margin-bottom:10px;">Class legend:</p>
                      <div style="display:flex; flex-direction:column; gap:8px;">
                        <div style="display:flex; align-items:center; gap:8px;">
                          <div style="width:20px; height:20px; background:#22c55e; border-radius:4px;"></div>
                          <span><b>Foot</b> — healthy tissue</span>
                        </div>
                        <div style="display:flex; align-items:center; gap:8px;">
                          <div style="width:20px; height:20px; background:#f97316; border-radius:4px;"></div>
                          <span><b>Perilesional</b> — periulcer area</span>
                        </div>
                        <div style="display:flex; align-items:center; gap:8px;">
                          <div style="width:20px; height:20px; background:#ef4444; border-radius:4px;"></div>
                          <span><b>Ulcer</b> — wound bed</span>
                        </div>
                      </div>
                    </div>""")

            # Step 3
            gr.HTML("""<div class="step-header" style="margin-top:12px;"><div class="step-number">3</div>
              <div class="step-title">Fitzpatrick / ITA Skin Type Estimation</div></div>""")
            output_fitz = gr.HTML()

            # Step 4
            gr.HTML("""<div class="step-header" style="margin-top:12px;"><div class="step-number">4</div>
              <div class="step-title">PWAT — Raw vs Fitzpatrick-Adjusted Scores</div></div>""")
            output_pwat = gr.HTML()

            # Download
            gr.HTML("""<div class="step-header" style="margin-top:16px;">
              <div class="step-number" style="background:#059669;">&#8681;</div>
              <div class="step-title">Download Clinical Report</div></div>
            <p style="font-size:0.88em; color:#6b7280; margin-bottom:8px;">
              Generates a PDF report with all visualizations and structured data.
              Run an analysis first.</p>""")
            download_btn = gr.Button("Download PDF Report", variant="secondary", size="lg")
            output_files = gr.File(label="Report Files (PDF + JSON)", file_count="multiple")

            with gr.Accordion("Full JSON (for integration)", open=False):
                output_json = gr.Code(label="JSON Output", language="json")

            analyze_btn.click(
                fn=analyze_image,
                inputs=[input_image],
                outputs=[output_dashboard, output_binary, output_multiclass, output_seg_stats,
                         output_fitz, output_pwat, output_json],
            )
            download_btn.click(fn=download_report, inputs=[], outputs=[output_files])

        # ══════════════════════════════════════════════════════════════════════
        # TAB 2: Guided Capture (webcam with foot guide overlay)
        # ══════════════════════════════════════════════════════════════════════
        with gr.Tab("Guided Capture"):
            gr.HTML("""
            <div style="padding:16px 0;">
              <h2 style="font-size:1.4em; margin:0 0 8px;">Guided Capture for Clinical Staff</h2>
              <p style="color:#6b7280; line-height:1.6;">
                Use the device camera to capture an image of the diabetic foot.
                The green silhouette guides correct foot positioning for optimal analysis.
              </p>
            </div>
            """)

            with gr.Row():
                with gr.Column(scale=3):
                    camera_input = gr.Image(
                        label="Camera — Position the foot inside the guide",
                        type="numpy",
                        sources=["webcam"],
                    )
                    camera_analyze_btn = gr.Button("Capture and Analyze", variant="primary", size="lg")

                with gr.Column(scale=2):
                    guide_image = gr.Image(
                        label="Positioning Guide",
                        value=generate_static_guide(),
                        interactive=False,
                    )

            gr.HTML("""
            <div style="background:#f0fdf4; border:1px solid #bbf7d0; border-radius:10px;
                        padding:16px; margin:12px 0;">
              <p style="font-weight:700; color:#166534; margin:0 0 8px;">
                Instructions for clinical staff:
              </p>
              <div style="display:grid; grid-template-columns:1fr 1fr; gap:8px 24px; font-size:0.92em; color:#15803d;">
                <div>1. Plantar surface facing the camera</div>
                <div>2. Distance: 30-40 cm from the lens</div>
                <div>3. Uniform lighting, no shadows</div>
                <div>4. Neutral background (white/blue sheet)</div>
                <div>5. Include the full ulcer + surrounding healthy skin</div>
                <div>6. Avoid direct flash (causes glare)</div>
                <div>7. Keep the device steady</div>
                <div>8. Clean the lens before capturing</div>
              </div>
            </div>
            """)

            # Integrated dashboard for camera capture
            gr.HTML("""<div class="step-header" style="margin-top:16px;">
              <div class="step-number" style="background:#0e7490;">&#9733;</div>
              <div class="step-title">Integrated Clinical Assessment Report</div></div>""")
            cam_dashboard = gr.Image(label="Integrated DFU Assessment Report",
                                     show_download_button=True, height=720)

            # Results from camera capture
            gr.HTML("""<div class="step-header" style="margin-top:16px;">
              <div class="step-number">1</div>
              <div class="step-title">Segmentation Result</div></div>""")
            with gr.Row():
                cam_binary = gr.Image(label="Binary Ulcer Mask")
                cam_multiclass = gr.Image(label="Multi-Class Overlay")

            with gr.Row():
                cam_seg_stats = gr.HTML()

            gr.HTML("""<div class="step-header"><div class="step-number">2</div>
              <div class="step-title">Fitzpatrick + PWAT</div></div>""")
            cam_fitz = gr.HTML()
            cam_pwat = gr.HTML()

            cam_download_btn = gr.Button("Download PDF Report", variant="secondary", size="lg")
            cam_files = gr.File(label="Report Files", file_count="multiple")

            with gr.Accordion("Full JSON", open=False):
                cam_json = gr.Code(label="JSON Output", language="json")

            camera_analyze_btn.click(
                fn=analyze_from_camera,
                inputs=[camera_input],
                outputs=[cam_dashboard, cam_binary, cam_multiclass, cam_seg_stats,
                         cam_fitz, cam_pwat, cam_json],
            )
            cam_download_btn.click(fn=download_report, inputs=[], outputs=[cam_files])

    gr.HTML("""
    <div style="text-align:center; padding:16px 0; font-size:0.82em; color:#9ca3af;
                border-top:1px solid #e5e7eb; margin-top:20px;">
      WoundNetB7 &bull; Doctoral Thesis &bull; Marcelo Marquez-Murillo &bull;
      Ulcer Dice 0.927 (95% CI: [0.917, 0.936]) &bull;
      Debiasing: 46.6% max group gap reduction (p &lt; 10<sup>-55</sup>)
    </div>
    """)

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