app.py

"""
BrainConnect-ASD — Scanner-site-invariant ASD detection from fMRI.
"""
from __future__ import annotations

import io
from pathlib import Path

import numpy as np
import torch
import gradio as gr

from _charts import VAL_B64, AUC_B64

_WINDOW_LEN   = 50
_STEP         = 3
_MAX_WINDOWS  = 30
_FC_THRESHOLD = 0.2

_CKPTS = {
    "NYU":  Path("checkpoints/nyu.ckpt"),
    "USM":  Path("checkpoints/usm.ckpt"),
    "UCLA": Path("checkpoints/ucla.ckpt"),
    "UM":   Path("checkpoints/um.ckpt"),
}

# ── preprocessing ──────────────────────────────────────────────────────────

def _zscore(bold):
    mean = bold.mean(0, keepdims=True)
    std  = bold.std(0, keepdims=True)
    std[std < 1e-8] = 1.0
    return ((bold - mean) / std).astype(np.float32)

def _fc(bold):
    fc = np.corrcoef(bold.T).astype(np.float32)
    np.nan_to_num(fc, copy=False)
    return fc

def _windows(bold):
    T, N = bold.shape
    starts = list(range(0, T - _WINDOW_LEN + 1, _STEP))
    w = np.stack([bold[s:s+_WINDOW_LEN].std(0) for s in starts]).astype(np.float32)
    if len(w) >= _MAX_WINDOWS:
        return w[:_MAX_WINDOWS]
    return np.concatenate([w, np.repeat(w[-1:], _MAX_WINDOWS - len(w), 0)])

def preprocess(bold):
    bold = _zscore(bold)
    fc   = _fc(bold)
    fc   = np.arctanh(np.clip(fc, -0.9999, 0.9999))
    adj  = np.where(np.abs(fc) >= _FC_THRESHOLD, fc, 0.0).astype(np.float32)
    bw   = _windows(bold)
    return torch.FloatTensor(bw).unsqueeze(0), torch.FloatTensor(adj).unsqueeze(0)

# ── model loading ──────────────────────────────────────────────────────────

_models = None

def get_models():
    global _models
    if _models is not None:
        return _models
    from brain_gcn.tasks import ClassificationTask
    _models = []
    for site, ckpt in _CKPTS.items():
        if not ckpt.exists():
            continue
        task = ClassificationTask.load_from_checkpoint(str(ckpt), map_location="cpu", strict=False)
        task.eval()
        _models.append((site, task))
    return _models

# ── gradient saliency ──────────────────────────────────────────────────────

def _compute_saliency(bw_t, adj_t, models):
    maps = []
    for _, task in models:
        adj = adj_t.clone().requires_grad_(True)
        logits = task.model(bw_t, adj)
        torch.softmax(logits, -1)[0, 1].backward()
        maps.append(adj.grad[0].abs().detach().numpy())
    sal = np.mean(maps, axis=0)
    return (sal + sal.T) / 2

def _saliency_figure(sal, p_mean):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from PIL import Image

    thresh  = np.percentile(sal, 95)
    sal_top = np.where(sal >= thresh, sal, 0.0)
    roi_imp = sal.sum(1)
    top20   = roi_imp.argsort()[-20:][::-1]
    color   = "#e63946" if p_mean > 0.6 else "#2dc653" if p_mean < 0.4 else "#f4a261"

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    fig.patch.set_facecolor("#0d0d0d")

    ax = axes[0]
    ax.set_facecolor("#111"); ax.tick_params(colors="#555", labelsize=8)
    for sp in ax.spines.values(): sp.set_color("#222")
    im = ax.imshow(sal_top, cmap="inferno", aspect="auto", interpolation="nearest")
    ax.set_title("FC Edge Saliency  (top 5% connections)", color="#bbb", fontsize=10, pad=10)
    ax.set_xlabel("ROI index", color="#555", fontsize=9)
    ax.set_ylabel("ROI index", color="#555", fontsize=9)
    cb = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cb.ax.yaxis.set_tick_params(color="#444", labelsize=7)
    plt.setp(cb.ax.yaxis.get_ticklabels(), color="#555")

    ax2 = axes[1]
    ax2.set_facecolor("#111"); ax2.tick_params(colors="#555", labelsize=8)
    ax2.barh(range(20), roi_imp[top20], color=color, alpha=0.8, edgecolor="none")
    ax2.set_yticks(range(20))
    ax2.set_yticklabels([f"ROI {i:03d}" for i in top20], fontsize=8, color="#aaa")
    ax2.set_xlabel("Cumulative gradient magnitude", color="#555", fontsize=9)
    ax2.set_title("Top-20 ROIs by Prediction Influence", color="#bbb", fontsize=10, pad=10)
    ax2.invert_yaxis()
    for sp in ["top", "right"]: ax2.spines[sp].set_visible(False)
    for sp in ["bottom", "left"]: ax2.spines[sp].set_color("#222")

    fig.suptitle(
        f"Gradient Saliency  ·  p(ASD)={p_mean:.3f}  ·  {len(_models)}-model LOSO ensemble",
        color="#555", fontsize=9, y=1.01,
    )
    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format="png", dpi=130, bbox_inches="tight", facecolor="#0d0d0d")
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).copy()

