SimurgArray/tests/unit/test_raw_xy_ingest.py

"""Phase P2: raw X/Y assembly rules (plan §7.1, IN-027) and ``simurg ingest`` behaviour (CLI-010/012/014)."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import pytest
from typer.testing import CliRunner

from simurg.cli.main import app
from simurg.core.errors import SimurgError
from simurg.formats.bpm import read_bpm
from simurg.formats.idat import read_idat
from simurg.signal.raw_xy import address_coverage, assemble_raw_xy, check_coverage, control_addresses
from simurg.synth.files import SynthIdat, SynthLocus, encode_bpm, encode_idat

runner = CliRunner()
# addresses 1..8 exist; per-address means differ by channel so X/Y sources are identifiable
ADDR = np.arange(1, 9, dtype=np.int32)
RED = np.array([101, 102, 103, 104, 105, 106, 107, 108], np.uint16)
GRN = np.array([201, 202, 203, 204, 205, 206, 207, 208], np.uint16)
BEADS = np.array([10, 10, 10, 10, 1, 10, 10, 10], np.uint8)  # address 5 has a single bead


def loci() -> list[SynthLocus]:
    return [
        SynthLocus("inf2", address_a=1, assay_type=0),
        SynthLocus("inf1_at", address_a=2, address_b=3, assay_type=1),
        SynthLocus("inf1_gc", address_a=4, address_b=6, assay_type=2),
        SynthLocus("low_beads", address_a=5, assay_type=0),
        SynthLocus("absent", address_a=99, assay_type=0),
        SynthLocus("inf1_absent_b", address_a=7, address_b=98, assay_type=1),
    ]


def write_sample(
    folder: Path, barcode: str, pos: str, *, addr=ADDR, red=RED, grn=GRN, beads=BEADS, sample_id: str = ""
) -> None:  # type: ignore[no-untyped-def]
    folder.mkdir(parents=True, exist_ok=True)
    for ch, mean in (("Red", red), ("Grn", grn)):
        s = SynthIdat(address_ids=addr, mean=mean, n_beads=beads, barcode=barcode, position=pos, sample_id=sample_id)
        (folder / f"{barcode}_{pos}_{ch}.idat").write_bytes(encode_idat(s))


def bpm(tmp_path: Path, control: str = "0000000008:0000000001,Staining,Red,DNP (High)\n") -> Path:
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(loci(), control_config=control))
    return p


def test_assembly_rules(tmp_path) -> None:  # type: ignore[no-untyped-def]
    write_sample(tmp_path / "i", "200000000001", "R01C01")
    m = read_bpm(bpm(tmp_path))
    g = read_idat(tmp_path / "i" / "200000000001_R01C01_Grn.idat")
    r = read_idat(tmp_path / "i" / "200000000001_R01C01_Red.idat")
    raw = assemble_raw_xy(g, r, m)
    # inf2: X=Red[1], Y=Grn[1] | A/T: X=Red[2], Y=Red[3] | G/C: X=Grn[4], Y=Grn[6]
    assert raw.raw_x.tolist() == [101, 102, 204, 0, 0, 0]
    assert raw.raw_y.tolist() == [201, 103, 206, 0, 0, 0]
    assert raw.missing.tolist() == [False, False, False, True, True, True]
    # controls: flattened addresses in config order, X=Red, Y=Grn, no bead filter
    assert control_addresses(m).tolist() == [8, 1]
    assert raw.controls_x.tolist() == [108, 101] and raw.controls_y.tolist() == [208, 201]


def test_address_order_independent(tmp_path) -> None:  # type: ignore[no-untyped-def]
    perm = np.array([7, 2, 5, 0, 3, 6, 1, 4])
    write_sample(tmp_path / "a", "200000000001", "R01C01")
    write_sample(
        tmp_path / "b", "200000000001", "R02C01", addr=ADDR[perm], red=RED[perm], grn=GRN[perm], beads=BEADS[perm]
    )
    m = read_bpm(bpm(tmp_path))
    out = []
    for d, pos in (("a", "R01C01"), ("b", "R02C01")):
        g = read_idat(tmp_path / d / f"200000000001_{pos}_Grn.idat")
        r = read_idat(tmp_path / d / f"200000000001_{pos}_Red.idat")
        out.append(assemble_raw_xy(g, r, m))
    assert np.array_equal(out[0].raw_x, out[1].raw_x) and np.array_equal(out[0].raw_y, out[1].raw_y)


def test_coverage_gate(tmp_path) -> None:  # type: ignore[no-untyped-def]  # IN-027
    write_sample(tmp_path / "i", "200000000001", "R01C01")
    g = read_idat(tmp_path / "i" / "200000000001_R01C01_Grn.idat")
    m = read_bpm(bpm(tmp_path))
    assert address_coverage(g, m) == pytest.approx(7 / 9)  # 99 and 98 absent
    with pytest.raises(SimurgError) as e:
        check_coverage(g, m)
    assert e.value.code == "E_IDAT_MANIFEST_MISMATCH"
    big = [SynthLocus(f"rs{i}", address_a=i + 1) for i in range(200)] + [SynthLocus("x", address_a=10_000)]
    (tmp_path / "big.bpm").write_bytes(encode_bpm(big))
    addr = np.arange(1, 201, dtype=np.int32)
    write_sample(
        tmp_path / "j",
        "200000000001",
        "R02C01",
        addr=addr,
        red=np.ones(200, np.uint16),
        grn=np.ones(200, np.uint16),
        beads=np.full(200, 5, np.uint8),
    )
    gj = read_idat(tmp_path / "j" / "200000000001_R02C01_Grn.idat")
    cov, notes = check_coverage(gj, read_bpm(tmp_path / "big.bpm"))
    assert cov == pytest.approx(200 / 201) and notes[0][0] == "W_IDAT_LOW_ADDRESS_COVERAGE"


def _full_bpm(tmp_path: Path) -> Path:
    p = tmp_path / "full.bpm"
    p.write_bytes(
        encode_bpm(
            [SynthLocus(f"rs{i}", address_a=int(a)) for i, a in enumerate(ADDR)],
            control_config="0000000001,Staining,Red,DNP (High)\n",
        )
    )
    return p


def run(*args: str):  # type: ignore[no-untyped-def]
    return runner.invoke(app, ["ingest", "--quiet", *args])


def test_ingest_end_to_end(tmp_path) -> None:  # type: ignore[no-untyped-def]
    idats = tmp_path / "idats"
    write_sample(idats / "c1", "200000000001", "R01C01", sample_id="NA12878")
    write_sample(idats / "c1", "200000000001", "R02C01")
    (idats / "c1" / "200000000001_R03C01_Grn.idat").write_bytes(b"IDAT garbage")  # corrupt, unpaired
    (idats / "c1" / "200000000001_R03C01_Red.idat").write_bytes(b"IDAT garbage")
    out = tmp_path / "out"
    r = run("--bpm", str(_full_bpm(tmp_path)), "--idats", str(idats), "--out", str(out), "--threads", "2")
    assert r.exit_code == 3, r.output  # partial failure
    rows = [ln.split("\t") for ln in (out / "ingest_summary.tsv").read_text().splitlines()]
    status = {row[0]: (row[1], row[2]) for row in rows[1:]}
    assert status["200000000001_R01C01"] == ("NA12878", "OK")  # IDAT name used (IN-010)
    assert status["200000000001_R02C01"] == ("200000000001_R02C01", "OK")
    assert status["200000000001_R03C01"][1] == "FAILED"
    t = pq.read_table(out / "samples" / "NA12878" / "NA12878.intensity.parquet")
    assert t.column("raw_x").to_pylist() == [*RED.tolist()[:4], 0, *RED.tolist()[5:]]
    md = {k.decode(): v.decode() for k, v in t.schema.metadata.items()}
    assert md["simurg.schema"] == "simurg.sample_intensity/1" and md["n_missing"] == "1"
    assert (out / "run_manifest.json").exists() and (out / "loci.parquet").exists()

    # rerun without --force -> refused before any work
    r = run("--bpm", str(_full_bpm(tmp_path)), "--idats", str(idats), "--out", str(out))
    assert r.exit_code == 2 and "E_OUTPUT_EXISTS" in r.output


def test_ingest_sheet_missing_and_all_failed(tmp_path) -> None:  # type: ignore[no-untyped-def]
    idats = tmp_path / "idats"
    write_sample(idats, "200000000001", "R01C01")
    sheet = tmp_path / "s.csv"
    sheet.write_text(
        "[Data]\nSentrixBarcode_A,SentrixPosition_A,Sample_Name\n200000000001,R01C01,alpha\n200000000001,R09C01,ghost\n"
    )
    out = tmp_path / "o"
    r = run("--bpm", str(_full_bpm(tmp_path)), "--idats", str(idats), "--sheet", str(sheet), "--out", str(out))
    assert r.exit_code == 3, r.output
    txt = (out / "ingest_summary.tsv").read_text()
    assert "alpha\tOK" in txt and "ghost\tMISSING_INPUT" in txt
    assert (out / "samples" / "alpha" / "alpha.intensity.parquet").exists()

    sheet.write_text("[Data]\nSentrixBarcode_A,SentrixPosition_A\n200000000001,R09C01\n")
    r = run(
        "--bpm", str(_full_bpm(tmp_path)), "--idats", str(idats), "--sheet", str(sheet), "--out", str(tmp_path / "o2")
    )
    assert r.exit_code == 4


def test_ingest_resume_force_flat_and_determinism(tmp_path) -> None:  # type: ignore[no-untyped-def]
    idats = tmp_path / "idats"
    for i in range(1, 5):
        write_sample(idats, "200000000001", f"R0{i}C01")
    b = str(_full_bpm(tmp_path))
    a_out, b_out = tmp_path / "a", tmp_path / "b"
    assert run("--bpm", b, "--idats", str(idats), "--out", str(a_out), "--threads", "1").exit_code == 0
    assert (
        run("--bpm", b, "--idats", str(idats), "--out", str(b_out), "--threads", "3", "--layout", "flat").exit_code == 0
    )
    for f in sorted((a_out / "samples").rglob("*.parquet")):
        assert f.read_bytes() == (b_out / f.name).read_bytes(), f.name  # threads/layout never change bytes
    victim = next((a_out / "samples").rglob("*.intensity.parquet"))
    victim.unlink()
    assert run("--bpm", b, "--idats", str(idats), "--out", str(a_out), "--resume").exit_code == 0
    import json

    assert json.loads((a_out / "run_manifest.json").read_text())["summary"]["resumed"] == 3
    assert victim.exists()
    assert run("--bpm", b, "--idats", str(idats), "--out", str(a_out), "--force").exit_code == 0


def test_ingest_usage_errors(tmp_path) -> None:  # type: ignore[no-untyped-def]
    r = run("--bpm", str(tmp_path / "nope.bpm"), "--idats", str(tmp_path), "--out", str(tmp_path / "o"))
    assert r.exit_code == 2 and "E_FILE_NOT_FOUND" in r.output
    r = run("--bpm", str(_full_bpm(tmp_path)), "--out", str(tmp_path / "o"))
    assert r.exit_code == 2 and "E_NO_INPUT" in r.output