"""
MyAbs — v0.2 product spine + developability flags.

Paste a heavy (VH) + light (VL) chain -> fold with ABodyBuilder2 -> view the 3D
structure with CDR loops highlighted -> scan for common developability liabilities
(PTM hotspots, glycosylation sequons, free cysteines, long CDR-H3) and paint the
offending residues onto the structure.

Runs locally (WSL 'base'/'myabs' env) and drops onto a free HF CPU Space unchanged.

Run:
    pip install gradio            # once, into the same env as ImmuneBuilder
    python app.py                 # opens http://127.0.0.1:7900

The liability flags are HEURISTIC screens, not disqualifiers. MyAbs yields in-silico
CANDIDATES, not patent-ready antibodies — wet-lab validation still required.
"""

import io
import itertools
import json
import os
import tempfile
import time
import gradio as gr
from ImmuneBuilder import ABodyBuilder2

try:
    from anarci import number as anarci_number
    HAVE_ANARCI = True
except Exception:
    HAVE_ANARCI = False

try:
    from Bio.PDB import PDBParser
    from Bio.PDB.SASA import ShrakeRupley
    HAVE_SASA = True
except Exception:
    HAVE_SASA = False

print("Loading ABodyBuilder2 ensemble...")
PREDICTOR = ABodyBuilder2()
print("Ready.")

# IMGT CDR position ranges (ABodyBuilder2 writes IMGT-numbered PDBs).
CDR_RANGES = {"1": (27, 38), "2": (56, 65), "3": (105, 117)}

# CDR cartoon colors: heavy = warm, light = cool. Framework stays grey.
CDR_COLORS = {
    "H": {"1": "#ffca28", "2": "#ff7043", "3": "#e53935"},
    "L": {"1": "#4dd0e1", "2": "#29b6f6", "3": "#1e88e5"},
}

# Illustrative demo VH / VL (verify before any real use).
DEMO_H = ("EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMSWVRQAPGKGLEWVSAISGSGGST"
          "YYADSVKGRFTISRDNSKNTLYLQMNSLRAEDTAVYYCAKDRGYYYGMDVWGQGTTVTVSS")
DEMO_L = ("DIQMTQSPSSLSASVGDRVTITCRASQSISSYLNWYQQKPGKAPKLLIYAASSLQSGVP"
          "SRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYSTPLTFGGGTKVEIK")

# Therapeutic library — variable domains only, self-curated from PUBLIC sources.
# Every sequence cross-checked against two independent authoritative sources.
PASTE = "— paste your own —"
LIBRARY = {
    PASTE: None,
    "Trastuzumab (Herceptin) · anti-HER2": {
        "H": "EVQLVESGGGLVQPGGSLRLSCAASGFNIKDTYIHWVRQAPGKGLEWVARIYPTNGYTRYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCSRWGGDGFYAMDYWGQGTLVTVSS",
        "L": "DIQMTQSPSSLSASVGDRVTITCRASQDVNTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSRSGTDFTLTISSLQPEDFATYYCQQHYTTPPTFGQGTKVEIK",
        "target": "HER2 (ERBB2)",
        "src": "PDB 1N8Z SEQRES + KEGG DRUG D03257 (identical)",
        "url": "https://www.rcsb.org/structure/1N8Z",
    },
    "Adalimumab (Humira) · anti-TNF-α": {
        "H": "EVQLVESGGGLVQPGRSLRLSCAASGFTFDDYAMHWVRQAPGKGLEWVSAITWNSGHIDYADSVEGRFTISRDNAKNSLYLQMNSLRAEDTAVYYCAKVSYLSTASSLDYWGQGTLVTVSS",
        "L": "DIQMTQSPSSLSASVGDRVTITCRASQGIRNYLAWYQQKPGKAPKLLIYAASTLQSGVPSRFSGSGSGTDFTLTISSLQPEDVATYYCQRYNRAPYTFGQGTKVEIK",
        "target": "TNF-α",
        "src": "PDB 6CR1 + DrugBank DB00051 (consensus; lone 3WD5 outlier residue rejected)",
        "url": "https://www.rcsb.org/structure/6CR1",
    },
    "Pembrolizumab (Keytruda) · anti-PD-1": {
        "H": "QVQLVQSGVEVKKPGASVKVSCKASGYTFTNYYMYWVRQAPGQGLEWMGGINPSNGGTNFNEKFKNRVTLTTDSSTTTAYMELKSLQFDDTAVYYCARRDYRFDMGFDYWGQGTTVTVSS",
        "L": "EIVLTQSPATLSLSPGERATLSCRASKGVSTSGYSYLHWYQQKPGQAPRLLIYLASYLESGVPARFSGSGSGTDFTLTISSLEPEDFAVYYCQHSRDLPLTFGGGTKLEIK",
        "target": "PD-1 (PDCD1)",
        "src": "PDB 5DK3 + 5GGS (identical)",
        "url": "https://www.rcsb.org/structure/5DK3",
    },
    "Rituximab (Rituxan) · anti-CD20": {
        "H": "QVQLQQPGAELVKPGASVKMSCKASGYTFTSYNMHWVKQTPGRGLEWIGAIYPGNGDTSYNQKFKGKATLTADKSSSTAYMQLSSLTSEDSAVYYCARSTYYGGDWYFNVWGAGTTVTVSA",
        "L": "QIVLSQSPAILSASPGEKVTMTCRASSSVSYIHWFQQKPGSSPKPWIYATSNLASGVPVRFSGSGSGTSYSLTISRVEAEDAATYYCQQWTSNPPTFGGGTKLEIK",
        "target": "CD20 (MS4A1)",
        "src": "PDB 2OSL + 6VJA (identical)",
        "url": "https://www.rcsb.org/structure/2OSL",
    },
    "Bevacizumab (Avastin) · anti-VEGF-A": {
        "H": "EVQLVESGGGLVQPGGSLRLSCAASGYTFTNYGMNWVRQAPGKGLEWVGWINTYTGEPTYAADFKRRFTFSLDTSKSTAYLQMNSLRAEDTAVYYCAKYPHYYGSSHWYFDVWGQGTLVTVSS",
        "L": "DIQMTQSPSSLSASVGDRVTITCSASQDISNYLNWYQQKPGKAPKVLIYFTSSLHSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQYSTVPWTFGQGTKVEIK",
        "target": "VEGF-A",
        "src": "PDB 1BJ1 + NIH GSRS (confirmed)",
        "url": "https://www.rcsb.org/structure/1BJ1",
    },
}

