app.py

"""
app.py
======
Mutation Explainability Intelligence System
Gradio Space — explanation-first clinical variant analysis

Three models:
  nileshhanotia/mutation-predictor-splice
  nileshhanotia/mutation-predictor-v4
  nileshhanotia/mutation-pathogenicity-predictor

Explanation ALWAYS precedes prediction panel.
"""

from __future__ import annotations
import io
import json
import logging
import os
import sys
import tempfile
import traceback

import gradio as gr
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LinearSegmentedColormap

def _fig_to_pil(fig):
    """Render matplotlib figure to PIL Image — required for gr.Image in Gradio 4.44."""
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110, bbox_inches="tight",
                facecolor=fig.get_facecolor())
    buf.seek(0)
    from PIL import Image as _PILImage
    img = _PILImage.open(buf).copy()
    plt.close(fig)
    return img


import requests
import time
from functools import lru_cache

# ── project imports ───────────────────────────────────────────────────────────
from model_loader import (
    ModelRegistry,
    encode_for_v2,
    find_mutation_pos,
)
from explainability_engine import (
    extract_splice_signals,
    extract_v4_signals,
    extract_classic_signals,
    compute_cross_model_analysis,
)
from decision_engine import build_decision, DecisionResult

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
)
logger = logging.getLogger("mutation_xai")

# ── Global registry (lazy) ────────────────────────────────────────────────────
REGISTRY = ModelRegistry(hf_token=os.environ.get("HF_TOKEN"))

# ── Ensembl fetch ─────────────────────────────────────────────────────────────
ENSEMBL_URL = "https://rest.ensembl.org/sequence/region/human"
WINDOW_HALF = 49   # 49 + 1 + 49 = 99 bp (matches all three models)


@lru_cache(maxsize=256)
def _fetch_ensembl(chrom: str, start: int, end: int) -> str:
    chrom = chrom.lstrip("chrCHR").strip()
    region = f"{chrom}:{start}..{end}:1"
    url    = f"{ENSEMBL_URL}/{region}"
    for attempt in range(3):
        try:
            r = requests.get(url, params={"content-type": "application/json"}, timeout=15)
            if r.status_code == 429:
                time.sleep(int(r.headers.get("Retry-After", 5)))
                continue
            r.raise_for_status()
            data = r.json()
            if isinstance(data, list): data = data[0]
            return data.get("seq", "").upper()
        except Exception as e:
            if attempt == 2:
                raise RuntimeError(f"Ensembl API failed: {e}")
            time.sleep(1.5 * (2 ** attempt))
    return ""


def fetch_window(chrom: str, pos: int) -> tuple[str, str, int]:
    """
    Returns (ref_seq_99bp, mut_seq_placeholder, mutation_pos_in_window).
    Caller must insert the alt base into mut_seq at mutation_pos.
    """
    chrom_clean = chrom.lstrip("chrCHR").strip()
    start = max(1, pos - WINDOW_HALF)
    end   = pos + WINDOW_HALF
    seq   = _fetch_ensembl(chrom_clean, start, end)
    if len(seq) < 1:
        raise ValueError(f"Empty sequence returned for chr{chrom}:{start}-{end}")
    # Pad/trim to 99
    seq = (seq + "N" * 99)[:99]
    mut_pos = pos - start  # 0-indexed position within window
    mut_pos = max(0, min(98, mut_pos))
    return seq, mut_pos


# ═══════════════════════════════════════════════════════════════════════════════
# Visualisation helpers
# ═══════════════════════════════════════════════════════════════════════════════

_BG    = "#0D1117"
_TEXT  = "#E6EDF3"
_MUTED = "#7D8590"
_BLUE  = "#58A6FF"
_GREEN = "#3FB950"
_RED   = "#F85149"
_ORG   = "#D29922"

_CMAP_ACTIVATION = LinearSegmentedColormap.from_list(
    "act", [(0.04,0.22,0.47),(0.96,0.96,0.96),(0.72,0.05,0.12)], N=256)
_CMAP_SPLICE = LinearSegmentedColormap.from_list(
    "splice", [(0.0,"#f7f7f7"),(0.3,"#fee08b"),(0.6,"#fc8d59"),(1.0,"#d73027")])


def _fig_base(w=15, h=2.8):
    fig, ax = plt.subplots(figsize=(w, h), facecolor=_BG)
    ax.set_facecolor(_BG)
    return fig, ax


