"""QC orchestration and I/O (spec QC-001 … QC-064, CLI-024, CLI-025, OUT-070, OUT-072; plan §9).

The science lives in ``simurg.qc``; this module reads IDATs/GTCs, streams cohort matrices through disk-backed
arrays (so thousands of samples fit in a small memory budget) and writes the QC tables.
"""

from __future__ import annotations

import datetime as _dt
import math
import multiprocessing as mp
import os
import shutil
import tempfile
from collections.abc import Iterator, Mapping, Sequence
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Any, cast

import numpy as np
import numpy.typing as npt

from simurg.core.atomic import atomic_write_text, check_outputs_absent
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter
from simurg.core.provenance import RunManifest
from simurg.formats.bpm import Manifest
from simurg.formats.idat import Idat
from simurg.formats.samplesheet import SampleSheet, fold
from simurg.genotype.baf_lrr import chrom_core
from simurg.qc.contamination import ContaminationResult, baf_regress
from simurg.qc.controls import (
    ASSAY_CONTROL_COLUMN,
    COMPAT_METRIC_NAMES,
    ControlBead,
    ControlQC,
    beads_from_definitions,
    canonical_id,
    compare_definitions,
    compute_control_metrics,
    control_intensities,
    metrics_for,
)
from simurg.qc.relatedness import DEGREE_CUTS, DUPLICATE_CONCORDANCE, degree, mendel_errors, pair_stats, select_sites
from simurg.qc.status import het_zscores
from simurg.qc.thresholds import Thresholds

CONTROLS_RAW = "controls_raw.tsv"
CONTROLS_QC = "controls_qc.tsv"
COMPAT_RAW = "controls.raw_metrics.csv"
COMPAT_QC = "controls.qc_metrics.csv"
MIN_CONTAMINATION_SAMPLES = 10


def fmt(v: Any, digits: int = 8) -> str:
    if v is None or v == "":
        return ""
    if isinstance(v, float | np.floating):
        f = float(v)
        if math.isinf(f):
            return "Inf" if f > 0 else "-Inf"
        return "NaN" if math.isnan(f) else f"{f:.{digits}f}"
    return str(v).replace("\t", " ").replace("\n", " ")


def write_tsv(path: Path, columns: Sequence[str], rows: Sequence[Mapping[str, Any]], sep: str = "\t") -> None:
    lines = [sep.join(columns)] + [sep.join(fmt(r.get(c, "")) for c in columns) for r in rows]
    atomic_write_text(path, "\n".join(lines) + "\n")


def autosome_mask(chroms: Sequence[str]) -> npt.NDArray[np.bool_]:
    """Chromosomes 1–22 only (het rate, relatedness, contamination sites)."""
    auto = {str(i) for i in range(1, 23)}
    return np.array([chrom_core(c) in auto for c in chroms], dtype=np.bool_)


# ================================================================================================ controls
def load_control_beads(manifest: Manifest | None, csv_path: Path | None, rep: Reporter) -> list[ControlBead]:
    """QC-001: BPM control config wins over the CSV ``[Controls]`` section; differences give one warning."""
    bpm_beads = beads_from_definitions(manifest.controls) if manifest is not None else []
    csv_beads: list[ControlBead] = []
    if csv_path is not None:
        from simurg.formats.csv_manifest import read_csv_manifest

        csv_beads = beads_from_definitions(read_csv_manifest(csv_path, load_sequences=False).controls)
    if bpm_beads and csv_beads and (diffs := compare_definitions(bpm_beads, csv_beads)):
        rep.warn("W_CTRL_DEFS_DIFFER", "BPM and CSV control definitions differ (using the BPM): " + "; ".join(diffs))
    return bpm_beads or csv_beads


@dataclass
class ControlSample:
    sample_id: str
    meta: dict[str, str]
    green: npt.NDArray[np.float64]
    red: npt.NDArray[np.float64]
    qc: ControlQC


def _param(pars: str, key: str) -> str:
    from simurg.genotype.sample import _param as p

    return p(pars, key)


def idat_metadata(grn: Idat, red: Idat, sheet_values: Mapping[str, str] | None = None) -> dict[str, str]:
    """QC-011 metadata: IDAT values, plate and well overridden by the sheet."""
    scan = red.scan_entry() or grn.scan_entry()
    sv = {fold(k): v for k, v in (sheet_values or {}).items()}
    return {
        "imaging_date": red.scan_time_iso or grn.scan_time_iso,
        "sample_plate": sv.get(fold("Sample_Plate")) or red.plate or grn.plate,
        "sample_well": sv.get(fold("Sample_Well")) or red.well or grn.well,
        "scanner_id": _param(scan.block_pars, "sherlockID") if scan else "",
        "scanner_version": scan.code_version if scan else "",
    }


def control_sample(
    sample_id: str,
    grn: Idat,
    red: Idat,
    beads: Sequence[ControlBead],
    bg_offset: float,
    sheet_values: Mapping[str, str] | None = None,
    array_type: str = "genotyping",
) -> ControlSample:
    g, r = control_intensities(grn, red, beads)
    qc = compute_control_metrics(beads, g, r, array_type, bg_offset)
    return ControlSample(sample_id, idat_metadata(grn, red, sheet_values), g, r, qc)


