#!/usr/bin/env python3
"""P8 scientific validation (plan §16.5) of SimurgArray calls on GSE165845 (GSA-MD v1, 120 1000 Genomes
individuals × 3 replicates) against the NYGC 30x WGS call set and GIAB HG001 v4.2.1.

Inputs are prepared by ``scripts/cloud/tasks/p8.sh``: every VCF/BCF is already split to biallelic records and
left-normalised against the same GRCh38 FASTA with Ensembl contig names (1…22, X):

    --array   cohort BCF from ``simurg gtc-to-vcf --merge`` (samples = <barcode>_<position>)
    --truth   NYGC records at the array positions, subset to the GSE165845 individuals
    --giab    GIAB HG001 records at the array positions; --giab-bed the v4.2.1 confident regions
    --map     samples_map.tsv from ``prepare.py``;  --gt-dir the ``simurg genotype`` output folder

Metrics: genotype concordance, non-reference discordance (NRD), no-call rate, sensitivity / specificity /
precision / F1 per genotype class and per truth MAF bin, SNV vs indel; replicate concordance; systematic
site discordance with the allele-swap (strand/allele error) pattern; per-chip call rate and LogR deviation;
duplicate detection by the cohort relatedness step. A site absent from a WGS call set is taken as homozygous
reference (reported separately from sites present in both).
"""

from __future__ import annotations

import argparse
import csv
import json
import statistics
import subprocess
import sys
import time
from collections import defaultdict
from itertools import combinations
from pathlib import Path

import numpy as np
import numpy.typing as npt

CHROMS = [str(i) for i in range(1, 23)] + ["X"]
GT_CODE = {
    "0/0": 0, "0|0": 0, "0": 0,
    "0/1": 1, "1/0": 1, "0|1": 1, "1|0": 1,
    "1/1": 2, "1|1": 2, "1": 2,
}  # fmt: skip
MAF_BINS = [0.0, 0.001, 0.01, 0.05, 0.1, 0.2, 0.5000001]
MAF_LABELS = ["<0.1%", "0.1-1%", "1-5%", "5-10%", "10-20%", "20-50%"]
CLASSES = ["hom_ref", "het", "hom_alt"]


def query(bcftools: str, path: Path, fmt: str) -> tuple[list[str], subprocess.Popen[str]]:
    samples = subprocess.run([bcftools, "query", "-l", str(path)], capture_output=True, text=True, check=True)
    p = subprocess.Popen([bcftools, "query", "-f", fmt, str(path)], stdout=subprocess.PIPE, text=True, bufsize=1 << 20)
    return samples.stdout.split(), p


def load(bcftools: str, path: Path, with_af: bool = False, chroms: list[str] | None = None) -> dict:
    """Records of chromosomes 1–22/X → keys, genotype matrix (int8, −1 = no call), AF (if asked)."""
    fmt = "%CHROM\t%POS\t%REF\t%ALT" + ("\t%INFO/AF" if with_af else "") + "[\t%GT]\n"
    samples, p = query(bcftools, path, fmt)
    nfix = 5 if with_af else 4
    keys: list[tuple[str, int, str, str]] = []
    rows: list[npt.NDArray[np.int8]] = []
    af: list[float] = []
    index: dict[tuple[str, int, str, str], int] = {}
    positions: set[tuple[str, int]] = set()
    dups = 0
    ok_chrom = set(chroms or CHROMS)
    get = GT_CODE.get
    assert p.stdout is not None
    for line in p.stdout:
        f = line.rstrip("\n").split("\t")
        if f[0] not in ok_chrom:
            continue
        key = (f[0], int(f[1]), f[2], f[3])
        if key in index:
            dups += 1
            continue
        index[key] = len(keys)
        keys.append(key)
        positions.add((f[0], key[1]))
        if with_af:
            try:
                af.append(float(f[4]))
            except ValueError:
                af.append(float("nan"))
        rows.append(np.fromiter((get(g, -1) for g in f[nfix:]), np.int8, len(f) - nfix))
    if p.wait() != 0:
        raise SystemExit(f"bcftools query failed on {path}")
    gt = np.vstack(rows) if rows else np.zeros((0, len(samples)), np.int8)
    return {
        "samples": samples, "keys": keys, "index": index, "positions": positions, "gt": gt,
        "af": np.array(af), "duplicate_keys": dups,
    }  # fmt: skip