def _style_ax(ax, title):
    ax.set_title(title, color=_TEXT, fontsize=9, loc="left", pad=4, fontweight="bold")
    for sp in ["top","right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    ax.tick_params(colors=_TEXT, labelsize=7)


def plot_activation_heatmap(profile: np.ndarray, mutation_pos: int,
                             label: str, prob: float):
    imp = profile.copy()
    if imp.max() > 0:
        imp /= imp.max()
    fig, ax = _fig_base(15, 2.5)
    im = ax.imshow(imp[np.newaxis,:], aspect="auto", cmap=_CMAP_ACTIVATION,
                   vmin=0, vmax=1, extent=[-0.5, 98.5, 0, 1])
    if mutation_pos >= 0:
        ax.axvline(x=mutation_pos, color=_GREEN, linewidth=2.0, linestyle="--",
                   label=f"Mutation pos {mutation_pos}")
        ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.6,
                  loc="upper right")
    cb = fig.colorbar(im, ax=ax, pad=0.01)
    cb.set_label("Activation intensity", color=_TEXT, fontsize=8)
    cb.ax.tick_params(colors=_TEXT, labelsize=7)
    ax.set_xlabel("Nucleotide position (99 bp window)", color=_TEXT, fontsize=9)
    ax.set_xticks(range(0, 99, 10))
    ax.set_yticks([])
    _style_ax(ax, f"CNN conv3 Activation — {label}  (prob={prob:.4f})")
    fig.tight_layout()
    return _fig_to_pil(fig)


def plot_splice_heatmap(ref_seq: str, mutation_pos: int):
    seq    = (ref_seq.upper() + "N" * 99)[:99]
    scores = np.zeros(99)
    donors, acceptors = [], []
    for i in range(len(seq)-1):
        if seq[i:i+2] == "GT": donors.append(i)
        if seq[i:i+2] == "AG": acceptors.append(i)
    for p in donors:
        for d in range(-8,9):
            if 0 <= p+d < 99: scores[p+d] = max(scores[p+d], 0.5)
    for p in acceptors:
        for d in range(-8,9):
            if 0 <= p+d < 99: scores[p+d] = max(scores[p+d], 0.5)
    for p in donors:
        if 0 <= p < 99: scores[p] = 1.0
    for p in acceptors:
        if 0 <= p < 99: scores[p] = max(scores[p], 0.8)

    fig, ax = _fig_base(15, 2.5)
    im = ax.imshow(scores[np.newaxis,:], aspect="auto", cmap=_CMAP_SPLICE,
                   vmin=0, vmax=1, extent=[-0.5, 98.5, 0, 1])
    if mutation_pos >= 0:
        ax.axvline(x=mutation_pos, color=_BLUE, linewidth=2.0, linestyle="--",
                   label=f"Mutation pos {mutation_pos}")
        ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.6,
                  loc="upper right")
    cb = fig.colorbar(im, ax=ax, pad=0.01)
    cb.set_label("Splice risk", color=_TEXT, fontsize=8)
    cb.ax.tick_params(colors=_TEXT, labelsize=7)
    ax.set_xlabel("Nucleotide position (99 bp window)", color=_TEXT, fontsize=9)
    ax.set_xticks(range(0, 99, 10))
    ax.set_yticks([])
    _style_ax(ax, "Splice Distance Risk — GT donor / AG acceptor signals")
    fig.tight_layout()
    return _fig_to_pil(fig)


def plot_gradient_heatmap(attr: np.ndarray, mutation_pos: int, label: str):
    fig, ax = _fig_base(15, 2.5)
    im = ax.imshow(attr[np.newaxis,:], aspect="auto", cmap="PuOr",
                   vmin=0, vmax=1, extent=[-0.5, 98.5, 0, 1])
    if mutation_pos >= 0:
        ax.axvline(x=mutation_pos, color=_GREEN, linewidth=2.0, linestyle="--",
                   label=f"Mutation pos {mutation_pos}")
        ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.6,
                  loc="upper right")
    cb = fig.colorbar(im, ax=ax, pad=0.01)
    cb.set_label("Gradient attribution", color=_TEXT, fontsize=8)
    cb.ax.tick_params(colors=_TEXT, labelsize=7)
    ax.set_xlabel("Nucleotide position", color=_TEXT, fontsize=9)
    ax.set_xticks(range(0, 99, 10))
    ax.set_yticks([])
    _style_ax(ax, f"Gradient Attribution Map — {label}")
    fig.tight_layout()
    return _fig_to_pil(fig)


