"""Nesso-1 — protein–ligand binding affinity prediction on ZeroGPU.

Mirrors the reference `nesso predict` CLI path (see
https://github.com/recursionpharma/nesso, docs/prediction.md): same
preprocessing (RDKit ETKDG conformer + CCD-backed protein tokenisation),
same ESM-2 650M embeddings, same defaults (5 recycling steps, two-stage
pocket refinement, bf16-mixed precision), same `predict_step`.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  — must precede any CUDA-touching import

import hashlib  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from rdkit import Chem, RDLogger  # noqa: E402
from rdkit.Chem import Draw  # noqa: E402
from safetensors.torch import save_file  # noqa: E402

from nesso.data import const  # noqa: E402
from nesso.data.esm import (  # noqa: E402
    DEFAULT_ESM2_MODEL,
    extract_esm_embedding,
    setup_esm_model,
)
from nesso.data.featurizer import NessoFeaturizer  # noqa: E402
from nesso.data.inference import (  # noqa: E402
    STANDARD_AA,
    InferenceDataset,
    inference_collate,
)
from nesso.data.types import Manifest  # noqa: E402
from nesso.data.yaml_input import (  # noqa: E402
    load_ccd_mol_dict,
    parse_schema,
    validate_schema,
)
from nesso.model.models.nesso1 import Nesso1  # noqa: E402

RDLogger.DisableLog("rdApp.*")

REPO_ID = "recursionpharma/nesso"
REVISION = "v1.0.0"
MAX_RESIDUES = 1200
DEFAULT_RECYCLING = 5

# --------------------------------------------------------------------------------------
# Load everything once, at module scope (ZeroGPU packs the weights at startup).
# --------------------------------------------------------------------------------------
print("Downloading Nesso-1 assets…", flush=True)
CCD_PATH = Path(hf_hub_download(REPO_ID, "ccd.pkl", revision=REVISION))
WEIGHTS_PATH = Path(
    hf_hub_download(REPO_ID, f"{REVISION}/model.safetensors", revision=REVISION)
)
hf_hub_download(REPO_ID, f"{REVISION}/hparams.json", revision=REVISION)

print("Loading CCD dictionary…", flush=True)
CCD_DICT = load_ccd_mol_dict(CCD_PATH)
STD_AA_MOLS = {aa: CCD_DICT.get(aa) for aa in STANDARD_AA}

print("Loading Nesso-1…", flush=True)
MODEL = Nesso1.from_pretrained(WEIGHTS_PATH.parent)
# Same predict_args the CLI sets (docs/prediction.md defaults).
MODEL.predict_args.update(
    {
        "pose_protein_cutoff": 15.0,
        "recycling_steps": DEFAULT_RECYCLING,
        "affinity_protein_cutoff": 15.0,
        "refine_protein_inference": True,
        "refine_protein_cutoff": 22.0,
        "refine_protein_tokens_budget": 256,
        "save_metadata": False,
    }
)
MODEL.eval()
MODEL.to("cuda")

print("Loading ESM-2 650M…", flush=True)
ESM_MODEL, ESM_TOKENIZER = setup_esm_model(DEFAULT_ESM2_MODEL, torch.device("cuda"))

torch.set_grad_enabled(False)
torch.set_float32_matmul_precision("highest")
print("Ready.", flush=True)

VALID_AA = set(const.prot_letter_to_token) - {"-"}


# --------------------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------------------
def _clean_sequence(raw: str) -> str:
    """Normalise a pasted protein sequence (accepts FASTA, whitespace, lowercase)."""
    lines = [ln for ln in (raw or "").splitlines() if not ln.strip().startswith(">")]
    seq = "".join("".join(lines).split()).upper()
    seq = "".join(ch for ch in seq if not ch.isdigit())
    return seq


def _format_affinity(value: float) -> str:
    """log10(IC50 / uM) -> a human-readable concentration."""
    ic50_um = 10.0**value
    if ic50_um < 1e-3:
        return f"{ic50_um * 1e6:.2f} pM"
    if ic50_um < 1.0:
        return f"{ic50_um * 1e3:.2f} nM"
    if ic50_um < 1e3:
        return f"{ic50_um:.2f} µM"
    return f"{ic50_um / 1e3:.2f} mM"


def _strength(value: float) -> str:
    if value <= -2.0:
        return "very strong (low-nM or better)"
    if value <= -1.0:
        return "strong"
    if value <= 0.0:
        return "moderate"
    if value <= 1.0:
        return "weak"
    return "very weak / likely non-binder"


def _estimate_duration(
    protein_sequence: str = "",
    ligand_smiles: str = "",
    recycling_steps: int = DEFAULT_RECYCLING,
    *args,
    **kwargs,
) -> int:
    try:
        n = len(_clean_sequence(protein_sequence)) or 400
    except Exception:
        n = 400
    try:
        steps = int(recycling_steps)
    except Exception:
        steps = DEFAULT_RECYCLING
    # Measured on ZeroGPU (H200): wall time scales ~cubically with the number of
    # protein tokens (8 s @ 384 aa, 12 s @ 600 aa, 39 s @ 1100 aa, 5 recycles).
    estimate = 5.0 + 2.5e-8 * (n**3) + 0.8 * steps
    return int(min(150, max(15, round(1.15 * estimate))))


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=_estimate_duration)
def predict_affinity(
    protein_sequence: str,
    ligand_smiles: str,
    recycling_steps: int = DEFAULT_RECYCLING,
    seed: int = 42,
    progress=gr.Progress(track_tqdm=True),
):
    """Predict the binding affinity between a protein and a small molecule.

    Args:
        protein_sequence: Target protein as a single-letter amino-acid sequence (FASTA accepted).
        ligand_smiles: Ligand as a SMILES string.
        recycling_steps: Number of trunk recycling iterations (Nesso-1 default is 5).
        seed: Random seed (controls RDKit conformer generation and featurisation).

    Returns:
        A 2D depiction of the ligand, a Markdown summary, the binder/non-binder
        probabilities, and the raw `affinity.json` scalars produced by Nesso-1.
    """
    seq = _clean_sequence(protein_sequence)
    if not seq:
        raise gr.Error("Please provide a protein amino-acid sequence.")
    bad = sorted(set(seq) - VALID_AA)
    if bad:
        raise gr.Error(f"Unsupported characters in the protein sequence: {bad}")
    if len(seq) > MAX_RESIDUES:
        raise gr.Error(
            f"Sequence has {len(seq)} residues; this demo is capped at {MAX_RESIDUES}. "
            "Paste the target domain (e.g. the kinase domain) instead of the full protein."
        )

    smiles = (ligand_smiles or "").strip()
    if not smiles:
        raise gr.Error("Please provide a ligand SMILES string.")
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise gr.Error(f"RDKit could not parse the SMILES string: {smiles!r}")

    steps = max(0, min(10, int(recycling_steps)))
    seed = int(seed)

    from lightning.pytorch import seed_everything

    seed_everything(seed, workers=True)

    ligand_png = Draw.MolToImage(mol, size=(420, 320))

    t0 = time.perf_counter()
    work = Path(tempfile.mkdtemp(prefix="nesso-"))
    processed = work / "processed"
    mol_dir = processed / "rdkit_conformers"
    structures_dir = processed / "structures"
    records_dir = processed / "records"
    esm_dir = processed / "esm_embeddings"
    for d in (mol_dir, structures_dir, records_dir, esm_dir):
        d.mkdir(parents=True, exist_ok=True)

    record_id = "complex"
    schema = {
        "sequences": [
            {"protein": {"id": "A", "sequence": seq}},
            {"ligand": {"id": "B", "smiles": smiles}},
        ],
        "properties": [{"affinity": {"binder": "B"}}],
    }
    validate_schema(schema)

    try:
        structure, record, entity_to_seq, _ = parse_schema(
            schema, mol_dir, ccd_dict=CCD_DICT, record_id=record_id
        )
    except Exception as exc:  # noqa: BLE001
        raise gr.Error(f"Could not build the complex: {exc}") from exc

    structure.dump(structures_dir / f"{record_id}.npz")
    record.dump(records_dir / f"{record_id}.json")

    # ESM-2 embeddings (same code path as the CLI's `run_esm`).
    for protein_seq in entity_to_seq.values():
        mid = hashlib.md5(protein_seq.encode("utf-8")).hexdigest()  # noqa: S324
        out_path = esm_dir / f"{mid}.safetensors"
        if not out_path.exists():
            emb = extract_esm_embedding(protein_seq, ESM_MODEL, ESM_TOKENIZER)
            save_file({"embeddings": emb}, out_path)

    featurizer = NessoFeaturizer(
        esm_emb_dir=esm_dir, esm_emb_dim=1280, esm_num_layers=33
    )
    dataset = InferenceDataset(
        manifest=Manifest([record]),
        target_dir=processed,
        featurizer=featurizer,
        ligand_dir=mol_dir,
        ccd_pkl=None,
        use_esm_all_layers=False,
    )
    # Reuse the CCD-backed standard residues loaded once at startup.
    dataset._standard_aa_mols = STD_AA_MOLS  # noqa: SLF001

    feats = dataset[0]
    if feats.get("exception"):
        raise gr.Error("Featurisation failed for this complex (see the Space logs).")

    batch = inference_collate([feats])
    batch = {
        k: (v.to("cuda", non_blocking=True) if torch.is_tensor(v) else v)
        for k, v in batch.items()
    }

    # `--precision bf16-mixed` equivalent.
    MODEL.predict_args["recycling_steps"] = steps
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        out = MODEL.predict_step(batch, 0)

    if out.get("exception"):
        raise gr.Error("Prediction failed for this complex (see the Space logs).")

    stats = {}
    for key, value in out.items():
        if not (key.startswith("affinity_") or key.startswith("entropy_")):
            continue
        if key == "entropy_pair":
            continue
        if torch.is_tensor(value) and value.numel() == 1:
            stats[key] = round(float(value.item()), 4)
        elif isinstance(value, (int, float)):
            stats[key] = round(float(value), 4)

    elapsed = time.perf_counter() - t0

    affinity = stats.get("affinity_pred_value")
    prob = stats.get("affinity_probability_binary", 0.0)
    entropy_pl = stats.get("entropy_crop_pl")

    if entropy_pl is not None and entropy_pl == 0.0:
        confidence = (
            "⚠️ **Low confidence** — `entropy_crop_pl` is 0.0, meaning the model could "
            "not confidently place the ligand. Do not trust this prediction."
        )
    else:
        confidence = (
            f"Interface distogram entropy (`entropy_crop_pl`): **{entropy_pl:.3f}** "
            "— lower is a more confident protein–ligand interface."
        )

    summary = f"""
