SimurgArray/tests/unit/test_vcf.py

"""Phase P6 unit tests: FASTA access (IN-060 … IN-064), BGZF/tabix (OUT-038), plus-strand mapping (OUT-015,
OUT-020), genotype combination (OUT-031, EDGE-011/012), ploidy (OUT-018), BAF orientation (EDGE-017), PLINK codes."""

from __future__ import annotations

import gzip
import struct
from pathlib import Path

import numpy as np
import pytest

from simurg.core.errors import SimurgError
from simurg.formats.bgzf import BgzfWriter, TabixBuilder, bgzf_compress, reg2bin
from simurg.formats.fasta import ContigMap, Genome, build_fai, build_gzi, detect_build, read_gzi
from simurg.harmonize.combine import AA, AB, BB, NC, failure_message, genotype, gt_text, merged
from simurg.harmonize.sites import (
    ManifestView,
    ProbeSite,
    Site,
    group_sites,
    left_shift,
    map_locus,
    normalize_indel,
    revcomp,
)
from simurg.pipeline.export import bed_code
from simurg.pipeline.vcf import CLASS_AUTO, CLASS_MT, CLASS_PAR, CLASS_X, CLASS_Y, _site_class, ploidy_for

SEQ = "ACGTACGTAAAACCCCGGGGTTTTACGTCAGCAGCAGTTT" * 3  # 120 bp


@pytest.fixture
def fasta(tmp_path: Path) -> Path:
    p = tmp_path / "g.fa"
    lines = [">1 test"] + [SEQ[i : i + 50] for i in range(0, len(SEQ), 50)] + [">X"] + [SEQ[:30].lower()]
    p.write_text("\n".join(lines) + "\n")
    return p


def test_fai_and_fetch_modes(fasta: Path, tmp_path: Path) -> None:  # IN-060, IN-064, EDGE-010
    fai = build_fai(fasta, False)
    assert [(e.name, e.length, e.line_bases, e.line_width) for e in fai] == [("1", 120, 50, 51), ("X", 30, 30, 31)]
    assert fai[0].offset == len(">1 test\n")
    g1 = Genome(fasta, use_cache=False, cache_dir=tmp_path / "c1")
    g2 = Genome(fasta, use_cache=True, cache_dir=tmp_path / "c2")
    for a, b in ((0, 120), (45, 55), (49, 51), (100, 200), (-5, 3)):
        assert g1.fetch("1", a, b) == g2.fetch("1", a, b) == SEQ[max(0, a) : min(120, b)]
    assert g1.fetch("X", 0, 5) == "ACGTA"  # upper-cased
    assert (tmp_path / "c1" / "index.fai").exists()  # built into the cache, not next to the input


def test_bgzip_fasta(fasta: Path, tmp_path: Path) -> None:
    bz = tmp_path / "g.fa.gz"
    bz.write_bytes(bgzf_compress(fasta.read_bytes()))
    assert read_gzi_roundtrip(bz)
    g = Genome(bz, use_cache=False, cache_dir=tmp_path / "c")
    assert g.bgzf and g.fetch("1", 40, 70) == SEQ[40:70]


def read_gzi_roundtrip(path: Path) -> bool:
    entries = build_gzi(path)
    gzi = path.with_suffix(".gzi")
    gzi.write_bytes(struct.pack("<Q", len(entries)) + b"".join(struct.pack("<QQ", c, u) for c, u in entries))
    return read_gzi(gzi) == entries


def test_fai_mismatch(fasta: Path, tmp_path: Path) -> None:  # IN-060
    Path(str(fasta) + ".fai").write_text("1\t120\t3\t50\t51\nX\t30\t140\t30\t31\n")
    with pytest.raises(SimurgError) as e:
        Genome(fasta, cache_dir=tmp_path / "c")
    assert e.value.code == "E_FAI_MISMATCH"


def test_build_and_contig_alias(fasta: Path, tmp_path: Path) -> None:  # IN-061, IN-063
    from simurg.formats.fasta import FaiEntry

    assert detect_build([FaiEntry("chr1", 248956422, 0, 60, 61)]) == "GRCh38"
    assert detect_build([FaiEntry("1", 249250621, 0, 60, 61)]) == "GRCh37"
    strict = ContigMap(["1", "X", "MT"], "strict")
    assert strict.get("XY") == "X" and strict.get("chr1") is None
    auto = ContigMap(["chr1", "chrX", "chrM", "chr22_KI270879v1_alt"], "auto")
    assert auto.get("1") == "chr1" and auto.get("MT") == "chrM" and auto.get("XY") == "chrX"
    assert auto.get("22_KI270879v1_alt") == "chr22_KI270879v1_alt"
    alias = tmp_path / "a.tsv"
    alias.write_text("1\tNC_000001\n")
    assert ContigMap(["NC_000001"], str(alias)).get("1") == "NC_000001"


