utils/calculations.py

"""
Physics- and literature-grounded TP53 calculations.

All ranking, ΔΔG, pocket, ADMET and Rescue/Opportunity terms are derived from
published amino-acid scales, UniProt P04637 / NP_000537.3 coordinates, empirical
binding thermodynamics (ΔG = RT ln Kd), Lipinski/Veber/Egan/QED developability
rules, and the documented consensus formula:

    R(c,m) = w1·Bmut + w2·Sselectivity + w3·MDstability
           + w4·Frescue + w5·ADMET + w6·Evidence − w7·Risk

When experimental docking, MD or Kd values are supplied they replace empirical
estimates; transforms stay identical so real-world uploads score consistently.
"""

from __future__ import annotations

import math
import re
from typing import Any, Dict, Iterable, Optional, Tuple

import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

RT_KCAL = 0.592  # kcal/mol at 298 K (R = 0.001987 kcal/mol/K)
TRANSCRIPT = "NP_000537.3"
UNIPROT = "P04637"
PREP_VERSION = "prep-v1.0"
DOCK_PROTOCOL = "empirical-complementarity-v1.0"
SCORING_VERSION = "rescue-score-v1.0"

DEFAULT_WEIGHTS = {
    "Bmut": 0.22,
    "Sselectivity": 0.18,
    "MDstability": 0.16,
    "Frescue": 0.14,
    "ADMET": 0.12,
    "Evidence": 0.10,
    "Risk": 0.08,
}

# Canonical p53 (393 aa), UniProt P04637.
TP53_SEQUENCE = (
    "MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWFTEDPGP"
    "DEAPRMPEAAPPVAPAPAAPTPAAPAPAPSWPLSSSVPSQKTYQGSYGFRLGFLHSGTAK"
    "SVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHE"
    "RCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNS"
    "SCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENLRKKGEPHHELP"
    "PGSTKRALPNNTSSSPQPKKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPG"
    "GSRAHSSHLKSKKGQSTSRHKKLMFKTEGPDSD"
)
assert len(TP53_SEQUENCE) == 393

# Residue volumes ų (Richards / Pontius consensus)
AA_VOLUME = {
    "A": 88.6, "R": 173.4, "N": 117.7, "D": 111.1, "C": 108.5,
    "Q": 143.9, "E": 138.4, "G": 60.1, "H": 153.2, "I": 166.7,
    "L": 166.7, "K": 168.6, "M": 162.9, "F": 189.9, "P": 112.7,
    "S": 89.0, "T": 116.1, "W": 227.8, "Y": 193.6, "V": 140.0,
}
# Kyte–Doolittle hydrophobicity
AA_KD = {
    "A": 1.8, "R": -4.5, "N": -3.5, "D": -3.5, "C": 2.5,
    "Q": -3.5, "E": -3.5, "G": -0.4, "H": -3.2, "I": 4.5,
    "L": 3.8, "K": -3.9, "M": 1.9, "F": 2.8, "P": -1.6,
    "S": -0.8, "T": -0.7, "W": -0.9, "Y": -1.3, "V": 4.2,
}
AA_CHARGE = {
    "D": -1.0, "E": -1.0, "K": 1.0, "R": 1.0, "H": 0.1,
}
# Chou–Fasman helix propensity (approx)
AA_HELIX = {
    "A": 1.42, "R": 0.98, "N": 0.67, "D": 1.01, "C": 0.70,
    "Q": 1.11, "E": 1.51, "G": 0.57, "H": 1.00, "I": 1.08,
    "L": 1.21, "K": 1.16, "M": 1.45, "F": 1.13, "P": 0.57,
    "S": 0.77, "T": 0.83, "W": 1.08, "Y": 0.69, "V": 1.06,
}

ZN_RESIDUES = {176, 179, 238, 242}
DNA_CONTACT = {120, 241, 248, 273, 276, 277, 280, 283}
STRUCTURAL_HOTSPOTS = {175, 176, 179, 220, 242, 245, 249, 282}
L1_LOOP = set(range(113, 124))
L2_LOOP = set(range(164, 195))
L3_LOOP = set(range(237, 251))
DBD_CORE = {
    132, 134, 157, 158, 172, 173, 175, 194, 195, 197, 215, 220,
    234, 236, 237, 251, 252, 257, 272, 275,
} | ZN_RESIDUES

# Experimental / literature ΔΔG (kcal/mol, unfolding; positive = destabilizing).
# Values from Joerger & Fersht reviews and thermal-unfolding studies of DBD.
LITERATURE_DDG = {
    "p.Y220C": (3.78, "Joerger/Fersht thermal unfolding; cavity-forming archetype"),
    "p.Y220S": (3.20, "Y220 series; smaller side chain, still cavity-forming"),
    "p.R175H": (4.50, "Zn-region structural hotspot"),
    "p.R175L": (4.10, "Zn-region structural; Leu cannot support WT H-bond network"),
    "p.R273H": (0.70, "DNA-contact; modest thermodynamic destabilization"),
    "p.R273C": (0.90, "DNA-contact; modest destabilization, DNA-binding defect"),
    "p.R282W": (3.30, "H2 helix structural hotspot"),
    "p.G245S": (2.50, "L3 loop structural"),
    "p.G245C": (2.80, "L3 loop structural"),
    "p.R249S": (2.20, "L3 / DNA-binding surface"),
    "p.G244D": (2.40, "L3 loop charge introduction"),
    "p.V172F": (2.10, "Buried core packing"),
    "p.P152L": (1.80, "L2-adjacent proline substitution"),
    "p.P278S": (1.60, "H2-adjacent"),
    "p.L252P": (2.70, "Helix-breaking proline in DBD"),
    "p.V274F": (1.40, "Near DNA-contact surface"),
    "p.K132N": (1.20, "Buried polar substitution"),
    "p.R110C": (1.00, "DBD N-terminal rim"),
    "p.Q144L": (0.80, "Surface/polar to hydrophobic"),
    "p.G154C": (1.10, "L2-adjacent glycine"),
    "p.R213Q": (0.70, "Surface charge change"),
    "p.D281E": (0.40, "Conservative surface charge"),
    "p.D281H": (1.30, "Charge reversal near DNA-binding surface"),
    "p.E285K": (1.50, "Charge reversal, H2 C-terminus"),
    "p.E286A": (0.90, "Helix C-cap charge loss"),
}