def is_snv(key: tuple[str, int, str, str]) -> bool:
    return len(key[2]) == 1 and len(key[3]) == 1


def confusion(t: npt.NDArray[np.int8], a: npt.NDArray[np.int8], group: npt.NDArray[np.int64], ngroups: int):
    """Counts (group, truth 0..2, array −1..2) over every cell; ``group`` has the shape of ``t``."""
    code = (t.astype(np.int64) * 4 + (a.astype(np.int64) + 1)) + 12 * group
    return np.bincount(code.ravel(), minlength=12 * ngroups).reshape(ngroups, 3, 4)


def metrics(c: npt.NDArray[np.int64]) -> dict:
    """``c`` is (truth 0..2) × (array no-call, 0, 1, 2)."""
    called = c[:, 1:]
    n = int(called.sum())
    nocall = int(c[:, 0].sum())
    out: dict = {"n_called": n, "n_nocall": nocall}
    if n == 0:
        return out
    trace = int(np.trace(called))
    out["concordance"] = trace / n
    out["nocall_rate"] = nocall / (n + nocall)
    nr = n - int(called[0, 0])
    out["nrd"] = (n - trace) / nr if nr else None
    for k, name in enumerate(CLASSES):
        tp = int(called[k, k])
        fn = int(called[k].sum()) - tp
        fp = int(called[:, k].sum()) - tp
        tn = n - tp - fn - fp
        sens = tp / (tp + fn) if tp + fn else None
        prec = tp / (tp + fp) if tp + fp else None
        out[name] = {
            "n_truth": tp + fn,
            "sensitivity": sens,
            "precision": prec,
            "specificity": tn / (tn + fp) if tn + fp else None,
            "f1": 2 * sens * prec / (sens + prec) if sens and prec else None,
            "nocall_rate": int(c[k, 0]) / (tp + fn + int(c[k, 0])) if tp + fn + int(c[k, 0]) else None,
        }
    # variant level (non-reference genotypes): a call is a true positive only with the exact genotype
    tp = int(called[1, 1] + called[2, 2])
    fn = int(called[1, 0] + called[2, 0] + called[1, 2] + called[2, 1])
    fp = int(called[0, 1] + called[0, 2] + called[1, 2] + called[2, 1])
    sens = tp / (tp + fn) if tp + fn else None
    prec = tp / (tp + fp) if tp + fp else None
    out["nonref"] = {
        "sensitivity": sens,
        "precision": prec,
        "f1": 2 * sens * prec / (sens + prec) if sens and prec else None,
    }
    out["matrix_truth_by_array_nc012"] = c.tolist()
    return out


def read_bed(path: Path) -> dict[str, tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]]:
    iv: dict[str, list[tuple[int, int]]] = defaultdict(list)
    with open(path) as fh:
        for line in fh:
            if line.startswith(("#", "track", "browser")):
                continue
            f = line.split("\t")
            iv[f[0].removeprefix("chr")].append((int(f[1]), int(f[2])))
    return {k: (np.array([s for s, _ in v]), np.array([e for _, e in v])) for k, v in iv.items() for v in [sorted(v)]}


def in_bed(bed, chrom: str, pos: int) -> bool:
    if chrom not in bed:
        return False
    s, e = bed[chrom]
    i = int(np.searchsorted(s, pos - 1, side="right")) - 1  # BED is 0-based half-open; pos is 1-based
    return i >= 0 and pos - 1 < e[i]


