app (7).py

"""
app.py
======
Mutation Explainability Intelligence System
Gradio Space — explanation ALWAYS precedes the prediction panel.

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

from __future__ import annotations
import io
import json
import logging
import os
import tempfile
import time
import traceback
from functools import lru_cache

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
import requests

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,
    V4Signals,
    ClassicSignals,
)
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")


# ═══════════════════════════════════════════════════════════════════════════════
# Model registry — loaded once at startup
# ═══════════════════════════════════════════════════════════════════════════════

REGISTRY = ModelRegistry(hf_token=os.environ.get("HF_TOKEN"))


# ═══════════════════════════════════════════════════════════════════════════════
# Ensembl sequence fetch
# ═══════════════════════════════════════════════════════════════════════════════

ENSEMBL_URL = "https://rest.ensembl.org/sequence/region/human"
WINDOW_HALF = 49   # 49 + 1 + 49 = 99 bp


@lru_cache(maxsize=512)
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:
                wait = int(r.headers.get("Retry-After", 5))
                logger.warning(f"Ensembl rate-limited — waiting {wait}s")
                time.sleep(wait)
                continue
            r.raise_for_status()
            data = r.json()
            if isinstance(data, list):
                data = data[0]
            return data.get("seq", "").upper()
        except Exception as exc:
            if attempt == 2:
                raise RuntimeError(
                    f"Ensembl API failed after 3 attempts: {exc}")
            time.sleep(1.5 * (2 ** attempt))
    return ""


def fetch_window(chrom: str, pos: int, ref: str, alt: str):
    """Fetch 99-bp window. Returns (ref_seq, mut_seq, mut_pos_in_window)."""
    chrom_clean = chrom.strip().lstrip("chrCHR")
    start       = max(1, pos - WINDOW_HALF)
    end         = pos + WINDOW_HALF
    raw         = _fetch_ensembl(chrom_clean, start, end)

    if not raw:
        raise ValueError(
            f"Empty sequence from Ensembl for chr{chrom}:{start}-{end}")

    seq     = (raw + "N" * 99)[:99]
    mut_pos = max(0, min(98, pos - start))

    genome_ref = seq[mut_pos] if mut_pos < len(seq) else "N"
    if genome_ref.upper() != ref.upper():
        logger.warning(
            f"Reference mismatch at chr{chrom}:{pos}: "
            f"Ensembl={genome_ref}, user={ref}. Using Ensembl sequence.")

    mut_list          = list(seq)
    mut_list[mut_pos] = alt.upper()
    mut_seq           = "".join(mut_list)

    return seq, mut_seq, mut_pos


# ═══════════════════════════════════════════════════════════════════════════════
# Colour palette & colour maps
# ═══════════════════════════════════════════════════════════════════════════════

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

_CMAP_ACT    = 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, "#f7f7f7"), (0.3, "#fee08b"), (0.6, "#fc8d59"), (1, "#d73027")])
_CMAP_GRAD   = matplotlib.colormaps.get_cmap("PuOr")


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

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


def _empty_pil():
    fig, ax = plt.subplots(figsize=(4, 2), facecolor=_BG)
    ax.set_facecolor(_BG)
    ax.axis("off")
    return _pil(fig)


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 _heatmap_pil(profile: np.ndarray, mutation_pos: int,
                  cmap, label: str, ylabel: str,
                  prob: float | None = None):
    imp = profile.copy()
    if imp.max() > 0:
        imp /= imp.max()
    fig, ax = plt.subplots(figsize=(15, 2.5), facecolor=_BG)
    ax.set_facecolor(_BG)
    im = ax.imshow(imp[np.newaxis, :], aspect="auto", cmap=cmap,
                   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(ylabel, 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([])
    title = label + (f"  (prob={prob:.4f})" if prob is not None else "")
    _style_ax(ax, title)
    fig.tight_layout()
    return _pil(fig)


def plot_splice_act(norm, pos, prob):
    return _heatmap_pil(norm, pos, _CMAP_ACT,
                         "Splice Model — conv3 Activation Norm",
                         "Activation", prob)


def plot_v4_act(norm, pos, prob):
    return _heatmap_pil(norm, pos, _CMAP_ACT,
                         "V4 Model — conv3 Activation Norm",
                         "Activation", prob)


def plot_classic_act(norm, pos, prob):
    return _heatmap_pil(norm, pos, _CMAP_ACT,
                         "Classic Model — conv3 Activation Norm",
                         "Activation", prob)


def plot_splice_distance(ref_seq: str, mut_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)
    return _heatmap_pil(scores, mut_pos, _CMAP_SPLICE,
                         "Splice Distance Risk Heatmap — GT/AG dinucleotides",
                         "Splice risk")


def plot_gradient(attr, pos, label):
    return _heatmap_pil(attr, pos, _CMAP_GRAD,
                         f"Gradient Attribution — {label}",
                         "Attribution")


def plot_counterfactual(cf: dict):
    table  = cf.get("table", [])
    orig_p = cf.get("original_probability", 0)
    if not table:
        fig, ax = plt.subplots(figsize=(8, 3), facecolor=_BG)
        ax.set_facecolor(_BG)
        ax.text(0.5, 0.5, "Counterfactual analysis not available",
                color=_TEXT, ha="center", va="center",
                transform=ax.transAxes, fontsize=11)
        ax.axis("off")
        return _pil(fig)

    labels = [r["mutation"] for r in table]
    probs  = [r["probability"] for r in table]
    p_max  = cf.get("max_probability", max(probs))
    p_min  = cf.get("min_probability", min(probs))
    colors = [_RED if r["probability"] == p_max
               else _BLUE if r["probability"] == p_min
               else "#74add1"
               for r in table]

    fig, ax = plt.subplots(figsize=(10, 3.5), facecolor=_BG)
    ax.set_facecolor(_SURF)
    bars = ax.bar(labels, probs, color=colors,
                  edgecolor="#30363D", linewidth=0.7)
    ax.axhline(0.5, color=_MUTED, linestyle="--",
               linewidth=1.0, label="Decision boundary (0.5)")
    ax.axhline(orig_p, color=_ORG, linestyle="-.", linewidth=1.5,
               label=f"Original mutation ({orig_p:.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)
    for sp in ["top", "right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    for bar, p in zip(bars, probs):
        ax.text(bar.get_x() + bar.get_width()/2,
                bar.get_height() + 0.015,
                f"{p:.3f}", ha="center", va="bottom",
                fontsize=8, color=_TEXT)
    ax.legend(fontsize=8, facecolor=_BG, labelcolor=_TEXT, framealpha=0.6)
    ax.set_title(
        f"Counterfactual Analysis  |  "
        f"Causal importance: {cf.get('probability_range', 0):.4f}  |  "
        f"Range: {p_min:.3f}–{p_max:.3f}",
        color=_TEXT, fontsize=9, loc="left", pad=4, fontweight="bold")
    fig.tight_layout()
    return _pil(fig)


def plot_ablation(abl: dict):
    keys   = ["splice_delta", "region_delta", "mutation_delta", "sequence_delta"]
    pkeys  = ["splice_pct",   "region_pct",   "mutation_pct",   "sequence_pct"]
    labels = [
        "Splice features\n(donor/acceptor/region)",
        "Region flags\n(exon/intron)",
        "Mutation type\n(one-hot)",
        "Sequence context\n(conv features)",
    ]
    colors = [_RED, _ORG, _BLUE, _GREEN]
    deltas = [abl.get(k, 0.0) for k in keys]
    pcts   = [abl.get(k, 0.0) for k in pkeys]

    fig, ax = plt.subplots(figsize=(10, 3.5), facecolor=_BG)
    ax.set_facecolor(_SURF)
    bars = ax.barh(labels, deltas, color=colors,
                   edgecolor="#30363D", linewidth=0.7)
    ax.set_xlabel("Probability delta (causal effect)",
                  color=_TEXT, fontsize=9)
    ax.tick_params(colors=_TEXT)
    for sp in ["top", "right"]:
        ax.spines[sp].set_visible(False)
    ax.spines["left"].set_color("#333")
    ax.spines["bottom"].set_color("#333")
    for bar, d, p in zip(bars, deltas, pcts):
        ax.text(bar.get_width() + 0.002,
                bar.get_y() + bar.get_height()/2,
                f" Δ{d:.4f} ({p}%)",
                va="center", color=_TEXT, fontsize=8)
    ax.set_xlim(0, max(deltas) * 1.65 + 0.02)
    ax.set_title(
        f"Feature Ablation Causal Analysis  |  "
        f"Baseline: {abl.get('baseline_probability', 0):.4f}",
        color=_TEXT, fontsize=9, loc="left", pad=4, fontweight="bold")
    fig.tight_layout()
    return _pil(fig)


def plot_xai_metrics(cross, sp_prob, v4_prob, cl_prob):
    """4-panel XAI metrics dashboard."""
    fig = plt.figure(figsize=(14, 7), facecolor=_BG)
    gs  = gridspec.GridSpec(2, 2, figure=fig, hspace=0.45, wspace=0.35)

    # ── TL: per-model probs ───────────────────────────────────────────────────
    ax0 = fig.add_subplot(gs[0, 0])
    ax0.set_facecolor(_SURF)
    names  = ["Splice", "V4", "Classic"]
    probs  = [sp_prob, v4_prob, cl_prob]
    col0   = [_RED if p >= 0.5 else _BLUE for p in probs]
    bars0  = ax0.bar(names, probs, color=col0,
                     edgecolor="#30363D", linewidth=0.7, width=0.5)
    ax0.axhline(0.5, color=_MUTED, linestyle="--", linewidth=1.0, alpha=0.7)
    ax0.set_ylim(0, 1.1)
    for bar, p in zip(bars0, probs):
        ax0.text(bar.get_x() + bar.get_width()/2,
                 bar.get_height() + 0.02,
                 f"{p:.4f}", ha="center", va="bottom",
                 color=_TEXT, fontsize=9)
    ax0.set_ylabel("Pathogenicity probability", color=_TEXT, fontsize=9)
    ax0.tick_params(colors=_TEXT)
    for sp in ["top", "right"]:
        ax0.spines[sp].set_visible(False)
    ax0.spines["left"].set_color("#333")
    ax0.spines["bottom"].set_color("#333")
    _style_ax(ax0, "Per-model Probability")

    # ── TR: XAI scores ────────────────────────────────────────────────────────
    ax1 = fig.add_subplot(gs[0, 1])
    ax1.set_facecolor(_SURF)
    xai_labels = [
        "Mut Peak Ratio\n(÷3 norm)",
        "CF Magnitude",
        "Cross-Model\nLocality",
        "Signal\nConcentration",
        "Explainability\nStrength",
    ]
    xai_raw = [
        cross["mutation_peak_ratio"],
        cross["counterfactual_magnitude"],
        cross["cross_model_locality_score"],
        cross["signal_concentration_index"],
        cross["explainability_strength_score"],
    ]
    xai_norm = [
        min(cross["mutation_peak_ratio"] / 3.0, 1.0),
        min(cross["counterfactual_magnitude"], 1.0),
        (cross["cross_model_locality_score"] + 1.0) / 2.0,
        cross["signal_concentration_index"],
        cross["explainability_strength_score"],
    ]
    col1 = [_GREEN if v >= 0.5 else _ORG if v >= 0.3 else _RED
             for v in xai_norm]
    bars1 = ax1.barh(xai_labels, xai_norm, color=col1,
                      edgecolor="#30363D", linewidth=0.7)
    ax1.set_xlim(0, 1.35)
    ax1.tick_params(colors=_TEXT, labelsize=7)
    for sp in ["top", "right"]:
        ax1.spines[sp].set_visible(False)
    ax1.spines["left"].set_color("#333")
    ax1.spines["bottom"].set_color("#333")
    for bar, raw in zip(bars1, xai_raw):
        ax1.text(bar.get_width() + 0.02,
                 bar.get_y() + bar.get_height()/2,
                 f"{raw:.3f}", va="center", color=_TEXT, fontsize=8)
    _style_ax(ax1, "XAI Engine Metrics (normalised 0–1)")

    # ── BL: cross-model activation overlap ────────────────────────────────────
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.set_facecolor(_SURF)
    x    = np.arange(99)
    sp_n = cross.get("_splice_norm",  np.zeros(99))
    v4_n = cross.get("_v4_norm",      np.zeros(99))
    cl_n = cross.get("_classic_norm", np.zeros(99))
    ax2.plot(x, sp_n, color=_RED,   linewidth=1.2, alpha=0.85, label="Splice")
    ax2.plot(x, v4_n, color=_BLUE,  linewidth=1.2, alpha=0.85, label="V4")
    ax2.plot(x, cl_n, color=_GREEN, linewidth=1.2, alpha=0.85, label="Classic")
    ax2.set_ylim(0, 1.15)
    ax2.set_xlabel("Position (99-bp window)", color=_TEXT, fontsize=8)
    ax2.set_ylabel("Norm. activation", color=_TEXT, fontsize=8)
    ax2.tick_params(colors=_TEXT, labelsize=7)
    for sp in ["top", "right"]:
        ax2.spines[sp].set_visible(False)
    ax2.spines["left"].set_color("#333")
    ax2.spines["bottom"].set_color("#333")
    ax2.legend(fontsize=7, facecolor=_BG, labelcolor=_TEXT,
               framealpha=0.6, loc="upper right")
    _style_ax(ax2, "Cross-model Activation Overlap")

    # ── BR: summary text ──────────────────────────────────────────────────────
    ax3 = fig.add_subplot(gs[1, 1])
    ax3.set_facecolor(_SURF)
    ax3.axis("off")
    summary = "\n".join([
        f"Activation pattern  :  {cross['activation_pattern_type']}",
        f"Model agreement     :  {cross['model_agreement']}",
        f"Probability std     :  {cross['prob_std']:.4f}",
        "",
        f"Splice prob         :  {sp_prob:.4f}",
        f"V4 prob             :  {v4_prob:.4f}",
        f"Classic prob        :  {cl_prob:.4f}",
        "",
        f"ESS score           :  {cross['explainability_strength_score']:.4f}",
        f"Cross-model loc.    :  {cross['cross_model_locality_score']:.4f}",
    ])
    ax3.text(0.05, 0.95, summary, transform=ax3.transAxes,
             color=_TEXT, fontsize=8, va="top",
             fontfamily="monospace",
             bbox=dict(facecolor="#21262D", edgecolor="#30363D",
                       alpha=0.8, boxstyle="round,pad=0.4"))
    _style_ax(ax3, "Summary")

    return _pil(fig)


# ═══════════════════════════════════════════════════════════════════════════════
# Pipeline
# ═══════════════════════════════════════════════════════════════════════════════

_EMPTY_PIL = _empty_pil()


def run_pipeline(chrom: str, pos_str: str,
                 ref: str, alt: str,
                 exon_flag: int, intron_flag: int):
    """
    Full XAI pipeline.  Returns 13 outputs for the Gradio UI.
    The ordering of computation enforces explanation-before-prediction:
      Step 3: extract all internal signals
      Step 4: run explainability engine
      Step 5: build unified decision (uses step-4 results)
    """

    def _err(msg):
        empty = _empty_pil()
        return (
            f"❌ **Error**\n\n{msg}",
            msg,
            empty, empty, empty, empty, empty,
            empty, empty, empty, empty,
            "{}", None,
        )

    # ── Validate input ────────────────────────────────────────────────────────
    try:
        pos = int(str(pos_str).strip())
    except ValueError:
        return _err(f"Invalid position: '{pos_str}'")

    ref = ref.strip().upper()
    alt = alt.strip().upper()
    if len(ref) != 1 or ref not in "ACGTN":
        return _err(f"Ref base must be a single nucleotide. Got: '{ref}'")
    if len(alt) != 1 or alt not in "ACGTN":
        return _err(f"Alt base must be a single nucleotide. Got: '{alt}'")
    if ref == alt:
        return _err("Reference and alternate bases are identical.")

    exon_flag   = int(exon_flag)
    intron_flag = int(intron_flag)

    # ── Step 1: Fetch 401-bp → trim to 99-bp window from Ensembl ─────────────
    logger.info(f"Fetching chr{chrom}:{pos} {ref}>{alt}")
    try:
        ref_seq, mut_seq, mut_win_pos = fetch_window(chrom, pos, ref, alt)
    except Exception as exc:
        logger.warning(f"Ensembl fetch failed ({exc}). Using synthetic window.")
        ref_seq      = "N" * 49 + ref + "N" * 49
        mut_seq      = "N" * 49 + alt + "N" * 49
        mut_win_pos  = 49

    # ── Step 2: Load models ───────────────────────────────────────────────────
    try:
        splice_model  = REGISTRY.splice
        v4_model      = REGISTRY.v4
        classic_model = REGISTRY.classic
    except Exception as exc:
        return _err(f"Model loading failed: {exc}")

    # ── Step 3: Extract internal signals ─────────────────────────────────────
    try:
        logger.info("Extracting splice signals …")
        splice_sig = extract_splice_signals(
            splice_model, ref_seq, mut_seq, exon_flag, intron_flag)
    except Exception as exc:
        return _err(f"Splice model failed: {exc}\n{traceback.format_exc()}")

    try:
        logger.info("Extracting V4 signals …")
        v4_sig = extract_v4_signals(
            v4_model, ref_seq, mut_seq, exon_flag, intron_flag)
    except Exception as exc:
        logger.warning(f"V4 model failed ({exc}), using fallback.")
        v4_sig = V4Signals(
            probability=0.5, conv3_norm=np.zeros(99),
            gradient_attribution=np.zeros(99),
            mutation_pos=mut_win_pos,
            mutation_peak_ratio=0.0, signal_concentration=0.0,
        )

    try:
        logger.info("Extracting classic signals …")
        classic_sig = extract_classic_signals(
            classic_model, ref_seq, mut_seq, exon_flag, intron_flag)
    except Exception as exc:
        logger.warning(f"Classic model failed ({exc}), using fallback.")
        classic_sig = ClassicSignals(
            probability=0.5, conv3_norm=np.zeros(99),
            importance_head=0.0, region_imp=np.zeros(2),
            mutation_pos=mut_win_pos,
            mutation_peak_ratio=0.0, signal_concentration=0.0,
        )

    # ── Step 4: Explainability engine — MANDATORY before decision ─────────────
    logger.info("Running explainability engine …")
    cross = compute_cross_model_analysis(splice_sig, v4_sig, classic_sig)

    # ── Step 5: Unified decision (uses cross results — ordering guaranteed) ───
    logger.info("Building unified decision …")
    result: DecisionResult = build_decision(
        chrom, pos, ref, alt,
        splice_sig, v4_sig, classic_sig, cross)

    # ── Step 6: Build visualisations ──────────────────────────────────────────
    try:
        xai_metrics  = plot_xai_metrics(
            cross, splice_sig.probability,
            v4_sig.probability, classic_sig.probability)
        splice_act   = plot_splice_act(
            splice_sig.conv3_norm, splice_sig.mutation_pos,
            splice_sig.probability)
        splice_dist  = plot_splice_distance(ref_seq, splice_sig.mutation_pos)
        v4_act       = plot_v4_act(
            v4_sig.conv3_norm, v4_sig.mutation_pos, v4_sig.probability)
        classic_act  = plot_classic_act(
            classic_sig.conv3_norm, classic_sig.mutation_pos,
            classic_sig.probability)
        v4_grad      = plot_gradient(
            v4_sig.gradient_attribution, v4_sig.mutation_pos, "V4")
        splice_grad  = plot_gradient(
            splice_sig.gradient_attribution, splice_sig.mutation_pos,
            "Splice")
        cf_plot      = plot_counterfactual(splice_sig.counterfactual)
        abl_plot     = plot_ablation(splice_sig.ablation)
    except Exception as exc:
        logger.error(f"Visualisation error: {exc}\n{traceback.format_exc()}")
        empty = _empty_pil()
        xai_metrics = splice_act = splice_dist = v4_act = empty
        classic_act = v4_grad = splice_grad = cf_plot = abl_plot = empty

    # ── Step 7: Downloadable JSON ─────────────────────────────────────────────
    json_str = result.report_json
    try:
        tmp = tempfile.NamedTemporaryFile(
            mode="w", suffix=".json",
            prefix=f"mutation_xai_{chrom}_{pos}_{ref}{alt}_",
            delete=False, encoding="utf-8")
        tmp.write(json_str)
        tmp.close()
        dl_path = tmp.name
    except Exception:
        dl_path = None

    # ── Step 8: Explanation-first summary markdown ────────────────────────────
    cf    = splice_sig.counterfactual
    abl   = splice_sig.ablation
    sp    = splice_sig
    cross_loc = cross["cross_model_locality_score"]

    prob_icon  = "🔴" if result.unified_probability >= 0.5 else "🟢"
    conf_icon  = {"High": "✅", "Moderate": "⚠️", "Low": "🔶"}.get(
        result.confidence, "❓")

    summary_md = f"""