# On-screen CDR + liability legend (static HTML).
def _swatch(label, color):
    return (f"<span style='display:inline-flex;align-items:center;gap:4px;margin-right:10px'>"
            f"<span style='width:13px;height:13px;background:{color};border-radius:2px;"
            f"display:inline-block;border:1px solid #999'></span>{label}</span>")

LEGEND_HTML = (
    "<div style='display:flex;flex-wrap:wrap;align-items:center;font-size:13px;margin:4px 0'>"
    "<b style='margin-right:10px'>Legend:</b>"
    + _swatch("CDR-H1", CDR_COLORS["H"]["1"]) + _swatch("H2", CDR_COLORS["H"]["2"])
    + _swatch("H3", CDR_COLORS["H"]["3"]) + _swatch("CDR-L1", CDR_COLORS["L"]["1"])
    + _swatch("L2", CDR_COLORS["L"]["2"]) + _swatch("L3", CDR_COLORS["L"]["3"])
    + _swatch("liability", "magenta") + _swatch("your pick", "lime")
    + _swatch("framework", "#cfd8dc")
    + "</div>"
)

SEV_ORDER = {"High": 0, "Medium": 1, "Low": 2, "Minimal": 3}
SEV_ICON = {"High": "🔴", "Medium": "🟠", "Low": "🟡", "Minimal": "⚪"}
SEV_LEVELS = ["High", "Medium", "Low", "Minimal"]

# Tien et al. 2013 theoretical max ASA (Ų), for relative solvent accessibility.
MAXASA = {
    "ALA": 129, "ARG": 274, "ASN": 195, "ASP": 193, "CYS": 167, "GLU": 223,
    "GLN": 225, "GLY": 104, "HIS": 224, "ILE": 197, "LEU": 201, "LYS": 236,
    "MET": 224, "PHE": 240, "PRO": 159, "SER": 155, "THR": 172, "TRP": 285,
    "TYR": 263, "VAL": 174,
}


def clean(seq: str) -> str:
    """Strip whitespace/newlines/numbers, uppercase — accept messy pasted input."""
    return "".join(c for c in seq.upper() if c.isalpha())


# ------------------------------------------------------------------ numbering
def number_chain(seq: str):
    """ANARCI IMGT-number a chain -> [(imgt_int, insertion, aa)] for present residues.
    Returns None if numbering fails / ANARCI unavailable."""
    if not HAVE_ANARCI:
        return None
    try:
        numbering, _chain_type = anarci_number(seq, scheme="imgt")
        if not numbering:
            return None
        out = []
        for (pos, ins), aa in numbering:
            if aa == "-":
                continue
            out.append((pos, ins, aa))
        return out
    except Exception:
        return None


def cdr_of(imgt: int):
    for name, (lo, hi) in CDR_RANGES.items():
        if lo <= imgt <= hi:
            return name
    return None