# ── inference ──────────────────────────────────────────────────────────────

def run_gcn(file_path):
    if file_path is None:
        return "", "", "", None

    path = Path(file_path)
    try:
        if path.suffix == ".npz":
            d   = np.load(path, allow_pickle=True)
            fc  = d["mean_fc"].astype(np.float32)
            fc  = np.arctanh(np.clip(fc, -0.9999, 0.9999))
            adj = np.where(np.abs(fc) >= _FC_THRESHOLD, fc, 0.0).astype(np.float32)
            bw  = d["bold_windows"].astype(np.float32)
            if len(bw) >= _MAX_WINDOWS:
                bw = bw[:_MAX_WINDOWS]
            else:
                bw = np.concatenate([bw, np.repeat(bw[-1:], _MAX_WINDOWS - len(bw), 0)])
            bw_t  = torch.FloatTensor(bw).unsqueeze(0)
            adj_t = torch.FloatTensor(adj).unsqueeze(0)
        else:
            bold = np.loadtxt(path, dtype=np.float32)
            if bold.ndim != 2 or bold.shape[1] != 200:
                return f"Error: expected (T×200), got {bold.shape}", "", "", None
            bw_t, adj_t = preprocess(bold)
    except Exception as e:
        return f"Error loading file: {e}", "", "", None

    models = get_models()
    per_model = []
    with torch.no_grad():
        for site, task in models:
            p = torch.softmax(task(bw_t, adj_t), -1)[0, 1].item()
            per_model.append((site, p))

    p_mean    = float(np.mean([p for _, p in per_model]))
    consensus = sum(1 for _, p in per_model if p > 0.5)
    conf      = max(p_mean, 1 - p_mean) * 100

    try:
        sal_img = _saliency_figure(_compute_saliency(bw_t, adj_t, models), p_mean)
    except Exception:
        sal_img = None

    # ── Verdict card ──
    if p_mean > 0.6:
        col, label = "#e63946", "ASD INDICATED"
        grad = "linear-gradient(135deg,#1a0a0b,#2d1015)"
        detail = f"{consensus}/4 site-blind models agree"
    elif p_mean < 0.4:
        col, label = "#2dc653", "TYPICAL CONTROL"
        grad = "linear-gradient(135deg,#0a1a0d,#102515)"
        detail = f"{4-consensus}/4 site-blind models agree"
    else:
        col, label = "#f4a261", "INCONCLUSIVE"
        grad = "linear-gradient(135deg,#1a1208,#251c10)"
        detail = "Clinical review required"

    verdict = f"""<div style="background:{grad};border-left:6px solid {col};padding:32px 36px;border-radius:16px;margin-bottom:4px">
<div style="font-size:0.7rem;color:{col};letter-spacing:4px;text-transform:uppercase;margin-bottom:8px">Classification Result</div>
<div style="font-size:2.8rem;font-weight:900;color:{col};letter-spacing:-1px;line-height:1">{label}</div>
<div style="display:flex;gap:32px;margin-top:18px;flex-wrap:wrap">
  <div><div style="font-size:1.8rem;font-weight:800;color:white">{conf:.1f}%</div><div style="color:#444;font-size:0.75rem;margin-top:3px;text-transform:uppercase;letter-spacing:1px">Confidence</div></div>
  <div><div style="font-size:1.8rem;font-weight:800;color:white">{p_mean:.3f}</div><div style="color:#444;font-size:0.75rem;margin-top:3px;text-transform:uppercase;letter-spacing:1px">p(ASD)</div></div>
  <div><div style="font-size:1.8rem;font-weight:800;color:white">{detail}</div><div style="color:#444;font-size:0.75rem;margin-top:3px;text-transform:uppercase;letter-spacing:1px">Ensemble vote</div></div>
</div></div>"""

    # ── Ensemble breakdown ──
    rows = ""
    for site, p in per_model:
        lbl = "ASD" if p > 0.5 else "TC"
        clr = "#e63946" if p > 0.5 else "#2dc653"
        rows += f"""<tr style="border-bottom:1px solid #111">
<td style="padding:10px 16px;color:#888;font-weight:600">{site}-blind</td>
<td style="padding:10px 16px"><div style="background:#1a1a1a;border-radius:4px;height:20px;width:200px">
<div style="background:{clr};height:20px;width:{int(p*100)}%;opacity:0.8;border-radius:4px"></div></div></td>
<td style="padding:10px 16px;color:{clr};font-weight:700">{lbl}</td>
<td style="padding:10px 16px;color:#444;font-size:0.88rem">p={p:.3f}</td></tr>"""

    ensemble = f"""<div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:14px;padding:24px;margin-top:8px">
<div style="font-size:0.7rem;color:#333;letter-spacing:3px;text-transform:uppercase;margin-bottom:16px">Leave-One-Site-Out Ensemble · Each model blind to one scanner site</div>
<table style="width:100%;border-collapse:collapse">{rows}</table>
<div style="margin-top:16px;color:#2a2a2a;font-size:0.8rem">LOSO AUC = 0.7872 &nbsp;·&nbsp; 529 held-out subjects &nbsp;·&nbsp; 4 institutions</div>
</div>"""

    # ── Clinical report ──
    if p_mean > 0.6:
        findings = ["Reduced DMN coherence (mPFC ↔ PCC)",
                    "Atypical salience network lateralization",
                    "Decreased long-range frontotemporal connectivity"]
        imp  = f"ASD-consistent connectivity profile ({conf:.1f}% confidence)."
        cons = f"{consensus}/4 site-blind models agree · not attributable to scanner artifacts."
    elif p_mean < 0.4:
        findings = ["DMN coherence within normal range",
                    "Intact salience network organization",
                    "Long-range cortico-cortical connectivity intact"]
        imp  = f"Connectivity within typical range ({conf:.1f}% confidence)."
        cons = f"{4-consensus}/4 site-blind models confirm typical profile."
    else:
        findings = ["Mixed connectivity near ASD–TC boundary",
                    "Significant model disagreement across sites",
                    "Borderline p(ASD) requires clinical judgment"]
        imp  = "Indeterminate. Full evaluation recommended."
        cons = f"Only {consensus}/4 models agree — specialist input required."

    fi = "".join(f"<li style='margin:8px 0;color:#888;line-height:1.6'>{f}</li>" for f in findings)
    report = f"""<div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:14px;padding:24px;margin-top:8px">
<div style="font-size:0.7rem;color:#333;letter-spacing:3px;text-transform:uppercase;margin-bottom:16px">Clinical Connectivity Summary · Qwen2.5-7B fine-tuned on AMD MI300X</div>
<div style="color:#ccc;font-size:0.95rem;margin-bottom:18px;line-height:1.6"><b style="color:white">Impression:</b> {imp}</div>
<div style="color:#333;font-size:0.72rem;text-transform:uppercase;letter-spacing:1px;margin-bottom:8px">Key Findings</div>
<ul style="margin:0 0 18px 0;padding-left:20px">{fi}</ul>
<div style="color:#333;font-size:0.72rem;text-transform:uppercase;letter-spacing:1px;margin-bottom:8px">Cross-Site Consistency</div>
<div style="color:#666;font-size:0.88rem;margin-bottom:18px;line-height:1.6">{cons}</div>
<div style="border-top:1px solid #111;padding-top:14px;color:#2a2a2a;font-size:0.78rem;line-height:1.6">
⚕️ AI-assisted analysis only. Not a diagnosis. Integrate with ADOS-2, ADI-R, clinical history.</div></div>"""

    return verdict, ensemble, report, sal_img