def _compat_date(iso: str) -> str:
    try:
        t = _dt.datetime.fromisoformat(iso)
    except ValueError:
        return iso
    from simurg.genotype.sample import invariant_datetime

    return invariant_datetime(t)


def write_control_tables(
    out: Path,
    samples: Sequence[ControlSample],
    beads: Sequence[ControlBead],
    bg_offset: float,
    rep: Reporter,
    *,
    array_type: str = "genotyping",
    assay_control: Mapping[str, float] | None = None,
    compat: bool = False,
) -> list[Path]:
    """QC-009 and QC-011 (or their compatibility CSVs). ``assay_control`` maps sample_id → assay-control norm R."""
    metrics = metrics_for(array_type)
    missing: dict[str, list[str]] = {}
    for s in samples:
        missing.update(s.qc.missing)
        for col in s.qc.zero_divisor:
            rep.warn("W_CTRL_ZERO_DIVISOR", f"{s.sample_id}: {col} has a zero denominator", sample_id=s.sample_id)
    if missing:
        ids = sorted({c for v in missing.values() for c in v})
        rep.warn(
            "W_CTRL_MISSING",
            f"manifest lacks control(s) {', '.join(ids)}; not evaluable: {', '.join(sorted(missing))}",
        )
    raw_rows = [
        {
            "sample_id": s.sample_id,
            "category": b.category,
            "bead_address": b.address,
            "control_name": b.name,
            "color_label": b.color,
            "green_intensity": "NaN" if math.isnan(gv) else int(gv),
            "red_intensity": "NaN" if math.isnan(rv) else int(rv),
        }
        for s in samples
        for b, gv, rv in zip(beads, s.green.tolist(), s.red.tolist(), strict=True)
    ]
    qc_rows = []
    for s in samples:
        row: dict[str, Any] = {
            "sample_id": s.sample_id,
            "array_type": array_type,
            **s.meta,
            "bg_offset": float(bg_offset),
            **{m.column: s.qc.values.get(m.column, math.nan) for m in metrics},
            "na_metrics": ";".join(s.qc.na_metrics),
        }
        if assay_control is not None:
            row[ASSAY_CONTROL_COLUMN] = assay_control.get(s.sample_id, math.nan)
        qc_rows.append(row)
    meta_cols = ["imaging_date", "sample_plate", "sample_well", "scanner_id", "scanner_version"]
    qc_cols = ["sample_id", "array_type", *meta_cols, "bg_offset", *(m.column for m in metrics)]
    if assay_control is not None:
        qc_cols.append(ASSAY_CONTROL_COLUMN)
    qc_cols.append("na_metrics")
    if not compat:
        raw_cols = list(raw_rows[0]) if raw_rows else ["sample_id"]
        write_tsv(out / CONTROLS_RAW, raw_cols, raw_rows)
        write_tsv(out / CONTROLS_QC, qc_cols, qc_rows)
        return [out / CONTROLS_RAW, out / CONTROLS_QC]
    ren_raw = {
        "sample_id": "SampleId",
        "category": "Category",
        "bead_address": "Beadtype",
        "control_name": "Control",
        "color_label": "Color",
        "green_intensity": "GreenIntensity",
        "red_intensity": "RedIntensity",
    }
    ren_qc = {
        "sample_id": "SampleId",
        "imaging_date": "ImagingDate",
        "sample_plate": "SamplePlate",
        "sample_well": "SampleWell",
        "scanner_id": "ScannerId",
        "scanner_version": "ScannerVersion",
        **COMPAT_METRIC_NAMES,
    }
    raw_c = [{ren_raw[k]: v for k, v in r.items()} for r in raw_rows]
    qc_keep = [c for c in qc_cols if c in ren_qc]
    qc_c = [
        {ren_qc[c]: (_compat_date(r[c]) if c == "imaging_date" else r.get(c, "")) for c in qc_keep} for r in qc_rows
    ]
    write_tsv(out / COMPAT_RAW, list(ren_raw.values()), raw_c, sep=",")
    write_tsv(out / COMPAT_QC, [ren_qc[c] for c in qc_keep], qc_c, sep=",")
    return [out / COMPAT_RAW, out / COMPAT_QC]


@dataclass
class QcControlsOptions:
    out: Path
    bpm: Path | None = None
    csv: Path | None = None
    idats: Path | None = None
    sheet: Path | None = None
    array_type: str = "genotyping"
    bg_offset: float = 3000.0
    compat: bool = False
    force: bool = False
    on_duplicate_idat: str = "error"
    argv: list[str] = field(default_factory=list)