# -------------------------------------------------------------- accessibility
def rsa_map(pdb: str):
    """Per-residue relative solvent accessibility from the folded Fv.
    Returns {(chain_id, imgt_resseq): rsa} or None if SASA is unavailable.
    Computed on the whole Fv, so the VH/VL interface counts as buried."""
    if not HAVE_SASA or not pdb:
        return None
    try:
        model = PDBParser(QUIET=True).get_structure("ab", io.StringIO(pdb))[0]
        ShrakeRupley().compute(model, level="R")  # sets .sasa on each residue
        out = {}
        for chain in model:
            for res in chain:
                if res.resname not in MAXASA:
                    continue
                rsa = res.sasa / MAXASA[res.resname]
                key = (chain.id, res.id[1])  # (chain, resseq); insertion code dropped
                out[key] = max(out.get(key, 0.0), rsa)  # keep max across insertions
        return out
    except Exception:
        return None


def exposure_tier(rsa):
    """buried (<15%), partial (15-30%), exposed (>=30%), or None if no RSA."""
    if rsa is None:
        return None
    if rsa < 0.15:
        return "buried"
    if rsa < 0.30:
        return "partial"
    return "exposed"


def adjust_severity(raw_sev: str, tier, desc: str) -> str:
    """Down-rank a sequence-motif flag by how buried the residue is: buried motifs
    can't undergo solvent-driven chemistry (oxidation, deamidation, glycosylation).
    A free cysteine is a covalent/structural concern beyond exposure, so it is
    never down-ranked more than one level."""
    if tier is None:
        return raw_sev
    steps = {"exposed": 0, "partial": 1, "buried": 2}[tier]
    if "cysteine" in desc.lower():
        steps = min(steps, 1)
    i = min(SEV_LEVELS.index(raw_sev) + steps, len(SEV_LEVELS) - 1)
    return SEV_LEVELS[i]


# ---------------------------------------------------------------- liabilities
def scan_chain(chain_label: str, residues):
    """Scan one numbered chain for sequence-liability motifs.
    Returns (flags, highlight_imgt_positions). Each flag: (sev, chain, imgt, desc, loc)."""
    flags, highlights = [], []
    aas = [r[2] for r in residues]
    imgts = [r[0] for r in residues]
    n = len(residues)

    for i in range(n):
        imgt, aa = imgts[i], aas[i]
        cdr = cdr_of(imgt)
        loc = f"CDR-{chain_label}{cdr}" if cdr else "framework"
        nxt = aas[i + 1] if i + 1 < n else ""
        nxt2 = aas[i + 2] if i + 2 < n else ""

        # N-glycosylation sequon  N-X-[S/T], X != P
        if aa == "N" and nxt and nxt != "P" and nxt2 in ("S", "T"):
            sev = "High" if cdr else "Medium"
            flags.append((sev, chain_label, imgt, f"N-glycosylation sequon (N{nxt}{nxt2})", loc))
            highlights.append(imgt)

        # Deamidation  NG (fast) ; NS/NT/NH (slow, only flag in CDR)
        if aa == "N" and nxt == "G":
            flags.append(("High" if cdr else "Medium", chain_label, imgt,
                          "Deamidation motif (NG)", loc))
            highlights.append(imgt)
        elif aa == "N" and nxt in ("S", "T", "H") and cdr:
            flags.append(("Low", chain_label, imgt, f"Deamidation motif (N{nxt})", loc))
            highlights.append(imgt)

        # Isomerization  DG (fast) ; DS/DT (slow, only flag in CDR)
        if aa == "D" and nxt == "G":
            flags.append(("High" if cdr else "Medium", chain_label, imgt,
                          "Isomerization motif (DG)", loc))
            highlights.append(imgt)
        elif aa == "D" and nxt in ("S", "T") and cdr:
            flags.append(("Low", chain_label, imgt, f"Isomerization motif (D{nxt})", loc))
            highlights.append(imgt)

        # Acid-labile peptide bond  DP
        if aa == "D" and nxt == "P":
            flags.append(("Low", chain_label, imgt, "Acid-labile bond (DP)", loc))
            highlights.append(imgt)

        # Oxidation-prone Met / Trp — only a liability when exposed (i.e. in a CDR)
        if aa in ("M", "W") and cdr:
            res = "Met" if aa == "M" else "Trp"
            flags.append(("Medium", chain_label, imgt, f"Oxidation-prone {res} in CDR", loc))
            highlights.append(imgt)

    # Free / non-canonical cysteine (canonical intradomain disulfide = IMGT 23 & 104)
    for i in range(n):
        if aas[i] == "C" and imgts[i] not in (23, 104):
            cdr = cdr_of(imgts[i])
            loc = f"CDR-{chain_label}{cdr}" if cdr else "framework"
            flags.append(("High", chain_label, imgts[i], "Unpaired / non-canonical cysteine", loc))
            highlights.append(imgts[i])

    return flags, highlights


