#!/usr/bin/env python3
"""Array genotypes vs 1000 Genomes 30x on one chromosome (ArrayTrain validation, plan §12).

    python scripts/train/concordance_chr.py --vcfs VCF_DIR --identity identity.tsv --truth chr22.vcf.gz \
        --bcftools BCFTOOLS --chrom 22 --out summary.json [--samples-glob '*.snv.vcf.gz']

Per array sample (identity table: sample_id → individual): biallelic SNVs of the chromosome whose (POS, REF, ALT)
is a record of the phased 3202-sample panel; genotype concordance and non-reference discordance (NRD = discordant
/ (discordant + concordant non-ref); plan §16.5) over sites called on the array, plus the array call rate there.
"""

from __future__ import annotations

import argparse
import csv
import json
import subprocess
from pathlib import Path

import numpy as np


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--vcfs", required=True, type=Path)
    ap.add_argument("--identity", required=True, type=Path)
    ap.add_argument("--truth", required=True, type=Path)
    ap.add_argument("--bcftools", default="bcftools")
    ap.add_argument("--chrom", default="22")
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    with a.identity.open() as fh:
        ident = {r["sample_id"]: r["individual"] for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")),
                                                                           delimiter="\t")}  # fmt: skip
    vcfs = {p.name.split(".")[0]: p for p in sorted(a.vcfs.rglob("*.snv.vcf.gz"))}
    vcfs = {k: v for k, v in vcfs.items() if k in ident}
    inds = sorted({ident[k] for k in vcfs})
    # array sites (from the first VCF) → truth genotypes of the individuals at those sites
    first = next(iter(vcfs.values()))
    q = subprocess.run([a.bcftools, "query", "-r", a.chrom, "-i", 'TYPE="snp" && N_ALT=1', "-f", "%POS\t%REF\t%ALT\n",
                        str(first)], capture_output=True, text=True, check=True).stdout  # fmt: skip
    sites = {(int(p), r, al) for p, r, al in (ln.split("\t") for ln in q.splitlines())}
    reg = a.out.with_suffix(".sites.tsv")
    chrom_t = f"chr{a.chrom}"
    reg.write_text("".join(f"{chrom_t}\t{p}\n" for p in sorted({s[0] for s in sites})))
    tq = subprocess.run([a.bcftools, "query", "-R", str(reg), "-s", ",".join(inds), "--force-samples", "-i",
                         'TYPE="snp"', "-f", "%POS\t%REF\t%ALT[\t%GT]\n", str(a.truth)], capture_output=True, text=True,
                        check=True).stdout  # fmt: skip
    head = subprocess.run([a.bcftools, "query", "-l", "-s", ",".join(inds), "--force-samples", str(a.truth)],
                          capture_output=True, text=True, check=True).stdout.split()  # fmt: skip
    truth: dict[tuple[int, str, str], dict[str, int]] = {}
    for ln in tq.splitlines():
        f = ln.split("\t")
        k = (int(f[0]), f[1], f[2])
        if k in sites:
            truth[k] = {
                s: sum(x == "1" for x in g.replace("|", "/").split("/")) for s, g in zip(head, f[3:], strict=True)
            }
    per = []
    for sid, path in sorted(vcfs.items()):
        ind = ident[sid]
        aq = subprocess.run([a.bcftools, "query", "-r", a.chrom, "-f", "%POS\t%REF\t%ALT[\t%GT]\n", str(path)],
                            capture_output=True, text=True, check=True).stdout  # fmt: skip
        n = called = conc = nrd_bad = nrd_ok = 0
        for ln in aq.splitlines():
            p, r, al, gt = ln.split("\t")
            k = (int(p), r, al)
            t = truth.get(k, {}).get(ind)
            if t is None:
                continue
            n += 1
            if "." in gt:
                continue
            called += 1
            d = sum(x == "1" for x in gt.replace("|", "/").split("/"))
            conc += d == t
            if d > 0 or t > 0:
                nrd_ok += d == t
                nrd_bad += d != t
        per.append({"sample": sid, "individual": ind, "sites": n, "call_rate": called / n if n else None,
                    "concordance": conc / called if called else None,
                    "nrd": nrd_bad / (nrd_bad + nrd_ok) if nrd_bad + nrd_ok else None})  # fmt: skip
    summ = {k: float(np.mean([p[k] for p in per if p[k] is not None])) for k in ("call_rate", "concordance", "nrd")}
    a.out.write_text(json.dumps({"chrom": a.chrom, "samples": len(per), "truth_sites": len(truth), **summ,
                                 "per_sample": per}, indent=1))  # fmt: skip
    print(json.dumps({"samples": len(per), "truth_sites": len(truth), **summ}))
    return 0


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