def plot_counterfactual(cf_table: list[dict], orig_prob: float, cf_delta: float):
    if not cf_table:
        fig, ax = plt.subplots(figsize=(8, 3), facecolor=_BG)
        ax.text(0.5, 0.5, "No counterfactual data", ha="center", va="center",
                color=_TEXT, fontsize=12)
        ax.axis("off")
        return _fig_to_pil(fig)

    labels = [r["mutation"] for r in cf_table]
    probs  = [r["probability"] for r in cf_table]
    max_p, min_p = max(probs), min(probs)
    colors = [_RED if p == max_p else (_BLUE if p == min_p else "#74add1") for p in probs]

    fig, ax = plt.subplots(figsize=(10, 3.5), facecolor=_BG)
    ax.set_facecolor(_BG)
    bars = ax.bar(labels, probs, color=colors, edgecolor="#444", linewidth=0.7)
    ax.axhline(0.5,    color=_MUTED, linestyle="--", linewidth=1.0, label="Decision boundary (0.5)")
    ax.axhline(orig_prob, color=_ORG,  linestyle="-.", linewidth=1.5,
               label=f"Original mutation ({orig_prob:.3f})")
    ax.set_ylim(0, 1.05)
    ax.set_xlabel("Alternative mutation", color=_TEXT, fontsize=10)
    ax.set_ylabel("Pathogenicity probability", color=_TEXT, fontsize=10)
    ax.tick_params(colors=_TEXT)
    ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.5)
    for b, p in zip(bars, probs):
        ax.text(b.get_x() + b.get_width()/2, b.get_height()+0.01,
                f"{p:.3f}", ha="center", va="bottom", fontsize=8, color=_TEXT)
    for sp in ["top","right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    ax.set_title(
        f"Counterfactual Analysis  |  Δ={cf_delta:.4f}  |  "
        f"range {min_p:.3f}–{max_p:.3f}",
        color=_TEXT, fontsize=10, loc="left")
    fig.tight_layout()
    return _fig_to_pil(fig)


