dee/core/features.py

"""The residues you should not mutate, from the people who curated them.

THE PROBLEM
-----------
The engine will happily rank a substitution at the catalytic serine of a
protease as promising. ESM-2 scores how UNUSUAL a residue is, not how load-
bearing it is, and an active-site residue is often perfectly ordinary in
sequence terms — it is the geometry that matters, and a language model over
sequence does not see geometry.

So the single cheapest improvement to design quality is not a better model. It
is asking UniProt what is already known about the positions being mutated:
active sites, binding sites, disulfides, metal ligands, domains. That is
curated, citable, human-reviewed annotation, and it turns "this substitution
looks tolerable" into "this substitution is at the catalytic residue".

WHAT THIS IS
------------
A read of UniProt's feature table, mapped onto positions the user cares about.
No prediction, no score, no model — every item returned carries the evidence
UniProt itself carries. When UniProt has nothing for a protein, that is
reported as "not annotated", never as "safe": absence of annotation is absence
of knowledge, and conflating the two is how a tool talks someone into an
experiment.

Only (accession) leaves the Space. Never a user sequence.
"""
from __future__ import annotations

import json
import re
import urllib.parse
import urllib.request
from typing import Any, Dict, Iterable, List, Optional

_UA = "TuringDNA/1.0 (https://turingdna.com)"
_ENTRY = "https://rest.uniprot.org/uniprotkb"

# Feature types worth interrupting a design for, in the order a designer cares.
# UniProt emits many more (VARIANT, CONFLICT, chains); those describe the
# record, not a reason to avoid a position, and including them would bury the
# signal.
CRITICAL = {
    "Active site":      "catalytic — mutating this typically abolishes activity",
    "Binding site":     "contacts the substrate or cofactor",
    "Site":             "functionally annotated position",
    "Metal binding":    "coordinates a metal ion the fold depends on",
    "Disulfide bond":   "forms a disulfide; losing one cysteine breaks both",
    "Modified residue": "post-translationally modified",
    "Glycosylation":    "glycosylation site",
    "Cross-link":       "covalently cross-linked",
}
CONTEXT = {
    "Domain":   "domain",
    "Region":   "region",
    "Motif":    "motif",
    "Repeat":   "repeat",
    "Transmembrane": "membrane-spanning",
    "Signal":   "signal peptide",
    "Propeptide": "propeptide — cleaved from the mature protein",
}


