explainability_engine.py

"""
explainability_engine.py
========================
Extracts ALL internal explainability signals from all three models.

Signal inventory (per model, as specified):

  Splice model (MutationPredictorCNN_v2):
    ✓ conv3 activation norm profile        (99,)
    ✓ mutation-centered activation peak    float
    ✓ splice aura distance                 (dist_donor, dist_acceptor)
    ✓ counterfactual delta                 |prob_mut - prob_ref|
    ✓ feature ablation response            {splice, region, mutation} Δprob
    ✓ risk tier classification

  v4 model (MutationPredictorCNN_v2):
    ✓ importance head vector               (256,) mutation-centered feature
    ✓ mutation-centered importance density float
    ✓ conv3 norm profile                   (99,)

  Classic model (MutationPredictorCNN):
    ✓ importance head output               float
    ✓ conv3 norm profile                   (99,)

  Cross-model:
    ✓ mutation peak ratio
    ✓ counterfactual magnitude
    ✓ cross-model locality score
    ✓ signal concentration index
    ✓ explainability strength score (0–1)
    ✓ activation pattern type
"""

from __future__ import annotations
import logging
from dataclasses import dataclass, field

import torch
import numpy as np

from model_loader import (
    MutationPredictorCNN_v2,
    MutationPredictorCNN,
    encode_for_v2,
    encode_for_classic,
    find_mutation_pos,
    MUT_TYPES,
)

logger = logging.getLogger(__name__)

ALL_BASES = ["A", "T", "C", "G"]

# ── Splice site helpers (from live splice app — exact logic) ──────────────────

def compute_splice_distances(mutation_pos: int, ref_seq: str):
    seq = ref_seq.upper()
    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)
    if mutation_pos < 0:
        return None, None, None, None
    d_donor = d_acceptor = nearest_d = nearest_a = None
    if donors:
        best = min(donors, key=lambda p: abs(mutation_pos - p))
        d_donor, nearest_d = abs(mutation_pos - best), best
    if acceptors:
        best = min(acceptors, key=lambda p: abs(mutation_pos - p))
        d_acceptor, nearest_a = abs(mutation_pos - best), best
    return d_donor, d_acceptor, nearest_d, nearest_a


def classify_splice_risk(distance) -> str:
    if distance is None: return "UNKNOWN"
    if distance <= 2:    return "CRITICAL SPLICE SITE"
    if distance <= 8:    return "SPLICE REGION"
    return "NON-SPLICE"


def classify_risk_tier(prob: float) -> tuple[str, str]:
    if prob >= 0.90: return "PATHOGENIC",          "Very high confidence"
    if prob >= 0.70: return "LIKELY PATHOGENIC",   "High confidence"
    if prob >= 0.50: return "POSSIBLY PATHOGENIC", "Moderate confidence"
    if prob >= 0.20: return "LIKELY BENIGN",        "Low pathogenic signal"
    return             "BENIGN",                   "Very low pathogenic signal"


# ═══════════════════════════════════════════════════════════════════════════════
# Data containers
# ═══════════════════════════════════════════════════════════════════════════════

@dataclass
class SpliceModelSignals:
    probability:            float
    imp_score:              float
    region_imp:             np.ndarray          # (2,) exon/intron
    splice_imp:             np.ndarray          # (3,) donor/acceptor/region
    conv3_profile:          np.ndarray          # (99,)
    mutation_peak:          float
    mutation_peak_ratio:    float
    splice_aura_donor:      int | None
    splice_aura_acceptor:   int | None
    nearest_donor_pos:      int | None
    nearest_acceptor_pos:   int | None
    splice_risk_donor:      str
    splice_risk_acceptor:   str
    risk_tier:              str
    risk_tier_desc:         str
    counterfactual_delta:   float
    counterfactual_table:   list[dict]
    ablation:               dict
    gradient_attribution:   np.ndarray          # (99,)


@dataclass
class V4ModelSignals:
    probability:            float
    imp_score:              float
    region_imp:             np.ndarray
    splice_imp:             np.ndarray
    conv3_profile:          np.ndarray
    mutation_peak:          float
    mutation_peak_ratio:    float
    importance_vector:      np.ndarray          # (256,)
    gradient_attribution:   np.ndarray


@dataclass
class ClassicModelSignals:
    probability:            float
    imp_score:              float
    conv3_profile:          np.ndarray
    mutation_peak:          float
    gradient_attribution:   np.ndarray