def developability(heavy: str, light: str, pdb: str = None):
    """Scan both chains + CDR-H3 length, then gate by solvent exposure using the
    folded structure. Returns (flags, highlights_by_chain, h3_len, ok).
    Each flag: (sev, chain, imgt, desc, loc, rsa, tier) where sev is the
    exposure-adjusted severity and rsa/tier are None when no structure is given."""
    raw_flags = []
    h3_len = None
    ok = True
    for label, seq in (("H", heavy), ("L", light)):
        residues = number_chain(seq)
        if residues is None:
            ok = False
            continue
        flags, _hi = scan_chain(label, residues)
        raw_flags += flags
        if label == "H":
            h3_len = sum(1 for (imgt, _ins, _aa) in residues if 105 <= imgt <= 117)
            if h3_len >= 18:
                raw_flags.append(("High", "H", 105,
                                  f"Long CDR-H3 ({h3_len} aa) — aggregation risk", "CDR-H3"))

    # Exposure-gate each flag against the folded structure (if available).
    rmap = rsa_map(pdb)
    enriched = []
    highlights = {"H": [], "L": []}
    for sev, chain, imgt, desc, loc in raw_flags:
        # CDR-H3 length is a whole-loop property, not single-residue exposure.
        is_h3len = desc.startswith("Long CDR-H3")
        rsa = None if is_h3len else (rmap.get((chain, imgt)) if rmap else None)
        tier = exposure_tier(rsa)
        adj = sev if is_h3len else adjust_severity(sev, tier, desc)
        enriched.append((adj, chain, imgt, desc, loc, rsa, tier))
        # Paint only residues whose flag survives exposure gating (High/Medium).
        if adj in ("High", "Medium") and not is_h3len and chain in highlights:
            highlights[chain].append(imgt)
    return enriched, highlights, h3_len, ok


def format_flags(flags, h3_len, ok):
    """Render the liability panel as Markdown."""
    if not ok:
        return ("**Developability:** could not IMGT-number the sequences (ANARCI). "
                "Flags unavailable — check that both chains are valid variable domains.")
    if not flags:
        return ("### ✅ No sequence liabilities flagged\n"
                "No glycosylation sequons, PTM hotspots, or free cysteines found in the CDRs "
                "or framework"
                + (f", and CDR-H3 length is normal ({h3_len} aa)" if h3_len else "")
                + ".")

    flags_sorted = sorted(flags, key=lambda f: (SEV_ORDER[f[0]], f[1], f[2]))
    highs = sum(1 for f in flags if f[0] == "High")
    meds = sum(1 for f in flags if f[0] == "Medium")
    lows = sum(1 for f in flags if f[0] == "Low")
    mins = sum(1 for f in flags if f[0] == "Minimal")
    have_exposure = any(f[6] is not None for f in flags)

    def _exp_cell(rsa, tier):
        if rsa is None:
            return "—"
        return f"{tier} ({rsa * 100:.0f}%)"

    tally = f"{highs} high · {meds} medium · {lows} low"
    if mins:
        tally += f" · {mins} minimal"
    lines = [
        f"### Developability flags — {tally}",
        ("Severity is **adjusted by solvent exposure** from the fold: buried motifs are "
         "down-ranked because they can't undergo solvent-driven chemistry. Surviving "
         "(high/medium) liabilities are drawn as **magenta sticks** on the structure."
         if have_exposure else
         "Flagged residues are drawn as **magenta sticks** on the structure."),
        "",
        "| Severity | Exposure | Chain | IMGT | Location | Liability |",
        "|---|---|---|---|---|---|",
    ]
    for sev, chain, imgt, desc, loc, rsa, tier in flags_sorted:
        lines.append(
            f"| {SEV_ICON[sev]} {sev} | {_exp_cell(rsa, tier)} | {chain} | {imgt} | {loc} | {desc} |"
        )
    footer = ("_Heuristic screens, not disqualifiers. Exposure (relative solvent accessibility) "
              "sharpens the ranking but is not the whole story: an exposed motif can still be "
              "fine (far from the paratope, slow kinetics, controlled by formulation). Confirm "
              "experimentally before acting._")
    lines += ["", footer]
    return "\n".join(lines)