def against(arr: dict, tru: dict, cols: list[int], tcols: list[int], restrict=None) -> dict:
    """Array samples ``cols`` vs truth columns ``tcols``: per-sample, per MAF bin × variant type, per site."""
    keys = arr["keys"]
    n = len(keys)
    rows_t = np.full(n, -1, np.int64)
    status = np.zeros(n, np.int8)  # 0 matched, 1 absent (→ hom ref), 2 allele mismatch, 3 outside region
    for i, k in enumerate(keys):
        if restrict is not None and not restrict(k):
            status[i] = 3
            continue
        j = tru["index"].get(k)
        if j is not None:
            rows_t[i] = j
        elif (k[0], k[1]) in tru["positions"]:
            status[i] = 2
        else:
            status[i] = 1
    use = np.flatnonzero(status <= 1)
    A = arr["gt"][np.ix_(use, cols)]
    T = np.zeros_like(A)
    m = rows_t[use] >= 0
    T[m] = tru["gt"][np.ix_(rows_t[use][m], tcols)]
    T[T < 0] = -9  # truth no-call → excluded below
    snv = np.array([is_snv(keys[i]) for i in use])
    if len(tru["af"]):
        af = np.where(m, tru["af"][np.maximum(rows_t[use], 0)], 0.0)
        af = np.nan_to_num(af, nan=0.0)
        maf = np.minimum(af, 1 - af)
        mbin = np.where(m, np.digitize(maf, MAF_BINS) - 1, len(MAF_LABELS))  # last bin = absent from truth
    else:
        mbin = np.where(m, 0, len(MAF_LABELS))
    matched_flag = m.astype(np.int64)

    S = len(cols)
    nb = len(MAF_LABELS) + 1
    per_sample = np.zeros((S, 3, 4), np.int64)
    by_group = np.zeros((2 * 2 * nb, 3, 4), np.int64)  # matched? × snv? × bin
    site_disc = np.zeros(len(use), np.int64)
    site_called = np.zeros(len(use), np.int64)
    site_swap = np.zeros(len(use), np.int64)
    step = 50_000
    for s0 in range(0, len(use), step):
        a = A[s0 : s0 + step]
        t = T[s0 : s0 + step]
        valid = t >= 0
        tt = np.where(valid, t, 0)
        sample_g = np.broadcast_to(np.arange(S), a.shape)
        per_sample += confusion(tt[valid], a[valid], sample_g[valid], S)
        g = (matched_flag[s0 : s0 + step] * 2 + snv[s0 : s0 + step]) * nb + mbin[s0 : s0 + step]
        gg = np.broadcast_to(g[:, None], a.shape)
        by_group += confusion(tt[valid], a[valid], gg[valid], 4 * nb)
        both = valid & (a >= 0)
        site_called[s0 : s0 + step] = both.sum(1)
        site_disc[s0 : s0 + step] = (both & (a != t)).sum(1)
        site_swap[s0 : s0 + step] = (both & (t != 1) & (a == 2 - t)).sum(1)

    def pooled(sel) -> dict:
        return metrics(by_group[np.atleast_1d(sel)].sum(0))

    idx = np.arange(4 * nb).reshape(2, 2, nb)  # [matched][snv][bin]
    res: dict = {
        "sites": {
            "array_records": n,
            "matched": int((status == 0).sum()),
            "absent_in_truth_as_hom_ref": int((status == 1).sum()),
            "allele_mismatch_excluded": int((status == 2).sum()),
            "outside_region": int((status == 3).sum()),
        },
        "all": pooled(idx.ravel()),
        "matched_only": pooled(idx[1].ravel()),
        "snv": pooled(idx[:, 1].ravel()),
        "indel": pooled(idx[:, 0].ravel()),
        "by_maf": {lab: pooled(idx[1, :, b].ravel()) for b, lab in enumerate(MAF_LABELS)},
        "by_maf_snv": {lab: pooled(idx[1, 1, b]) for b, lab in enumerate(MAF_LABELS)},
        "absent_in_truth": pooled(idx[:, :, nb - 1].ravel()),
    }
    ps = [metrics(per_sample[i]) for i in range(S)]
    res["per_sample"] = ps
    conc = [p["concordance"] for p in ps if "concordance" in p]
    nrd = [p["nrd"] for p in ps if p.get("nrd") is not None]
    res["per_sample_summary"] = {
        "n": len(conc),
        "concordance_median": statistics.median(conc) if conc else None,
        "concordance_min": min(conc) if conc else None,
        "nrd_median": statistics.median(nrd) if nrd else None,
        "nrd_max": max(nrd) if nrd else None,
    }
    # systematic discordance: sites wrong in at least half of the called samples (≥ 10 called)
    frac = np.divide(site_disc, site_called, out=np.zeros(len(use)), where=site_called > 0)
    sysm = (site_called >= 10) & (frac >= 0.5)
    swap = sysm & (site_swap >= 0.8 * np.maximum(site_disc, 1))
    ambiguous = np.array([{keys[i][2], keys[i][3]} in ({"A", "T"}, {"C", "G"}) for i in use])
    evaluable = int((site_called >= 10).sum())
    res["site_errors"] = {
        "evaluable_sites": evaluable,
        "systematic_discordant_sites": int(sysm.sum()),
        "systematic_rate": sysm.sum() / evaluable if evaluable else None,
        "allele_swap_sites": int(swap.sum()),
        "allele_swap_rate": swap.sum() / evaluable if evaluable else None,
        "allele_swap_at_cg_ambiguous": int((swap & ambiguous).sum()),
        "systematic_indels": int((sysm & ~snv).sum()),
        "systematic_absent_in_truth": int((sysm & ~m).sum()),
    }
    res["_systematic_rows"] = [
        (keys[use[i]], int(site_called[i]), int(site_disc[i]), int(site_swap[i]), bool(m[i]))
        for i in np.flatnonzero(sysm)
    ]
    return res


