SimurgArray/tests/oracle/test_p5_qc.py

"""Phase P5 acceptance on P0 data (plan §18 P5).

* BAF-regression contamination equals gtc2vcf ``+BAFregress`` (MIT, built by scripts/make_oracle_gtcs.sh) at its
  printed precision for all 36 GSA oracle samples, with identical homozygous-site counts.
* Locus summary equals BeadArrayFiles ``examples/locus_summary.py`` (BSD-2; scripts/make_oracle_qc.sh) exactly for
  every column except ChiTest_P100 (≤ 1e-13 relative; scipy's normal CDF vs libm, known_unknowns K-018).
* ``qc controls`` and ``genotype`` write the QC-009/QC-011 tables and the QC columns of the summary.
"""

from __future__ import annotations

import csv
import math
import os
import subprocess
from pathlib import Path

import numpy as np
import pytest
from tests.conftest import DATA, GSA_BPM, GSA_EGT, GSA_IDATS, ORACLE_GTCS, REFS, need

from simurg.formats.gtc import read_gtc
from simurg.qc.contamination import baf_regress
from simurg.qc.locus_summary import summarize_block

pytestmark = pytest.mark.data
BCFTOOLS = REFS / "tools" / "bcftools-1.22"
LOCUS_ORACLE = DATA / "oracle" / "locus_summary" / "gsa36_beadarrayfiles.csv"


@pytest.fixture(scope="module")
def cohort():  # type: ignore[no-untyped-def]
    need(ORACLE_GTCS)
    files = sorted(ORACLE_GTCS.glob("*.gtc"))
    gs = [read_gtc(f) for f in files]
    names = [f.stem for f in files]
    gt = np.stack([g.genotypes for g in gs], axis=1)  # loci × samples
    baf = np.stack([g.b_allele_freqs for g in gs], axis=1).astype(np.float32)
    scores = np.stack([g.genotype_scores for g in gs], axis=0)  # samples × loci
    return names, gt, baf, scores


@pytest.mark.oracle
def test_baf_regress_matches_gtc2vcf(cohort, tmp_path: Path) -> None:  # type: ignore[no-untyped-def]
    if not (BCFTOOLS / "plugins" / "BAFregress.so").exists():
        pytest.skip("BAFregress plugin not built (run scripts/make_oracle_gtcs.sh)")
    names, gt, baf, _ = cohort
    vcf = tmp_path / "t.vcf"
    gtmap = {0: "./.", 1: "0/0", 2: "0/1", 3: "1/1"}
    with vcf.open("w") as fh:
        fh.write(
            "##fileformat=VCFv4.2\n##contig=<ID=1>\n"
            '##INFO=<ID=ALLELE_A,Number=1,Type=Integer,Description="A">\n'
            '##INFO=<ID=ALLELE_B,Number=1,Type=Integer,Description="B">\n'
            '##FORMAT=<ID=GT,Number=1,Type=String,Description="GT">\n'
            '##FORMAT=<ID=BAF,Number=1,Type=Float,Description="BAF">\n'
            "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t" + "\t".join(names) + "\n"
        )
        for i in range(gt.shape[0]):
            cells = "\t".join(
                f"{gtmap[int(g)]}:{f'{b:.9g}' if np.isfinite(b) else '.'}" for g, b in zip(gt[i], baf[i], strict=True)
            )
            fh.write(f"1\t{i + 1}\t.\tA\tC\t.\t.\tALLELE_A=0;ALLELE_B=1\tGT:BAF\t{cells}\n")
    env = dict(os.environ, BCFTOOLS_PLUGINS=str(BCFTOOLS / "plugins"))
    out = subprocess.run(
        [str(BCFTOOLS / "bcftools"), "+BAFregress", str(vcf)], capture_output=True, text=True, env=env, check=True
    ).stdout
    ref = {ln.split("\t")[0]: ln.split("\t")[1:] for ln in out.strip().splitlines()[1:]}
    res = baf_regress(gt, baf)
    for j, nm in enumerate(names):
        assert f"{res.estimate[j]:.4f}" == ref[nm][0], nm
        assert int(res.n_hom[j]) == int(ref[nm][1]), nm