def run_qc_controls(opts: QcControlsOptions, rep: Reporter) -> int:
    """``simurg qc controls`` — control tables from IDATs without genotyping (QC-012, QC-014)."""
    from simurg.formats.bpm import read_bpm
    from simurg.formats.idat import check_pair, read_idat
    from simurg.pipeline.inputs import check_paths, resolve_sample_set, sheet_values

    rm = RunManifest(command="qc controls", argv=opts.argv)
    rep.run_id = rm.run_id
    if opts.array_type == "methylation" and opts.csv is None:
        raise SimurgError("E_NO_INPUT", "methylation control QC needs the methylation CSV manifest (--csv)")
    if opts.bpm is None and opts.csv is None:
        raise SimurgError("E_NO_INPUT", "give --bpm or --csv for the control definitions (QC-001)")
    if opts.bg_offset < 0:
        raise SimurgError("E_QC_THRESHOLD_RANGE", f"--bg-offset must be >= 0, got {opts.bg_offset}")
    check_paths([(opts.bpm, "BPM"), (opts.csv, "CSV manifest"), (opts.idats, "IDAT folder"), (opts.sheet, "sheet")])
    sset = resolve_sample_set(opts.idats, opts.sheet, opts.on_duplicate_idat, rep, rm)
    names = (COMPAT_RAW, COMPAT_QC) if opts.compat else (CONTROLS_RAW, CONTROLS_QC)
    targets = [opts.out / n for n in names]
    check_outputs_absent(targets, force=opts.force)  # QC-014: other files in the folder are fine
    if opts.array_type == "methylation":
        from simurg.formats.methyl_manifest import read_methyl_manifest

        assert opts.csv is not None
        mman = read_methyl_manifest(opts.csv)
        beads = [ControlBead(c.type, c.color_channel, c.name, c.address, canonical_id(c.name)) for c in mman.controls]
    else:
        manifest = read_bpm(opts.bpm) if opts.bpm else None
        beads = load_control_beads(manifest, opts.csv, rep)
    if not beads:
        raise SimurgError("E_BPM_FORMAT", "the manifest has no control definitions")
    samples: list[ControlSample] = []
    for s in sset.samples:
        if s.pair is None or not s.pair.complete:
            rep.warn("W_IDAT_UNPAIRED", f"{s.sample_id}: no complete IDAT pair; skipped", sample_id=s.sample_id)
            continue
        assert s.pair.grn is not None and s.pair.red is not None
        grn, red = read_idat(s.pair.grn), read_idat(s.pair.red)
        check_pair(grn, red)
        samples.append(
            control_sample(s.sample_id, grn, red, beads, opts.bg_offset, sheet_values(s, sset.sheet), opts.array_type)
        )
    # QC-014: sample IDs of an existing summary must match
    summary = opts.out / "sample_summary.tsv"
    if summary.exists():
        ids = {ln.split("\t", 1)[0] for ln in summary.read_text(encoding="utf-8").splitlines()[1:] if ln}
        mine = {s.sample_id for s in samples}
        if ids ^ mine:
            rep.warn("W_QC_UNMATCHED_SAMPLES", f"{len(ids ^ mine)} sample ID(s) differ from {summary}")
    opts.out.mkdir(parents=True, exist_ok=True)
    written = write_control_tables(
        opts.out, samples, beads, opts.bg_offset, rep, array_type=opts.array_type, compat=opts.compat
    )
    for p in written:
        rm.add_output(p, "simurg.controls/1")
    rm.summary = {"samples": len(samples)}
    rm.status = "completed"
    rm.write(opts.out / "run_manifest.qc_controls.json")
    rep.info(f"control QC for {len(samples)} sample(s) written to {opts.out}")
    return ExitCode.OK


# ================================================================================================ cohort store
@dataclass
class SampleInfo:
    sample_id: str
    path: str
    sample_name: str = ""
    call_rate: float = math.nan
    logr_dev: float = math.nan
    sex_est: str = "U"
    has_baf: bool = False
    error: str = ""


class CohortStore:
    """Disk-backed (samples × loci) matrices: genotypes (uint8), BAF and GenCall scores (float32, optional)."""

    def __init__(
        self,
        root: Path,
        n_samples: int,
        n_loci: int,
        *,
        baf: bool,
        scores: bool,
        lrr: bool = False,
        intensities: bool = False,
    ) -> None:
        self.root, self.shape = root, (n_samples, n_loci)
        root.mkdir(parents=True, exist_ok=True)
        self.gt = np.lib.format.open_memmap(root / "gt.npy", "w+", np.uint8, self.shape)
        self.baf = np.lib.format.open_memmap(root / "baf.npy", "w+", np.float32, self.shape) if baf else None
        self.score = np.lib.format.open_memmap(root / "score.npy", "w+", np.float32, self.shape) if scores else None
        self.lrr = np.lib.format.open_memmap(root / "lrr.npy", "w+", np.float32, self.shape) if lrr else None
        # raw and normalized intensities (needs root/lookup.npy: normalization-transform index per locus)
        self.rx = self.ry = self.nx = self.ny = None
        if intensities:
            self.rx = np.lib.format.open_memmap(root / "rx.npy", "w+", np.uint16, self.shape)
            self.ry = np.lib.format.open_memmap(root / "ry.npy", "w+", np.uint16, self.shape)
            self.nx = np.lib.format.open_memmap(root / "nx.npy", "w+", np.float32, self.shape)
            self.ny = np.lib.format.open_memmap(root / "ny.npy", "w+", np.float32, self.shape)

    def close(self) -> None:
        for name in ("gt", "baf", "score", "lrr", "rx", "ry", "nx", "ny"):
            a = getattr(self, name)
            if a is not None:
                a.flush()
        self.gt = self.baf = self.score = self.lrr = None  # type: ignore[assignment]
        self.rx = self.ry = self.nx = self.ny = None
        shutil.rmtree(self.root, ignore_errors=True)