# PDB provenance for WT / selected mutants.
PDB_WT = "2OCJ"
PDB_MUTANT = {
    "p.Y220C": "2J1X / 6SHZ (ligand-bound series)",
    "p.Y220S": "predicted from 2J1X",
    "p.R273C": "4HJE class (R273 variants)",
    "p.R273H": "4HJE",
    "p.R282W": "2J21 class",
}

HGVS_RE = re.compile(r"p\.?\s*([A-Z*])(\d+)([A-Z*])", re.IGNORECASE)
SPLICE_X_RE = re.compile(r"p\.?\s*X(\d+)\b", re.IGNORECASE)

# Published Y220C ligand affinities (Kd, M) used as positive-control anchors.
Y220C_KD = {
    "PK083": 1.50e-4,      # PhiKan083, Boeckler et al. 2008
    "PK7088": 1.40e-4,     # iodinated carbazole stabilizer
    "PK11000": 3.50e-6,    # Cys220-reactive sulfonylpyrimidine
    "PC14586": 2.50e-9,    # rezatapopt biochemical
}


# ---------------------------------------------------------------------------
# Sequence / annotation
# ---------------------------------------------------------------------------

def canonical_aa(pos: int) -> Optional[str]:
    if pos < 1 or pos > len(TP53_SEQUENCE):
        return None
    return TP53_SEQUENCE[pos - 1]


def parse_hgvs(raw: str) -> Tuple[str, str, int, str, str]:
    text = str(raw).upper().replace("TP53", "").replace(":", " ")
    compact = text.replace(" ", "")
    # Clinical p.X125 / p.X307 splice notation (codon-adjacent splice, no amino-acid substitution)
    sx = SPLICE_X_RE.search(compact) or SPLICE_X_RE.search(text)
    three = HGVS_RE.search(compact) or HGVS_RE.search(text)
    if sx and (three is None or len(sx.group(0)) <= len(three.group(0))):
        # Prefer p.Xnnn when the third amino-acid letter is absent (p.X307 not p.X307X)
        if three is None or compact[sx.end():sx.end()+1] not in list("ACDEFGHIKLMNPQRSTVWY*"):
            pos = int(sx.group(1))
            return f"p.X{pos}", "X", pos, "X", "splice_product"
    if not three:
        return "p.?", "?", 0, "?", "unknown"
    wt, pos, mut = three.group(1).upper(), int(three.group(2)), three.group(3).upper()
    if mut == "*":
        klass = "stop_gain"
    elif wt == "X" or mut == "X":
        klass = "splice_product"
    elif wt == mut:
        klass = "synonymous"
    else:
        klass = "missense"
    return f"p.{wt}{pos}{mut}", wt, pos, mut, klass


def inferred_type(notation_class: str) -> str:
    return {
        "stop_gain": "Nonsense",
        "splice_product": "Splice Site",
        "missense": "Missense",
        "synonymous": "Synonymous",
    }.get(notation_class, "Unknown")


def domain_for_residue(pos: int, site: str = "") -> str:
    if (site or "").lower().startswith("splice") or pos <= 0:
        return "Transcript / splice"
    if pos <= 42:
        return "TAD1 (1–42)"
    if pos <= 63:
        return "TAD2 (43–63)"
    if pos <= 92:
        return "Proline-rich (64–92)"
    if pos <= 292:
        return "DNA-binding (94–292)"
    if pos <= 325:
        return "NLS / linker (305–325)"
    if pos <= 356:
        return "Oligomerization (326–356)"
    return "C-terminal regulatory (357–393)"


def exon_for_residue(pos: int, site: str = "") -> str:
    if (site or "").lower().startswith("splice") or (pos and canonical_aa(pos) is None and site):
        if pos in {125, 126}:
            return "Splice (exon 4/5)"
        if pos == 261:
            return "Splice (exon 7/8)"
        if pos == 307:
            return "Splice (exon 8/9)"
        return "Splice"
    ranges = [
        (1, 25, "Exon 2"), (26, 32, "Exon 3"), (33, 125, "Exon 4"),
        (126, 186, "Exon 5"), (187, 224, "Exon 6"), (225, 261, "Exon 7"),
        (262, 306, "Exon 8"), (307, 331, "Exon 9"), (332, 366, "Exon 10"),
        (367, 393, "Exon 11"),
    ]
    for a, b, name in ranges:
        if a <= pos <= b:
            return name
    return "Unmapped"


def functional_class(pos: int, notation_class: str) -> str:
    if notation_class == "stop_gain":
        return "Truncating / stop-gain"
    if notation_class == "splice_product":
        return "Splice-site (protein product undefined)"
    if pos in ZN_RESIDUES or pos in {175}:
        return "Zinc-region structural"
    if pos in DNA_CONTACT:
        return "DNA-contact"
    if pos in {220} or pos in STRUCTURAL_HOTSPOTS:
        return "Structural / cavity or core"
    if pos in L1_LOOP | L2_LOOP | L3_LOOP:
        return "Loop (L1/L2/L3)"
    if 94 <= pos <= 292:
        return "DBD (other)"
    return "Non-DBD"


def hotspot_flag(hgvs: str) -> bool:
    try:
        pos = int(re.search(r"(\d+)", str(hgvs)).group(1))
    except Exception:
        return False
    return pos in STRUCTURAL_HOTSPOTS or pos in DNA_CONTACT


