tests/test_annotate.py

"""Scanning a construct for what nothing has annotated yet."""
import random

import pytest

from dee.core.annotate import (CURATED, MOTIF, PREDICTED, merge_annotations,
                               orf_chance_expectation, scan_construct,
                               _chance_model)
from dee.core.plasmid import _MOTIFS, find_orfs
from dee.core.promoters import PROMOTER_LIBRARY

NO_STOP = ["GCT", "TGC", "GAT", "GAA", "TTC", "GGT", "CAT", "ATC", "AAA",
           "CTG", "AAC", "CCG", "CAG", "CGT", "TCT", "ACC", "GTT", "TGG"]


def _filler(n, seed=0):
    r = random.Random(seed)
    return "".join(r.choices("ACGT", k=n))


def _orf(n_codons, seed=1):
    r = random.Random(seed)
    return "ATG" + "".join(r.choices(NO_STOP, k=n_codons)) + "TAA"


def _motif(name):
    return next(sq for m, sq, _t in _MOTIFS if m == name)


@pytest.fixture
def planted():
    """A construct with known elements at known coordinates."""
    ara = next(p for p in PROMOTER_LIBRARY["ecoli"] if p.name == "araBAD")
    parts = [
        (None, _filler(600, 1)),
        ("araBAD promoter", ara.sequence),
        (None, _filler(30, 2)),
        ("CDS", _orf(400)),
        (None, _filler(200, 3)),
        ("6×His tag", _motif("6×His tag")),
        (None, _filler(400, 4)),
        ("loxP", _motif("loxP")),
        (None, _filler(900, 5)),
    ]
    seq, truth = "", {}
    for name, part in parts:
        if name:
            truth[name] = (len(seq), len(seq) + len(part))
        seq += part
    return seq, truth


def test_finds_planted_elements_in_an_unannotated_construct(planted):
    seq, truth = planted
    r = scan_construct(seq, features=[], topology="circular")
    at = {(f["start"], f["end"]): f for f in r["new"]}
    for name in ("araBAD promoter", "6×His tag", "loxP"):
        assert truth[name] in at, f"{name} not recovered at {truth[name]}"
        assert at[truth[name]]["name"] == name
    s, e = truth["CDS"]
    # find_orfs opens at the first ATG since the last stop, which can sit
    # upstream of the planted one, so the call covers the CDS rather than
    # matching it base for base.
    assert any(f["tier"] == PREDICTED and "ORF" in f["name"]
               and f["start"] <= s and f["end"] >= e
               for f in r["new"]), "planted ORF not recovered"


def test_an_exact_match_is_not_displaced_by_a_longer_predicted_orf(planted):
    """Regression: a His tag inside a CDS was dropped because the ORF was longer.

    Biology nests — tags sit inside coding sequence, operators inside
    promoters. An exact sequence match must never lose to a guess.
    """
    seq, truth = planted
    r = scan_construct(seq, features=[], topology="circular")
    names = {f["name"] for f in r["new"]}
    assert "6×His tag" in names


def test_the_same_element_from_two_sources_is_reported_once():
    """T7 is in both the curated library and the motif list."""
    seq = _filler(400, 9) + _motif("T7 promoter") + _filler(400, 10)
    r = scan_construct(seq, features=[], topology="linear")
    t7 = [f for f in r["new"] if "T7" in f["name"]]
    assert len(t7) == 1, [f["name"] for f in t7]
    assert t7[0]["tier"] == CURATED       # stronger provenance wins
    assert not t7[0]["name"].endswith("promoter promoter")


# ── the point of the whole thing: partial records ────────────────────────
def test_existing_annotations_are_never_touched(planted):
    seq, truth = planted
    s, e = truth["CDS"]
    existing = [{"name": "my gene", "type": "CDS", "start": s, "end": e,
                 "strand": 1, "color": "#abc", "notes": "curated"}]
    r = scan_construct(seq, features=existing, topology="circular")
    assert r["existing"] == [dict(existing[0])]
    assert all(f["name"] != "my gene" for f in r["new"])


def test_gaps_are_filled_around_what_is_already_annotated(planted):
    """auto_annotate bails on any annotated record; this must not."""
    seq, truth = planted
    s, e = truth["CDS"]
    existing = [{"name": "my gene", "type": "CDS", "start": s, "end": e,
                 "strand": 1}]
    r = scan_construct(seq, features=existing, topology="circular")
    found = {f["name"] for f in r["new"]}
    assert "araBAD promoter" in found and "loxP" in found
    assert r["coverage"]["after_pct"] > r["coverage"]["before_pct"]


def test_a_prediction_over_annotated_sequence_is_reported_separately(planted):
    seq, truth = planted
    s, e = truth["CDS"]
    r = scan_construct(seq, features=[{"name": "known", "type": "CDS",
                                       "start": s, "end": e, "strand": 1}],
                       topology="circular")
    over = [f for f in r["confirms"] if f["tier"] == PREDICTED]
    assert over, "the ORF under the existing annotation should be set aside"
    assert over[0]["overlaps"] == "known"