def _load_one(task: tuple[int, str, str, str, bool, bool, bool, bool]) -> SampleInfo:
    from simurg.formats.gtc import read_gtc

    i, sid, path, root, want_baf, want_scores, want_lrr, want_int = task
    info = SampleInfo(sid, path)
    try:
        g = read_gtc(path)
        info.sample_name, info.call_rate, info.logr_dev, info.sex_est = (
            g.sample_name,
            float(g.call_rate),
            float(g.logr_dev) if g.version >= 4 else math.nan,
            g.gender_char,
        )
        gt = np.load(Path(root) / "gt.npy", mmap_mode="r+")
        if len(g.genotypes) != gt.shape[1]:
            raise SimurgError("E_GTC_UNSUITABLE", f"{path}: {len(g.genotypes)} loci, manifest has {gt.shape[1]}")
        gt[i] = g.genotypes
        gt.flush()
        if want_baf:
            baf = np.load(Path(root) / "baf.npy", mmap_mode="r+")
            info.has_baf = len(g.b_allele_freqs) == gt.shape[1]
            baf[i] = g.b_allele_freqs if info.has_baf else np.nan
            baf.flush()
        if want_lrr:
            lr = np.load(Path(root) / "lrr.npy", mmap_mode="r+")
            lr[i] = g.logr_ratios if len(g.logr_ratios) == gt.shape[1] else np.nan
            lr.flush()
        if want_scores:
            sc = np.load(Path(root) / "score.npy", mmap_mode="r+")
            sc[i] = g.genotype_scores
            sc.flush()
        if want_int:
            from simurg.signal.normalization.apply import normalize

            lookup = np.load(Path(root) / "lookup.npy")
            nx, ny = normalize(g.raw_x, g.raw_y, lookup, g.normalization_transforms)
            for name, arr in (("rx", g.raw_x), ("ry", g.raw_y), ("nx", nx), ("ny", ny)):
                mm = np.load(Path(root) / f"{name}.npy", mmap_mode="r+")
                mm[i] = arr
                mm.flush()
    except SimurgError as e:
        info.error = f"{e.code}: {e.message}"
    except Exception as e:
        info.error = f"E_INTERNAL: {type(e).__name__}: {e}"
    return info


def load_cohort(
    samples: Sequence[tuple[str, str]],
    n_loci: int,
    tmp: Path,
    *,
    baf: bool,
    scores: bool,
    workers: int,
    lrr: bool = False,
    lookup: npt.NDArray[np.int32] | None = None,
) -> tuple[CohortStore, list[SampleInfo]]:
    """Read GTCs (sample_id, path) into a :class:`CohortStore` with a process pool. With ``lookup`` (the manifest's
    normalization-transform index per locus) raw and normalized intensities are stored too."""
    store = CohortStore(tmp, len(samples), n_loci, baf=baf, scores=scores, lrr=lrr, intensities=lookup is not None)
    if lookup is not None:
        np.save(tmp / "lookup.npy", lookup)
    store.gt.flush()
    tasks = [(i, sid, p, str(tmp), baf, scores, lrr, lookup is not None) for i, (sid, p) in enumerate(samples)]
    if workers <= 1 or len(tasks) <= 1:
        infos = [_load_one(t) for t in tasks]
    else:
        with ProcessPoolExecutor(workers, mp_context=mp.get_context("spawn")) as ex:
            infos = list(ex.map(_load_one, tasks, chunksize=4))
    return store, infos


class SiteView:
    """(sites × samples) row access onto a (samples × loci) matrix, optionally restricted to some samples."""

    def __init__(self, m: npt.NDArray[Any], samples: npt.NDArray[np.intp] | None = None) -> None:
        self.m, self.samples = m, samples
        self.shape = (m.shape[1], m.shape[0] if samples is None else len(samples))

    def __getitem__(self, rows: Any) -> npt.NDArray[Any]:
        block = np.asarray(self.m[:, rows])
        return (block if self.samples is None else block[self.samples]).T


def _blocks(n: int, size: int) -> Iterator[slice]:
    for s in range(0, n, size):
        yield slice(s, min(n, s + size))


def genotype_counts(
    gt: npt.NDArray[np.uint8], samples: npt.NDArray[np.intp] | None = None, block: int = 65536
) -> npt.NDArray[np.int64]:
    """(loci × 4) counts of no-call/AA/AB/BB over (a subset of) samples, streamed in locus blocks."""
    out = np.zeros((gt.shape[1], 4), np.int64)
    for sl in _blocks(gt.shape[1], block):
        g = np.asarray(gt[:, sl])
        if samples is not None:
            g = g[samples]
        for k in range(4):
            out[sl, k] = (g == k).sum(axis=0)
    return out


def het_rates(gt: npt.NDArray[np.uint8], loci: npt.NDArray[np.intp], block: int = 65536) -> npt.NDArray[np.float64]:
    het = np.zeros(gt.shape[0], np.int64)
    called = np.zeros(gt.shape[0], np.int64)
    for s in range(0, len(loci), block):
        g = np.asarray(gt[:, loci[s : s + block]])
        het += (g == 2).sum(axis=1)
        called += (g > 0).sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(called > 0, het / called, np.nan)


# ================================================================================================ pedigree
@dataclass
class Pedigree:
    parents: dict[str, list[str]]  # child sample_id -> genotyped parents (0–2)
    replicates: dict[str, str]  # sample_id -> declared replicate sample_id

    def expected(self, a: str, b: str) -> str:
        if b in self.parents.get(a, []) or a in self.parents.get(b, []):
            return "parent-offspring"
        if self.replicates.get(a) == b or self.replicates.get(b) == a:
            return "replicate"
        pa, pb = set(self.parents.get(a, [])), set(self.parents.get(b, []))
        if pa & pb:
            return "sibling"
        return ""