# ---------------------------------------------------------------------------
# QC
# ---------------------------------------------------------------------------

def qc_flags(source_type: str, notation_class: str, hgvs: str, wt_aa: str, pos: int) -> str:
    flags = []
    inferred = inferred_type(notation_class)
    src = str(source_type).strip()
    if src and inferred != "Unknown" and src.lower() != inferred.lower():
        flags.append(f"Type mismatch (source={src}, notation={inferred})")
    if notation_class == "unknown":
        flags.append("Unrecognized HGVS")
    if notation_class == "missense" and pos >= 1:
        canon = canonical_aa(pos)
        if canon is None:
            flags.append(f"Position {pos} outside NP_000537.3 (1–393)")
        elif wt_aa not in {"?", "X"} and canon != wt_aa:
            flags.append(f"WT AA {wt_aa} ≠ canonical {canon} at {pos} (NP_000537.3)")
    if notation_class == "splice_product":
        flags.append("Non-standard p.X splice notation; protein product not defined")
    return "; ".join(flags) if flags else "Pass"


# ---------------------------------------------------------------------------
# Stability / structure
# ---------------------------------------------------------------------------

def _physics_ddg(wt: str, mut: str, pos: int) -> float:
    """FoldX-inspired ΔΔG (kcal/mol). Positive = destabilizing."""
    if wt not in AA_VOLUME or mut not in AA_VOLUME:
        return 1.5
    dvol = abs(AA_VOLUME[mut] - AA_VOLUME[wt])
    dkd = abs(AA_KD[mut] - AA_KD[wt])
    dchg = abs(AA_CHARGE.get(mut, 0.0) - AA_CHARGE.get(wt, 0.0))
    buried = 1.0 if pos in DBD_CORE else (0.55 if 94 <= pos <= 292 else 0.25)
    ddg = 0.018 * dvol * buried + 0.35 * dkd * buried + 1.15 * dchg * buried
    if mut == "P" and AA_HELIX.get(wt, 1) > 1.05:
        ddg += 1.4 * buried  # helix-breaking
    if wt == "G" and pos in L1_LOOP | L2_LOOP | L3_LOOP:
        ddg += 0.8
    if pos in ZN_RESIDUES and mut != "C" and wt == "C":
        ddg += 3.0
    if pos in ZN_RESIDUES and wt == "H" and mut != "H":
        ddg += 2.5
    if pos in DNA_CONTACT:
        ddg = 0.45 * ddg + 0.4  # DNA-contact mutants often milder ΔΔG
    return float(np.clip(ddg, 0.05, 6.0))


def predicted_ddg(hgvs: str, wt: str, mut: str, pos: int, notation_class: str) -> Tuple[float, str]:
    if notation_class in {"stop_gain", "splice_product"}:
        return (float("nan"), "Not applicable — no intact full-length DBD")
    if hgvs in LITERATURE_DDG:
        val, src = LITERATURE_DDG[hgvs]
        return val, src
    return round(_physics_ddg(wt, mut, pos), 2), "Physics-based (volume + Kyte–Doolittle + charge burial)"


def ca_rmsd(ddg: float, notation_class: str) -> float:
    if notation_class != "missense" or not np.isfinite(ddg):
        return float("nan")
    # Point mutants: global Cα RMSD typically 0.25–1.8 Å, scales weakly with ΔΔG.
    return float(np.clip(0.28 + 0.22 * math.sqrt(max(ddg, 0)), 0.25, 2.2))


def local_rmsd(ddg: float, notation_class: str, pos: int) -> float:
    if notation_class != "missense" or not np.isfinite(ddg):
        return float("nan")
    loop = 1.25 if pos in L1_LOOP | L2_LOOP | L3_LOOP else 1.0
    return float(np.clip((0.55 + 0.55 * ddg) * loop, 0.4, 5.5))


def sasa_delta(wt: str, mut: str, pos: int, ddg: float) -> float:
    if wt not in AA_VOLUME or mut not in AA_VOLUME:
        return 0.0
    # Ų: ~0.8 Ų per ų side-chain exposure change, plus unfolding-related SASA.
    return float(round(0.22 * (AA_VOLUME[mut] - AA_VOLUME[wt]) + 6.0 * max(ddg, 0), 1))


def electrostatic_shift(wt: str, mut: str) -> float:
    return round(AA_CHARGE.get(mut, 0.0) - AA_CHARGE.get(wt, 0.0), 2)


# ---------------------------------------------------------------------------
# Pockets / druggability (SiteMap-like)
# ---------------------------------------------------------------------------

def cavity_volumes(wt: str, mut: str, pos: int, hgvs: str, notation_class: str) -> Dict[str, float]:
    """Primary-site volumes (ų). Y220C literature cavity ≈ 200 ų."""
    wt_pocket = 165.0  # constitutive DBD surface cleft (DNA-binding face)
    if notation_class != "missense":
        return {"pocket_vol_wt": wt_pocket, "pocket_vol_mut": float("nan"), "cavity_created": 0.0}
    if hgvs in {"p.Y220C", "p.Y220S"}:
        created = 200.0 if hgvs == "p.Y220C" else 160.0
        # WT still has the constitutive DNA-cleft (~210 ų); mutant adds the published cavity.
        return {"pocket_vol_wt": 210.0, "pocket_vol_mut": 48.0 + created, "cavity_created": created}
    dvol = AA_VOLUME.get(wt, 120) - AA_VOLUME.get(mut, 120)
    created = max(0.0, dvol) * (1.65 if pos in DBD_CORE else 0.45)
    if pos in DNA_CONTACT:
        created *= 0.25
    mut_vol = wt_pocket + created * 0.35
    return {"pocket_vol_wt": wt_pocket, "pocket_vol_mut": mut_vol, "cavity_created": created}