# -------------------------------------------------------------------- viewer
# Per-residue colors for the clickable sequence tracks (CDRs reuse the 3D colors;
# "liab" is the same magenta as the liability sticks on the structure).
TRACK_COLORS = {
    "H1": CDR_COLORS["H"]["1"], "H2": CDR_COLORS["H"]["2"], "H3": CDR_COLORS["H"]["3"],
    "L1": CDR_COLORS["L"]["1"], "L2": CDR_COLORS["L"]["2"], "L3": CDR_COLORS["L"]["3"],
    "liab": "#ff00ff", "fw": "#eceff1",
}


def build_track(chain_label: str, residues, liab_imgts=()):
    """Numbered residues -> (HighlightedText tokens, index->imgt map).
    Each residue is its own token (combine_adjacent=False) so it clicks individually.
    Liability residues (same set shown as magenta sticks in 3D) are tagged 'liab'
    so they read magenta in the sequence, matching the structure."""
    tokens, idxmap = [], []
    liab = set(liab_imgts)
    for imgt, _ins, aa in residues:
        cdr = cdr_of(imgt)
        if imgt in liab:
            cat = "liab"
        elif cdr:
            cat = f"{chain_label}{cdr}"
        else:
            cat = "fw"
        tokens.append((aa, cat))
        idxmap.append(imgt)
    return tokens, idxmap


FOCUS_CHOICES = ["Both chains", "Heavy only (VH)", "Light only (VL)"]
MODE_CHOICES = ["Grey out others", "Hide others"]

_NONCE = itertools.count(1)


def build_payload(pdb: str, highlights) -> str:
    """JSON handed to the client-side viewer: the structure + liability positions.
    The nonce guarantees the value changes each fold so the .change bridge fires."""
    highlights = highlights or {}
    return json.dumps({
        "pdb": pdb,
        "highlights": {"H": highlights.get("H", []), "L": highlights.get("L", [])},
        "n": next(_NONCE),
    })


def toggle_payload(chain_label: str, evt, idxmap) -> str:
    """JSON telling the client to toggle a lime stick on one clicked residue."""
    idx = getattr(evt, "index", None)
    if isinstance(idx, (list, tuple)):
        idx = idx[0] if idx else None
    imap = (idxmap or {}).get(chain_label, [])
    imgt = imap[idx] if isinstance(idx, int) and 0 <= idx < len(imap) else None
    return json.dumps({"chain": chain_label, "imgt": imgt, "n": next(_NONCE)})