@pytest.mark.oracle
def test_locus_summary_matches_beadarrayfiles(cohort) -> None:  # type: ignore[no-untyped-def]
    if not LOCUS_ORACLE.exists():
        pytest.skip("locus-summary oracle missing (run scripts/make_oracle_qc.sh)")
    _, gt, _, scores = cohort
    st = summarize_block(gt.T, scores)
    exact = {
        "#No_Calls": "n_nocalls",
        "#Calls": "n_calls",
        "Call_Freq": "call_freq",
        "A/A_Freq": "aa_freq",
        "A/B_Freq": "ab_freq",
        "B/B_Freq": "bb_freq",
        "Minor_Freq": "minor_freq",
        "50%_GC_Score": "gc50",
        "10%_GC_Score": "gc10",
        "Het_Excess_Freq": "het_excess",
    }
    worst = 0.0
    with LOCUS_ORACLE.open() as fh:
        r = csv.reader(fh)
        next(r)
        next(r)
        hdr = next(r)
        for i, row in enumerate(r):
            d = dict(zip(hdr, row, strict=True))
            for col, key in exact.items():
                ref, ours = float(d[col]), float(st[key][i])
                assert ref == ours or (math.isnan(ref) and math.isnan(ours)), (i, col, d[col], ours)
            ref, ours = float(d["ChiTest_P100"]), float(st["chi_test_p100"][i])
            if ref != ours:
                worst = max(worst, abs(ours - ref) / abs(ref))
    assert worst < 1e-13


def test_qc_controls_and_genotype_qc_columns(tmp_path: Path) -> None:
    need(GSA_BPM, GSA_EGT, GSA_IDATS)
    from typer.testing import CliRunner

    from simurg.cli.main import app

    sub = tmp_path / "idats"
    chip = sorted(p for p in GSA_IDATS.iterdir() if p.is_dir())[0]
    (sub / chip.name).mkdir(parents=True)
    for f in sorted(chip.glob("*.idat"))[:6]:  # three samples
        (sub / chip.name / f.name).symlink_to(f)
    runner = CliRunner()
    out1 = tmp_path / "ctl"
    res = runner.invoke(
        app,
        ["qc", "controls", "--out", str(out1), "--array-type", "genotyping", "--bpm", str(GSA_BPM),
         "--idats", str(sub), "--quiet"],
    )  # fmt: skip
    assert res.exit_code == 0, res.output
    raw = (out1 / "controls_raw.tsv").read_text().splitlines()
    assert raw[0].split("\t") == [
        "sample_id", "category", "bead_address", "control_name", "color_label", "green_intensity", "red_intensity"
    ]  # fmt: skip
    assert len(raw) == 1 + 3 * 23
    qc = list(csv.DictReader((out1 / "controls_qc.tsv").open(), delimiter="\t"))
    assert len(qc) == 3 and float(qc[0]["staining_green"]) > 5 and qc[0]["na_metrics"] == ""

    out2 = tmp_path / "gt"
    res = runner.invoke(
        app,
        ["genotype", "--bpm", str(GSA_BPM), "--egt", str(GSA_EGT), "--idats", str(sub), "--out", str(out2),
         "--threads", "1", "--qc-profile", "genotyping", "--quiet"],
    )  # fmt: skip
    assert res.exit_code == 0, res.output
    rows = list(csv.DictReader((out2 / "sample_summary.tsv").open(), delimiter="\t"))
    for r in rows:
        assert r["functional_status"] == "PASS" and r["control_status"] in ("CLEAR", "FLAGGED")
        assert 0.1 < float(r["het_rate"]) < 0.3
        assert r["contamination"] == ""  # < 10 samples: not estimated (W_QC_TOO_FEW)
    gqc = list(csv.DictReader((out2 / "controls_qc.tsv").open(), delimiter="\t"))
    assert [r["staining_green"] for r in gqc] == [r["staining_green"] for r in qc]
    assert "assay_control_norm_r" in gqc[0]
    assert (out2 / "relatedness.tsv").exists()