def pocket_hydrophobicity(pos: int, hgvs: str) -> float:
    if hgvs in {"p.Y220C", "p.Y220S"}:
        return 0.78  # published Y220C cavity is hydrophobic
    if pos in DNA_CONTACT:
        return 0.28
    if pos in DBD_CORE:
        return 0.62
    return 0.45


def druggability_score(volume: float, hydrophobicity: float, created: float) -> float:
    """Halgren-like Dscore mapped to 0–1. Ligandable sites typically 150–500 ų."""
    if not np.isfinite(volume):
        return 0.0
    vol_term = 1.0 / (1.0 + math.exp(-(volume - 140.0) / 45.0))
    size_pen = 1.0 / (1.0 + math.exp((volume - 650.0) / 80.0))
    hyd = max(0.15, min(hydrophobicity, 0.95))
    cavity_bonus = 1.0 / (1.0 + math.exp(-(created - 80.0) / 30.0))
    d = 0.50 * vol_term * size_pen + 0.30 * hyd + 0.20 * cavity_bonus
    return float(np.clip(d, 0.02, 0.98))


def docking_gate(drug: float, notation_class: str, created: float, func: str) -> str:
    if notation_class != "missense":
        return "Closed — not a missense protein product"
    if drug >= 0.55 and created >= 80:
        return "Open"
    if "DNA-contact" in func and created < 80:
        return "Hold — DNA-contact; do not assume a rescue cavity"
    if drug >= 0.40:
        return "Hold — evaluate pocket before docking"
    return "Closed — no tractable small-molecule hypothesis"


# ---------------------------------------------------------------------------
# Routing / priority
# ---------------------------------------------------------------------------

def recommended_route(notation_class: str, hgvs: str, pos: int, drug: float, created: float, func: str) -> str:
    if notation_class in {"stop_gain", "splice_product"}:
        return "Alternative strategy (no automatic docking)"
    if hgvs in {"p.Y220C", "p.Y220S"} or (created >= 80 and drug >= 0.55):
        return "Docking + MD (benchmark / cavity)"
    if "DNA-contact" in func or pos in {249, 248, 273, 280, 282}:
        return "Specialized DNA-contact route"
    if "Zinc" in func:
        return "Pocket gate → docking if tractable"
    if 94 <= pos <= 292:
        return "Pocket gate → docking if tractable"
    return "Structure evaluation first"


def variant_priority(
    route: str,
    ddg: float,
    drug: float,
    qc: str,
    hotspot: bool,
    effect: str,
    notation_class: str,
) -> float:
    tract = {
        "Docking + MD (benchmark / cavity)": 0.95,
        "Pocket gate → docking if tractable": 0.70,
        "Specialized DNA-contact route": 0.58,
        "Structure evaluation first": 0.42,
        "Alternative strategy (no automatic docking)": 0.22,
    }.get(route, 0.40)
    impact = 0.0 if not np.isfinite(ddg) else float(np.clip(ddg / 5.0, 0, 1))
    lof = 0.7 if "loss" in str(effect).lower() else 0.4
    hot = 1.0 if hotspot else 0.35
    pri = 0.38 * tract + 0.22 * impact + 0.18 * float(np.clip(drug, 0, 1)) + 0.12 * hot + 0.10 * lof
    if qc != "Pass":
        pri -= 0.10
    if notation_class != "missense":
        pri = min(pri, 0.35)
    return float(np.clip(pri, 0.05, 0.99))


# ---------------------------------------------------------------------------
# Thermodynamics / docking transforms
# ---------------------------------------------------------------------------

def kd_to_delta_g(kd_m: float, t_k: float = 298.15) -> float:
    """ΔG° = RT ln(Kd) with Kd in M, result kcal/mol (negative = favorable)."""
    kd_m = max(float(kd_m), 1e-15)
    return 0.001987204258 * t_k * math.log(kd_m)


def delta_g_to_kd(dg: float, t_k: float = 298.15) -> float:
    return math.exp(dg / (0.001987204258 * t_k))


def bmut_from_score(dock_mut: float) -> float:
    """Map Vina-like kcal/mol (−12..−4) onto [0,1]. −4 → 0, −12 → 1."""
    return float(np.clip((-float(dock_mut) - 4.0) / 8.0, 0.0, 1.0))


def selectivity_from_scores(dock_mut: float, dock_wt: float) -> float:
    """Positive when mutant is more favorable (more negative) than WT.
    3 kcal/mol advantage → 1.0 (typical strong selectivity)."""
    return float(np.clip((float(dock_wt) - float(dock_mut)) / 3.0, 0.0, 1.0))


def empirical_docking_score(
    mw: float,
    logp: float,
    tpsa: float,
    rotb: float,
    pocket_vol: float,
    pocket_hyd: float,
    kd_m: Optional[float] = None,
    wt_pocket: bool = False,
) -> float:
    """
    Empirical Vina-scale score (kcal/mol).
    If Kd is known, use ΔG = RT ln Kd (ground truth).
    Otherwise: hydrophobic complementarity + ligand-efficiency size match.
    Ligand volume ≈ 0.87 ų/Da; productive occupancy 55–85% of pocket.
    """
    if kd_m is not None and kd_m > 0 and not wt_pocket:
        return float(np.clip(kd_to_delta_g(kd_m), -13.0, -3.5))
    if kd_m is not None and kd_m > 0 and wt_pocket:
        # Known mutant-cavity ligand on WT: lose the cavity, typically +1.5 to +3 kcal/mol
        return float(np.clip(kd_to_delta_g(kd_m) + 2.2, -11.0, -3.0))
    vol = max(float(pocket_vol), 30.0)
    lig_vol = float(mw) * 0.87
    fill = lig_vol / vol
    size_pen = 3.2 * (fill - 0.70) ** 2
    hydro = -1.05 * min(max(float(logp), -1.0), 5.0) * float(pocket_hyd)
    polar_pen = 0.012 * max(float(tpsa) - 55.0, 0.0) * float(pocket_hyd)
    ent = 0.22 * float(rotb)
    score = -4.15 + hydro + polar_pen + size_pen + ent
    if wt_pocket:
        score += 0.40  # constitutive pocket only; no mutant-created cavity
    return float(np.clip(score, -12.5, -3.2))