# Client-side 3Dmol controller. One persistent viewer; every interaction updates
# it IN PLACE (no re-render, no zoomTo) so the camera / orientation is preserved.
# Loaded once via demo.load(js=...); it also injects 3Dmol.js from the CDN.
_CDR_JS = json.dumps(CDR_COLORS)
_RANGES_JS = json.dumps({k: list(v) for k, v in CDR_RANGES.items()})
CONTROLLER_JS = """
() => {
  if (window.__myabsReady) return;
  window.__myabsReady = true;
  const CDR = __CDR__;
  const RANGES = __RANGES__;
  const rlist = (lo,hi) => { let a=[]; for(let i=lo;i<=hi;i++) a.push(i); return a; };
  const S = window.myabsState = {viewer:null, highlights:{H:[],L:[]}, picks:{H:[],L:[]},
                                 focus:["H","L"], mode:"grey", spinning:false, colorMode:"cdr"};

  // Per-residue predicted error (PDB B-factor, Å) -> color. 0 = confident (blue),
  // >= CONF_MAX = uncertain (red), through yellow. Matches AlphaFold-style intuition.
  const CONF_MAX = 1.5;
  const confColor = (atom) => {
    const t = Math.max(0, Math.min(1, (atom.b || 0) / CONF_MAX));
    let r,g,b;
    if (t < 0.5){ const u=t/0.5; r=Math.round(43+(240-43)*u); g=Math.round(131+(200-131)*u); b=Math.round(186+(50-186)*u); }
    else { const u=(t-0.5)/0.5; r=Math.round(240+(215-240)*u); g=Math.round(200+(48-200)*u); b=Math.round(50+(39-50)*u); }
    return "rgb("+r+","+g+","+b+")";
  };

  window.myabsApply = () => {
    const v = S.viewer; if(!v) return;
    v.setStyle({}, {});
    ["H","L"].forEach(ch => {
      if (S.focus.includes(ch)) {
        if (S.colorMode === "confidence") {
          v.addStyle({chain:ch}, {cartoon:{colorfunc: confColor}});   // color by predicted error
        } else {
          v.addStyle({chain:ch}, {cartoon:{color:"#cfd8dc"}});
          for (const c in RANGES){ const r=RANGES[c];
            v.addStyle({chain:ch, resi:rlist(r[0],r[1])}, {cartoon:{color:CDR[ch][c]}}); }
        }
        const hl=S.highlights[ch]||[]; if(hl.length) v.addStyle({chain:ch, resi:hl},{stick:{color:"magenta",radius:0.3}});
        const pk=S.picks[ch]||[]; if(pk.length) v.addStyle({chain:ch, resi:pk},{stick:{color:"lime",radius:0.3}});
      } else if (S.mode === "grey") {
        v.addStyle({chain:ch}, {cartoon:{color:"#c7ccd1", opacity:0.35}});
      }
    });
    v.render();
  };

  window.myabsInit = () => {
    if (S.viewer) return true;
    const el = document.getElementById("myabs-viewer");
    if (!el || !window.$3Dmol) return false;
    S.viewer = $3Dmol.createViewer(el, {backgroundColor:"white"});
    return true;
  };

  window.myabsLoad = (payload) => {
    if (!payload) return;
    let p; try { p = JSON.parse(payload); } catch(e){ return; }
    if (!p.pdb) return;
    if (!window.myabsInit()) { setTimeout(()=>window.myabsLoad(payload), 200); return; }
    const v = S.viewer;
    S.highlights = p.highlights || {H:[],L:[]}; S.picks = {H:[],L:[]}; S.spinning = false;
    v.removeAllModels(); v.addModel(p.pdb, "pdb");
    window.myabsApply();
    v.zoomTo(); v.render();   // new molecule: recentering here is expected
  };

  window.myabsSetFocus = (choice) => {
    const M = {"Both chains":["H","L"], "Heavy only (VH)":["H"], "Light only (VL)":["L"]};
    S.focus = M[choice] || ["H","L"]; window.myabsApply();   // no zoomTo -> view kept
  };
  window.myabsSetMode = (m) => {
    S.mode = (m && m.indexOf("Hide")>=0) ? "hide" : "grey"; window.myabsApply();
  };
  window.myabsSetColorMode = (m) => {
    S.colorMode = (m && m.indexOf("onfidence")>=0) ? "confidence" : "cdr"; window.myabsApply();
  };
  window.myabsToggleResidue = (payload) => {
    let p; try { p = JSON.parse(payload); } catch(e){ return; }
    if (p.imgt == null || !p.chain) return;
    const arr = S.picks[p.chain] || (S.picks[p.chain]=[]);
    const i = arr.indexOf(p.imgt); if (i>=0) arr.splice(i,1); else arr.push(p.imgt);
    window.myabsApply();
  };
  window.myabsClear = () => { S.picks = {H:[],L:[]}; window.myabsApply(); };
  window.myabsToggleSpin = () => {
    const v = S.viewer; if (!v) return "▶ Spin";
    S.spinning = !S.spinning; v.spin(S.spinning ? "y" : false);
    return S.spinning ? "⏸ Stop" : "▶ Spin";
  };

  if (!window.$3Dmol) {
    const s = document.createElement("script");
    // Pinned, immutable version + Subresource Integrity so a compromised CDN
    // cannot inject arbitrary JS into users' browsers.
    s.src = "https://cdn.jsdelivr.net/npm/3dmol@2.5.5/build/3Dmol-min.js";
    s.integrity = "sha384-OsczYbldvrHgslr9fFp/i4GiLSeuw9l+QIlv99ITw8soOwXcoGeflFMLg+CU/X1d";
    s.crossOrigin = "anonymous";
    s.onload = () => window.myabsInit();
    document.head.appendChild(s);
  } else { window.myabsInit(); }
}
""".replace("__CDR__", _CDR_JS).replace("__RANGES__", _RANGES_JS)


# ---------------------------------------------------------------------- fold
def fold_confidence(pdb: str) -> str:
    """Readout of ABodyBuilder2's per-residue predicted error (PDB B-factor = RMS
    spread across the 4-model ensemble, Å; lower = models agree = more confident)."""
    if not HAVE_SASA or not pdb:
        return ""
    try:
        model = PDBParser(QUIET=True).get_structure("ab", io.StringIO(pdb))[0]
    except Exception:
        return ""
    per, h3 = [], []
    for chain in model:
        for res in chain:
            bs = [a.get_bfactor() for a in res]
            if not bs:
                continue
            b = sum(bs) / len(bs)
            per.append(b)
            if chain.id == "H" and 105 <= res.id[1] <= 117:
                h3.append(b)
    if not per:
        return ""

    def band(x):
        return "high" if x < 0.5 else ("moderate" if x < 1.5 else "low")

    mean = sum(per) / len(per)
    out = [f"**Fold confidence:** mean predicted error **{mean:.2f} Å** "
           f"({band(mean)} confidence)."]
    if h3:
        h3m = sum(h3) / len(h3)
        note = ""
        if h3m >= 1.5:
            note = " — the least certain region, so treat its liability flags with extra caution"
        elif h3m >= 0.8:
            note = " — moderate certainty"
        out.append(f"CDR-H3 **{h3m:.2f} Å**{note}.")
    out.append("_Predicted error = spread across the 4-model ensemble (lower = models agree). "
               "Set the viewer's \"Color by\" to **Fold confidence** to see it on the structure "
               "(blue = confident → red = uncertain)._")
    return " ".join(out)