def replicate_concordance(arr: dict, groups: dict[str, list[int]]) -> dict:
    gt = arr["gt"]
    pairs = []
    for ind, cols in sorted(groups.items()):
        for x, y in combinations(cols, 2):
            a, b = gt[:, x], gt[:, y]
            both = (a >= 0) & (b >= 0)
            one = (a >= 0) ^ (b >= 0)
            nb = int(both.sum())
            eq = int((both & (a == b)).sum())
            nonref = both & ((a > 0) | (b > 0))
            pairs.append(
                {
                    "individual": ind,
                    "a": arr["samples"][x],
                    "b": arr["samples"][y],
                    "both_called": nb,
                    "concordance": eq / nb if nb else None,
                    "nonref_discordance": int((nonref & (a != b)).sum()) / int(nonref.sum()) if nonref.any() else None,
                    "one_side_nocall": int(one.sum()),
                }
            )
    c = [p["concordance"] for p in pairs if p["concordance"] is not None]
    d = [p["nonref_discordance"] for p in pairs if p["nonref_discordance"] is not None]
    tot_both = sum(p["both_called"] for p in pairs)
    tot_eq = sum(round(p["concordance"] * p["both_called"]) for p in pairs if p["concordance"] is not None)
    return {
        "n_pairs": len(pairs),
        "pooled_concordance": tot_eq / tot_both if tot_both else None,
        "median_concordance": statistics.median(c) if c else None,
        "min_concordance": min(c) if c else None,
        "median_nonref_discordance": statistics.median(d) if d else None,
        "pairs": pairs,
    }