def plot_ablation(ablation: dict):
    labels = [
        "Splice features\n(donor/acceptor/region)",
        "Region features\n(exon/intron flags)",
        "Mutation type\n(one-hot)",
    ]
    deltas = [ablation["splice_causal_effect"],
              ablation["region_causal_effect"],
              ablation["mutation_causal_effect"]]
    pcts   = [ablation["splice_pct"], ablation["region_pct"], ablation["mutation_pct"]]
    colors = [_RED, _ORG, _BLUE]

    fig, ax = plt.subplots(figsize=(9, 3.0), facecolor=_BG)
    ax.set_facecolor(_BG)
    bars = ax.barh(labels, deltas, color=colors, edgecolor="#444", linewidth=0.6)
    ax.set_xlabel("Probability delta when ablated (causal effect)", color=_TEXT, fontsize=9)
    ax.tick_params(colors=_TEXT, labelsize=8)
    ax.set_title(
        f"Feature Ablation  |  baseline prob={ablation['baseline_probability']:.4f}",
        color=_TEXT, fontsize=10, loc="left")
    for b, d, p in zip(bars, deltas, pcts):
        ax.text(b.get_width()+0.002, b.get_y()+b.get_height()/2,
                f"  Δ{d:.4f}  ({p}%)", va="center", fontsize=9, color=_TEXT)
    ax.set_xlim(0, max(deltas+[0.01]) * 1.6)
    for sp in ["top","right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    fig.tight_layout()
    return _fig_to_pil(fig)


def plot_xai_metrics(xai):
    """Radar-style bar chart of explainability metrics."""
    labels  = ["Model\nAgreement", "XAI\nStrength", "CF\nMagnitude",
               "Locality\nScore", "Concentration\nIndex"]
    values  = [
        xai.model_agreement,
        xai.explainability_strength,
        min(xai.counterfactual_magnitude / 0.4, 1.0),
        xai.cross_model_locality_score,
        xai.signal_concentration_index,
    ]
    colors  = [_GREEN if v >= 0.65 else (_ORG if v >= 0.40 else _RED) for v in values]

    fig, ax = plt.subplots(figsize=(10, 3.0), facecolor=_BG)
    ax.set_facecolor(_BG)
    bars = ax.bar(labels, values, color=colors, edgecolor="#444", linewidth=0.6, width=0.5)
    ax.axhline(0.65, color=_GREEN, linestyle="--", linewidth=0.8, alpha=0.6, label="High (≥0.65)")
    ax.axhline(0.40, color=_ORG,   linestyle="--", linewidth=0.8, alpha=0.6, label="Moderate (≥0.40)")
    ax.set_ylim(0, 1.1)
    ax.set_ylabel("Score (0–1)", color=_TEXT, fontsize=9)
    ax.tick_params(colors=_TEXT, labelsize=8)
    ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.4, loc="upper right")
    for b, v in zip(bars, values):
        ax.text(b.get_x()+b.get_width()/2, b.get_height()+0.02,
                f"{v:.3f}", ha="center", fontsize=9, color=_TEXT)
    for sp in ["top","right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    ax.set_title("Explainability Metrics Panel", color=_TEXT, fontsize=10, loc="left")
    fig.tight_layout()
    return _fig_to_pil(fig)


# ═══════════════════════════════════════════════════════════════════════════════
# Core pipeline
# ═══════════════════════════════════════════════════════════════════════════════

def run_pipeline(
    chrom: str,
    position: str,
    ref_base: str,
    alt_base: str,
    exon_flag: int,
    intron_flag: int,
):
    """Main Gradio callback. Returns all outputs."""
    chrom    = chrom.strip()
    ref_base = ref_base.strip().upper()
    alt_base = alt_base.strip().upper()

    try:
        pos = int(position.strip().replace(",",""))
    except ValueError:
        return _error(f"Invalid position: '{position}'")

    for b, name in [(ref_base,"Reference"),(alt_base,"Alternate")]:
        if b not in "ACGT" or len(b) != 1:
            return _error(f"{name} base must be A, C, G, or T. Got: '{b}'")
    if ref_base == alt_base:
        return _error("Reference and alternate bases are identical.")

    try:
        ref_seq, mutation_pos = fetch_window(chrom, pos)

        # Validate reference base
        actual_ref = ref_seq[mutation_pos].upper()
        if actual_ref != ref_base:
            return _error(
                f"Reference mismatch at chr{chrom}:{pos}: "
                f"genome has '{actual_ref}', you entered '{ref_base}'."
            )

        # Build mutated sequence
        mut_seq = ref_seq[:mutation_pos] + alt_base + ref_seq[mutation_pos+1:]

        splice_sig = extract_splice_signals(
            REGISTRY.splice, ref_seq, mut_seq, exon_flag, intron_flag)

        v4_sig = extract_v4_signals(
            REGISTRY.v4, ref_seq, mut_seq, exon_flag, intron_flag)

        classic_sig = extract_classic_signals(
            REGISTRY.classic, ref_seq, mut_seq)

        xai = compute_cross_model_analysis(splice_sig, v4_sig, classic_sig, mutation_pos)

        result = build_decision(
            chrom=chrom, pos=pos, ref=ref_base, alt=alt_base,
            ref_seq=ref_seq, mut_seq=mut_seq, mutation_pos=mutation_pos,
            splice=splice_sig, v4=v4_sig, classic=classic_sig, xai=xai,
        )

        plots = _build_all_plots(result)

        json_str  = result.to_json()
        json_file = _write_json_file(json_str)

        demo_banner = (
            "\n> ⚠️ **DEMO MODE** — models are running with random weights. "
            "Place real checkpoints or ensure HF_TOKEN is set.\n"
            if REGISTRY.demo_mode else ""
        )

        summary_md = _build_summary_md(result, demo_banner)

        return (
            summary_md,                       # 0: explanation summary (FIRST)
            result.final_explanation,         # 1: final explanation text
            plots["xai_metrics"],             # 2: XAI metrics panel
            plots["splice_activation"],       # 3: splice conv3 heatmap
            plots["splice_heatmap"],          # 4: splice distance heatmap
            plots["v4_activation"],           # 5: v4 conv3 heatmap
            plots["classic_activation"],      # 6: classic conv3 heatmap
            plots["v4_gradient"],             # 7: v4 gradient attribution
            plots["splice_gradient"],         # 8: splice gradient attribution
            plots["counterfactual"],          # 9: counterfactual chart
            plots["ablation"],                # 10: ablation chart
            json_str,                         # 11: JSON report
            json_file,                        # 12: download file
        )

    except Exception as exc:
        logger.error("Pipeline error: %s\n%s", exc, traceback.format_exc())
        return _error(f"Error: {exc}\n\n```\n{traceback.format_exc()}\n```")


def _build_all_plots(r: DecisionResult) -> dict:
    mp = r.mutation_pos
    return {
        "xai_metrics":        plot_xai_metrics(r.xai),
        "splice_activation":  plot_activation_heatmap(
            r.splice.conv3_profile, mp, "Splice Model", r.splice.probability),
        "splice_heatmap":     plot_splice_heatmap(r.ref_seq, mp),
        "v4_activation":      plot_activation_heatmap(
            r.v4.conv3_profile, mp, "V4 Model", r.v4.probability),
        "classic_activation": plot_activation_heatmap(
            r.classic.conv3_profile, mp, "Classic Model", r.classic.probability),
        "v4_gradient":        plot_gradient_heatmap(
            r.v4.gradient_attribution, mp, "V4 Model"),
        "splice_gradient":    plot_gradient_heatmap(
            r.splice.gradient_attribution, mp, "Splice Model"),
        "counterfactual":     plot_counterfactual(
            r.splice.counterfactual_table,
            r.splice.probability,
            r.splice.counterfactual_delta),
        "ablation":           plot_ablation(r.splice.ablation),
    }


def _write_json_file(json_str: str) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=".json", delete=False,
                                     mode="w", encoding="utf-8")
    tmp.write(json_str)
    tmp.close()
    return tmp.name