def test_bgzf_blocks_and_tabix(tmp_path: Path) -> None:  # OUT-038
    out = tmp_path / "t.gz"
    tbx = TabixBuilder()
    with out.open("wb") as fh:
        w = BgzfWriter(fh)
        w.write_record(b"#header\n")
        for i in range(5000):
            line = f"1\t{i * 10 + 1}\trs{i}\tA\tC\n".encode()
            s, e = w.write_record(line)
            tbx.add("1", i * 10, i * 10 + 1, s, e)
        w.close()
    text = gzip.decompress(out.read_bytes()).decode()
    assert text.count("\n") == 5001
    idx = gzip.decompress(tbx.encode())
    assert idx[:4] == b"TBI\x01" and struct.unpack_from("<i", idx, 4)[0] == 1
    assert reg2bin(0, 1) == 4681 and reg2bin(0, 1 << 20) == 73 and reg2bin(0, 1 << 29) == 0


# ------------------------------------------------------------------------------------------------ mapping
def fetcher(seq: str):  # type: ignore[no-untyped-def]
    return lambda c, a, b: seq[max(0, a) : max(0, b)]


def mv(**kw: object) -> ManifestView:
    base = dict(names=["p"], snps=["[A/G]"], chroms=["1"], positions=[5], ilmn_strands=["TOP"],
                source_strands=["TOP"], ref_strands=["+"], source_seqs=[""], assay_types=[0],
                intensity_only=[False])  # fmt: skip
    base.update({k: [v] for k, v in kw.items()})
    return ManifestView(**base)  # type: ignore[arg-type]


def test_snv_plus_and_minus() -> None:  # plan §8.1, OUT-015
    genome = "TTTTAGGGG"  # position 5 = A
    p, a = map_locus(0, mv(), fetcher(genome), lambda c: c)
    assert p is not None and p.ref == "A" and p.alleles == ("A", "G") and not a.swapped
    p, a = map_locus(0, mv(snps="[T/C]", ref_strands="-"), fetcher(genome), lambda c: c)
    assert p is not None and p.alleles == ("A", "G") and a.flipped
    p, a = map_locus(0, mv(snps="[C/G]"), fetcher(genome), lambda c: c)
    assert a.reason == "W_REF_NOT_QUERIED" and p is not None
    assert map_locus(0, mv(chroms="0"), fetcher(genome), lambda c: c)[1].reason == "W_SKIP_UNMAPPED"
    assert map_locus(0, mv(intensity_only=True), fetcher(genome), lambda c: c)[1].reason == "W_SKIP_INTENSITY_ONLY"
    assert map_locus(0, mv(chroms="Z"), fetcher(genome), lambda c: None)[1].reason == "W_SKIP_CONTIG"


def test_indel_deletion_and_insertion_normalised() -> None:  # OUT-020, IN-046
    #          1234567890123
    genome = "GGCATCATCATTT"
    # CAT copies at 3-5, 6-8, 9-11: deleting the last copy normalises to the G anchor at position 2
    pos, ref, alt = normalize_indel(fetcher(genome), "1", 9, "CAT", True)
    assert (pos, ref, alt) == (2, "GCAT", "G")
    pos, ref, alt = normalize_indel(fetcher(genome), "1", 11, "CAT", False)  # insertion after position 11
    assert (pos, ref, alt) == (2, "G", "GCAT")
    assert left_shift("GGCATCAT", "CAT", "TT") == ("GG", "CATCATTT")
    assert revcomp("ACGTN") == "NACGT"


def test_indel_mapping_from_source_seq() -> None:
    genome = "AAAAACCCCGATTACAGGGGGTTTTT"
    src = "AAAAACCCCG[-/ATTACA]GGGGGTTTTT"  # the reference carries ATTACA at 11-16 → deletion
    view = mv(snps="[I/D]", positions=11, ilmn_strands="PLUS", source_strands="PLUS", source_seqs=src)
    p, a = map_locus(0, view, fetcher(genome), lambda c: c)
    assert p is not None and p.kind == "indel"
    assert (p.pos, p.ref, p.alleles) == (10, "GATTACA", ("GATTACA", "G"))  # I = REF, D = ALT
    assert a.swapped is False