def batch_table(gt_dir: Path, sample_map: dict[str, dict], per_sample_conc: dict[str, float]) -> dict:
    path = gt_dir / "sample_summary.tsv"
    if not path.exists():
        return {}
    chips: dict[str, list[dict]] = defaultdict(list)
    for r in csv.DictReader(path.open(), delimiter="\t"):
        chips[r["sample_folder"] or r["sample_id"].split("_")[0]].append(r)

    def f(x: str) -> float | None:
        try:
            return float(x)
        except ValueError:
            return None

    table = []
    for chip, rs in sorted(chips.items()):
        cr = [v for r in rs if (v := f(r["call_rate"])) is not None]
        ld = [v for r in rs if (v := f(r["logr_dev"])) is not None]
        cc = [per_sample_conc[r["sample_id"]] for r in rs if r["sample_id"] in per_sample_conc]
        table.append(
            {
                "chip": chip,
                "n": len(rs),
                "call_rate_mean": statistics.fmean(cr) if cr else None,
                "call_rate_min": min(cr) if cr else None,
                "logr_dev_median": statistics.median(ld) if ld else None,
                "wgs_concordance_mean": statistics.fmean(cc) if cc else None,
                "status_fail": sum(r["status"] != "OK" for r in rs),
            }
        )
    means = [t["call_rate_mean"] for t in table if t["call_rate_mean"] is not None]
    med = statistics.median(means)
    mad = statistics.median(abs(x - med) for x in means) or 1e-9
    for t in table:
        t["call_rate_outlier"] = t["call_rate_mean"] is not None and abs(t["call_rate_mean"] - med) > 3 * 1.4826 * mad
    all_cr = [f(r["call_rate"]) for rs in chips.values() for r in rs]
    all_ld = [f(r["logr_dev"]) for rs in chips.values() for r in rs]
    return {
        "n_chips": len(table),
        "n_samples": sum(t["n"] for t in table),
        "call_rate_median": statistics.median(x for x in all_cr if x is not None),
        "logr_dev_median": statistics.median(x for x in all_ld if x is not None),
        "chips_call_rate_outliers": [t["chip"] for t in table if t["call_rate_outlier"]],
        "chips": table,
    }


