#!/usr/bin/env python3
"""Phase P5 cohort-scale validation on the 1000 Genomes Omni2.5 set (2141 samples).

Input is a ``simurg genotype`` output folder (GTC v5 + sample_summary.tsv + controls_qc.tsv + relatedness.tsv).
The GTCs carry lab sample names, not 1000G IDs, so each sample is first identified by its heterozygosity pattern
(orientation-free) against the public 1000G Omni2.5 VCF. Then:

* sex estimate vs pedigree sex;
* parent–offspring pairs of the pedigree are detected as first degree (KING-robust, relatedness.tsv);
* Mendelian error rates of complete trios;
* distributions of contamination, control metrics, statuses and flags.

Usage: validate_p5_large.py --gt-dir DIR --bpm F --vcf F.vcf.gz --ped F.ped --out summary.json [--workers N]
"""

from __future__ import annotations

import argparse
import csv
import gzip
import json
import multiprocessing as mp
import os
import sys
import time
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any

import numpy as np

_W: dict[str, Any] = {}


def _init(idx: np.ndarray) -> None:  # type: ignore[type-arg]
    _W["idx"] = idx


def _sel(path: str) -> np.ndarray:  # type: ignore[type-arg]
    from simurg.formats.gtc import read_gtc

    return read_gtc(path).genotypes[_W["idx"]]


def _trio(t: tuple[str, str, str, str, str, str]) -> dict[str, Any]:
    from simurg.qc.relatedness import mendel_errors

    k, f, m, pk, pf, pm = t
    e, n = mendel_errors(_sel(pk), _sel(pf), _sel(pm))
    return {"child": k, "father": f, "mother": m, "errors": e, "informative": n, "rate": e / n if n else None}