def test_grouping_and_order() -> None:  # OUT-030, OUT-033, OUT-034, OUT-014
    def ps(i: int, name: str, alleles: tuple[str, str], pos: int = 5) -> ProbeSite:
        return ProbeSite(i, name, "1", pos, "A", alleles, "snv", 2, "1")

    probes = [ps(0, "b", ("A", "G")), ps(1, "a", ("A", "C")), ps(2, "c", ("A", "G"), pos=3)]
    sites = group_sites(probes, {"1": 0})
    assert [s.id for s in sites] == ["c", "b,a"]
    assert sites[1].alts == ["G", "C"]  # BPM order, then allele order
    dup = [ps(0, "x", ("A", "G")), ps(1, "y", ("A", "G"))]
    assert len(group_sites(dup, {"1": 0}, unsquash=True)) == 2
    assert len(group_sites(probes[:2], {"1": 0}, unsquash=True)) == 1  # multi-allelic stays merged


# ------------------------------------------------------------------------------------------------ combination
def site_of(*probes: ProbeSite) -> Site:
    ref = probes[0].ref
    alts: list[str] = []
    for p in probes:
        alts += [x for x in p.alleles if x != ref and x not in alts]
    return Site("1", 5, ref, alts, list(probes))


def P(name: str, a: str, b: str, inf: int, ref: str = "A") -> ProbeSite:
    return ProbeSite(0, name, "1", 5, ref, (a, b), "snv", inf, "1")


def test_edge_011_scenarios() -> None:  # OUT-031
    # Infinium II [A/G] = AA with Infinium I [T/G] = TT → ambiguous between A/T and T/T
    s = site_of(P("ii", "A", "G", 2), P("i", "T", "G", 1))
    assert merged(s, (AA, AA), 2).reason == "W_MERGE_AMBIGUOUS"
    # Infinium I [T/A] homozygous T with Infinium II [A/G] no call → ambiguous (not a T/T call)
    s = site_of(P("i", "T", "A", 1), P("ii", "A", "G", 2))
    assert merged(s, (AA, NC), 2).reason == "W_MERGE_AMBIGUOUS"
    # inconsistent pair
    s = site_of(P("p1", "A", "G", 1), P("p2", "A", "G", 1))
    assert merged(s, (AA, BB), 2).reason == "W_MERGE_INCONSISTENT"
    assert merged(s, (NC, NC), 2).reason == "W_MERGE_ALL_NC"
    # consistent duplicates → a call
    o = merged(s, (AB, AB), 2)
    assert o.alleles == (0, 1) and gt_text(o, 2) == "0/1"
    # tri-allelic: A/G and A/C Infinium I probes
    s = site_of(P("g", "A", "G", 1), P("c", "A", "C", 1))
    assert merged(s, (AB, AB), 2).reason == "W_MERGE_INCONSISTENT"
    assert merged(s, (BB, AB), 2).reason == "W_MERGE_INCONSISTENT"
    assert merged(s, (AB, AA), 2).alleles == (0, 1)  # A/G
    assert "InfiniumI" in failure_message(s, (AB, BB), "W_MERGE_INCONSISTENT")


def test_inf1_nocall_option() -> None:  # OUT-037
    s = site_of(P("ii", "A", "G", 2), P("i", "T", "G", 1))
    assert merged(s, (AA, NC), 2).reason == "W_MERGE_AMBIGUOUS"
    o = merged(s, (AA, NC), 2, inf1_nc_informative=True)
    assert o.alleles == (0, 0)  # the T/G probe saw nothing → no T and no G → A/A


def test_edge_012_opposite_strand_indels() -> None:
    a = ProbeSite(0, "fwd", "1", 5, "ACT", ("ACT", "A"), "indel", 1, "1")  # [I/D]
    b = ProbeSite(1, "rev", "1", 5, "ACT", ("A", "ACT"), "indel", 1, "1")  # [D/I]
    s = Site("1", 5, "ACT", ["A"], [a, b], "indel")
    assert gt_text(genotype(s, (AB, AB), 2), 2) == "0/1"


def test_ploidy_rules() -> None:  # OUT-018, EDGE-038
    cls = np.array([CLASS_AUTO, CLASS_X, CLASS_Y, CLASS_MT, CLASS_PAR], np.int8)
    assert ploidy_for(cls, "M", "sex-aware").tolist() == [2, 1, 1, 1, 2]
    assert ploidy_for(cls, "F", "sex-aware").tolist() == [2, 2, 0, 1, 2]
    assert ploidy_for(cls, "U", "sex-aware").tolist() == [2, 2, 1, 1, 2]
    assert ploidy_for(cls, "M", "diploid").tolist() == [2, 2, 2, 1, 2]
    assert _site_class("X", 100_000, False, "GRCh38") == CLASS_PAR
    assert _site_class("X", 5_000_000, False, "GRCh38") == CLASS_X
    assert _site_class("chrM", 5, False, "GRCh38") == CLASS_MT
    s = site_of(P("p", "A", "G", 2))
    assert gt_text(genotype(s, (AB,), 1), 1) == "."  # het in haploid context
    assert gt_text(genotype(s, (BB,), 1), 1) == "1"


