#!/usr/bin/env python3
"""Derive a PGx probe mask from array-vs-WGS concordance (spec PGX-010) for the individuals of a GSA PGx batch.

    python scripts/pgx/derive_masks.py --db DB.sqlite --vcfs VCF_DIR --identity demo36_identity.tsv \
        --exclude-individuals getrm.txt --manifest NAME --out MASK.tsv [--min-concordance 0.95]

Truth: the 1000 Genomes 30x phased panel (NYGC, 3202 samples), fetched per gene window over HTTPS with pysam (MIT;
a script-only dependency: ``uv pip install pysam``). For every PGx-DB defining variant the array record(s) at that
site are compared with the truth alt dosage; records whose concordance is below the threshold (≥ 5 comparisons)
are masked by their probe names. Individuals listed in ``--exclude-individuals`` (the validation set) are not used,
so the mask stays independent of the star-allele validation. The spec asks for ≥ 100 samples; this batch has fewer
(recorded as K-029).
"""

from __future__ import annotations

import argparse
import csv
import gzip
import sqlite3
from collections import defaultdict
from pathlib import Path

PANEL = (
    "https://ftp.1000genomes.ebi.ac.uk/vol1/ftp/data_collections/1000G_2504_high_coverage/working/"
    "20220422_3202_phased_SNV_INDEL_SV/1kGP_high_coverage_Illumina.{c}.filtered.SNV_INDEL_SV_phased_panel.vcf.gz"
)