def pedigree_from_sheet(sheet: SampleSheet | None, sample_ids: Sequence[str], rep: Reporter) -> Pedigree:
    """IN-018: Parent1/Parent2/Replicate values name Sample_ID, Sample_Name or sample_id (in that order)."""
    ped = Pedigree({}, {})
    if sheet is None:
        return ped
    ids = set(sample_ids)
    by_col: dict[str, dict[str, str]] = {"sample_id_col": {}, "name": {}}
    rows = [r for r in sheet.rows if r.sample_id in ids]
    cols = {fold(c): c for c in sheet.columns}
    for r in rows:
        if (c := cols.get(fold("Sample_ID"))) and r.values.get(c):
            by_col["sample_id_col"].setdefault(r.values[c], r.sample_id)
        if r.sample_name:
            by_col["name"].setdefault(r.sample_name, r.sample_id)

    def resolve(v: str) -> str | None:
        v = v.strip()
        if v in ("", "0", "NA", "na", "-"):
            return None
        for table in (by_col["sample_id_col"], by_col["name"]):
            if v in table:
                return table[v]
        if v in ids:
            return v
        rep.warn("W_QC_PEDIGREE", f"{v!r} does not name a sample in this set")
        return None

    for r in rows:
        for key in ("Parent1", "Parent2"):
            if (c := cols.get(fold(key))) and (p := resolve(r.values.get(c, ""))) and p != r.sample_id:
                ped.parents.setdefault(r.sample_id, []).append(p)
        if (c := cols.get(fold("Replicate"))) and (p := resolve(r.values.get(c, ""))) and p != r.sample_id:
            ped.replicates[r.sample_id] = p
    return ped


# ================================================================================================ cohort QC
@dataclass
class CohortQC:
    het_rate: npt.NDArray[np.float64]
    het_z: list[float]
    contamination: ContaminationResult | None = None
    pairs: list[dict[str, Any]] = field(default_factory=list)
    duplicates: dict[str, list[str]] = field(default_factory=dict)
    unexpected: set[str] = field(default_factory=set)
    mendel: list[dict[str, Any]] = field(default_factory=list)


def cohort_qc(
    store: CohortStore,
    infos: Sequence[SampleInfo],
    manifest: Manifest,
    rep: Reporter,
    *,
    pedigree: Pedigree | None = None,
    contamination: bool = True,
    relatedness: bool = True,
    all_pairs: bool = False,
    median: str = "reference",
    workers: int = 1,
) -> CohortQC:
    ids = [s.sample_id for s in infos]
    ok = np.array([not s.error for s in infos])
    auto = np.flatnonzero(autosome_mask(manifest.chroms) & ~manifest.intensity_only)
    het = het_rates(store.gt, auto)
    het[~ok] = np.nan
    res = CohortQC(het_rate=het, het_z=het_zscores(het.tolist()))
    if contamination and store.baf is not None:
        usable = ok & np.array([s.has_baf for s in infos])
        for s in infos:
            if not s.error and not s.has_baf:
                rep.warn("W_QC_NO_BAF", f"{s.sample_id}: GTC has no BAF; contamination not estimated")
        if usable.sum() < MIN_CONTAMINATION_SAMPLES:
            rep.warn(
                "W_QC_TOO_FEW",
                f"contamination needs >= {MIN_CONTAMINATION_SAMPLES} samples with BAF (have {int(usable.sum())})",
            )
        else:
            rows_idx = np.flatnonzero(usable)
            sel = None if usable.all() else rows_idx
            c = genotype_counts(store.gt, sel)[:, 1:].astype(np.float64)
            an = 2 * c.sum(axis=1)
            site_af = np.full(manifest.num_loci, np.nan)
            with np.errstate(invalid="ignore", divide="ignore"):
                af = np.where(an > 0, (c[:, 1] + 2 * c[:, 2]) / an, np.nan)
            site_af[auto] = af[auto]  # contamination sites: autosomes only
            with _pool(workers) as ex:
                cr = baf_regress(
                    SiteView(store.gt, sel),  # type: ignore[arg-type]
                    SiteView(store.baf, sel),  # type: ignore[arg-type]
                    site_af,
                    median=median,
                    map_fn=ex.map if ex is not None else map,
                )
            est = np.full(len(ids), np.nan)
            nh = np.zeros(len(ids), np.int64)
            est[rows_idx], nh[rows_idx] = cr.estimate, cr.n_hom
            res.contamination = ContaminationResult(est, nh, cr.n_sites)
    if relatedness and ok.sum() >= 2:
        usable_loci = np.zeros(manifest.num_loci, np.bool_)
        usable_loci[auto] = True
        g_ok = np.flatnonzero(ok)
        sites = select_sites(genotype_counts(store.gt, None if ok.all() else g_ok), usable_loci)
        G = SiteView(store.gt, None if ok.all() else g_ok)[sites]
        ps = pair_stats(G)
        ped = pedigree or Pedigree({}, {})
        sub = [ids[i] for i in g_ok]
        for a in range(len(sub)):
            for b in range(a + 1, len(sub)):
                phi, conc = float(ps.kinship[a, b]), float(ps.concordance[a, b])
                exp = ped.expected(sub[a], sub[b])
                dup = conc >= DUPLICATE_CONCORDANCE
                if dup:
                    res.duplicates.setdefault(sub[a], []).append(sub[b])
                    res.duplicates.setdefault(sub[b], []).append(sub[a])
                rel = math.isfinite(phi) and phi > DEGREE_CUTS[2][0]  # second degree or closer
                if rel and not exp:
                    res.unexpected.update((sub[a], sub[b]))
                if all_pairs or dup or exp or (math.isfinite(phi) and phi > DEGREE_CUTS[3][0]):
                    res.pairs.append(
                        {
                            "sample_a": sub[a],
                            "sample_b": sub[b],
                            "n_sites": int(ps.n_sites[a, b]),
                            "kinship": phi,
                            "concordance": conc,
                            "degree": degree(phi) if math.isfinite(phi) else "",
                            "expected": exp,
                        }
                    )
        for child, parents in (ped.parents or {}).items():
            if child not in ids:
                continue
            ci = ids.index(child)
            prow = [np.asarray(store.gt[ids.index(p)])[auto] for p in parents[:2]]
            e, n = mendel_errors(
                np.asarray(store.gt[ci])[auto], prow[0] if prow else None, prow[1] if len(prow) > 1 else None
            )
            res.mendel.append(
                {
                    "child": child,
                    "parent1": parents[0] if parents else "",
                    "parent2": parents[1] if len(parents) > 1 else "",
                    "n_errors": e,
                    "n_informative": n,
                    "error_rate": e / n if n else math.nan,
                }
            )
    return res


