"""
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 os
import time
import gradio as gr
import py3Dmol
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("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
def render_structure(pdb: str, spin: bool = False, highlights=None) -> str:
    """py3Dmol view -> self-contained HTML in an iframe (renders inside Gradio).
    Framework grey, CDR loops colored, liability residues as magenta sticks."""
    view = py3Dmol.view(width=760, height=560)
    view.addModel(pdb, "pdb")
    view.setStyle({"cartoon": {"color": "#cfd8dc"}})
    for chain in ("H", "L"):
        for cdr, (lo, hi) in CDR_RANGES.items():
            view.addStyle(
                {"chain": chain, "resi": list(range(lo, hi + 1))},
                {"cartoon": {"color": CDR_COLORS[chain][cdr]}},
            )
    if highlights:
        for chain in ("H", "L"):
            pos = highlights.get(chain, [])
            if pos:
                view.addStyle(
                    {"chain": chain, "resi": pos},
                    {"stick": {"color": "magenta", "radius": 0.3}},
                )
    view.zoomTo()
    if spin:
        view.spin(True)
    html = view._make_html()
    esc = html.replace("&", "&amp;").replace('"', "&quot;")
    return f'<iframe style="width:100%;height:580px;border:none;border-radius:8px" srcdoc="{esc}"></iframe>'


# ---------------------------------------------------------------------- fold
def fold(heavy: str, light: str):
    """Fold VH+VL, scan liabilities, render, and reset the spin toggle."""
    heavy, light = clean(heavy), clean(light)
    reset_btn = gr.update(value="▶ Spin")
    if not heavy or not light:
        return ("<p style='padding:1em'>Enter both a heavy and a light chain.</p>",
                "Need both chains.", "", None, "", {}, False, reset_btn)
    try:
        t0 = time.time()
        antibody = PREDICTOR.predict({"H": heavy, "L": light})
        dt = time.time() - t0
        antibody.save("myabs_fold.pdb")
    except Exception as e:  # OpenMM refinement / numbering can occasionally fail
        return (f"<p style='padding:1em;color:#c00'>Fold failed: {e}</p>",
                f"Error: {e}", "", None, "", {}, False, reset_btn)

    pdb = open("myabs_fold.pdb").read()
    flags, highlights, h3_len, ok = developability(heavy, light, pdb)
    viewer = render_structure(pdb, spin=False, highlights=highlights)
    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).")
    flags_md = format_flags(flags, h3_len, ok)
    return viewer, status, flags_md, "myabs_fold.pdb", pdb, highlights, False, 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 toggle_spin(spin: bool, pdb: str, highlights):
    """Flip spin on/off and re-render the stored structure (no re-fold)."""
    spin = not spin
    label = "⏸ Stop" if spin else "▶ Spin"
    if not pdb:
        return spin, gr.update(value=label), gr.update()
    return spin, gr.update(value=label), render_structure(pdb, spin, highlights)


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()
            pdb_file = gr.File(label="Download structure (.pdb)")
        with gr.Column(scale=3):
            viewer = gr.HTML()
            gr.HTML(LEGEND_HTML)
            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("---")
    flags_md = gr.Markdown()

    pdb_state = gr.State("")       # last folded PDB, for re-render without re-folding
    hi_state = gr.State({})        # liability highlight positions by chain
    spin_state = gr.State(False)   # is the structure currently spinning?

    fold_outputs = [viewer, status, flags_md, pdb_file, pdb_state, hi_state, spin_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
    )

    spin_btn.click(
        toggle_spin, inputs=[spin_state, pdb_state, hi_state],
        outputs=[spin_state, spin_btn, viewer],
    )

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())