def _get(url: str, timeout: float = 15.0) -> Optional[dict]:
    try:
        req = urllib.request.Request(url, headers={"User-Agent": _UA,
                                                   "Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return json.loads(r.read().decode("utf-8"))
    except Exception:  # noqa: BLE001
        return None


def _span(f: Dict[str, Any]) -> Optional[tuple]:
    loc = (f.get("location") or {})
    try:
        s = loc["start"]["value"]
        e = loc["end"]["value"]
    except (KeyError, TypeError):
        return None
    if s is None or e is None:
        return None
    return (int(s), int(e))


def fetch(accession: str) -> Dict[str, Any]:
    """UniProt's curated feature table for one accession."""
    acc = re.sub(r"[^A-Za-z0-9\-]", "", accession or "").upper()
    if not acc:
        return {"ok": False, "error": "No accession given."}
    # The whole entry, filtered here. A ?fields= list looked tidier and is
    # not valid on the entry endpoint — it returns an error, not a subset, so
    # the "tidy" version fetched nothing at all.
    data = _get(f"{_ENTRY}/{urllib.parse.quote(acc)}.json")
    if not data:
        return {"ok": False, "kind": "unreachable",
                "error": f"Couldn't reach UniProt for {acc}.",
                "next": "Retry; this is usually momentary."}

    feats = []
    for f in (data.get("features") or []):
        sp = _span(f)
        if not sp:
            continue
        ftype = str(f.get("type") or "")
        tier = "critical" if ftype in CRITICAL else (
            "context" if ftype in CONTEXT else None)
        if tier is None:
            continue
        feats.append({
            "type": ftype,
            "start": sp[0], "end": sp[1],
            "tier": tier,
            "description": str(f.get("description") or "")[:180],
            "means": CRITICAL.get(ftype) or CONTEXT.get(ftype, ""),
            # UniProt's own evidence codes travel with the claim. ECO:0000269
            # is experimental; ECO:0000250 is inferred by similarity, and a
            # designer should weight those differently.
            "evidence": sorted({str(e.get("evidenceCode") or "")
                                for e in (f.get("evidences") or [])
                                if e.get("evidenceCode")}),
        })

    name = ""
    try:
        name = (data["proteinDescription"]["recommendedName"]
                    ["fullName"]["value"])
    except (KeyError, TypeError):
        pass
    return {
        "ok": True,
        "accession": data.get("primaryAccession") or acc,
        "protein": name,
        "length": ((data.get("sequence") or {}).get("length")),
        "features": feats,
        "critical_count": sum(1 for f in feats if f["tier"] == "critical"),
        "source": "UniProt",
    }


def annotate(accession: str, positions: Iterable) -> Dict[str, Any]:
    """Which of `positions` land on annotated features.

    Positions may be integers or substitution labels ('R175H'); the label form
    is what the rest of the engine speaks, so accepting it saves the caller a
    parsing step and a chance to get the offset wrong.
    """
    table = fetch(accession)
    if not table.get("ok"):
        return table

    wanted: List[Dict[str, Any]] = []
    for p in positions or []:
        if isinstance(p, int):
            wanted.append({"label": str(p), "pos": p})
            continue
        m = re.match(r"^([A-Za-z])?(\d+)([A-Za-z*])?$", str(p).strip())
        if m:
            wanted.append({"label": str(p).strip(), "pos": int(m.group(2))})

    hits = []
    for w in wanted:
        on = [f for f in table["features"] if f["start"] <= w["pos"] <= f["end"]]
        crit = [f for f in on if f["tier"] == "critical"]
        hits.append({
            "label": w["label"], "position": w["pos"],
            "critical": crit, "context": [f for f in on if f["tier"] == "context"],
            # The whole point of the tool, stated so the agent cannot miss it.
            "verdict": ("AVOID — annotated functional residue" if crit else
                        "no functional annotation at this position"),
        })

    flagged = [h for h in hits if h["critical"]]
    return {
        "ok": True,
        "accession": table["accession"], "protein": table["protein"],
        "length": table["length"],
        "positions": hits,
        "flagged_count": len(flagged),
        "verdict": (f"{len(flagged)} of {len(hits)} positions sit on annotated "
                    f"functional residues." if flagged else
                    f"None of the {len(hits)} positions are annotated as "
                    f"functional."),
        # The failure that would make this tool actively harmful. UniProt
        # numbers the FULL PRECURSOR — bovine trypsin's catalytic triad is at
        # 63/107/200 here, not the classic His57/Asp102/Ser195, because the
        # signal peptide and propeptide are counted. A user thinking in mature
        # or chymotrypsin numbering gets a confident flag on the wrong residue,
        # which is worse than no flag at all.
        "numbering": (f"Positions are UniProt's, numbered from residue 1 of the "
                      f"full precursor ({table['length']} aa) including any "
                      f"signal peptide or propeptide. If you are using mature-"
                      f"protein or a classic numbering scheme, these will be "
                      f"offset — check one known residue before trusting the "
                      f"rest."),
        "caveat": ("Absence of annotation is absence of KNOWLEDGE, not evidence "
                   "that a position is safe — most proteins are annotated "
                   "sparsely, and an unannotated residue may still be "
                   "essential. Evidence codes are UniProt's own: ECO:0000269 "
                   "is experimental, ECO:0000250 is inferred by similarity."),
        "source": "UniProt",
    }