def _build_summary_md(r: DecisionResult, demo_banner: str) -> str:
    mech_icon = {
        "Splice-driven":  "🔀",
        "Protein-driven": "🧬",
        "Consensus":      "✅",
        "Ambiguous":      "⚠️",
    }.get(r.dominant_mechanism, "❓")

    tier_icon = {
        "PATHOGENIC":          "🔴",
        "LIKELY PATHOGENIC":   "🟠",
        "POSSIBLY PATHOGENIC": "🟡",
        "LIKELY BENIGN":       "🟢",
        "BENIGN":              "🟢",
    }.get(r.risk_tier, "⚪")

    conf_icon = {"High": "🔵", "Moderate": "🟡", "Low": "🔴"}.get(r.confidence, "⚪")

    return f"""{demo_banner}
## {tier_icon} Risk Tier: **{r.risk_tier}**

| Field | Value |
|---|---|
| **Variant** | `chr{r.chrom}:g.{r.pos}{r.ref}>{r.alt}` |
| **Unified Probability** | `{r.unified_probability:.4f}` |
| **Dominant Mechanism** | {mech_icon} {r.dominant_mechanism} |
| **Confidence** | {conf_icon} {r.confidence} |
| **Splice Model** | `{r.splice.probability:.4f}` — {r.splice.risk_tier} |
| **V4 Model** | `{r.v4.probability:.4f}` |
| **Classic Model** | `{r.classic.probability:.4f}` |

---

### Explainability Metrics

| Metric | Value |
|---|---|
| **Mutation Peak Ratio** | `{r.xai.mutation_peak_ratio:.4f}` |
| **Counterfactual Magnitude** | `{r.xai.counterfactual_magnitude:.4f}` |
| **Cross-Model Locality** | `{r.xai.cross_model_locality_score:.4f}` |
| **Signal Concentration** | `{r.xai.signal_concentration_index:.4f}` |
| **XAI Strength Score** | `{r.xai.explainability_strength:.4f}` |
| **Activation Pattern** | `{r.xai.activation_pattern_type}` |
| **Model Agreement** | `{r.xai.model_agreement:.4f}` |

---

### Interpretation Briefs

**Splice:** {r.splice_analysis[:300]}{'…' if len(r.splice_analysis)>300 else ''}

**Protein:** {r.protein_analysis[:250]}{'…' if len(r.protein_analysis)>250 else ''}

**Agreement:** {r.agreement_analysis[:250]}{'…' if len(r.agreement_analysis)>250 else ''}
"""