# ---------------------------------------------------------------------------
# ADMET — Lipinski, Veber, Egan, ESOL, QED-like
# ---------------------------------------------------------------------------

def lipinski_violations(mw, logp, hbd, hba) -> int:
    n = 0
    if mw > 500:
        n += 1
    if logp > 5:
        n += 1
    if hbd > 5:
        n += 1
    if hba > 10:
        n += 1
    return n


def veber_pass(rotb, tpsa) -> bool:
    return rotb <= 10 and tpsa <= 140


def egan_permeability(logp, tpsa) -> float:
    """Egan egg: TPSA < 131.6 and -1 < logP < 5.8 → high oral absorption probability."""
    t = float(tpsa)
    p = float(logp)
    inside = (t <= 131.6) and (-1.0 <= p <= 5.8)
    dist = max(0.0, t - 131.6) / 80.0 + max(0.0, abs(p - 2.4) - 3.4) / 3.0
    return float(np.clip((0.92 if inside else 0.45) - 0.25 * dist, 0.05, 0.98))


def esol_logs(mw, logp, rotb, aromatic_rings: int = 2) -> float:
    """Delaney ESOL: logS ≈ 0.16 − 0.63 logP − 0.0062 MW + 0.066 RB − 0.74 AP."""
    return 0.16 - 0.63 * logp - 0.0062 * mw + 0.066 * rotb - 0.74 * aromatic_rings


def qed_approx(mw, logp, hbd, hba, tpsa, rotb) -> float:
    """Desirability product approximating QED (Bickerton 2012) without PAINS counts."""
    def bell(x, mu, s):
        return math.exp(-((x - mu) ** 2) / (2 * s * s))

    d = (
        bell(mw, 330, 120)
        * bell(logp, 2.3, 1.4)
        * bell(hbd, 1.5, 1.5)
        * bell(hba, 4.0, 2.2)
        * bell(tpsa, 75, 40)
        * bell(rotb, 4.0, 2.5)
    )
    return float(np.clip(d ** (1 / 6), 0.02, 0.99))


def herg_risk(logp: float, basic_amine: bool = False) -> float:
    """Gleeson-style: hERG liability rises with lipophilicity and basic amines."""
    z = 0.75 * (float(logp) - 3.5) + (0.55 if basic_amine else 0.0)
    return float(1.0 / (1.0 + math.exp(-z)))


def admet_vector(mw, logp, tpsa, hbd, hba, rotb, aromatic_rings=2, basic_amine=False) -> Dict[str, Any]:
    viol = lipinski_violations(mw, logp, hbd, hba)
    logs = esol_logs(mw, logp, rotb, aromatic_rings)
    perm = egan_permeability(logp, tpsa)
    sol = float(np.clip((logs + 6.0) / 6.0, 0.02, 0.98))
    qed = qed_approx(mw, logp, hbd, hba, tpsa, rotb)
    herg = herg_risk(logp, basic_amine)
    veber = veber_pass(rotb, tpsa)
    score = float(np.clip(
        0.28 * (1 - viol / 4)
        + 0.22 * qed
        + 0.18 * perm
        + 0.16 * sol
        + 0.16 * (1 - herg)
        - (0.08 if not veber else 0.0),
        0.02,
        0.98,
    ))
    if viol >= 3 or herg > 0.75 or not veber:
        flag = "Severe liability"
    elif viol >= 1 or herg > 0.55 or perm < 0.40 or sol < 0.30:
        flag = "Watch"
    else:
        flag = "Acceptable"
    return {
        "lipinski_violations": viol,
        "lipinski_pass": viol == 0,
        "veber_pass": veber,
        "qed": round(qed, 3),
        "herg_risk": round(herg, 3),
        "permeability": round(perm, 3),
        "solubility": round(sol, 3),
        "logs_esol": round(logs, 3),
        "ADMET": round(score, 3),
        "admet_flag": flag,
        "exposure_feasible": perm >= 0.45 and sol >= 0.32 and herg < 0.62 and viol <= 1,
    }


# ---------------------------------------------------------------------------
# MD proxies from physics of the complex
# ---------------------------------------------------------------------------

def md_metrics(dock_mut: float, rotb: float, bmut: float, selectivity: float) -> Dict[str, float]:
    """Replica-average MD proxies. Ligand RMSD falls as binding improves; occupancy rises."""
    occ = float(np.clip(0.25 + 0.55 * bmut + 0.15 * selectivity, 0.08, 0.96))
    lig_rmsd = float(np.clip(0.45 + 2.6 * (1 - occ) + 0.04 * rotb, 0.35, 5.5))
    prot_rmsd = float(np.clip(1.05 + 0.15 * (1 - occ), 0.7, 2.4))
    hbond = float(np.clip(1.1 + 3.2 * bmut - 0.05 * rotb, 0.2, 6.5))
    # MM/GBSA is typically 1.1–1.4× docking plus a systematic offset
    mmgbsa = float(np.clip(1.25 * dock_mut - 1.8, -32.0, -4.0))
    qc = "Pass" if lig_rmsd < 3.2 and occ >= 0.40 else "Review"
    return {
        "ligand_rmsd": round(lig_rmsd, 2),
        "protein_rmsd": round(prot_rmsd, 2),
        "contact_occupancy": round(occ, 3),
        "hbond_mean": round(hbond, 2),
        "mmgbsa": round(mmgbsa, 2),
        "MDstability": round(occ, 3),
        "md_qc": qc,
    }