@dataclass
class CrossModelAnalysis:
    mutation_peak_ratio:        float   # mean of per-model peak ratios
    counterfactual_magnitude:   float   # |prob_with_mut - prob_without_mut|
    cross_model_locality_score: float   # cosine similarity of conv3 profiles
    signal_concentration_index: float   # how focused activation is at mutation
    explainability_strength:    float   # 0–1 composite
    activation_pattern_type:    str     # Sharp / Broad / Flat
    model_agreement:            float   # 1 - std of the 3 probabilities


# ═══════════════════════════════════════════════════════════════════════════════
# Per-model signal extractors
# ═══════════════════════════════════════════════════════════════════════════════

def _run_v2(model: MutationPredictorCNN_v2,
            ref_seq: str, mut_seq: str,
            exon_flag: int, intron_flag: int,
            donor_flag: int = 0, acceptor_flag: int = 0,
            region_flag: int = 0) -> tuple:
    enc = encode_for_v2(ref_seq, mut_seq, exon_flag, intron_flag,
                        donor_flag, acceptor_flag, region_flag)
    with torch.no_grad():
        logit, imp, r_imp, s_imp = model(enc.unsqueeze(0))
    prob = torch.sigmoid(logit).item()
    return (prob,
            float(imp.item()),
            r_imp.squeeze(0).numpy(),
            s_imp.squeeze(0).numpy(),
            enc)


def extract_splice_signals(
    model: MutationPredictorCNN_v2,
    ref_seq: str,
    mut_seq: str,
    exon_flag: int = 0,
    intron_flag: int = 0,
) -> SpliceModelSignals:

    # ── Forward pass ──────────────────────────────────────────
    prob, imp, r_imp, s_imp, enc = _run_v2(
        model, ref_seq, mut_seq, exon_flag, intron_flag)

    mutation_pos   = find_mutation_pos(ref_seq, mut_seq)
    conv3_profile  = model.conv3_norm_profile() or np.zeros(99)
    mut_peak       = float(conv3_profile[mutation_pos]) if mutation_pos >= 0 else 0.0
    peak_ratio     = float(mut_peak / (conv3_profile.mean() + 1e-9))

    # ── Splice distances ───────────────────────────────────────
    d_don, d_acc, n_don, n_acc = compute_splice_distances(mutation_pos, ref_seq)
    risk_don = classify_splice_risk(d_don)
    risk_acc = classify_splice_risk(d_acc)
    tier, tier_desc = classify_risk_tier(prob)

    # ── Counterfactual delta ───────────────────────────────────
    cf_table, cf_delta = _counterfactual(
        model, ref_seq, mut_seq, mutation_pos, exon_flag, intron_flag, prob)

    # ── Feature ablation ──────────────────────────────────────
    ablation = _feature_ablation(model, enc, prob)

    # ── Gradient attribution ──────────────────────────────────
    grad_attr = _gradient_attribution_v2(model, enc)

    return SpliceModelSignals(
        probability=round(prob, 4),
        imp_score=round(imp, 4),
        region_imp=r_imp,
        splice_imp=s_imp,
        conv3_profile=conv3_profile,
        mutation_peak=round(mut_peak, 4),
        mutation_peak_ratio=round(peak_ratio, 4),
        splice_aura_donor=d_don,
        splice_aura_acceptor=d_acc,
        nearest_donor_pos=n_don,
        nearest_acceptor_pos=n_acc,
        splice_risk_donor=risk_don,
        splice_risk_acceptor=risk_acc,
        risk_tier=tier,
        risk_tier_desc=tier_desc,
        counterfactual_delta=round(cf_delta, 4),
        counterfactual_table=cf_table,
        ablation=ablation,
        gradient_attribution=grad_attr,
    )


def extract_v4_signals(
    model: MutationPredictorCNN_v2,
    ref_seq: str,
    mut_seq: str,
    exon_flag: int = 0,
    intron_flag: int = 0,
) -> V4ModelSignals:

    prob, imp, r_imp, s_imp, enc = _run_v2(
        model, ref_seq, mut_seq, exon_flag, intron_flag)

    mutation_pos  = find_mutation_pos(ref_seq, mut_seq)
    conv3_profile = model.conv3_norm_profile() or np.zeros(99)
    mut_peak      = float(conv3_profile[mutation_pos]) if mutation_pos >= 0 else 0.0
    peak_ratio    = float(mut_peak / (conv3_profile.mean() + 1e-9))
    imp_vector    = model.importance_head_vector() or np.zeros(256)
    grad_attr     = _gradient_attribution_v2(model, enc)

    return V4ModelSignals(
        probability=round(prob, 4),
        imp_score=round(imp, 4),
        region_imp=r_imp,
        splice_imp=s_imp,
        conv3_profile=conv3_profile,
        mutation_peak=round(mut_peak, 4),
        mutation_peak_ratio=round(peak_ratio, 4),
        importance_vector=imp_vector,
        gradient_attribution=grad_attr,
    )