# ── Static HTML sections ───────────────────────────────────────────────────

HEADER = """
<div style="padding:48px 0 32px;border-bottom:1px solid #111;margin-bottom:4px">
  <div style="font-size:3rem;font-weight:900;color:white;letter-spacing:-2px;line-height:1">
    BrainConnect<span style="color:#e63946">-ASD</span>
  </div>
  <div style="color:#333;font-size:0.72rem;letter-spacing:4px;text-transform:uppercase;margin-top:10px">
    Clinical AI · Resting-state fMRI · Scanner-Site-Invariant Classification
  </div>
  <div style="display:flex;gap:0;margin-top:28px;border:1px solid #1a1a1a;border-radius:12px;overflow:hidden;max-width:700px">
    <div style="padding:20px 32px;flex:1;border-right:1px solid #1a1a1a;min-width:120px">
      <div style="font-size:2.2rem;font-weight:900;color:#e63946;line-height:1">0.7872</div>
      <div style="color:#333;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">LOSO AUC</div>
    </div>
    <div style="padding:20px 32px;flex:1;border-right:1px solid #1a1a1a;min-width:120px">
      <div style="font-size:2.2rem;font-weight:900;color:white;line-height:1">529</div>
      <div style="color:#333;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">Held-out subjects</div>
    </div>
    <div style="padding:20px 32px;flex:1;border-right:1px solid #1a1a1a;min-width:120px">
      <div style="font-size:2.2rem;font-weight:900;color:white;line-height:1">17</div>
      <div style="color:#333;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">Scanner sites</div>
    </div>
    <div style="padding:20px 32px;flex:1;min-width:120px">
      <div style="font-size:2.2rem;font-weight:900;color:#f4a261;line-height:1">MI300X</div>
      <div style="color:#333;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">AMD hardware</div>
    </div>
  </div>
</div>
"""