def main() -> int:
    import pysam

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--db", required=True, type=Path)
    ap.add_argument("--vcfs", required=True, type=Path)
    ap.add_argument("--identity", required=True, type=Path)
    ap.add_argument("--exclude-individuals", type=Path)
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--min-concordance", type=float, default=0.95)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    ident = {}
    with a.identity.open() as fh:
        for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")), delimiter="\t"):
            ident[r["sample_id"]] = r["individual"]
    excl = set(a.exclude_individuals.read_text().split()) if a.exclude_individuals else set()
    use = {s: i for s, i in ident.items() if i not in excl}
    con = sqlite3.connect(a.db)
    # the highest population frequency among the alleles that carry each variant (CPIC frequencies)
    vfreq: dict[tuple[str, int, str, str], float] = defaultdict(float)
    for c, p, r, al, f in con.execute(
        "SELECT v.chrom, v.pos, v.ref, v.alt, MAX(f.freq) FROM allele_variants v JOIN frequencies f "
        "ON f.gene = v.gene AND f.allele = v.allele GROUP BY v.chrom, v.pos, v.ref, v.alt"
    ):
        vfreq[(c, p, r, al)] = float(f or 0.0)
    keys = sorted({(c, p, r, al, g) for g, c, p, r, al in con.execute(
        "SELECT gene, chrom, pos, ref, alt FROM allele_variants")})  # fmt: skip
    # array dosages per (chrom, pos, ref, alt) and record ID
    wanted = defaultdict(set)
    for c, p, *_ in keys:
        wanted[c].add(p)
    arr: dict[tuple[str, int, str, str], dict[str, dict[str, int]]] = defaultdict(lambda: defaultdict(dict))
    gsd: dict[tuple[str, int, str, str, str, str], float] = {}
    for sid in sorted(use):
        path = a.vcfs / f"{sid}.snv.vcf.gz"
        if not path.exists():
            continue
        with gzip.open(path, "rt") as fh:
            for ln in fh:
                if ln[0] == "#":
                    continue
                c, p, rid, ref, alts = ln.split("\t", 5)[:5]
                if int(p) not in wanted.get(c, ()):
                    continue
                sample = ln.rstrip("\n").split("\t")[9].split(":")
                gt = sample[0].replace("|", "/").split("/")
                if "." in gt:
                    continue
                gs_v = float(sample[1]) if len(sample) > 1 and sample[1] not in (".", "") else float("nan")
                for ai, alt in enumerate(alts.split(","), start=1):
                    arr[(c, int(p), ref, alt)][rid][sid] = sum(int(x) == ai for x in gt)
                    gsd[(c, int(p), ref, alt, rid, sid)] = gs_v
    # truth per gene window
    truth: dict[tuple[str, int, str, str], dict[str, int]] = {}
    by_gene = defaultdict(list)
    for c, p, r, al, g in keys:
        by_gene[(g, c)].append((p, r, al))
    inds = sorted(set(use.values()))
    for (g, c), vs in sorted(by_gene.items()):
        core = c.removeprefix("chr")
        chrom = "chrM" if core in ("MT", "M") else f"chr{core}"
        if chrom == "chrM":
            continue
        try:
            vf = pysam.VariantFile(PANEL.format(c=chrom))
        except (OSError, ValueError):
            print(f"{g} {chrom}: panel file not available, skipped")
            continue
        samples = [s for s in inds if s in vf.header.samples]
        vf.subset_samples(samples)
        lo, hi = min(p for p, _, _ in vs) - 10, max(p + len(r) for p, r, _ in vs) + 10
        want = {(p, r, al) for p, r, al in vs}
        seen = set()
        for rec in vf.fetch(chrom, lo, hi):
            for ai, alt in enumerate(rec.alts or (), start=1):
                k = (rec.pos, rec.ref, alt)
                if k in want:
                    truth[(c, *k)] = {s: sum(x == ai for x in rec.samples[s]["GT"] if x is not None) for s in samples}
                    seen.add(k)
        for p, r, al in want - seen:
            # absent from the (filtered) panel: homozygous reference only for a rare variant; a common one was
            # most likely filtered out (e.g. CYP2B6 rs2279343 next to CYP2B7) and is not assessed
            if vfreq.get((c, p, r, al), 0.0) < 0.05:
                truth[(c, p, r, al)] = dict.fromkeys(samples, 0)
        print(f"{g} {chrom}: {len(want)} variant(s), {len(seen)} in the panel")
    rows = []
    stats = [0, 0]
    bins = [0.0, 0.15, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 1.01]
    calib = [[0, 0] for _ in bins[:-1]]  # [discordant, n] per GS bin over unmasked records (error-model calibration)
    for k, recs in sorted(arr.items()):
        t = truth.get(k)
        if t is None:
            continue
        for rid, dos in sorted(recs.items()):
            pairs = [(d, t[use[s]]) for s, d in dos.items() if use[s] in t]
            if len(pairs) < 5:
                continue
            conc = sum(x == y for x, y in pairs) / len(pairs)
            stats[1] += 1
            if conc >= a.min_concordance:
                for s_, d in dos.items():
                    if use[s_] not in t:
                        continue
                    gv = gsd.get((*k, rid, s_), float("nan"))
                    for bi in range(len(bins) - 1):
                        if bins[bi] <= gv < bins[bi + 1]:
                            calib[bi][0] += d != t[use[s_]]
                            calib[bi][1] += 1
            if conc < a.min_concordance:
                stats[0] += 1
                for probe in rid.split(","):
                    rows.append((probe, f"concordance {conc:.2f} vs 1000G 30x at {k[0]}:{k[1]} {k[2]}>{k[3]} "
                                        f"(n={len(pairs)})", a.manifest))  # fmt: skip
    with a.out.open("w") as fh:
        fh.write(f"# PGx probe mask for {a.manifest} (PGX-010): array vs 1000G 30x, {len(inds)} individuals, "
                 f"threshold {a.min_concordance}\n")  # fmt: skip
        for r in sorted(set(rows)):
            fh.write("\t".join(r) + "\n")
    print(f"masked {stats[0]} of {stats[1]} assessed record(s); {len(set(rows))} probe name(s) -> {a.out}")
    for bi, (dsc, n) in enumerate(calib):
        if n:
            print(f"GS [{bins[bi]:.2f}, {bins[bi + 1]:.2f}): discordance {dsc}/{n} = {dsc / n:.4f} (unmasked records)")
    return 0


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