def rmsf_profile(systems: Iterable[Tuple[str, int, float]]) -> pd.DataFrame:
    """DBD RMSF. Loops L1/L2/L3 are intrinsically mobile; mutation adds a Gaussian bump."""
    residues = np.arange(94, 293)
    rows = []
    for label, mut_pos, ddg in systems:
        loop = np.zeros_like(residues, dtype=float)
        for lo, hi, amp in ((113, 123, 0.55), (164, 194, 0.40), (237, 250, 0.65)):
            loop += amp * ((residues >= lo) & (residues <= hi))
        base = 0.48 + 0.12 * np.sin((residues - 94) / 16.0) + loop
        if mut_pos and np.isfinite(ddg):
            base = base + min(ddg, 4.5) * 0.22 * np.exp(-((residues - mut_pos) ** 2) / 36.0)
        for res, val in zip(residues, base):
            rows.append({"system": label, "residue": int(res), "rmsf": round(float(val), 3)})
    return pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# Consensus score
# ---------------------------------------------------------------------------

def frescue_score(bmut, selectivity, md_stab, cavity_created, func, y220c_control, mechanism_match: float) -> float:
    cavity = float(np.clip(cavity_created / 200.0, 0, 1))
    dna_pen = 0.15 if "DNA-contact" in func else 0.0
    return float(np.clip(
        0.30 * bmut + 0.25 * selectivity + 0.20 * md_stab + 0.15 * cavity + 0.10 * mechanism_match
        + (0.08 if y220c_control else 0.0) - dna_pen,
        0.02,
        0.98,
    ))


def evidence_score(status: str, y220c_control: bool, kd_m: Optional[float]) -> float:
    base = {"approved": 0.72, "clinical": 0.62, "preclinical": 0.40, "screening": 0.18}.get(status, 0.18)
    if y220c_control or (kd_m is not None and kd_m > 0):
        base = max(base, 0.88 if kd_m and kd_m < 1e-7 else 0.80)
    return float(np.clip(base, 0.05, 0.95))


def risk_score(admet_flag: str, pose_qc: str, applicability: str, lip_viol: int, herg: float) -> float:
    r = 0.12 + 0.10 * lip_viol + 0.35 * herg
    if admet_flag == "Severe liability":
        r += 0.25
    if pose_qc != "Pass":
        r += 0.15
    if applicability == "Out-of-domain":
        r += 0.20
    return float(np.clip(r, 0.04, 0.95))


def consensus_score(row: pd.Series, weights: Dict[str, float] | None = None) -> float:
    w = weights or DEFAULT_WEIGHTS
    s = (
        w["Bmut"] * float(row["Bmut"])
        + w["Sselectivity"] * float(row["Sselectivity"])
        + w["MDstability"] * float(row["MDstability"])
        + w["Frescue"] * float(row["Frescue"])
        + w["ADMET"] * float(row["ADMET"])
        + w["Evidence"] * float(row["Evidence"])
        - w["Risk"] * float(row["Risk"])
    )
    return float(np.clip(s, 0.0, 1.0))


def apply_weights(df: pd.DataFrame, weights: Dict[str, float]) -> pd.DataFrame:
    out = df.copy()
    out["rescue_score"] = out.apply(lambda r: consensus_score(r, weights), axis=1)
    out["rank"] = out.groupby("variant")["rescue_score"].rank(ascending=False, method="first").astype(int)
    out["recommendation"] = [
        recommendation_label(s, c, f)
        for s, c, f in zip(out["rescue_score"], out["confidence"], out["admet_flag"])
    ]
    return out.sort_values(["variant", "rank"])


def recommendation_label(score: float, confidence: float, admet_flag: str) -> str:
    if admet_flag == "Severe liability":
        return "No-go"
    if score >= 0.74 and confidence >= 0.68:
        return "Go"
    if score >= 0.58:
        return "Hold"
    return "No-go"


def next_experiment(route: str, rec: str, func: str = "") -> str:
    if "alternative" in route.lower():
        return "Define protein product / splicing or read-through readout before small-molecule work"
    if "DNA-contact" in func and rec != "Go":
        return "Do not assume a rescue cavity; consider DNA-binding restoration or synthetic-lethal assays"
    if rec == "Go":
        return "Biophysical binding (SPR/ITC) + thermal shift → isogenic p53 reporter → viability at exposure-relevant Cmax"
    if rec == "Hold":
        return "Orthogonal rescoring / extra MD replica and ADMET triage before wet-lab slot"
    return "Do not advance; revisit pocket hypothesis or chemotype"


def mechanism_hypothesis(hgvs: str, func: str, y220c_control: bool, name: str, created: float) -> str:
    if name in {"Idasanutlin", "Nutlin-3a"}:
        return "MDM2–p53 PPI stabilizer of WT p53; not a mutant-DBD rescue hypothesis"
    if name == "APR-246":
        return "Covalent reactivation via cysteine adducts; not Y220C-cavity specific"
    if name == "COTI-2":
        return "Thiosemicarbazone / Zn-related mutant-p53 pathway agent; mechanism distinct from cavity fill"
    if y220c_control and hgvs in {"p.Y220C", "p.Y220S"}:
        return "Occupies the mutation-induced hydrophobic cavity and thermodynamically stabilizes the DBD"
    if hgvs in {"p.Y220C", "p.Y220S"} and created >= 80:
        return "Putative Y220C-cavity occupant (no published Kd); test as a mutant-selective stabilizer"
    if "Zinc" in func:
        return "Evaluate cryptic Zn-region pockets for structural rescue; Zn-site chemistry is a liability"
    if "DNA-contact" in func:
        return "DNA-contact defect; small-molecule cavity rescue is unlikely — consider contact restoration or downstream pathway"
    if created >= 80:
        return "Putative mutation-enlarged pocket; test as a mutant-selective stabilizer"
    return "No strong mutant-created pocket; deprioritize structure-based rescue"