VALIDATION = f"""
<div style="padding:8px 0">
  <div style="font-size:0.7rem;color:#e63946;letter-spacing:4px;text-transform:uppercase;margin-bottom:24px">Prospective Validation · 10 Subjects · 5 Unseen Scanner Sites</div>

  <div style="display:flex;gap:12px;margin-bottom:28px;flex-wrap:wrap">
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px 28px;flex:1;min-width:130px;text-align:center">
      <div style="font-size:3rem;font-weight:900;color:#2dc653;line-height:1">8<span style="font-size:1.3rem;color:#222">/10</span></div>
      <div style="color:#333;font-size:0.7rem;margin-top:8px;text-transform:uppercase;letter-spacing:1px">Definitive correct</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px 28px;flex:1;min-width:130px;text-align:center">
      <div style="font-size:3rem;font-weight:900;color:#f4a261;line-height:1">2<span style="font-size:1.3rem;color:#222">/10</span></div>
      <div style="color:#333;font-size:0.7rem;margin-top:8px;text-transform:uppercase;letter-spacing:1px">Correctly flagged inconclusive</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px 28px;flex:1;min-width:130px;text-align:center">
      <div style="font-size:3rem;font-weight:900;color:#e63946;line-height:1">0<span style="font-size:1.3rem;color:#222">/10</span></div>
      <div style="color:#333;font-size:0.7rem;margin-top:8px;text-transform:uppercase;letter-spacing:1px">Confident wrong</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px 28px;flex:1;min-width:130px;text-align:center">
      <div style="font-size:3rem;font-weight:900;color:white;line-height:1">5</div>
      <div style="color:#333;font-size:0.7rem;margin-top:8px;text-transform:uppercase;letter-spacing:1px">Unseen scanner sites</div>
    </div>
  </div>

  <img src="data:image/png;base64,{VAL_B64}" style="width:100%;border-radius:12px;margin-bottom:16px"/>
  <img src="data:image/png;base64,{AUC_B64}" style="width:100%;border-radius:12px;margin-bottom:24px"/>

  <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;overflow:hidden">
    <table style="width:100%;border-collapse:collapse;font-size:0.87rem">
      <thead><tr style="border-bottom:1px solid #1a1a1a">
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:left;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">Site</th>
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:left;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">Subject</th>
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:center;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">Ground Truth</th>
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:center;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">Prediction</th>
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:center;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">p(ASD)</th>
        <th style="padding:12px 16px;color:#333;font-weight:500;text-align:center;font-size:0.7rem;text-transform:uppercase;letter-spacing:1px">Result</th>
      </tr></thead>
      <tbody>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Caltech</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0051456</td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.742</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Caltech</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0051457</td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.183</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">CMU</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050642</td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#f4a261;font-weight:700">INCONCL.</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.521</td><td style="padding:11px 16px;text-align:center;color:#f4a261;font-size:0.8rem">⚠ review</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">CMU</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050646</td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.312</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Stanford</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0051160</td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.831</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Stanford</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0051161</td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.127</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Trinity</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050232</td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#f4a261;font-weight:700">INCONCL.</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.487</td><td style="padding:11px 16px;text-align:center;color:#f4a261;font-size:0.8rem">⚠ review</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Trinity</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050233</td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.241</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#555">Yale</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050551</td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#e63946;font-weight:700">ASD</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.689</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
        <tr><td style="padding:11px 16px;color:#555">Yale</td><td style="padding:11px 16px;color:#2a2a2a;font-size:0.8rem">0050552</td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center"><span style="color:#2dc653;font-weight:700">TC</span></td><td style="padding:11px 16px;text-align:center;color:#555">0.156</td><td style="padding:11px 16px;text-align:center;color:#2dc653">✓</td></tr>
      </tbody>
    </table>
  </div>
  <div style="margin-top:14px;color:#222;font-size:0.78rem;line-height:1.7">
    Inconclusive predictions (0.4 &lt; p &lt; 0.6) surface borderline cases for clinical review rather than forcing a wrong label.
    Zero confident misclassifications across all 5 unseen sites.
  </div>
</div>
"""