def duplicate_check(gt_dir: Path, sample_map: dict[str, dict]) -> dict:
    path = gt_dir / "relatedness.tsv"
    if not path.exists():
        return {}
    ind = {s: m["individual"] for s, m in sample_map.items()}
    found: set[frozenset[str]] = set()
    false_dup = []
    for r in csv.DictReader(path.open(), delimiter="\t"):
        a, b = r["sample_a"], r["sample_b"]
        if r["degree"] == "duplicate":
            if ind.get(a) == ind.get(b):
                found.add(frozenset((a, b)))
            else:
                false_dup.append((a, b, r["kinship"]))
    groups: dict[str, list[str]] = defaultdict(list)
    for s, i in ind.items():
        groups[i].append(s)
    expected = {frozenset(p) for g in groups.values() for p in combinations(g, 2)}
    return {
        "replicate_pairs": len(expected),
        "detected_as_duplicate": len(found & expected),
        "missed": sorted(" ".join(sorted(p)) for p in expected - found),
        "duplicates_between_different_individuals": false_dup,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--array", required=True, type=Path)
    ap.add_argument("--truth", required=True, type=Path)
    ap.add_argument("--giab", type=Path)
    ap.add_argument("--giab-bed", type=Path)
    ap.add_argument("--giab-sample", default="NA12878")
    ap.add_argument("--map", required=True, type=Path)
    ap.add_argument("--gt-dir", type=Path)
    ap.add_argument("--bcftools", default="bcftools")
    ap.add_argument("--chroms", default="", help="comma-separated subset of 1..22,X (tests)")
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    t0 = time.time()
    a.out.mkdir(parents=True, exist_ok=True)

    smap = {r["sample_id"]: r for r in csv.DictReader(a.map.open(), delimiter="\t")}
    chroms = a.chroms.split(",") if a.chroms else None
    arr = load(a.bcftools, a.array, chroms=chroms)
    tru = load(a.bcftools, a.truth, with_af=True)
    print(f"array {arr['gt'].shape}, truth {tru['gt'].shape} ({time.time() - t0:.0f} s)", file=sys.stderr)

    tcol = {s: i for i, s in enumerate(tru["samples"])}
    cols, tcols, missing = [], [], []
    for i, s in enumerate(arr["samples"]):
        ind = smap.get(s, {}).get("individual")
        if ind in tcol:
            cols.append(i)
            tcols.append(tcol[ind])
        else:
            missing.append(s)
    summary: dict = {
        "array_samples": len(arr["samples"]),
        "individuals": len({smap[s]["individual"] for s in arr["samples"] if s in smap}),
        "samples_without_truth": missing,
        "array_duplicate_keys": arr["duplicate_keys"],
    }
    wgs = against(arr, tru, cols, tcols)
    sysrows = wgs.pop("_systematic_rows")
    per = wgs.pop("per_sample")
    per_conc = {}
    with open(a.out / "per_sample.tsv", "w") as fh:
        fh.write("sample_id\tindividual\treplicate\tn_called\tn_nocall\tconcordance\tnrd\n")
        for c, p in zip(cols, per, strict=True):
            s = arr["samples"][c]
            per_conc[s] = p.get("concordance")
            fh.write(
                f"{s}\t{smap[s]['individual']}\t{smap[s]['replicate']}\t{p['n_called']}\t{p['n_nocall']}\t"
                f"{p.get('concordance', '')}\t{p.get('nrd', '')}\n"
            )
    with open(a.out / "systematic_sites.tsv", "w") as fh:
        fh.write("chrom\tpos\tref\talt\tcalled\tdiscordant\tallele_swap\tin_truth\n")
        for k, nc, nd, ns, inb in sysrows:
            fh.write(f"{k[0]}\t{k[1]}\t{k[2]}\t{k[3]}\t{nc}\t{nd}\t{ns}\t{int(inb)}\n")
    summary["nygc_30x"] = wgs

    groups: dict[str, list[int]] = defaultdict(list)
    for i, s in enumerate(arr["samples"]):
        if s in smap:
            groups[smap[s]["individual"]].append(i)
    rep = replicate_concordance(arr, groups)
    with open(a.out / "replicate_pairs.tsv", "w") as fh:
        fh.write("individual\ta\tb\tboth_called\tconcordance\tnonref_discordance\tone_side_nocall\n")
        for p in rep.pop("pairs"):
            fh.write("\t".join(str(p[k]) for k in p) + "\n")
    summary["replicates"] = rep

    if a.giab and a.giab_bed:
        gia = load(a.bcftools, a.giab)
        bed = read_bed(a.giab_bed)
        gcols = [i for i, s in enumerate(arr["samples"]) if smap.get(s, {}).get("individual") == a.giab_sample]
        if gcols and gia["samples"]:
            g = against(arr, gia, gcols, [0] * len(gcols), restrict=lambda k: k[0] != "X" and in_bed(bed, k[0], k[1]))
            g.pop("by_maf"), g.pop("by_maf_snv")
            g["systematic_sites"] = [f"{k[0]}:{k[1]}:{k[2]}:{k[3]}" for k, *_ in g.pop("_systematic_rows")]
            g["per_replicate"] = [
                {"sample": arr["samples"][c], **{k: p.get(k) for k in ("n_called", "concordance", "nrd")}}
                for c, p in zip(gcols, g.pop("per_sample"), strict=True)
            ]
            summary["giab_hg001"] = g
    if a.gt_dir:
        summary["batch"] = batch_table(a.gt_dir, smap, per_conc)
        summary["duplicate_detection"] = duplicate_check(a.gt_dir, smap)
    summary["seconds"] = round(time.time() - t0, 1)
    (a.out / "summary.json").write_text(json.dumps(summary, indent=1, default=float))
    w = summary["nygc_30x"]
    print(
        f"NYGC: concordance {w['all'].get('concordance'):.5f}  NRD {w['all'].get('nrd'):.5f}  "
        f"no-call {w['all'].get('nocall_rate'):.5f}  replicate {rep['pooled_concordance']:.5f}"
    )
    if "giab_hg001" in summary:
        gg = summary["giab_hg001"]["all"]
        print(f"GIAB HG001: concordance {gg.get('concordance'):.5f}  NRD {gg.get('nrd')}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())