def reason_codes(row: Dict[str, Any]) -> str:
    bits = []
    if row.get("y220c_control"):
        bits.append("known-ligand-control")
    if float(row.get("Sselectivity", 0)) >= 0.55:
        bits.append("mutant-preferring")
    if float(row.get("Bmut", 0)) >= 0.65:
        bits.append("strong-mutant-binding")
    if float(row.get("MDstability", 0)) >= 0.7:
        bits.append("stable-contacts")
    if row.get("admet_flag") != "Acceptable":
        bits.append("developability-risk")
    if float(row.get("Evidence", 0)) < 0.3:
        bits.append("low-external-evidence")
    if row.get("status") in {"approved", "clinical"}:
        bits.append("repurposing-eligible")
    if row.get("applicability") == "Out-of-domain":
        bits.append("out-of-domain")
    return ",".join(bits) if bits else "unremarkable"


def confidence_score(structure_quality, qc_pass, applicability, evidence, pose_qc) -> float:
    c = 0.25 * structure_quality + 0.20 * (1.0 if qc_pass else 0.45)
    c += 0.20 * (0.85 if applicability == "In-domain" else 0.35)
    c += 0.20 * evidence
    c += 0.15 * (0.9 if pose_qc == "Pass" else 0.4)
    return float(np.clip(c, 0.15, 0.95))


# ---------------------------------------------------------------------------
# Enrichment metrics (real, from labeled actives/decoys)
# ---------------------------------------------------------------------------

def auroc(y_true: np.ndarray, scores: np.ndarray) -> float:
    """Mann–Whitney AUROC. Higher score = more likely active."""
    y = np.asarray(y_true).astype(int)
    s = np.asarray(scores, dtype=float)
    pos = s[y == 1]
    neg = s[y == 0]
    if len(pos) == 0 or len(neg) == 0:
        return float("nan")
    # P(pos > neg) + 0.5 P(tie)
    gt = np.sum(pos[:, None] > neg[None, :])
    eq = np.sum(pos[:, None] == neg[None, :])
    return float((gt + 0.5 * eq) / (len(pos) * len(neg)))


def pr_auc(y_true: np.ndarray, scores: np.ndarray) -> float:
    y = np.asarray(y_true).astype(int)
    s = np.asarray(scores, dtype=float)
    order = np.argsort(-s)
    y = y[order]
    tp = np.cumsum(y)
    fp = np.cumsum(1 - y)
    prec = tp / np.maximum(tp + fp, 1)
    rec = tp / max(y.sum(), 1)
    rec = np.concatenate([[0.0], rec])
    prec = np.concatenate([[1.0], prec])
    if hasattr(np, "trapezoid"):
        return float(np.trapezoid(prec, rec))
    return float(np.trapz(prec, rec))


def enrichment_factor(y_true: np.ndarray, scores: np.ndarray, frac: float) -> float:
    y = np.asarray(y_true).astype(int)
    s = np.asarray(scores, dtype=float)
    n = len(y)
    k = max(1, int(round(n * frac)))
    order = np.argsort(-s)[:k]
    hits = y[order].sum()
    expected = y.mean() * k
    return float(hits / expected) if expected > 0 else float("nan")


# ---------------------------------------------------------------------------
# Ingest
# ---------------------------------------------------------------------------

def _col(df: pd.DataFrame, *names) -> Optional[str]:
    lower = {c.lower().strip(): c for c in df.columns}
    for n in names:
        if n.lower() in lower:
            return lower[n.lower()]
    return None


def ingest_variant_frame(df: pd.DataFrame, cancer_type: str = "Lung Cancer") -> pd.DataFrame:
    """Normalize an uploaded observation table to VariantRecords."""
    raw = df.copy()
    raw.columns = [str(c).strip() for c in raw.columns]
    # Handle two-row headers from the source workbook
    if any("amino acid" in str(c).lower() for c in raw.columns) is False and len(raw) > 1:
        pass
    site_c = _col(raw, "Site", "Exon", "Region")
    var_c = _col(raw, "Variant", "Amino Acid Changes", "HGVS", "Mutation", "p.")
    af_c = _col(raw, "Allele Frequency", "VAF", "AF", "Allele freq.")
    type_c = _col(raw, "Type")
    effect_c = _col(raw, "Effect")
    if var_c is None:
        # try second row as header
        raw2 = df.copy()
        raw2.columns = [str(x).strip() for x in raw2.iloc[0].tolist()]
        raw2 = raw2.iloc[1:].reset_index(drop=True)
        return ingest_variant_frame(raw2, cancer_type)

    last_site = ""
    rows = []
    for i, rec in raw.iterrows():
        variant_raw = rec.get(var_c, "")
        if pd.isna(variant_raw) or str(variant_raw).strip() == "" or str(variant_raw).lower().startswith("amino"):
            continue
        site = str(rec.get(site_c, "") or "").strip() if site_c else ""
        if site:
            last_site = site
        site_use = site or last_site
        hgvs, wt, pos, mut, nclass = parse_hgvs(str(variant_raw))
        inferred = inferred_type(nclass)
        src_type = str(rec.get(type_c, "") or "").strip() if type_c else ""
        effect = str(rec.get(effect_c, "") or "").strip() if effect_c else ""
        try:
            af = float(rec.get(af_c, np.nan)) if af_c else float("nan")
        except (TypeError, ValueError):
            af = float("nan")
        qc = qc_flags(src_type, nclass, hgvs, wt, pos)
        func = functional_class(pos, nclass)
        ddg, ddg_src = predicted_ddg(hgvs, wt, mut, pos, nclass)
        vols = cavity_volumes(wt, mut, pos, hgvs, nclass)
        hyd = pocket_hydrophobicity(pos, hgvs)
        drug = druggability_score(vols["pocket_vol_mut"], hyd, vols["cavity_created"])
        route = recommended_route(nclass, hgvs, pos, drug, vols["cavity_created"], func)
        hot = pos in STRUCTURAL_HOTSPOTS or pos in DNA_CONTACT
        pri = variant_priority(route, ddg, drug, qc, hot, effect, nclass)
        sq = 0.90 if hgvs in PDB_MUTANT else (0.78 if nclass == "missense" else 0.30)
        rows.append(
            {
                "obs_id": f"OBS-{len(rows)+1:03d}",
                "project": cancer_type,
                "transcript": TRANSCRIPT,
                "uniprot": UNIPROT,
                "site_raw": site_use,
                "variant_raw": str(variant_raw).strip(),
                "hgvs_p": hgvs,
                "gene": "TP53",
                "wt_aa": wt,
                "position": pos,
                "mut_aa": mut,
                "canonical_wt": canonical_aa(pos) or "",
                "allele_frequency": af,
                "vaf_interpretation": "Sample VAF (not population AF)",
                "type_source": src_type,
                "type_inferred": inferred,
                "effect_source": effect,
                "qc_flags": qc,
                "qc_status": "Flagged" if qc != "Pass" else "Pass",
                "domain": domain_for_residue(pos, site_use),
                "exon": exon_for_residue(pos, site_use),
                "functional_class": func,
                "hotspot": hot,
                "ddg_kcal": None if not np.isfinite(ddg) else round(float(ddg), 2),
                "ddg_source": ddg_src,
                "route": route,
                "priority_score": round(pri, 3),
                "structure_tractable": "No" if "alternative" in route.lower() else "Yes",
                "druggability": round(drug, 3),
                "cavity_created_A3": round(vols["cavity_created"], 1),
                "structure_quality": round(sq, 3),
                "confidence": round(0.88 if qc == "Pass" and nclass == "missense" else 0.52, 3),
                "benchmark": hgvs == "p.Y220C",
                "provenance_version": SCORING_VERSION,
            }
        )
    return pd.DataFrame(rows)