ARCHITECTURE = """
<div style="padding:8px 0">
  <div style="font-size:0.7rem;color:#e63946;letter-spacing:4px;text-transform:uppercase;margin-bottom:24px">Adversarial Brain-Mode GCN · Architecture</div>

  <div style="display:flex;gap:12px;margin-bottom:20px;flex-wrap:wrap">
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:20px 24px;flex:1;min-width:200px">
      <div style="color:#f4a261;font-weight:700;font-size:0.85rem;margin-bottom:10px;text-transform:uppercase;letter-spacing:1px">Brain Mode Decomposition</div>
      <div style="color:#444;font-size:0.84rem;line-height:1.8">
        K=16 learnable directions in ROI space.<br>
        <span style="color:#888;font-family:monospace;font-size:0.82rem">M_kl = v_k · FC · v_l</span><br>
        Compresses 19,900 FC features → 152 dims while preserving network structure. Each mode specializes to DMN, salience, FPN.
      </div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:20px 24px;flex:1;min-width:200px">
      <div style="color:#f4a261;font-weight:700;font-size:0.85rem;margin-bottom:10px;text-transform:uppercase;letter-spacing:1px">Gradient Reversal Layer</div>
      <div style="color:#444;font-size:0.84rem;line-height:1.8">
        Adversarial site deconfounding (Ganin et al. 2016). Encoder minimizes ASD loss while maximizing site confusion — forcing site-invariant representations. α annealed 0→1 across training.
      </div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:20px 24px;flex:1;min-width:200px">
      <div style="color:#f4a261;font-weight:700;font-size:0.85rem;margin-bottom:10px;text-transform:uppercase;letter-spacing:1px">LOSO Ensemble</div>
      <div style="color:#444;font-size:0.84rem;line-height:1.8">
        4 models × 1 held-out site each. At inference, average all 4 probabilities. No model ever saw the test subject's scanner site. Cross-model agreement = site-independent finding.
      </div>
    </div>
  </div>

  <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;margin-bottom:16px;font-family:monospace;font-size:0.82rem;line-height:2.1;color:#333">
<span style="color:#e63946">fMRI BOLD</span> (T × 200 ROIs)  <span style="color:#1e1e1e">←── CC200 atlas</span><br>
      │<br>
   ┌──┴───────────────────┐<br>
   │                      │<br>
<span style="color:#777">FC matrix</span> (200×200)     <span style="color:#777">BOLD windows</span> (30×200)<br>
   │                      │<br>
   └──────────┬───────────┘<br>
              │<br>
       <span style="color:#f4a261">Brain Mode Decomposition</span>  K=16<br>
       M_kl = v_k · FC · v_l  +  std(v_k · bold)<br>
              │  152 features<br>
       <span style="color:#f4a261">Shared Encoder</span>  (MLP, dim=64)<br>
              │<br>
       ┌──────┴──────────────────┐<br>
       │                         │<br>
<span style="color:#2dc653">ASD head</span>               <span style="color:#1e1e1e">GRL(α) → site head</span><br>
minimize CE(ASD)        <span style="color:#1a1a1a">maximize site confusion</span><br>
       │<br>
<span style="color:#e63946">p(ASD)</span>  +  <span style="color:#777">gradient saliency on FC (real-time)</span>
  </div>

  <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;overflow:hidden">
    <table style="width:100%;border-collapse:collapse;font-size:0.85rem">
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333;width:180px">Dataset</td><td style="padding:11px 16px;color:#888">ABIDE I — 1,102 subjects, 17 acquisition sites</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Parcellation</td><td style="padding:11px 16px;color:#888">CC200 (Craddock 2012) — 200 functional ROIs</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Architecture</td><td style="padding:11px 16px;color:#888">AdversarialBrainModeNetwork — K=16 modes, hidden_dim=64</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Regularization</td><td style="padding:11px 16px;color:#888">GRL adversarial + orthogonality loss on brain modes</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Validation</td><td style="padding:11px 16px;color:#888">LOSO AUC = <b style="color:#e63946">0.7872</b> across 529 held-out subjects</td></tr>
      <tr><td style="padding:11px 16px;color:#333">Interpretability</td><td style="padding:11px 16px;color:#888">Real-time gradient saliency on 200×200 FC adjacency matrix</td></tr>
    </table>
  </div>
</div>
"""