### {prob_icon} `{result.variant}`

| Field | Value |
|---|---|
| **Risk Tier** | `{result.risk_tier}` · {result.tier_desc} |
| **Unified Probability** | `{result.unified_probability:.4f}` |
| **Dominant Mechanism** | `{result.dominant_mechanism}` |
| **Confidence** | {conf_icon} `{result.confidence}` |

---

#### 🔬 Explainability Engine Output

| Metric | Raw Value | Interpretation |
|---|---|---|
| Mutation Peak Ratio | `{cross["mutation_peak_ratio"]:.4f}` | {"Strongly localised to mutation site" if cross["mutation_peak_ratio"] > 2 else "Above-average localisation" if cross["mutation_peak_ratio"] > 1 else "Diffuse — signal not mutation-centred"} |
| Counterfactual Magnitude | `{cross["counterfactual_magnitude"]:.4f}` | {"Strong position-level causality" if cross["counterfactual_magnitude"] > 0.25 else "Moderate causality" if cross["counterfactual_magnitude"] > 0.10 else "Weak positional causality"} |
| Cross-model Locality | `{cross["cross_model_locality_score"]:.4f}` | {"Models align on same region" if cross_loc > 0.5 else "Partial alignment" if cross_loc > 0 else "Models attend to different regions"} |
| Signal Concentration Index | `{cross["signal_concentration_index"]:.4f}` | Fraction of activation energy at mutation site |
| **Explainability Strength (ESS)** | **`{cross["explainability_strength_score"]:.4f}`** | 0–1 composite quality score |
| Activation Pattern | `{cross["activation_pattern_type"]}` | Shape of conv3 profile |
| Model Agreement | `{cross["model_agreement"]}` | std={cross["prob_std"]:.4f} across models |