# ── circular origin ──────────────────────────────────────────────────────
def test_an_orf_across_the_origin_is_found_only_when_circular():
    cds = _orf(300, seed=4)
    seq = cds[400:] + _filler(2500, 6) + cds[:400]
    assert not find_orfs(seq, min_aa=250), "fixture must hide it from a linear scan"
    circ = scan_construct(seq, features=[], topology="circular", min_aa=250)
    assert any(f.get("wraps") for f in circ["new"])
    lin = scan_construct(seq, features=[], topology="linear", min_aa=250)
    assert not [f for f in lin["new"] if "ORF" in f["name"]]


# ── the chance model ─────────────────────────────────────────────────────
def test_chance_expectation_falls_as_the_orf_grows():
    p_atg, p_stop = _chance_model(_filler(5000, 7))
    vals = [orf_chance_expectation(10000, aa, p_atg, p_stop)
            for aa in (40, 80, 160, 320)]
    assert vals == sorted(vals, reverse=True)
    assert vals[0] > 1.0 > vals[-1]


def test_the_chance_model_tracks_real_counts_within_an_order_of_magnitude():
    """It is an estimate, not a bound.

    An earlier version of this claimed to be an upper bound on the strength of
    three samples; across 36 conditions real counts run above it about a fifth
    of the time, by up to 1.8x. What it must actually do is stay in the right
    order of magnitude, because the figure is quoted to scientists.
    """
    ratios = []
    for gc in (0.35, 0.45, 0.55, 0.65):
        for seed in range(3):
            r = random.Random(1000 + seed)
            w = [(1 - gc) / 2, gc / 2, gc / 2, (1 - gc) / 2]
            seq = "".join(r.choices("ACGT", weights=w, k=30000))
            p_atg, p_stop = _chance_model(seq)
            for aa in (40, 60, 80):
                pred = orf_chance_expectation(30000, aa, p_atg, p_stop)
                if pred > 0.5:
                    obs = len(find_orfs(seq, min_aa=aa, max_orfs=100000))
                    ratios.append(obs / pred)
    assert ratios
    assert max(ratios) < 3.0, f"over-confident: max obs/pred {max(ratios):.2f}"
    mid = sorted(ratios)[len(ratios) // 2]
    assert 0.4 < mid < 1.6, f"median obs/pred drifted to {mid:.2f}"


def test_a_real_length_orf_is_separated_from_noise(planted):
    seq, truth = planted
    r = scan_construct(seq, features=[], topology="circular")
    s, e = truth["CDS"]
    planted_orf = next(f for f in r["new"]
                       if "ORF" in f["name"] and f["start"] <= s and f["end"] >= e)
    assert planted_orf["chance"] < 0.01
    assert "unidentified" in planted_orf["name"]      # never a named gene


def test_pure_noise_orfs_are_not_reported():
    """A construct of random DNA should not come back full of confident calls."""
    r = scan_construct(_filler(12000, 8), features=[], topology="circular",
                       min_aa=30)
    for f in r["new"]:
        if f["tier"] == PREDICTED and "ORF" in f["name"]:
            assert f["chance"] <= 100.0


# ── coverage ─────────────────────────────────────────────────────────────
def test_coverage_and_dark_regions_add_up(planted):
    seq, _ = planted
    r = scan_construct(seq, features=[], topology="circular")
    cov = r["coverage"]
    assert 0 <= cov["after_pct"] <= 100
    assert cov["after_pct"] >= cov["before_pct"]
    for d in cov["dark_regions"]:
        assert d["end"] - d["start"] == d["length"] >= 150
        assert 0 <= d["start"] < d["end"] <= r["length"]


def test_empty_sequence_is_refused():
    with pytest.raises(ValueError):
        scan_construct("")


def test_merge_keeps_existing_and_marks_the_rest(planted):
    seq, truth = planted
    s, e = truth["CDS"]
    existing = [{"name": "my gene", "type": "CDS", "start": s, "end": e,
                 "strand": 1}]
    r = scan_construct(seq, features=existing, topology="circular")
    merged = merge_annotations(existing, r["new"])
    assert len(merged) == len(existing) + len(r["new"])
    assert merged == sorted(merged, key=lambda f: (f["start"], f["end"]))
    mine = next(f for f in merged if f["name"] == "my gene")
    assert not mine.get("auto")
    assert all(f.get("auto") for f in merged if f["name"] != "my gene")


# ── the tool ─────────────────────────────────────────────────────────────
def test_the_tool_is_registered_and_takes_the_bound_construct():
    from dee.core import agent_tools as t
    assert t.needs_target("scan_construct"), \
        "it must act on the loaded construct, not on sequence the model retypes"
    names = [f["function"]["name"] for f in t.TOOL_SPECS]
    assert "scan_construct" in names
    spec = next(f for f in t.TOOL_SPECS if f["function"]["name"] == "scan_construct")
    props = spec["function"]["parameters"]["properties"]
    assert "sequence" not in props, "the sequence is injected, never passed in"
    assert not spec["function"]["parameters"]["required"]


def test_the_tool_reports_without_touching_the_construct(planted):
    from dee.core.agent_tools import _tool_scan_construct
    seq, truth = planted
    s, e = truth["CDS"]
    existing = [{"name": "my gene", "type": "CDS", "start": s, "end": e,
                 "strand": 1}]
    before = [dict(f) for f in existing]
    res = _tool_scan_construct({"sequence": seq, "features": existing,
                                "topology": "circular"})
    assert res["ok"] and res["found_in_gaps"] > 0
    assert existing == before, "the tool must not mutate what it was given"
    assert res["_ui"]["existing"] == before


def test_the_model_copy_stays_inside_the_truncation_budget():
    """A tool result is stored as json.dumps(result)[:8000]. A scan of a dense
    construct must not blow that, or the findings arrive cut in half."""
    import json
    from dee.core.agent_tools import _tool_scan_construct
    seq = _filler(40000, 12)
    res = _tool_scan_construct({"sequence": seq, "features": [],
                                "topology": "circular", "min_aa": 30})
    model_copy = {k: v for k, v in res.items() if k != "_ui"}
    assert len(json.dumps(model_copy)) < 8000, len(json.dumps(model_copy))


def test_the_tool_refuses_when_nothing_is_loaded():
    from dee.core.agent_tools import _tool_scan_construct
    res = _tool_scan_construct({"sequence": "", "features": []})
    assert not res["ok"] and "loaded" in res["error"]


def test_evidence_tiers_reach_the_model(planted):
    """The model has to be able to tell an exact match from a guess."""
    from dee.core.agent_tools import _tool_scan_construct
    seq, _ = planted
    res = _tool_scan_construct({"sequence": seq, "features": [],
                                "topology": "circular"})
    tiers = {f["evidence_tier"] for f in res["findings"]}
    assert {CURATED, MOTIF, PREDICTED} & tiers
    assert "unidentified ORF" in res["note"] or "unidentified" in res["note"]
    assert res["by_evidence"]["curated"] >= 1


def test_the_scan_is_not_gated_behind_a_confirm():
    """It writes nothing, so asking permission would be asking about nothing."""
    from dee.core import agent_tools as t
    assert not t._TOOLS["scan_construct"].get("requires_confirm")


# ── the overlay ──────────────────────────────────────────────────────────
def _static(p):
    with open(p, encoding="utf-8") as fh:
        return fh.read()


def test_the_scan_paints_but_never_writes():
    """The painter holds findings unapplied, like propose_edit — it must not
    call the path that commits a construct."""
    js = _static("dee/static/cockpit.js")
    import re
    m = re.search(r"scan_construct: function \(panel\) \{(.*?)\},", js, re.S)
    assert m, "scan_construct painter missing"
    assert "paintScan" in m.group(1)
    assert "paintAgentRun" not in m.group(1), "paintAgentRun COMMITS; a scan must not"
    assert "scan_construct: \"plasmid\"" in js, "the scan has a canvas now"


def test_the_painter_refuses_findings_from_another_construct():
    """Findings are bare coordinates. Painted onto the wrong construct they
    are plausible and meaningless — the exact failure this feature exists to
    prevent."""
    js = _static("dee/static/app.js")
    import re
    m = re.search(r"window\.TDPlasmid\.paintScan = function \(panel\) \{(.*?)\n    \};",
                  js, re.S)
    assert m, "paintScan missing"
    b = m.group(1)
    assert "panel.length" in b and "plasmid.length" in b, "no length guard"
    assert "return false" in b


def test_applying_a_scan_carries_the_evidence_into_the_record():
    js = _static("dee/static/app.js")
    import re
    m = re.search(r"function _applyScan\(which\) \{(.*?)\n    \}", js, re.S)
    assert m, "_applyScan missing"
    b = m.group(1)
    assert "f.evidence" in b, "a finding's basis must survive into the feature"
    assert "evidence_tier" in b
    assert "'exact'" in b and "curated" in b and "motif" in b, \
        "exact matches and predictions must be separately acceptable"


def test_the_ring_draws_what_nothing_explains():
    """The dark regions are the part no other map shows; if the ring ever
    stops drawing them this feature is just another feature renderer."""
    js = _static("dee/static/app.js")
    assert "dark_regions" in js and "pm-scan-dark" in js
    css = _static("dee/static/app.css")
    assert ".pm-scan-dark" in css


def test_evidence_tier_is_carried_by_stroke_not_a_second_colour_ramp():
    """The axis is confidence, not category. A predicted finding reads as
    provisional without a legend."""
    css = _static("dee/static/app.css")
    import re
    m = re.search(r'\.pm-scan-find\[data-tier="predicted"\]\s*\{([^}]*)\}', css)
    assert m, "predicted tier has no style"
    assert "stroke-dasharray" in m.group(1)