MAX_CHAIN_LEN = 250  # antibody variable domains are ~110-130 aa; a generous DoS cap


def fold(heavy: str, light: str):
    """Fold VH+VL, scan liabilities, and hand the structure to the client viewer."""
    heavy, light = clean(heavy), clean(light)
    reset_btn = gr.update(value="▶ Spin")
    # trailing 8 outputs (everything after `status`) for the early-return paths
    tail = ("", "", None, "", [], [], {"H": [], "L": []}, reset_btn)
    if not heavy or not light:
        return ("Enter both a heavy and a light chain.",) + tail
    if len(heavy) > MAX_CHAIN_LEN or len(light) > MAX_CHAIN_LEN:
        return (f"Sequence too long (VH {len(heavy)}, VL {len(light)} aa; max "
                f"{MAX_CHAIN_LEN} per chain). Paste one antibody variable domain per box.",) + tail
    try:
        t0 = time.time()
        antibody = PREDICTOR.predict({"H": heavy, "L": light})
        dt = time.time() - t0
        # Per-fold unique dir so concurrent public users never share the output
        # file (the download basename stays a clean "myabs_fold.pdb").
        out_path = os.path.join(tempfile.mkdtemp(prefix="myabs_"), "myabs_fold.pdb")
        antibody.save(out_path)
    except Exception:  # OpenMM refinement / numbering can occasionally fail
        return ("Fold failed. Check that both inputs are valid antibody "
                "variable-domain sequences.",) + tail

    pdb = open(out_path).read()
    flags, highlights, h3_len, ok = developability(heavy, light, pdb)
    status = (f"Folded in {dt:.1f} s  ·  VH {len(heavy)} aa / VL {len(light)} aa  ·  "
              f"CDRs highlighted (H: yellow/orange/red, L: cyan/blue).")
    conf_md = fold_confidence(pdb)
    flags_md = format_flags(flags, h3_len, ok)

    h_tokens, h_idx = build_track("H", number_chain(heavy) or [], highlights.get("H", []))
    l_tokens, l_idx = build_track("L", number_chain(light) or [], highlights.get("L", []))
    idxmap = {"H": h_idx, "L": l_idx}
    payload = build_payload(pdb, highlights)  # -> client-side viewer via .change bridge
    return (status, conf_md, flags_md, out_path, payload,
            h_tokens, l_tokens, idxmap, reset_btn)


def load_library(name: str):
    """Populate VH/VL from the therapeutic library + show provenance."""
    entry = LIBRARY.get(name)
    if not entry:  # "paste your own"
        return (DEMO_H, DEMO_L,
                "_Built-in demo Fv (illustrative, not a real drug). "
                "Pick a therapeutic above, or paste your own sequences._")
    prov = (f"**{name.split(' · ')[0]}**  ·  Target: **{entry['target']}**  ·  "
            f"variable domains from a public source: {entry['src']} "
            f"([reference]({entry['url']})).")
    return entry["H"], entry["L"], prov


def click_heavy(idxmap, evt: gr.SelectData):
    """Map a heavy-chain track click to its IMGT position for the client to toggle."""
    return toggle_payload("H", evt, idxmap)


def click_light(idxmap, evt: gr.SelectData):
    return toggle_payload("L", evt, idxmap)