def _error(msg: str):
    empties = [None] * 9
    return (
        f"❌ **Error**\n\n{msg}",
        "", empty, empty, empty, empty, empty,
        empty, empty, empty, empty,
        "{}", None,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Gradio UI
# ═══════════════════════════════════════════════════════════════════════════════

CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&family=Inter:wght@300;400;600;700&display=swap');
:root {
    --bg:      #0D1117; --surface: #161B22; --border: #30363D;
    --text:    #E6EDF3; --muted:   #7D8590;
    --blue:    #58A6FF; --green:   #3FB950; --red: #F85149; --orange: #D29922;
    --font:    'Inter', system-ui; --mono: 'JetBrains Mono', monospace;
}
body, .gradio-container { background: var(--bg) !important; color: var(--text) !important; font-family: var(--font) !important; }
.xai-header { background: linear-gradient(135deg, #0D1117 0%, #161B22 60%, #1a2332 100%);
    border-bottom: 1px solid var(--border); padding: 2rem 2.5rem 1.5rem; margin-bottom: 1.5rem; }
.xai-header h1 { font-size: 1.7rem; font-weight: 700; letter-spacing: -0.03em; margin: 0 0 0.3rem; }
.xai-header h1 em { color: var(--blue); font-style: normal; }
.xai-header p { color: var(--muted); font-size: 0.82rem; margin: 0; }
.section-title { font-size: 0.68rem; font-weight: 600; letter-spacing: 0.12em;
    text-transform: uppercase; color: var(--muted); border-bottom: 1px solid var(--border);
    padding-bottom: 0.4rem; margin-bottom: 1rem; }
.gradio-textbox input, .gradio-textbox textarea, .gradio-number input {
    background: #161B22 !important; border: 1px solid var(--border) !important;
    color: var(--text) !important; border-radius: 6px !important;
    font-family: var(--mono) !important; font-size: 0.88rem !important; }
label span { color: var(--muted) !important; font-size: 0.76rem !important; font-weight: 500 !important; }
.run-btn { background: linear-gradient(135deg, #1f6feb 0%, #388bfd 100%) !important;
    border: none !important; color: white !important; font-weight: 700 !important;
    font-size: 0.92rem !important; border-radius: 6px !important; letter-spacing: 0.04em !important; }
.run-btn:hover { transform: translateY(-1px) !important; box-shadow: 0 4px 14px rgba(88,166,255,.35) !important; }
.explanation-first { border: 1px solid var(--blue) !important;
    border-radius: 8px !important; background: rgba(88,166,255,0.04) !important; padding: 1rem !important; }
.gradio-markdown table { border-collapse: collapse; width: 100%; font-size: 0.83rem; }
.gradio-markdown th { background: #161B22; color: var(--muted); font-size: 0.68rem;
    letter-spacing: 0.08em; text-transform: uppercase; padding: 0.45rem 0.7rem;
    border: 1px solid var(--border); }
.gradio-markdown td { padding: 0.42rem 0.7rem; border: 1px solid var(--border);
    font-family: var(--mono); font-size: 0.80rem; }
.gradio-markdown code { background: #161B22; padding: 1px 5px; border-radius: 3px;
    font-family: var(--mono); color: var(--blue); font-size: 0.85em; }
.gradio-image img { border-radius: 6px; border: 1px solid var(--border); }
.gradio-tabs button { font-size: 0.80rem !important; color: var(--muted) !important;
    border-bottom: 2px solid transparent !important; background: transparent !important; }
.gradio-tabs button[aria-selected=true] { color: var(--blue) !important;
    border-bottom-color: var(--blue) !important; }
.gradio-textbox textarea { font-family: var(--mono) !important; font-size: 0.76rem !important; line-height: 1.5 !important; }
"""

HEADER_HTML = """
<div class="xai-header">
  <h1>Mutation <em>Explainability</em> Intelligence System</h1>
  <p>
    Three-model ensemble ·
    Explanation before prediction ·
    conv3 activations · gradient attribution · counterfactual analysis ·
    feature ablation · splice distance · cross-model locality
  </p>
</div>
"""

EXAMPLES = [
    ["17", "43071077", "G", "A", 1, 0],   # BRCA1 region
    ["11", "5226929",  "T", "C", 1, 0],   # HBB region
    ["7",  "117548628","T", "A", 1, 0],   # CFTR region
    ["3",  "37053577", "A", "C", 0, 1],   # intronic
    ["19", "44908684", "G", "T", 1, 0],   # APOE region
]


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Mutation Explainability Intelligence System",
                   css=CSS) as demo:

        gr.HTML(HEADER_HTML)

        with gr.Row(equal_height=False):
            # ─── INPUT PANEL ──────────────────────────────────────────
            with gr.Column(scale=1, min_width=280):
                gr.HTML('<div class="section-title">Variant Input</div>')
                chrom_in   = gr.Textbox(label="Chromosome", value="17", max_lines=1)
                pos_in     = gr.Textbox(label="Position (hg38, 1-based)", value="43071077", max_lines=1)
                with gr.Row():
                    ref_in = gr.Textbox(label="Ref Base", value="G", max_lines=1)
                    alt_in = gr.Textbox(label="Alt Base", value="A", max_lines=1)
                with gr.Row():
                    exon_in   = gr.Radio([0,1], label="Exon flag",   value=1)
                    intron_in = gr.Radio([0,1], label="Intron flag",  value=0)

                run_btn = gr.Button("▶  Analyse Variant", variant="primary",
                                    elem_classes="run-btn")

                gr.HTML('<div class="section-title" style="margin-top:1rem">Examples</div>')
                gr.Examples(
                    examples=EXAMPLES,
                    inputs=[chrom_in, pos_in, ref_in, alt_in, exon_in, intron_in],
                    label="",
                    examples_per_page=5,
                )

            # ─── OUTPUT PANEL ─────────────────────────────────────────
            with gr.Column(scale=3, min_width=640):
                # ══════════════════════════════════════════════════════
                # EXPLANATION FIRST — always shown before prediction
                # ══════════════════════════════════════════════════════
                gr.HTML('<div class="section-title">① Explanation & Signal Analysis</div>')

                summary_out = gr.Markdown(
                    value="*Run an analysis to see the full explanation.*",
                    elem_classes="explanation-first",
                )
                final_exp_out = gr.Textbox(
                    label="Final Explanation (grounded in internal signals)",
                    lines=10, max_lines=18, show_copy_button=True,
                )

                gr.HTML('<div class="section-title" style="margin-top:1.5rem">② Explainability Metrics Panel</div>')
                xai_metrics_plot = gr.Image(label="Explainability Metrics")

                # ── Tabs for visualisations ────────────────────────────
                gr.HTML('<div class="section-title" style="margin-top:1.5rem">③ Internal Model Signals</div>')

                with gr.Tabs():
                    with gr.TabItem("🔬 Splice Model"):
                        splice_act_plot  = gr.Image(label="conv3 Activation Heatmap — Splice")
                        splice_dist_plot = gr.Image(label="Splice Distance Risk Heatmap")
                        splice_grad_plot = gr.Image(label="Gradient Attribution — Splice")

                    with gr.TabItem("🧬 V4 Model"):
                        v4_act_plot  = gr.Image(label="conv3 Activation Heatmap — V4")
                        v4_grad_plot = gr.Image(label="Gradient Attribution — V4")

                    with gr.TabItem("📊 Classic Model"):
                        classic_act_plot = gr.Image(label="conv3 Activation Heatmap — Classic")

                    with gr.TabItem("⚗️ Causal Analysis"):
                        cf_plot  = gr.Image(label="Counterfactual Mutation Analysis")
                        abl_plot = gr.Image(label="Feature Ablation Causal Chart")

                    with gr.TabItem("📋 JSON Report"):
                        json_out  = gr.Textbox(label="Structured JSON Report",
                                               lines=30, max_lines=60,
                                               show_copy_button=True)
                        dl_btn    = gr.File(label="⬇ Download JSON Report")

        # ── Wire outputs ──────────────────────────────────────────────
        all_outputs = [
            summary_out, final_exp_out,
            xai_metrics_plot,
            splice_act_plot, splice_dist_plot, v4_act_plot,
            classic_act_plot, v4_grad_plot, splice_grad_plot,
            cf_plot, abl_plot,
            json_out, dl_btn,
        ]
        run_btn.click(
            fn=run_pipeline,
            inputs=[chrom_in, pos_in, ref_in, alt_in, exon_in, intron_in],
            outputs=all_outputs,
            show_progress=True,
        )

        gr.HTML("""
        <div style="text-align:center; color:#7D8590; font-size:0.70rem;
                    padding:1rem; margin-top:1rem; border-top:1px solid #30363D;">
            Mutation Explainability Intelligence System
            · Models: nileshhanotia/{mutation-predictor-splice, mutation-predictor-v4, mutation-pathogenicity-predictor}
            · For Research Use Only · Not for Clinical Diagnosis
        </div>
        """)

    return demo


demo = build_ui()

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,
    share=False,
)