SimurgArray/tests/oracle/test_p6_vcf.py

"""Phase P6 acceptance on P0 data (plan §18 P6): VCFs validate with bcftools and index correctly; genotypes agree
with gtc2vcf apart from documented categories; BAF is ALT-oriented (EDGE-017); bedGraph equals GTC values (EDGE-023).

Needs the GRCh38 FASTA (``scripts/fetch_data.sh p6``) and, for the comparison, the gtc2vcf oracle VCF
(``scripts/make_oracle_vcf.sh``).
"""

from __future__ import annotations

import gzip
import subprocess
from pathlib import Path

import numpy as np
import pytest
from tests.conftest import GSA_BPM, GSA_CSV, ORACLE_GTCS, REFS, ROOT, need

pytestmark = [pytest.mark.data, pytest.mark.slow]
FASTA = ROOT / "data" / "p6" / "GRCh38.primary.fa"
BCFTOOLS = REFS / "tools" / "bcftools-1.22" / "bcftools"
ORACLE_VCF = ROOT / "data" / "p6" / "oracle" / "gsa36_gtc2vcf.vcf.gz"


@pytest.fixture(scope="module")
def vcfs(tmp_path_factory: pytest.TempPathFactory) -> Path:
    need(GSA_BPM, GSA_CSV, FASTA, ORACLE_GTCS)
    from typer.testing import CliRunner

    from simurg.cli.main import app

    gtcs = tmp_path_factory.mktemp("gtc")
    for f in sorted(ORACLE_GTCS.glob("*.gtc"))[:3]:
        (gtcs / f.name).symlink_to(f)
    out = tmp_path_factory.mktemp("vcf")
    res = CliRunner().invoke(
        app,
        ["gtc-to-vcf", "--bpm", str(GSA_BPM), "--csv", str(GSA_CSV), "--fasta", str(FASTA), "--gtcs", str(gtcs),
         "--out", str(out), "--threads", "1", "--unsquash-duplicates", "--quiet", "--force"],
    )  # fmt: skip
    assert res.exit_code == 0, res.output
    return out


def test_bcftools_validates_and_index_works(vcfs: Path) -> None:  # OUT-038, OUT-039, EDGE-020, EDGE-021
    if not BCFTOOLS.exists():
        pytest.skip("bcftools not built (scripts/make_oracle_gtcs.sh)")
    f = next(vcfs.glob("*.snv.vcf.gz"))
    n = int(subprocess.run([str(BCFTOOLS), "view", "-H", str(f)], capture_output=True, check=True).stdout.count(b"\n"))
    nrec = subprocess.run([str(BCFTOOLS), "index", "--nrecords", str(f)], capture_output=True, text=True, check=True)
    assert int(nrec.stdout.strip()) == n > 600_000
    for region in ("1:1000000-2000000", "X:150000000-156000000", "MT", "Y"):
        idx = subprocess.run([str(BCFTOOLS), "view", "-H", "-r", region, str(f)], capture_output=True, check=True)
        chrom, _, span = region.partition(":")
        lo, _, hi = span.partition("-")
        scan = 0
        with gzip.open(f, "rt") as fh:
            for line in fh:
                if line.startswith("#"):
                    continue
                c, p = line.split("\t", 2)[:2]
                if c == chrom and (not lo or int(lo) <= int(p) <= int(hi)):
                    scan += 1
        assert idx.stdout.count(b"\n") == scan


def test_matches_gtc2vcf(vcfs: Path, tmp_path: Path) -> None:  # plan §18 P6 gate
    if not ORACLE_VCF.exists() or not BCFTOOLS.exists():
        pytest.skip("gtc2vcf oracle VCF missing (scripts/make_oracle_vcf.sh)")
    names = ",".join(p.name.removesuffix(".snv.vcf.gz") for p in sorted(vcfs.glob("*.snv.vcf.gz")))
    sub = tmp_path / "oracle.vcf.gz"
    subprocess.run([str(BCFTOOLS), "view", "-s", names, "-Oz", "-o", str(sub), str(ORACLE_VCF)], check=True)
    out = tmp_path / "cmp.json"
    r = subprocess.run(
        [str(ROOT / ".venv" / "bin" / "python"), str(ROOT / "scripts" / "validate_p6_vcf.py"), "--ours", str(vcfs),
         "--oracle", str(sub), "--fasta", str(FASTA), "--bcftools", str(BCFTOOLS), "--out", str(out)],
        capture_output=True, text=True,
    )  # fmt: skip
    assert r.returncode == 0, r.stdout + r.stderr
    import json

    d = json.loads(out.read_text())
    assert d["unexplained"] == 0 and d["categories"]["match_called"] > 1_500_000