with gr.Blocks(title="MyAbs") as demo:  # theme moved to launch() in Gradio 6
    gr.Markdown(
        "# MyAbs\n"
        "**Look at an antibody candidate, fold it, see its CDR loops, and flag its "
        "developability liabilities — in your browser.**\n\n"
        "_In-silico candidates only. Not patent-ready antibodies; wet-lab validation required._"
    )
    with gr.Row():
        with gr.Column(scale=2):
            lib_dd = gr.Dropdown(
                choices=list(LIBRARY.keys()), value=PASTE,
                label="Load a known therapeutic (public sequences) — or paste your own",
            )
            provenance = gr.Markdown()
            h_in = gr.Textbox(label="Heavy chain (VH)", value=DEMO_H, lines=4)
            l_in = gr.Textbox(label="Light chain (VL)", value=DEMO_L, lines=4)
            with gr.Row():
                go = gr.Button("Fold", variant="primary")
                spin_btn = gr.Button("▶ Spin")
            status = gr.Markdown()
            conf_md = gr.Markdown()
            pdb_file = gr.File(label="Download structure (.pdb)")
        with gr.Column(scale=3):
            with gr.Row():
                focus_dd = gr.Radio(choices=FOCUS_CHOICES, value="Both chains",
                                    label="Focus chain")
                mode_dd = gr.Radio(choices=MODE_CHOICES, value="Grey out others",
                                   label="The non-focused chain is…")
                color_dd = gr.Radio(choices=["CDR regions", "Fold confidence"],
                                    value="CDR regions", label="Color by")
            # Persistent 3Dmol viewer container. Updated in place by the client-side
            # controller (CONTROLLER_JS) so interactions never reset the camera.
            viewer = gr.HTML(
                '<div id="myabs-viewer" style="width:100%;height:580px;position:relative;'
                'border:1px solid #eee;border-radius:8px"></div>'
            )
            gr.HTML(LEGEND_HTML)
            gr.HTML(
                "<div style='font-size:13px;margin:2px 0'><b>Confidence scale</b> "
                "(when Color by → Fold confidence): "
                "<span style='color:rgb(43,131,186)'>■&nbsp;confident</span> → "
                "<span style='color:rgb(240,200,50)'>■</span> → "
                "<span style='color:rgb(215,48,39)'>■&nbsp;uncertain</span> "
                "&nbsp;(predicted error, Å)</div>"
            )
            gr.Markdown(
                "**Rotate it yourself** — _Touchscreen:_ one-finger drag = rotate · "
                "pinch = zoom · two-finger drag = pan.  _Mouse:_ drag = rotate · "
                "scroll = zoom · right-drag = pan."
            )
            gr.Markdown(
                "**Click a residue below to highlight it on the structure (lime stick).** "
                "Click it again to remove it. CDR residues are pre-colored; "
                "**liability residues are magenta** (same as the sticks in 3D)."
            )
            h_track = gr.HighlightedText(label="Heavy chain (VH)", combine_adjacent=False,
                                         show_legend=False, color_map=TRACK_COLORS)
            l_track = gr.HighlightedText(label="Light chain (VL)", combine_adjacent=False,
                                         show_legend=False, color_map=TRACK_COLORS)
            clear_btn = gr.Button("Clear clicked highlights", size="sm")

    gr.Markdown("---")
    flags_md = gr.Markdown()

    idxmap_state = gr.State({"H": [], "L": []})  # track index -> imgt, per chain
    # Hidden bridges: server writes JSON here, a .change(js=...) hands it to the viewer.
    load_box = gr.Textbox(visible=False)     # fold -> myabsLoad
    toggle_box = gr.Textbox(visible=False)   # residue click -> myabsToggleResidue

    fold_outputs = [status, conf_md, flags_md, pdb_file, load_box, h_track, l_track,
                    idxmap_state, spin_btn]

    go.click(fold, inputs=[h_in, l_in], outputs=fold_outputs)

    # Pick a therapeutic -> load its sequences + provenance -> auto-fold.
    lib_dd.change(load_library, inputs=lib_dd, outputs=[h_in, l_in, provenance]).then(
        fold, inputs=[h_in, l_in], outputs=fold_outputs
    )

    # --- client-side viewer controls (no server round-trip, camera preserved) ---
    load_box.change(None, inputs=[load_box], js="(p) => window.myabsLoad(p)")
    toggle_box.change(None, inputs=[toggle_box], js="(p) => window.myabsToggleResidue(p)")
    focus_dd.change(None, inputs=[focus_dd], js="(c) => window.myabsSetFocus(c)")
    mode_dd.change(None, inputs=[mode_dd], js="(m) => window.myabsSetMode(m)")
    color_dd.change(None, inputs=[color_dd], js="(m) => window.myabsSetColorMode(m)")
    spin_btn.click(None, outputs=[spin_btn], js="() => window.myabsToggleSpin()")
    clear_btn.click(None, js="() => window.myabsClear()")

    # Residue click: server maps track index -> IMGT, client applies the lime stick.
    h_track.select(click_heavy, inputs=[idxmap_state], outputs=[toggle_box])
    l_track.select(click_light, inputs=[idxmap_state], outputs=[toggle_box])

    # Load 3Dmol.js + define the controller once, on page load.
    demo.load(None, js=CONTROLLER_JS)

if __name__ == "__main__":
    # Local dev defaults to 127.0.0.1:7900. On an HF Docker Space the Dockerfile
    # sets GRADIO_SERVER_NAME=0.0.0.0 and GRADIO_SERVER_PORT=7860 (app_port).
    host = os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1")
    port = int(os.environ.get("GRADIO_SERVER_PORT", "7900"))
    demo.launch(server_name=host, server_port=port, theme=gr.themes.Soft())