---

#### 📊 Per-model Probabilities

| Model | Probability |
|---|---|
| `mutation-predictor-splice` | `{sp.probability:.4f}` · {sp.risk_tier} |
| `mutation-predictor-v4` | `{v4_sig.probability:.4f}` |
| `mutation-pathogenicity-predictor` | `{classic_sig.probability:.4f}` |

---

#### ⚗️ Splice Signals

| Signal | Value |
|---|---|
| Splice aura score | `{sp.splice_aura_score:.4f}` |
| Donor importance | `{float(sp.splice_imp[0]):.4f}` |
| Acceptor importance | `{float(sp.splice_imp[1]):.4f}` |
| Nearest GT donor | `{sp.dist_donor if sp.dist_donor is not None else "N/A"} bp` — {sp.splice_risk_donor} |
| Nearest AG acceptor | `{sp.dist_acceptor if sp.dist_acceptor is not None else "N/A"} bp` — {sp.splice_risk_acceptor} |
| Counterfactual delta | `{cf.get("probability_range", 0):.4f}` |
| Dominant ablation feature | `{abl.get("dominant_feature", "—")}` |
"""

    logger.info("Pipeline complete.")

    return (
        summary_md,                   # ①  explanation summary — FIRST
        result.final_explanation,     # ②  final human-readable explanation
        xai_metrics,                  # ③  XAI metrics dashboard
        splice_act,                   # ④  splice conv3 heatmap
        splice_dist,                  # ⑤  splice distance heatmap
        v4_act,                       # ⑥  v4 conv3 heatmap
        classic_act,                  # ⑦  classic conv3 heatmap
        v4_grad,                      # ⑧  v4 gradient attribution
        splice_grad,                  # ⑨  splice gradient attribution
        cf_plot,                      # ⑩  counterfactual chart
        abl_plot,                     # ⑪  feature ablation chart
        json_str,                     # ⑫  JSON report text
        dl_path,                      # ⑬  downloadable file
    )


# ═══════════════════════════════════════════════════════════════════════════════
# 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:-.03em;margin:0 0 .3rem;}
.xai-header h1 em{color:var(--blue);font-style:normal;}
.xai-header p{color:var(--muted);font-size:.82rem;margin:0;}
.section-title{font-size:.68rem;font-weight:600;letter-spacing:.12em;text-transform:uppercase;
    color:var(--muted);border-bottom:1px solid var(--border);padding-bottom:.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:.88rem!important;}
label span{color:var(--muted)!important;font-size:.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:.92rem!important;border-radius:6px!important;letter-spacing:.04em!important;}
.run-btn:hover{transform:translateY(-1px)!important;box-shadow:0 4px 14px rgba(88,166,255,.35)!important;}
.explanation-panel{border:1px solid var(--blue)!important;border-radius:8px!important;
    background:rgba(88,166,255,.04)!important;padding:1rem!important;}
.gradio-markdown table{border-collapse:collapse;width:100%;font-size:.83rem;}
.gradio-markdown th{background:#161B22;color:var(--muted);font-size:.68rem;
    letter-spacing:.08em;text-transform:uppercase;padding:.45rem .7rem;border:1px solid var(--border);}
.gradio-markdown td{padding:.42rem .7rem;border:1px solid var(--border);
    font-family:var(--mono);font-size:.80rem;}
.gradio-markdown code{background:#161B22;padding:1px 5px;border-radius:3px;
    font-family:var(--mono);color:var(--blue);font-size:.85em;}
.gradio-image img{border-radius:6px;border:1px solid var(--border);}
.gradio-tabs button{font-size:.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:.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 &nbsp;·&nbsp; Explanation always before prediction &nbsp;·&nbsp;
    conv3 activations &nbsp;·&nbsp; gradient attribution &nbsp;·&nbsp;
    counterfactual analysis &nbsp;·&nbsp; feature ablation &nbsp;·&nbsp;
    splice distance &nbsp;·&nbsp; cross-model locality
  </p>
</div>
"""