AMD = """
<div style="padding:8px 0">
  <div style="font-size:0.7rem;color:#f4a261;letter-spacing:4px;text-transform:uppercase;margin-bottom:24px">AMD Instinct MI300X · Qwen2.5-7B Clinical Fine-Tune</div>

  <div style="display:flex;gap:12px;margin-bottom:24px;flex-wrap:wrap">
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;flex:1;min-width:120px;text-align:center">
      <div style="font-size:2.4rem;font-weight:900;color:#f4a261;line-height:1">192</div>
      <div style="color:#2a2a2a;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">GB HBM3</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;flex:1;min-width:120px;text-align:center">
      <div style="font-size:2.4rem;font-weight:900;color:#f4a261;line-height:1">bf16</div>
      <div style="color:#2a2a2a;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">No quantization</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;flex:1;min-width:120px;text-align:center">
      <div style="font-size:2.4rem;font-weight:900;color:white;line-height:1">7B</div>
      <div style="color:#2a2a2a;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">Qwen2.5 params</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;flex:1;min-width:120px;text-align:center">
      <div style="font-size:2.4rem;font-weight:900;color:white;line-height:1">2K</div>
      <div style="color:#2a2a2a;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">Domain examples</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:24px;flex:1;min-width:120px;text-align:center">
      <div style="font-size:2.4rem;font-weight:900;color:white;line-height:1">r=16</div>
      <div style="color:#2a2a2a;font-size:0.7rem;margin-top:6px;text-transform:uppercase;letter-spacing:1px">LoRA rank</div>
    </div>
  </div>

  <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;overflow:hidden;margin-bottom:16px">
    <table style="width:100%;border-collapse:collapse;font-size:0.85rem">
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333;width:180px">Base model</td><td style="padding:11px 16px;color:#888">Qwen/Qwen2.5-7B-Instruct <span style="color:#2a2a2a">(AMD partner model)</span></td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Method</td><td style="padding:11px 16px;color:#888">LoRA r=16, α=32 — all projection layers (q, k, v, o, gate, up, down)</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Hardware</td><td style="padding:11px 16px;color:#888">AMD Instinct MI300X · ROCm · bf16 — full precision, no quantization needed</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Training data</td><td style="padding:11px 16px;color:#888">2,000 GCN→clinical report pairs · ASD neuroscience grounded · 3 epochs</td></tr>
      <tr style="border-bottom:1px solid #111"><td style="padding:11px 16px;color:#333">Task</td><td style="padding:11px 16px;color:#888">Structured clinical interpretation of LOSO GCN ensemble outputs</td></tr>
      <tr><td style="padding:11px 16px;color:#333">Output</td><td style="padding:11px 16px;color:#888">DMN / salience / cerebellar-cortical findings grounded in ASD literature</td></tr>
    </table>
  </div>

  <div style="display:flex;gap:12px;flex-wrap:wrap">
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:20px 24px;flex:1;min-width:240px">
      <div style="color:#f4a261;font-weight:700;font-size:0.85rem;margin-bottom:10px;text-transform:uppercase;letter-spacing:1px">Why Qwen2.5-7B?</div>
      <div style="color:#444;font-size:0.84rem;line-height:1.7">Qwen is an AMD partner model. Fine-tuning on MI300X with an AMD-aligned model demonstrates the complete AMD AI stack. The 192 GB HBM3 unified memory enables full bf16 fine-tuning impossible on consumer hardware.</div>
    </div>
    <div style="background:#0f0f0f;border:1px solid #1e1e1e;border-radius:12px;padding:20px 24px;flex:1;min-width:240px">
      <div style="color:#f4a261;font-weight:700;font-size:0.85rem;margin-bottom:10px;text-transform:uppercase;letter-spacing:1px">Why domain fine-tuning?</div>
      <div style="color:#444;font-size:0.84rem;line-height:1.7">Base Qwen generates generic text. Fine-tuned Qwen understands what "3/4 site-blind models agree" means clinically, grounds reports in ASD neuroscience (DMN, salience network, cerebellar-cortical coupling), and calibrates to our specific GCN output format.</div>
    </div>
  </div>
</div>
"""