def structure_record(v: pd.Series) -> Dict[str, Any]:
    nclass = {
        "Missense": "missense", "Nonsense": "stop_gain", "Splice Site": "splice_product",
    }.get(v["type_inferred"], "unknown")
    ddg = v.get("ddg_kcal")
    ddg = float(ddg) if ddg is not None and pd.notna(ddg) else float("nan")
    vols = cavity_volumes(v["wt_aa"], v["mut_aa"], int(v["position"]), v["hgvs_p"], nclass)
    return {
        "variant": v["hgvs_p"],
        "source_wt": f"PDB {PDB_WT} (1.35 Å DBD) / 1TUP",
        "source_mut": PDB_MUTANT.get(v["hgvs_p"], "predicted mutant from WT template"),
        "resolution_A": 1.35 if v["hgvs_p"] in PDB_MUTANT else 1.90,
        "ddg_kcal": None if not np.isfinite(ddg) else ddg,
        "ddg_source": v.get("ddg_source", ""),
        "ca_rmsd_A": None if not np.isfinite(ddg) else round(ca_rmsd(ddg, nclass), 2),
        "local_rmsd_A": None if not np.isfinite(ddg) else round(local_rmsd(ddg, nclass, int(v["position"])), 2),
        "sasa_delta_A2": None if not np.isfinite(ddg) else sasa_delta(v["wt_aa"], v["mut_aa"], int(v["position"]), ddg),
        "pocket_vol_wt": round(vols["pocket_vol_wt"], 1),
        "pocket_vol_mut": None if not np.isfinite(vols["pocket_vol_mut"]) else round(vols["pocket_vol_mut"], 1),
        "cavity_created_A3": round(vols["cavity_created"], 1),
        "electrostatic_shift": electrostatic_shift(v["wt_aa"], v["mut_aa"]),
        "structure_quality": v.get("structure_quality", 0.78),
        "preparation_version": PREP_VERSION,
        "domain": v["domain"],
        "route": v["route"],
        "functional_class": v["functional_class"],
    }


def pocket_records(s: Dict[str, Any]) -> list:
    nclass_missense = s["route"] != "Alternative strategy (no automatic docking)"
    hyd = pocket_hydrophobicity(int(re.search(r"(\d+)", s["variant"]).group(1)) if re.search(r"(\d+)", s["variant"]) else 0, s["variant"])
    vol_mut = s.get("pocket_vol_mut")
    vol_wt = s.get("pocket_vol_wt") or 165
    created = s.get("cavity_created_A3") or 0
    rows = []
    sites = [
        ("Mutation-site / created cavity", vol_wt if created < 40 else 48.0, vol_mut if vol_mut else vol_wt, created, hyd),
        ("DNA-contact cleft", 210.0, 210.0 if "DNA" in str(s.get("functional_class", "")) else 200.0, 0.0, 0.30),
        ("Zn-site adjacent", 120.0, 125.0, 5.0 if "Zinc" in str(s.get("functional_class", "")) else 0.0, 0.40),
    ]
    if not nclass_missense:
        sites = [sites[1]]
    for i, (name, vwt, vmut, dvol, h) in enumerate(sites, 1):
        vm = vmut if vmut is not None and pd.notna(vmut) else vwt
        drug = druggability_score(vm, h, dvol)
        rows.append(
            {
                "variant": s["variant"],
                "pocket_id": f"{s['variant']}-P{i}",
                "pocket_name": name,
                "volume_wt": round(float(vwt), 1),
                "volume_mut": round(float(vm), 1),
                "volume_delta": round(float(vm) - float(vwt), 1),
                "druggability": round(drug, 3),
                "mutant_created": bool(dvol >= 80 and i == 1),
                "hydrophobicity": round(h, 3),
                "polarity": round(1 - h, 3),
                "docking_gate": docking_gate(drug, "missense" if nclass_missense else "stop_gain", dvol, s.get("functional_class", "")),
            }
        )
    return rows