EXAMPLES = [
    ["17", "43071077", "G", "A", 1, 0],
    ["11", "5226929",  "T", "C", 1, 0],
    ["7",  "117548628","T", "A", 1, 0],
    ["3",  "37053577", "A", "C", 0, 1],
    ["19", "44908684", "G", "T", 1, 0],
]


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 PANEL — ALWAYS RENDERED FIRST
                #    Prediction score does not appear without this panel
                # ══════════════════════════════════════════════════════════════
                gr.HTML('<div class="section-title">'
                         '① Explanation &amp; Signal Analysis</div>')

                summary_out = gr.Markdown(
                    value=(
                        "*Run an analysis to see the full explanation.*\n\n"
                        "*This panel always renders **before** the prediction score.*"
                    ),
                    elem_classes="explanation-panel",
                )

                final_exp_out = gr.Textbox(
                    label="Final Explanation (grounded in internal signals)",
                    lines=10, max_lines=18,
                    show_copy_button=True,
                )

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

                # ── ③ Internal model signal tabs ──────────────────────────────
                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 all outputs ──────────────────────────────────────────────────
        all_outputs = [
            summary_out,       # ①  explanation summary (always first)
            final_exp_out,     # ①  detailed explanation text
            xai_metrics_plot,  # ②  XAI dashboard
            splice_act_plot,   # ③  splice tab
            splice_dist_plot,
            v4_act_plot,       # ③  v4 tab
            classic_act_plot,  # ③  classic tab
            v4_grad_plot,
            splice_grad_plot,
            cf_plot,           # ③  causal tab
            abl_plot,
            json_out,          # ③  JSON tab
            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:.70rem;
                    padding:1rem;margin-top:1rem;border-top:1px solid #30363D;">
            Mutation Explainability Intelligence System
            &nbsp;·&nbsp;
            Models: nileshhanotia/{mutation-predictor-splice,
            mutation-predictor-v4, mutation-pathogenicity-predictor}
            &nbsp;·&nbsp; For Research Use Only &nbsp;·&nbsp;
            Not for Clinical Diagnosis
        </div>
        """)

    return demo


demo = build_ui()

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