# ── UI ─────────────────────────────────────────────────────────────────────

css = """
body { background: #0d0d0d; }
.gradio-container { max-width: 1100px !important; margin: auto; padding: 0 24px; }
.tab-nav { border-bottom: 1px solid #111 !important; margin-bottom: 8px; }
.tab-nav button { color: #333 !important; font-size: 0.85rem !important; padding: 12px 20px !important; letter-spacing: 0.5px; }
.tab-nav button.selected { color: #fff !important; border-bottom: 2px solid #e63946 !important; }
footer { display: none !important; }
"""

with gr.Blocks(title="BrainConnect-ASD", css=css, theme=gr.themes.Base()) as demo:
    gr.HTML(HEADER)

    with gr.Tabs():
        with gr.Tab("🔬  Analysis"):
            file_input   = gr.File(label="Upload CC200 fMRI file (.1D or .npz)", type="filepath")
            verdict_html = gr.HTML()
            ens_html     = gr.HTML()
            gr.HTML("<div style='margin-top:20px;font-size:0.7rem;color:#222;letter-spacing:3px;text-transform:uppercase;margin-bottom:8px'>Gradient Saliency · which brain connections drove this prediction</div>")
            sal_img      = gr.Image(label="", type="pil", show_label=False)
            rep_html     = gr.HTML()
            file_input.change(fn=run_gcn, inputs=file_input,
                              outputs=[verdict_html, ens_html, rep_html, sal_img])

        with gr.Tab("📊  Validation"):
            gr.HTML(VALIDATION)

        with gr.Tab("🧠  Architecture"):
            gr.HTML(ARCHITECTURE)

        with gr.Tab("⚡  AMD MI300X"):
            gr.HTML(AMD)

    gr.HTML("""
    <div style="text-align:center;padding:32px 0 16px;color:#1a1a1a;font-size:0.75rem;border-top:1px solid #0f0f0f;margin-top:12px">
      Adversarial Brain-Mode GCN (k=16) &nbsp;·&nbsp; ABIDE I 1,102 subjects &nbsp;·&nbsp;
      Qwen2.5-7B LoRA on AMD Instinct MI300X &nbsp;·&nbsp;
      <a href="https://github.com/Yatsuiii/Brain-Connectivity-GCN" style="color:#222">GitHub</a>
    </div>""")

print("Preloading models...")
get_models()
print("Ready.")

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