class _NullPool:
    def __enter__(self) -> None:
        return None

    def __exit__(self, *a: object) -> None:
        return None


def _pool(workers: int) -> Any:
    if workers <= 1:
        return _NullPool()
    return ProcessPoolExecutor(workers, mp_context=mp.get_context("spawn"))


def default_tmp(out: Path) -> Path:
    base = Path(os.environ.get("SIMURG_TMP", out))
    base.mkdir(parents=True, exist_ok=True)
    return Path(tempfile.mkdtemp(prefix=".simurg_cohort_", dir=base))


def contamination_rows(ids: Sequence[str], cr: ContaminationResult | None) -> dict[str, dict[str, Any]]:
    if cr is None:
        return {}
    return {sid: {"contamination": float(cr.estimate[i]), "n_hom_sites": int(cr.n_hom[i])} for i, sid in enumerate(ids)}


def status_from_thresholds(th: Thresholds) -> list[dict[str, Any]]:
    return th.as_records()


# ================================================================================================ locus summary
def write_locus_summary(
    path: Path, store: CohortStore, manifest: Manifest, egt: Any, samples: npt.NDArray[np.intp] | None = None
) -> None:
    """OUT-072, streamed in locus blocks. ``egt`` is a parsed cluster file (records joined by locus name)."""
    from simurg.core.atomic import atomic_open
    from simurg.qc.locus_summary import COLUMNS, summarize_block

    assert store.score is not None
    idx = egt.name_index()
    rec = np.array([idx.get(nm, -1) for nm in manifest.names], dtype=np.int64)
    has = rec >= 0
    safe = np.where(has, rec, 0)

    def take(a: npt.NDArray[Any]) -> npt.NDArray[np.float64]:
        v = np.asarray(a)[safe].astype(np.float64)
        v[~has] = np.nan
        return v

    cl = {
        "locus_id": np.where(has, np.asarray(egt.addresses)[safe], 0),
        "gentrain_score": take(egt.gentrain_score),
        "cluster_sep": take(egt.cluster_sep),
    }
    for j, gname in enumerate(("aa", "ab", "bb")):
        cl[f"{gname}_t_mean"] = take(egt.theta_mean[:, j])
        cl[f"{gname}_t_sd"] = take(egt.theta_dev[:, j])
        cl[f"{gname}_r_mean"] = take(egt.r_mean[:, j])
        cl[f"{gname}_r_sd"] = take(egt.r_dev[:, j])
    strands = [s if s in ("+", "-") else "U" for s in (manifest.ref_strands or [""] * manifest.num_loci)]

    def f(v: Any) -> str:
        if isinstance(v, float | np.floating):
            return "NaN" if math.isnan(v) else f"{float(v):.8g}"
        return str(v)

    with atomic_open(path, "w", encoding="utf-8") as raw:
        fh = cast(IO[str], raw)
        fh.write("\t".join(COLUMNS) + "\n")
        for sl in _blocks(manifest.num_loci, 65536):
            g = np.asarray(store.gt[:, sl])
            s = np.asarray(store.score[:, sl])
            if samples is not None:
                g, s = g[samples], s[samples]
            stats = summarize_block(g, s)
            for k in range(sl.stop - sl.start):
                i = sl.start + k
                row = {"locus_name": manifest.names[i], "plus_minus_strand": strands[i]}
                vals = [row[c] if c in row else (cl[c][i] if c in cl else stats[c][k]) for c in COLUMNS]
                fh.write("\t".join(f(v) for v in vals) + "\n")