def test_plink_codes() -> None:  # OUT-043
    assert [bed_code(g, 1) for g in ("0/0", "0/1", "1/1", "./.", "1", "0")] == [3, 2, 0, 1, 0, 3]
    assert bed_code("0/2", 1) == 3 and bed_code("0/2", 2) == 2  # split multi-allelic record


# ------------------------------------------------------------------------------------------------ P6 options
def test_merge_snv_indel_edge_013() -> None:  # OUT-021
    snv = ProbeSite(0, "snv", "1", 5, "A", ("A", "G"), "snv", 2, "1")
    dele = ProbeSite(1, "del", "1", 5, "ACT", ("ACT", "A"), "indel", 1, "1")
    sep = group_sites([snv, dele], {"1": 0})
    assert len(sep) == 2  # default: separate records
    (s,) = group_sites([snv, dele], {"1": 0}, merge_snv_indel=True)
    assert s.ref == "ACT" and s.alts == ["GCT", "A"] and s.id == "snv;del"
    ext_snv, ext_del = s.probes
    assert ext_del.alleles == ("ACT", "A")
    assert ext_snv.alleles == ("ACT", "GCT")
    # SNV probe homozygous A (sees the anchor base of the deletion too), deletion probe het → REF/DEL
    assert gt_text(genotype(s, (AA, AB), 2), 2) == "0/2"
    # SNV het A/G, deletion hom REF → REF/SNV
    assert gt_text(genotype(s, (AB, AA), 2), 2) == "0/1"
    # incompatible REFs (not prefixes of each other) stay separate
    other = ProbeSite(2, "x", "1", 5, "GT", ("GT", "G"), "indel", 1, "1")
    assert len(group_sites([snv, other], {"1": 0}, merge_snv_indel=True)) == 2


def _aux_probe(alleles: tuple[str, str] = ("A", "C"), idx: int = 0, name: str = "mnv") -> ProbeSite:
    return ProbeSite(idx, name, "1", 5, "A", alleles, "snv", 2, "1")


def test_auxiliary_loci_rules() -> None:  # OUT-035, EDGE-015
    from simurg.harmonize.sites import AuxRecord, apply_auxiliary

    genome = "TTTTATGCC"  # positions 5-6 = AT
    f = fetcher(genome)
    plus = lambda i: "+"  # noqa: E731
    ok = AuxRecord("mnv", "1", 5, "AT", ["CG"])
    probes, errs = apply_auxiliary([_aux_probe()], [ok], f, lambda c: c, plus)
    assert not errs and probes[0].ref == "AT" and probes[0].alleles == ("AT", "CG") and probes[0].reason == "aux_locus"
    cases = {
        "(a)": AuxRecord("mnv", "1", 5, "AT", ["CG", "GG"]),
        "(b)": AuxRecord("mnv", "1", 5, "AT", ["C"]),
        "(d)": AuxRecord("mnv", "1", 5, "AT", ["TG"]),
        "(e)": AuxRecord("mnv", "1", 5, "AG", ["CG"]),
    }
    for rule, rec in cases.items():
        _, errs = apply_auxiliary([_aux_probe()], [rec], f, lambda c: c, plus)
        assert len(errs) == 1 and rule in errs[0], (rule, errs)
    dup = [_aux_probe(), _aux_probe(idx=1, name="twin")]
    _, errs = apply_auxiliary(dup, [ok], f, lambda c: c, plus)
    assert "(c)" in errs[0]
    _, errs = apply_auxiliary([_aux_probe()], [AuxRecord("nope", "1", 5, "AT", ["CG"])], f, lambda c: c, plus)
    assert "not a probe" in errs[0]
    # minus-strand design: the interrogated base is the MNV's last base on the plus strand
    probes, errs = apply_auxiliary([_aux_probe(("T", "G"))], [AuxRecord("mnv", "1", 5, "AT", ["CG"])], f,
                                   lambda c: c, lambda i: "-")  # fmt: skip
    assert not errs and probes[0].alleles == ("AT", "CG")


def test_pgen_codes() -> None:  # plan §8.4
    from simurg.pipeline.export import PGEN_FROM_BED, _pack

    # PLINK 1 (A1 = ALT): 0 hom ALT, 1 missing, 2 het, 3 hom REF → PLINK 2: 2, 3, 1, 0
    assert list(PGEN_FROM_BED) == [2, 3, 1, 0]
    assert _pack([0, 1, 2, 3, 0], 2) == bytes([0b11100100, 0b00000000])