### Predicted binding affinity

| | |
|---|---|
| **log₁₀(IC₅₀ / µM)** | **{affinity:.2f}** ({_strength(affinity)}) |
| Estimated IC₅₀ | **{_format_affinity(affinity)}** |
| pIC₅₀ (= 6 − value) | {6.0 - affinity:.2f} |
| Binder probability | {prob * 100:.1f}% |
| Ensemble members | {stats.get("affinity_pred_value1", float("nan")):.2f} / {stats.get("affinity_pred_value2", float("nan")):.2f} |

{confidence}

<sub>{len(seq)} residues · {mol.GetNumAtoms()} heavy atoms · {steps} recycling steps · {elapsed:.1f}s</sub>
"""

    label = {"binder": float(prob), "non-binder": float(1.0 - prob)}
    return ligand_png, summary, label, stats


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
TUTORIAL_PROTEIN = (
    "MVTPEGNVSLVDESLLVGVTDEDRAVRSAHQFYERLIGLWAPAVMEAAHELGVFAALAEAPADSGELARRLDCDARAMRVL"
    "LDALYAYDVIDRIHDTNGFRYLLSAEARECLLPGTLFSLVGKFMHDINVAWPAWRNLAEVVRHGARDTSGAESPNGIAQED"
    "YESLVGGINFWAPPIVTTLSRKLRASGRSGDATASVLDVGCGTGLYSQLLLREFPRWTATGLDVERIATLANAQALRLGVE"
    "ERFATRAGDFWRGGWGTGYDLVLFANIFHLQTPASAVRLMRHAAACLAPDGLVAVVDQIVDADREPKTPQDRFALLFAASM"
    "TNTGGGDAYTFQEYEEWFTAAGLQRIETLDTPMHRILLARRATEPSAVPEGQASENLYFQ"
)
ABL1_KINASE = (
    "ITMKHKLGGGQYGEVYEGVWKKYSLTVAVKTLKEDTMEVEEFLKEAAVMKEIKHPNLVQLLGVCTREPPFYIITEFMTYGN"
    "LLDYLRECNRQEVNAVVLLYMATQISSAMEYLEKKNFIHRDLAARNCLVGENHLVKVADFGLSRLMTGDTYTAHAGAKFPI"
    "KWTAPESLAYNKFSIKSDVWAFGVLLWEIATYGMSPYPGIDLSQVYELLEKDYRMERPEGCPEKVYELMRACWQWNPSDRP"
    "SFAEIHQAF"
)
EGFR_KINASE = (
    "FKKIKVLGSGAFGTVYKGLWIPEGEKVKIPVAIKELREATSPKANKEILDEAYVMASVDNPHVCRLLGICLTSTVQLITQL"
    "MPFGCLLDYVREHKDNIGSQYLLNWCVQIAKGMNYLEDRRLVHRDLAARNVLVKTPQHVKITDFGLAKLLGAEEKEYHAEG"
    "GKVPIKWMALESILHRIYTHQSDVWSYGVTVWELMTFGSKPYDGIPASEISSILEKGERLPQPPICTIDVYMIMVKCWMID"
    "ADSRPKFRELIIEFSKMARDPQRYL"
)
CDK2 = (
    "MENFQKVEKIGEGTYGVVYKARNKLTGEVVALKKIRLDTETEGVPSTAIREISLLKELNHPNIVKLLDVIHTENKLYLVFE"
    "FLHQDLKKFMDASALTGIPLPLIKSYLFQLLQGLAFCHSHRVLHRDLKPQNLLINTEGAIKLADFGLARAFGVPVRTYTHE"
    "VVTLWYRAPEILLGCKYYSTAVDIWSLGCIFAEMVTRRALFPGDSEIDQLFRIFRTLGTPDEVVWPGVTSMPDYKPSFPKW"
    "ARQDFSKVVPPLDEDGRSLLSQMLHYDPNKRISAKAALAHPFFQDVTKPVPHLRL"
)

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Nesso-1") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# 🧬 Nesso-1 — binding affinity prediction

Predict how strongly a small molecule binds a protein, from **sequence + SMILES only**
(no MSA, no structure). [Nesso-1](https://huggingface.co/recursionpharma/nesso) is a
coarse-grained cofolding model from Valence Labs (Recursion) —
[code](https://github.com/recursionpharma/nesso) ·
[technical report](https://www.biorxiv.org/content/10.64898/2026.08.01.742196v1).
"""
        )
        with gr.Row():
            with gr.Column(scale=1):
                protein = gr.Textbox(
                    label="Protein sequence",
                    placeholder="Single-letter amino-acid sequence (FASTA is fine)…",
                    lines=8,
                    max_lines=12,
                )
                ligand = gr.Textbox(
                    label="Ligand SMILES",
                    placeholder="CC1=C(C=C(C=C1)NC(=O)…",
                    lines=2,
                )
                run = gr.Button("Predict affinity", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    recycling = gr.Slider(
                        1,
                        8,
                        value=DEFAULT_RECYCLING,
                        step=1,
                        label="Recycling steps",
                        info="Nesso-1 was evaluated with 5. More steps = slower.",
                    )
                    seed = gr.Number(value=42, precision=0, label="Seed")
            with gr.Column(scale=1):
                summary_out = gr.Markdown(label="Prediction")
                binder_out = gr.Label(label="Binder classification", num_top_classes=2)
                ligand_out = gr.Image(label="Ligand", height=260)
                with gr.Accordion("Raw output (affinity.json)", open=False):
                    json_out = gr.JSON(label="Nesso-1 scalars")

        gr.Markdown(
            "**Reading the output** — `affinity_pred_value` is log₁₀(IC₅₀ / µM): "
            "**−3 ≈ 1 nM** (strong), **0 ≈ 1 µM** (moderate), **+2 ≈ 100 µM** (weak). "
            "`entropy_crop_pl` measures confidence in the predicted protein–ligand "
            "interface; **0.0 means the prediction should not be trusted**. "
            "Research use only — not for clinical or diagnostic decisions."
        )

        gr.Examples(
            examples=[
                [TUTORIAL_PROTEIN, "N[C@@H](Cc1ccc(O)cc1)C(=O)O"],
                [
                    ABL1_KINASE,
                    "CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=CC(=N4)C5=CN=CC=C5",
                ],
                [
                    EGFR_KINASE,
                    "COC1=C(C=C2C(=C1)N=CN=C2NC3=CC(=C(C=C3)F)Cl)OCCCN4CCOCC4",
                ],
                [
                    CDK2,
                    "C[C@@]12[C@@H]([C@@H](C[C@@H](O1)N3C4=CC=CC=C4C5=C6C(=C7C8=CC=CC=C8N2C7=C53)CNC6=O)NC)OC",
                ],
                [ABL1_KINASE, "CN1C=NC2=C1C(=O)N(C(=O)N2C)C"],
            ],
            example_labels=[
                "Nesso tutorial complex + L-tyrosine",
                "ABL1 kinase domain + imatinib",
                "EGFR kinase domain + gefitinib",
                "CDK2 + staurosporine",
                "ABL1 kinase domain + caffeine (negative control)",
            ],
            inputs=[protein, ligand],
            outputs=[ligand_out, summary_out, binder_out, json_out],
            fn=predict_affinity,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=predict_affinity,
        inputs=[protein, ligand, recycling, seed],
        outputs=[ligand_out, summary_out, binder_out, json_out],
        api_name="predict",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)