# ================================================================================================ GTC-based commands
SAMPLE_QC_COLUMNS = [
    "sample_id",
    "sample_name",
    "status",
    "call_rate",
    "autosomal_call_rate",
    "logr_dev",
    "sex_est",
    "sex_reported",
    "sex_check",
    "het_rate",
    "het_z",
    "contamination",
    "n_hom_sites",
    "functional_status",
    "flags",
]
RELATEDNESS_COLUMNS = ["sample_a", "sample_b", "n_sites", "kinship", "concordance", "degree", "expected"]
MENDEL_COLUMNS = ["child", "parent1", "parent2", "n_errors", "n_informative", "error_rate"]


@dataclass
class QcCohortOptions:
    """Options shared by ``qc samples | contamination | relatedness | loci`` (CLI-025)."""

    command: str
    out: Path
    bpm: Path
    gtcs: Path | None = None
    sheet: Path | None = None
    egt: Path | None = None
    threads: int = 0
    force: bool = False
    all_pairs: bool = False
    pedigree_from_sheet: bool = True
    median: str = "reference"
    thresholds: Thresholds | None = None
    argv: list[str] = field(default_factory=list)


def _gtc_samples(opts: QcCohortOptions, rep: Reporter) -> tuple[list[tuple[str, str]], SampleSheet | None]:
    """(sample_id, path) in sheet order (sheet given) or discovery order (IN-030, IN-016)."""
    from simurg.formats.discovery import discover_gtcs
    from simurg.formats.samplesheet import parse_samplesheet

    sheet = parse_samplesheet(opts.sheet) if opts.sheet else None
    root = opts.gtcs or (sheet.root_folder if sheet else None)
    if root is None:
        raise SimurgError("E_NO_INPUT", "give --gtcs or a sample sheet with a RootFolder")
    found = discover_gtcs(root)
    if sheet is None:
        return [(k, str(v)) for k, v in found.items()], None
    out = []
    for r in sheet.rows:
        if r.sample_id in found:
            out.append((r.sample_id, str(found[r.sample_id])))
        else:
            rep.warn("W_QC_UNMATCHED_SAMPLES", f"{r.sample_id}: no GTC under {root}", sample_id=r.sample_id)
    if not out:
        raise SimurgError("E_NO_INPUT", f"none of the sheet's samples has a GTC under {root}")
    return out, sheet


def run_qc_cohort(opts: QcCohortOptions, rep: Reporter) -> int:
    """``simurg qc samples | contamination | relatedness | loci`` (CLI-025)."""
    from simurg.formats.bpm import read_bpm
    from simurg.pipeline.inputs import check_paths
    from simurg.qc.status import functional_status, sample_flags
    from simurg.qc.thresholds import resolve_thresholds

    rm = RunManifest(command=f"qc {opts.command}", argv=opts.argv)
    rep.run_id = rm.run_id
    check_paths([(opts.bpm, "BPM"), (opts.gtcs, "GTC folder"), (opts.sheet, "sample sheet"), (opts.egt, "EGT")])
    if opts.command == "loci" and opts.egt is None:
        raise SimurgError("E_NO_INPUT", "qc loci needs --egt (cluster statistics, OUT-072)")
    names = {
        "samples": ["sample_qc.tsv", "contamination.tsv", "relatedness.tsv"],
        "contamination": ["contamination.tsv"],
        "relatedness": ["relatedness.tsv"],
        "loci": ["locus_summary.tsv"],
    }[opts.command]
    outputs = [opts.out / n for n in names]
    check_outputs_absent(outputs, force=opts.force)
    samples, sheet = _gtc_samples(opts, rep)
    rep.info(f"reading manifest {opts.bpm}")
    manifest = read_bpm(opts.bpm)
    workers = max(1, opts.threads or (os.cpu_count() or 1))
    opts.out.mkdir(parents=True, exist_ok=True)
    tmp = default_tmp(opts.out)
    want_baf = opts.command in ("samples", "contamination")
    rep.info(f"loading {len(samples)} GTC(s) with {workers} worker(s)")
    store, infos = load_cohort(
        samples, manifest.num_loci, tmp, baf=want_baf, scores=opts.command == "loci", workers=workers
    )
    try:
        for s in infos:
            if s.error:
                rep.error(s.error.split(":", 1)[0], f"{s.sample_id}: {s.error}", sample_id=s.sample_id)
        ids = [s.sample_id for s in infos]
        if opts.command == "loci":
            egt_path = opts.egt
            assert egt_path is not None
            from simurg.formats.egt import read_egt

            ok = np.flatnonzero([not s.error for s in infos])
            write_locus_summary(opts.out / "locus_summary.tsv", store, manifest, read_egt(egt_path), ok)
        else:
            ped = pedigree_from_sheet(sheet, ids, rep) if opts.pedigree_from_sheet else None
            res = cohort_qc(
                store,
                infos,
                manifest,
                rep,
                pedigree=ped,
                contamination=opts.command in ("samples", "contamination"),
                relatedness=opts.command in ("samples", "relatedness"),
                all_pairs=opts.all_pairs,
                median=opts.median,
                workers=workers,
            )
            crow = contamination_rows(ids, res.contamination)
            if opts.command in ("samples", "contamination"):
                write_tsv(
                    opts.out / "contamination.tsv",
                    ["sample_id", "contamination", "n_hom_sites"],
                    [{"sample_id": sid, **crow.get(sid, {})} for sid in ids],
                )
            if opts.command in ("samples", "relatedness"):
                write_tsv(opts.out / "relatedness.tsv", RELATEDNESS_COLUMNS, res.pairs)
                if res.mendel:
                    write_tsv(opts.out / "mendel.tsv", MENDEL_COLUMNS, res.mendel)
                    outputs.append(opts.out / "mendel.tsv")
            if opts.command == "samples":
                th = opts.thresholds or resolve_thresholds()
                auto = autosome_mask(manifest.chroms) & ~manifest.intensity_only
                auto_idx = np.flatnonzero(auto)
                sex_rep = {r.sample_id: r.reported_sex for r in sheet.rows} if sheet else {}
                rows = []
                for i, s in enumerate(infos):
                    g = np.asarray(store.gt[i])[auto_idx]
                    row: dict[str, Any] = {
                        "sample_id": s.sample_id,
                        "sample_name": s.sample_name,
                        "status": "FAILED" if s.error else "OK",
                        "call_rate": s.call_rate,
                        "autosomal_call_rate": float((g > 0).mean()) if len(g) and not s.error else math.nan,
                        "logr_dev": s.logr_dev,
                        "sex_est": s.sex_est,
                        "sex_reported": sex_rep.get(s.sample_id, "U"),
                        "het_rate": float(res.het_rate[i]),
                        "het_z": res.het_z[i],
                        **crow.get(s.sample_id, {}),
                    }
                    known = {"M", "F"}
                    row["sex_check"] = (
                        ("match" if row["sex_est"] == row["sex_reported"] else "mismatch")
                        if row["sex_est"] in known and row["sex_reported"] in known
                        else "unknown"
                    )
                    row["functional_status"] = functional_status(row, th)
                    row["flags"] = ";".join(
                        sample_flags(
                            row,
                            th,
                            het_z=res.het_z[i],
                            duplicates=res.duplicates.get(s.sample_id, []),
                            unexpected_relatedness=s.sample_id in res.unexpected,
                        )
                    )
                    rows.append(row)
                write_tsv(opts.out / "sample_qc.tsv", SAMPLE_QC_COLUMNS, rows)
    finally:
        store.close()
    for p in outputs:
        if p.exists():
            rm.add_output(p, f"simurg.{p.stem}/1")
    rm.summary = {"samples": len(infos), "failed": sum(bool(s.error) for s in infos)}
    rm.status = "completed"
    rm.write(opts.out / f"run_manifest.qc_{opts.command}.json")
    failed = sum(bool(s.error) for s in infos)
    rep.info(f"done: qc {opts.command} on {len(infos) - failed}/{len(infos)} sample(s); outputs in {opts.out}")
    if failed == 0:
        return ExitCode.OK
    return ExitCode.PARTIAL if failed < len(infos) else ExitCode.ALL_FAILED