def test_baf_is_alt_oriented(vcfs: Path) -> None:  # OUT-016, EDGE-017
    from simurg.formats.bpm import read_bpm
    from simurg.formats.gtc import read_gtc

    f = sorted(vcfs.glob("*.snv.vcf.gz"))[0]
    g = read_gtc(ORACLE_GTCS / f.name.replace(".snv.vcf.gz", ".gtc"))
    m = read_bpm(GSA_BPM)
    idx = m.name_index()
    flipped = same = 0
    with gzip.open(f, "rt") as fh:
        for line in fh:
            if line.startswith("#"):
                continue
            c = line.rstrip("\n").split("\t")
            if "," in c[2] or ";" in c[2] or c[8] != "GT:GS:BAF:LRR" or "," in c[4]:
                continue
            baf = c[9].split(":")[2]
            if baf == ".":
                continue
            i = idx[c[2]]
            gtc_baf = float(g.b_allele_freqs[i])
            gt = c[9].split(":")[0]
            call = int(g.genotypes[i])
            if gt in ("0/0", "1/1") and call in (1, 3):
                a_is_ref = (gt == "0/0") == (call == 1)
                want = gtc_baf if a_is_ref else float(np.float32(1.0) - np.float32(gtc_baf))
                assert abs(float(baf) - want) < 1e-6, (c[2], baf, gtc_baf)
                flipped += not a_is_ref
                same += a_is_ref
            if flipped > 500 and same > 500:
                break
    assert flipped > 500 and same > 500  # both orientations exercised


def test_bedgraph_equals_gtc(tmp_path: Path) -> None:  # OUT-050, OUT-052, EDGE-023
    need(GSA_BPM, ORACLE_GTCS)
    from simurg.core.log import Reporter
    from simurg.formats.bpm import read_bpm
    from simurg.formats.gtc import read_gtc
    from simurg.pipeline.export import BedgraphOptions, run_export_bedgraph

    gtcs = tmp_path / "g"
    gtcs.mkdir()
    src = sorted(ORACLE_GTCS.glob("*.gtc"))[0]
    (gtcs / src.name).symlink_to(src)
    out = tmp_path / "bg"
    assert run_export_bedgraph(BedgraphOptions(out=out, bpm=GSA_BPM, gtcs=gtcs), Reporter(command="t", quiet=True)) == 0
    m = read_bpm(GSA_BPM)
    g = read_gtc(src)
    single = {}
    seen: dict[tuple[str, int], int] = {}
    for i, (c, p) in enumerate(zip(m.chroms, m.map_infos.tolist(), strict=True)):
        key = ("X" if c == "XY" else c, p)
        seen[key] = seen.get(key, 0) + 1
        single[key] = i
    lines = gzip.decompress((out / f"{src.stem}.lrr.bedgraph.gz").read_bytes()).decode().splitlines()
    assert lines[0].startswith('track type=bedGraph name="') and "viewLimits=-2:2" in lines[0]
    checked = 0
    for ln in lines[1:]:
        c, s, e, v = ln.split("\t")
        key = (c, int(e))
        assert int(e) - int(s) == 1
        if seen[key] == 1:
            assert v == f"{float(g.logr_ratios[single[key]]):.6g}"
            checked += 1
    assert checked > 500_000


