SimurgArray/tests/unit/test_bpm_egt.py

"""BPM and EGT readers on synthetic files (spec IN-040, IN-044; plan §6.2 normalization-bin rule)."""

from __future__ import annotations

import numpy as np
import pytest

from simurg.core.errors import FormatError
from simurg.formats.bpm import infer_assay_type, parse_control_config, read_bpm
from simurg.formats.egt import read_egt
from simurg.synth.files import SynthLocus, encode_bpm, encode_egt


def _loci() -> list[SynthLocus]:
    return [
        SynthLocus("rs1", "1", 100, "[A/G]", 11, 0, 0, norm_id=1),
        SynthLocus("rs2", "2", 200, "[A/T]", 12, 13, 1, norm_id=1, ilmn_strand="BOT", ref_strand="-"),
        SynthLocus("rs3", "X", 300, "[C/G]", 14, 15, 2, norm_id=99),
        SynthLocus("rs4", "Y", 400, "[T/C]", 16, 0, 0, norm_id=2, intensity_only=1),
        SynthLocus("rs5", "MT", 500, "[A/C]", 17, 0, 0, norm_id=99),
    ]


def test_bpm_roundtrip(tmp_path) -> None:  # type: ignore[no-untyped-def]
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(_loci(), control_config="0000000001,Staining,Red,DNP (High)\n"))
    m = read_bpm(p)
    assert m.names == ["rs1", "rs2", "rs3", "rs4", "rs5"]
    assert m.chroms == ["1", "2", "X", "Y", "MT"]
    assert m.map_infos.tolist() == [100, 200, 300, 400, 500]
    assert m.address_a.tolist() == [11, 12, 14, 16, 17]
    assert m.address_b.tolist() == [0, 13, 15, 0, 0]
    assert m.assay_types.tolist() == [0, 1, 2, 0, 0]
    assert m.intensity_only.tolist() == [False, False, False, True, False]
    assert m.ref_strands == ["+", "-", "+", "+", "+"]
    assert m.ilmn_strands[1] == "BOT"
    assert m.controls == [{"address": "0000000001", "category": "Staining", "color": "Red", "name": "DNP (High)"}]
    assert m.header_lines == ["Illumina, Inc."]


def test_norm_bins_mod256_and_rank(tmp_path) -> None:  # type: ignore[no-untyped-def]
    # rs3: 99 + 100*2 = 299 -> 43 (byte wrap) ; rs2: 1 + 100 = 101
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(_loci()))
    m = read_bpm(p)
    assert m.norm_bin_ids().tolist() == [1, 101, 43, 2, 99]
    # sorted distinct: [1, 2, 43, 99, 101] -> ranks
    assert m.normalization_lookups().tolist() == [0, 4, 2, 1, 3]


def test_entries_mapped_by_name(tmp_path) -> None:  # type: ignore[no-untyped-def]
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(_loci(), entry_order=[4, 3, 2, 1, 0]))
    m = read_bpm(p)
    assert m.map_infos.tolist() == [100, 200, 300, 400, 500]


@pytest.mark.parametrize("entry_version", [4, 6, 7, 8])
def test_entry_versions(tmp_path, entry_version) -> None:  # type: ignore[no-untyped-def]
    loci = [
        SynthLocus("rs1", source_seq="ACGT[A/G]TTGA", probe_a="ACGTA"),
        SynthLocus("rs2", address_b=5, assay_type=1, source_seq="AAAT[A/T]AAAC", probe_a="ACGTA", probe_b="ACGTT"),
    ]
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(loci, entry_version=entry_version))
    m = read_bpm(p)
    assert m.assay_types.tolist() == [0, 1]
    assert (m.base_fractions is not None) == (entry_version >= 7)
    assert m.has_ref_strand == (entry_version >= 8)


def test_bpm_errors(tmp_path) -> None:  # type: ignore[no-untyped-def]
    good = encode_bpm(_loci())
    cases = {
        "magic": b"XPM" + good[3:],
        "truncated": good[: len(good) // 2],
        "bad_assay": encode_bpm([SynthLocus("rs1", address_b=0, assay_type=1)]),
        "bad_norm": encode_bpm([SynthLocus("rs1", norm_id=150)]),
    }
    for label, data in cases.items():
        p = tmp_path / f"{label}.bpm"
        p.write_bytes(data)
        with pytest.raises(FormatError) as e:
            read_bpm(p)
        assert e.value.code == "E_BPM_FORMAT", label


def test_long_manifest_name(tmp_path) -> None:  # type: ignore[no-untyped-def]  # EDGE-006
    name = "X" * 200 + "_ünicode"
    p = tmp_path / "long name (1).bpm"
    p.write_bytes(encode_bpm(_loci(), manifest_name=name))
    assert read_bpm(p).manifest_name == name


def test_infer_assay_type() -> None:
    assert infer_assay_type("ACGT", "", "AC[A/G]TT") == 0
    assert infer_assay_type("ACGA", "ACGT", "AAA[A/T]TTT") == 1
    assert infer_assay_type("ACGA", "ACGT", "CCC[C/G]GGG") == 2
    assert infer_assay_type("", "", "") == 255


def test_control_config_parse() -> None:
    txt = "0027630314:0027630314,Staining,Red,DNP (High)\n\nbad line\n1,Hyb,Green,Hyb (High), extra\n"
    assert parse_control_config(txt)[1]["name"] == "Hyb (High),extra"
    assert len(parse_control_config(txt)) == 2


@pytest.mark.parametrize("dbv", [9, 8, 7, 6, 5, 4])
def test_egt_versions(tmp_path, dbv) -> None:  # type: ignore[no-untyped-def]
    names = ["rs1", "rs2"]
    tm = [[0.05, 0.5, 0.95], [0.1, 0.45, 0.9]]
    rm = [[1.0, 1.1, 0.9], [0.8, 0.9, 1.0]]
    p = tmp_path / "c.egt"
    p.write_bytes(encode_egt(names, theta_mean=tm, r_mean=rm, intensity_threshold=[0.2, 0.0], data_block_version=dbv))
    e = read_egt(p)
    assert e.names == names
    np.testing.assert_allclose(e.theta_mean, tm, rtol=1e-6)
    np.testing.assert_allclose(e.r_mean, rm, rtol=1e-6)
    assert e.n.tolist() == [[100, 100, 100]] * 2
    if dbv >= 7:
        assert e.intensity_threshold.tolist()[0] == pytest.approx(0.2)
    else:
        assert np.isnan(e.intensity_threshold).all()
    np.testing.assert_allclose(e.mean_r(), np.mean(rm, axis=1), rtol=1e-6)


def test_egt_count_mismatch(tmp_path) -> None:  # type: ignore[no-untyped-def]
    p = tmp_path / "c.egt"
    p.write_bytes(encode_egt(["rs1"], theta_mean=[[0.1, 0.5, 0.9]], r_mean=[[1, 1, 1]], corrupt_counts=True))
    with pytest.raises(FormatError) as e:
        read_egt(p)
    assert e.value.code == "E_EGT_LOAD"


def test_egt_bad_version(tmp_path) -> None:  # type: ignore[no-untyped-def]
    p = tmp_path / "c.egt"
    p.write_bytes(b"\x04\x00\x00\x00" + b"\x00" * 50)
    with pytest.raises(FormatError):
        read_egt(p)