def read_vcf_hets(
    path: Path, names: dict[str, int], n_snps: int
) -> tuple[list[str], list[int], np.ndarray, np.ndarray]:  # type: ignore[type-arg]
    """First ``n_snps`` common SNPs (0.2 <= AF <= 0.8) that the manifest has: het and called indicators."""
    samples: list[str] = []
    loci: list[int] = []
    het_rows, call_rows = [], []
    with gzip.open(path, "rt") as fh:
        for line in fh:
            if line.startswith("##"):
                continue
            f = line.rstrip("\n").split("\t")
            if line.startswith("#"):
                samples = f[9:]
                continue
            if f[2] not in names:
                continue
            info = dict(kv.split("=", 1) for kv in f[7].split(";") if "=" in kv)
            try:
                af = float(info.get("AF", "nan").split(",")[0])
            except ValueError:
                continue
            if not 0.2 <= af <= 0.8:
                continue
            gts = f[9:]
            het_rows.append(np.array([g[:3] in ("0/1", "1/0", "0|1", "1|0") for g in gts]))
            call_rows.append(np.array([g[:1] != "." for g in gts]))
            loci.append(names[f[2]])
            if len(loci) >= n_snps:
                break
    return samples, loci, np.array(het_rows), np.array(call_rows)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--gt-dir", required=True, type=Path)
    ap.add_argument("--bpm", required=True, type=Path)
    ap.add_argument("--vcf", required=True, type=Path)
    ap.add_argument("--ped", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--workers", type=int, default=os.cpu_count() or 1)
    ap.add_argument("--n-snps", type=int, default=20000)
    a = ap.parse_args()
    t0 = time.time()

    from simurg.formats.bpm import read_bpm
    from simurg.formats.samplesheet import output_stem
    from simurg.pipeline.qc import autosome_mask

    manifest = read_bpm(a.bpm)
    names = manifest.name_index()
    summary = list(csv.DictReader((a.gt_dir / "sample_summary.tsv").open(), delimiter="\t"))
    ok_rows = [r for r in summary if r["status"] == "OK"]
    paths = [
        str(a.gt_dir / "samples" / output_stem(r["sample_name"]) / f"{output_stem(r['sample_name'])}.gtc")
        for r in ok_rows
    ]
    missing = [p for p in paths if not Path(p).exists()]
    if missing:
        print(f"GTC paths not found, e.g. {missing[:3]}", file=sys.stderr)
        return 2

    # 1. identify samples ------------------------------------------------------------------------------
    vs, loci, vhet, vcall = read_vcf_hets(a.vcf, names, a.n_snps)
    idx = np.array(loci, dtype=np.int64)
    with ProcessPoolExecutor(a.workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(idx,)) as ex:
        ours = np.stack(list(ex.map(_sel, paths, chunksize=8)))  # samples × snps
    H = (ours == 2).astype(np.float32)
    C = (ours > 0).astype(np.float32)
    V = (vhet.T & vcall.T).astype(np.float32)  # vcf samples × snps
    W = vcall.T.astype(np.float32)
    joint = C @ W.T
    agree = H @ V.T + (C - H) @ (W - V).T
    conc = agree / np.maximum(joint, 1)
    best = conc.argmax(axis=1)
    order = np.sort(conc, axis=1)
    top, second = order[:, -1], order[:, -2]
    sid2k = {}
    ambiguous = []
    for i, r in enumerate(ok_rows):
        if top[i] >= 0.98 and top[i] - second[i] >= 0.05:
            sid2k[r["sample_id"]] = vs[best[i]]
        else:
            ambiguous.append({"sample_id": r["sample_id"], "top": float(top[i]), "second": float(second[i])})
    dup_targets = [k for k, c in Counter(sid2k.values()).items() if c > 1]

    # 2. pedigree --------------------------------------------------------------------------------------
    ped: dict[str, dict[str, str]] = {}
    with a.ped.open() as fh:
        for r in csv.DictReader(fh, delimiter="\t"):
            ped[r["Individual ID"]] = r
    k2sid = {k: s for s, k in sid2k.items()}

    # 3. sex -------------------------------------------------------------------------------------------
    sex_cmp = Counter()
    sex_mismatch = []
    for r in ok_rows:
        k = sid2k.get(r["sample_id"])
        if not k or k not in ped:
            sex_cmp["no_pedigree"] += 1
            continue
        want = {"1": "M", "2": "F"}.get(ped[k]["Gender"], "U")
        got = r["sex_est"]
        sex_cmp["match" if got == want else "mismatch"] += 1
        if got != want:
            sex_mismatch.append({"sample_id": r["sample_id"], "id": k, "est": got, "ped": want})

    # 4. relatedness -----------------------------------------------------------------------------------
    rel = list(csv.DictReader((a.gt_dir / "relatedness.tsv").open(), delimiter="\t"))
    found = {frozenset((p["sample_a"], p["sample_b"])): p for p in rel}
    po_pairs = set()
    for k, r in ped.items():
        for par in (r["Paternal ID"], r["Maternal ID"]):
            if k in k2sid and par in k2sid:
                po_pairs.add(frozenset((k2sid[k], k2sid[par])))
    po_detected = [found[p] for p in po_pairs if p in found and found[p]["degree"] == "first"]
    po_kin = [float(found[p]["kinship"]) for p in po_pairs if p in found]
    po_missed = [sorted(p) for p in po_pairs if p not in found or found[p]["degree"] != "first"]
    sib = set()
    for k, r in ped.items():
        for s in (r.get("Siblings") or "").replace(",", " ").split():
            if k in k2sid and s in k2sid:
                sib.add(frozenset((k2sid[k], k2sid[s])))
    close = [p for p in rel if p["degree"] in ("duplicate", "first", "second")]
    unexplained = [
        {**p, "id_a": sid2k.get(p["sample_a"], ""), "id_b": sid2k.get(p["sample_b"], "")}
        for p in close
        if frozenset((p["sample_a"], p["sample_b"])) not in po_pairs | sib
    ]

    # 5. Mendelian errors -------------------------------------------------------------------------------
    auto = np.flatnonzero(autosome_mask(manifest.chroms) & ~manifest.intensity_only)
    path_of = {r["sample_id"]: p for r, p in zip(ok_rows, paths, strict=True)}
    trios = []
    for k, r in ped.items():
        f, m = r["Paternal ID"], r["Maternal ID"]
        if k in k2sid and f in k2sid and m in k2sid:
            trios.append((k, f, m))
    tasks = [(k, f, m, path_of[k2sid[k]], path_of[k2sid[f]], path_of[k2sid[m]]) for k, f, m in trios]
    with ProcessPoolExecutor(a.workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(auto,)) as ex:
        mendel = list(ex.map(_trio, tasks, chunksize=2))
    rates = np.array([x["rate"] for x in mendel if x["rate"] is not None])

    # 6. distributions ---------------------------------------------------------------------------------
    def q(vals: list[float]) -> dict[str, float]:
        v = np.array([x for x in vals if np.isfinite(x)])
        if not len(v):
            return {}
        return {
            k: float(np.percentile(v, p))
            for k, p in (("min", 0), ("p01", 1), ("median", 50), ("p99", 99), ("max", 100))
        } | {"n": len(v)}

    def num(x: str) -> float:
        try:
            return float(x)
        except ValueError:
            return float("nan")

    cqc = list(csv.DictReader((a.gt_dir / "controls_qc.tsv").open(), delimiter="\t"))
    metric_cols = [
        c
        for c in (cqc[0] if cqc else {})
        if c
        not in (
            "sample_id",
            "array_type",
            "imaging_date",
            "sample_plate",
            "sample_well",
            "scanner_id",
            "scanner_version",
            "bg_offset",
            "na_metrics",
        )
    ]
    contam = [num(r["contamination"]) for r in ok_rows]
    top_contam = sorted(
        (
            {"sample_id": r["sample_id"], "id": sid2k.get(r["sample_id"], ""), "c": num(r["contamination"])}
            for r in ok_rows
        ),
        key=lambda x: -x["c"] if np.isfinite(x["c"]) else 0,
    )[:15]
    flags = Counter(f.split(":")[0] for r in summary for f in r["flags"].split(";") if f)
    codes = Counter(c for r in summary for c in r["control_flag_codes"].split(";") if c)

    out: dict[str, Any] = {
        "n_samples": len(summary),
        "n_ok": len(ok_rows),
        "identification": {
            "snps": len(loci),
            "identified": len(sid2k),
            "ambiguous": ambiguous[:20],
            "n_ambiguous": len(ambiguous),
            "ids_matched_twice": dup_targets,
            "top_concordance": q(top.tolist()),
            "margin": q((top - second).tolist()),
        },
        "sex_vs_pedigree": dict(sex_cmp),
        "sex_mismatches": sex_mismatch[:30],
        "parent_offspring": {
            "pairs_in_pedigree": len(po_pairs),
            "detected_first_degree": len(po_detected),
            "kinship": q(po_kin),
            "missed": po_missed[:20],
        },
        "close_pairs_total": len(close),
        "close_pairs_unexplained_by_pedigree": len(unexplained),
        "unexplained_examples": unexplained[:25],
        "mendel": {
            "trios": len(mendel),
            "error_rate": q(rates.tolist()),
            "worst": sorted(mendel, key=lambda x: -(x["rate"] or 0))[:10],
        },
        "contamination": q(contam) | {"above_0.02": int(sum(c > 0.02 for c in contam if np.isfinite(c)))},
        "contamination_top": top_contam,
        "functional_status": dict(Counter(r["functional_status"] for r in summary)),
        "control_status": dict(Counter(r["control_status"] for r in summary)),
        "control_flag_codes": dict(codes),
        "flags": dict(flags),
        "control_metrics": {c: q([num(r[c]) for r in cqc]) for c in metric_cols},
        "het_rate": q([num(r["het_rate"]) for r in ok_rows]),
        "seconds": round(time.time() - t0, 1),
    }
    po_rate = len(po_detected) / max(len(po_pairs), 1)
    sex_rate = sex_cmp["match"] / max(sex_cmp["match"] + sex_cmp["mismatch"], 1)
    out["checks"] = {
        "identified_ge_99pct": len(sid2k) >= 0.99 * len(ok_rows),
        "sex_agreement_ge_99pct": sex_rate >= 0.99,
        "parent_offspring_detected_ge_99pct": po_rate >= 0.99,
        "median_mendel_rate_lt_0.001": bool(len(rates) and float(np.median(rates)) < 0.001),
    }
    out["passed"] = all(out["checks"].values())
    a.out.parent.mkdir(parents=True, exist_ok=True)
    a.out.write_text(json.dumps(out, indent=2, default=str) + "\n")
    print(json.dumps({k: out[k] for k in ("n_ok", "sex_vs_pedigree", "checks", "passed", "seconds")}, default=str))
    print(
        json.dumps(
            {
                "identification": {k: out["identification"][k] for k in ("identified", "n_ambiguous")},
                "po": {k: out["parent_offspring"][k] for k in ("pairs_in_pedigree", "detected_first_degree")},
                "mendel_trios": out["mendel"]["trios"],
            }
        )
    )
    return 0 if out["passed"] else 1


if __name__ == "__main__":
    sys.exit(main())