# ================================================================================================ suggest-thresholds
def _read_table(path: Path) -> list[dict[str, str]]:
    import csv

    text = path.read_text(encoding="utf-8-sig")
    delim = "\t" if "\t" in text.split("\n", 1)[0] else ","
    return list(csv.DictReader(text.splitlines(), delimiter=delim))


def dataset_values(folder: Path) -> dict[str, list[float]]:
    """Per-key metric values of one QC dataset folder (QC-040 file contract)."""
    from simurg.qc.thresholds import KEY_TO_COLUMN

    def num(v: str | None) -> float:
        try:
            return float(v) if v not in (None, "") else math.nan
        except ValueError:
            return math.nan

    out: dict[str, list[float]] = {}
    qc = folder / CONTROLS_QC
    if not qc.exists():
        raise SimurgError(
            "E_QC_DATASET_INCOMPLETE", f"Dataset {folder.name} is missing required file(s): {CONTROLS_QC}"
        )
    rows = _read_table(qc)
    for key, col in KEY_TO_COLUMN.items():
        if rows and col in rows[0]:
            out[key] = [num(r.get(col)) for r in rows]
    summary = folder / "sample_summary.tsv"
    if summary.exists():
        srows = [r for r in _read_table(summary) if r.get("status", "OK") == "OK"]
        out["callRate"] = [num(r.get("autosomal_call_rate")) for r in srows]
        out["logRDev"] = [num(r.get("logr_dev")) for r in srows]
    return out


def run_suggest_thresholds(
    data: Sequence[Path], out_yaml: Path, th: Thresholds, rep: Reporter, *, k: float = 4.0, min_samples: int = 30
) -> int:
    from simurg.qc.suggest import render_yaml, suggest

    if not data:
        raise SimurgError("E_QC_NO_DATASETS", "No valid dataset folders detected.")
    merged: dict[str, list[float]] = {}
    for d in data:
        if not d.exists():
            raise SimurgError("E_QC_NO_DATASETS", f"Input data folder does not exist: {d}", path=str(d))
        for key, vals in dataset_values(d).items():
            merged.setdefault(key, []).extend(vals)
    sugs = suggest(merged, th, k=k, min_samples=min_samples)
    for s in sugs:
        if s.note:
            rep.warn("W_QC_TOO_FEW", f"{s.key}: {s.note}; left unset")
    out_yaml.parent.mkdir(parents=True, exist_ok=True)
    atomic_write_text(out_yaml, render_yaml(sugs, k))
    write_tsv(
        out_yaml.with_suffix(".tsv"),
        ["key", "n", "centre", "spread", "value", "pct_flagged_current", "pct_flagged_suggested", "note"],
        [s.__dict__ for s in sugs],
    )
    rep.info(f"suggested thresholds written to {out_yaml}")
    return ExitCode.OK