def extract_classic_signals(
    model: MutationPredictorCNN,
    ref_seq: str,
    mut_seq: str,
) -> ClassicModelSignals:

    x_tensor = encode_for_classic(ref_seq, mut_seq)
    with torch.no_grad():
        logit, imp = model(x_tensor)
    prob = torch.sigmoid(logit).item()
    imp_val = float(imp.item())

    mutation_pos  = find_mutation_pos(ref_seq, mut_seq)
    conv3_profile = model.conv3_norm_profile() or np.zeros(99)
    mut_peak      = float(conv3_profile[mutation_pos]) if mutation_pos >= 0 else 0.0
    grad_attr     = _gradient_attribution_classic(model, x_tensor.squeeze(0))

    return ClassicModelSignals(
        probability=round(prob, 4),
        imp_score=round(imp_val, 4),
        conv3_profile=conv3_profile,
        mutation_peak=round(mut_peak, 4),
        gradient_attribution=grad_attr,
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Cross-model analysis
# ═══════════════════════════════════════════════════════════════════════════════

def compute_cross_model_analysis(
    splice: SpliceModelSignals,
    v4:     V4ModelSignals,
    classic: ClassicModelSignals,
    mutation_pos: int,
) -> CrossModelAnalysis:

    # ── 1. Mutation peak ratio (mean across models) ────────────
    peak_ratios = [splice.mutation_peak_ratio, v4.mutation_peak_ratio]
    mean_peak_ratio = float(np.mean(peak_ratios))

    # ── 2. Counterfactual magnitude ───────────────────────────
    cf_mag = splice.counterfactual_delta   # primary signal from splice model

    # ── 3. Cross-model locality score ─────────────────────────
    # Cosine similarity between splice and v4 conv3 profiles
    p1 = splice.conv3_profile.astype(np.float32)
    p2 = v4.conv3_profile.astype(np.float32)
    p3 = classic.conv3_profile.astype(np.float32)
    cos_12 = _cosine(p1, p2)
    cos_13 = _cosine(p1, p3)
    cos_23 = _cosine(p2, p3)
    locality_score = float(np.mean([cos_12, cos_13, cos_23]))

    # ── 4. Signal concentration index ─────────────────────────
    # What fraction of total activation is within ±5 positions of mutation
    concentration = _concentration_index(p1, mutation_pos, window=5)

    # ── 5. Explainability strength score ─────────────────────
    # Weighted composite: peak_ratio (norm), cf_mag, locality
    pr_norm  = min(mean_peak_ratio / 5.0, 1.0)   # 5x mean = saturated
    cf_norm  = min(cf_mag / 0.4, 1.0)            # 0.4 delta = saturated
    xai_score = round(0.35 * pr_norm + 0.35 * cf_norm + 0.30 * locality_score, 4)

    # ── 6. Activation pattern type ────────────────────────────
    pattern = _activation_pattern(p1, mutation_pos)

    # ── 7. Model agreement ───────────────────────────────────
    probs = [splice.probability, v4.probability, classic.probability]
    agreement = round(1.0 - float(np.std(probs)), 4)

    return CrossModelAnalysis(
        mutation_peak_ratio=round(mean_peak_ratio, 4),
        counterfactual_magnitude=cf_mag,
        cross_model_locality_score=round(locality_score, 4),
        signal_concentration_index=round(concentration, 4),
        explainability_strength=xai_score,
        activation_pattern_type=pattern,
        model_agreement=max(0.0, min(1.0, agreement)),
    )


# ═══════════════════════════════════════════════════════════════════════════════
# Gradient attribution (exact logic from live splice app)
# ═══════════════════════════════════════════════════════════════════════════════

def _gradient_attribution_v2(model: MutationPredictorCNN_v2,
                              enc: torch.Tensor) -> np.ndarray:
    model.eval()
    leaf = enc.clone().detach()
    leaf.requires_grad_(True)
    logit, _, _, _ = model(leaf.unsqueeze(0))
    model.zero_grad()
    logit.backward()
    grad = leaf.grad  # (1106,)
    seq_grad = grad[:1089].view(99, 11)
    attr = seq_grad.abs().norm(dim=1).detach().numpy()
    mx = attr.max()
    return attr / (mx + 1e-9)


def _gradient_attribution_classic(model: MutationPredictorCNN,
                                   enc: torch.Tensor) -> np.ndarray:
    model.eval()
    leaf = enc.clone().detach()          # (8, 99)
    leaf.requires_grad_(True)
    logit, _ = model(leaf.unsqueeze(0))
    model.zero_grad()
    logit.backward()
    grad = leaf.grad  # (8, 99)
    attr = grad.abs().norm(dim=0).detach().numpy()
    mx = attr.max()
    return attr / (mx + 1e-9)


# ═══════════════════════════════════════════════════════════════════════════════
# Counterfactual analysis (from live splice app — exact logic preserved)
# ═══════════════════════════════════════════════════════════════════════════════

def _counterfactual(model: MutationPredictorCNN_v2,
                    ref_seq: str, mut_seq: str,
                    mutation_pos: int,
                    exon_flag: int, intron_flag: int,
                    orig_prob: float) -> tuple[list[dict], float]:
    if mutation_pos < 0 or mutation_pos >= len(ref_seq):
        return [], 0.0
    ref_base = ref_seq[mutation_pos].upper()
    results  = []
    for alt in ALL_BASES:
        if alt == ref_base:
            continue
        alt_seq = ref_seq[:mutation_pos] + alt + ref_seq[mutation_pos+1:]
        enc_cf  = encode_for_v2(ref_seq, alt_seq, exon_flag, intron_flag)
        with torch.no_grad():
            logit, _, _, _ = model(enc_cf.unsqueeze(0))
            p = torch.sigmoid(logit).item()
        results.append({"mutation": f"{ref_base}>{alt}", "alt_base": alt,
                        "probability": round(p, 4)})
    results_sorted = sorted(results, key=lambda x: x["probability"], reverse=True)
    all_probs  = [r["probability"] for r in results] + [orig_prob]
    cf_delta   = round(max(all_probs) - min(all_probs), 4)
    return results_sorted, cf_delta


# ═══════════════════════════════════════════════════════════════════════════════
# Feature ablation (from live splice app — exact logic preserved)
# ═══════════════════════════════════════════════════════════════════════════════

def _feature_ablation(model: MutationPredictorCNN_v2,
                      enc: torch.Tensor, base_prob: float) -> dict:
    def _p(e):
        with torch.no_grad():
            logit, _, _, _ = model(e.unsqueeze(0))
            return torch.sigmoid(logit).item()

    e_no_splice = enc.clone(); e_no_splice[1103:1106] = 0.0
    e_no_region = enc.clone(); e_no_region[1101:1103] = 0.0
    e_no_mut    = enc.clone(); e_no_mut[1089:1101]    = 0.0

    ds = round(abs(base_prob - _p(e_no_splice)), 4)
    dr = round(abs(base_prob - _p(e_no_region)), 4)
    dm = round(abs(base_prob - _p(e_no_mut)),    4)
    total = ds + dr + dm + 1e-9

    return {
        "baseline_probability":   round(base_prob, 4),
        "splice_causal_effect":   ds,
        "region_causal_effect":   dr,
        "mutation_causal_effect": dm,
        "splice_pct":   round(ds / total * 100, 1),
        "region_pct":   round(dr / total * 100, 1),
        "mutation_pct": round(dm / total * 100, 1),
    }


# ═══════════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════════

def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    denom = (np.linalg.norm(a) * np.linalg.norm(b)) + 1e-9
    return float(np.dot(a, b) / denom)


def _concentration_index(profile: np.ndarray, center: int, window: int = 5) -> float:
    if center < 0 or len(profile) == 0:
        return 0.0
    lo  = max(0, center - window)
    hi  = min(len(profile), center + window + 1)
    local = profile[lo:hi].sum()
    total = profile.sum() + 1e-9
    return float(local / total)


def _activation_pattern(profile: np.ndarray, mutation_pos: int) -> str:
    if mutation_pos < 0 or profile.max() < 1e-6:
        return "Flat"
    peak_val   = profile[mutation_pos]
    mean_val   = profile.mean()
    # count positions above 70% of peak
    high_count = int((profile >= 0.7 * peak_val).sum())
    if peak_val > 2.5 * mean_val and high_count <= 8:
        return "Sharp"
    if peak_val > 1.5 * mean_val:
        return "Broad"
    return "Flat"