def test_info_and_extra_format_match_gtc2vcf(tmp_path: Path) -> None:  # OUT-012, plan §8.3
    """INFO ASSAY_TYPE/GENTRAIN/CLUSTER_SEP and FORMAT X/Y/NORMX/NORMY equal gtc2vcf's; THETA/R too, except where the
    normalized intensities are both 0 (SimurgArray: missing, plan §7.4; gtc2vcf: 0)."""
    from tests.conftest import GSA_EGT

    plugin = BCFTOOLS.parent / "plugins" / "gtc2vcf.so"
    need(GSA_BPM, GSA_CSV, GSA_EGT, FASTA, ORACLE_GTCS)
    if not plugin.exists():
        pytest.skip("gtc2vcf plugin not built")
    import os

    from typer.testing import CliRunner

    from simurg.cli.main import app

    gtcs = tmp_path / "g"
    gtcs.mkdir()
    src = sorted(ORACLE_GTCS.glob("*.gtc"))[0]
    (gtcs / src.name).symlink_to(src)
    out = tmp_path / "v"
    res = CliRunner().invoke(
        app,
        ["gtc-to-vcf", "--bpm", str(GSA_BPM), "--csv", str(GSA_CSV), "--fasta", str(FASTA), "--gtcs", str(gtcs),
         "--out", str(out), "--threads", "1", "--unsquash-duplicates", "--quiet", "--egt", str(GSA_EGT),
         "--info-fields", "ASSAY_TYPE,GENTRAIN,CLUSTER_SEP", "--extra-format", "X,Y,NORMX,NORMY,THETA,R"],
    )  # fmt: skip
    assert res.exit_code == 0, res.output
    env = dict(os.environ, BCFTOOLS_PLUGINS=str(plugin.parent))
    g2v = tmp_path / "g2v.vcf"
    with g2v.open("w") as fh:  # streamed: the text of ~700 000 records would cost GBs held as dicts
        subprocess.run(
            [str(BCFTOOLS), "+gtc2vcf", "--no-version", "-t", "GT,X,Y,NORMX,NORMY,THETA,R", "-b", str(GSA_BPM), "-c",
             str(GSA_CSV), "-e", str(GSA_EGT), "-f", str(FASTA), "-g", str(gtcs), "-O", "v"],
            stdout=fh, env=env, check=True,
        )  # fmt: skip
    keys = ("X", "Y", "NORMX", "NORMY", "THETA", "R")
    ref: dict[str, tuple[str, ...]] = {}  # id → (ASSAY_TYPE, GenTrain_Score, Cluster_Sep, *keys)
    with g2v.open() as fh:
        for line in fh:
            if line.startswith("#"):
                continue
            c = line.rstrip("\n").split("\t")
            info = dict(kv.split("=", 1) for kv in c[7].split(";") if "=" in kv)
            fmt = dict(zip(c[8].split(":"), c[9].split(":"), strict=True))
            ref[c[2]] = (info["ASSAY_TYPE"], info["GenTrain_Score"], info["Cluster_Sep"], *(fmt[k] for k in keys))
    g2v.unlink()

    def close(a: str, b: str) -> bool:
        na, nb = a in (".", "") or "nan" in a.lower(), b in (".", "") or "nan" in b.lower()
        if na or nb:
            return na and nb
        return abs(float(a) - float(b)) <= 1e-5 * max(1.0, abs(float(b)))

    checked = zero_norm = 0
    with gzip.open(out / f"{src.stem}.snv.vcf.gz", "rt") as fh:
        for line in fh:
            if line.startswith("#"):
                continue
            c = line.rstrip("\n").split("\t")
            if "," in c[2] or ";" in c[2] or c[2] not in ref:
                continue
            info = dict(kv.split("=", 1) for kv in c[7].split(";"))
            fmt = dict(zip(c[8].split(":"), c[9].split(":"), strict=True))
            at, gt_score, sep, *vals = ref[c[2]]
            rf = dict(zip(keys, vals, strict=True))
            assert info["ASSAY_TYPE"] == at
            assert close(info["GENTRAIN"], gt_score) and close(info["CLUSTER_SEP"], sep)
            for k in ("X", "Y", "NORMX", "NORMY"):
                assert close(fmt[k], rf[k]), (c[2], k, fmt[k], rf[k])
            if fmt["NORMX"] == "0" and fmt["NORMY"] == "0":
                zero_norm += 1
                assert fmt["THETA"] == "." and float(rf["THETA"]) == 0.0
            else:
                assert close(fmt["THETA"], rf["THETA"]) and close(fmt["R"], rf["R"]), (c[2], fmt, rf)
            checked += 1
    assert checked > 600_000 and zero_norm < 1000