#!/usr/bin/env python3
"""Validate `simurg cn call` on GSE165845 (120 1000 Genomes individuals × 3 GSA arrays) against the NYGC 30x
structural-variant callset (1KGP_3202 GATK-SV/svtools freeze V3) and between replicates.

    python scripts/v2/cn_validate_truth.py --calls CN_OUT --map samples_map.tsv --truth truth.tsv --bpm M.bpm --out DIR

``truth.tsv``: ``chrom start end svtype sample`` for non-reference DEL/DUP genotypes of the 120 individuals
(``bcftools query``; chr prefix removed). Matching: reciprocal overlap ≥ 50 % and the same class (DEL ↔ losses
incl. mosaic, DUP ↔ gains). Only truth events with ≥ 10 array probes are counted as detectable.
"""

from __future__ import annotations

import argparse
import csv
import json
import statistics
from collections import defaultdict
from pathlib import Path

import numpy as np

from simurg.formats.bpm import read_bpm

BINS = [(25_000, 50_000), (50_000, 100_000), (100_000, 500_000), (500_000, 1_000_000), (1_000_000, 10**9)]
CLASS = {"DEL": "DEL", "MDEL": "DEL", "DUP": "DUP", "MDUP": "DUP"}


def ro(a: tuple[int, int], b: tuple[int, int]) -> float:
    ov = min(a[1], b[1]) - max(a[0], b[0]) + 1
    if ov <= 0:
        return 0.0
    return min(ov / (a[1] - a[0] + 1), ov / (b[1] - b[0] + 1))


def bin_of(size: int) -> str:
    for lo, hi in BINS:
        if lo <= size < hi:
            return f"{lo // 1000}-{hi // 1000 if hi < 10**9 else 'inf'}kb"
    return "<25kb"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--calls", required=True, type=Path)
    ap.add_argument("--map", required=True, type=Path)
    ap.add_argument("--truth", required=True, type=Path)
    ap.add_argument("--bpm", required=True)
    ap.add_argument("--min-qual", type=float, default=20)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    a.out.mkdir(parents=True, exist_ok=True)
    man = read_bpm(a.bpm)
    chroms = np.array(["X" if c == "XY" else c for c in man.chroms])
    pos = man.map_infos.astype(np.int64)
    probe_pos = {c: np.sort(pos[(chroms == c) & (pos > 0)]) for c in set(chroms.tolist())}

    def n_probes(c: str, s: int, e: int) -> int:
        p = probe_pos.get(c)
        return 0 if p is None else int(np.searchsorted(p, e, "right") - np.searchsorted(p, s, "left"))

    with a.map.open() as fh:
        smap = {r["sample_id"]: r for r in csv.DictReader(fh, delimiter="\t")}
    side = json.loads((a.calls / "cn_call_manifest.json").read_text())
    calls: dict[str, list[dict]] = {}
    for s in side["samples"]:
        if s.get("status") != "OK":
            continue
        calls[s["sample_id"]] = [e for e in s["events"] if e["quality"] >= a.min_qual and e["filters"] == ["PASS"]]
    truth: dict[str, list[tuple[str, int, int, str]]] = defaultdict(list)
    with a.truth.open() as fh:
        for line in fh:
            c, s_, e_, t, ind = line.rstrip("\n").split("\t")
            c = c[3:] if c.startswith("chr") else c
            if t in ("DEL", "DUP") and int(e_) - int(s_) + 1 >= 25_000:
                truth[ind].append((c, int(s_), int(e_), t))

    # sensitivity per class and size bin
    sens: dict[str, dict[str, list[int]]] = defaultdict(lambda: defaultdict(lambda: [0, 0]))
    prec = defaultdict(lambda: [0, 0])
    per_sample = []
    for sid, evs in calls.items():
        ind = smap.get(sid, {}).get("individual")
        if ind is None:
            continue
        tr = [t for t in truth.get(ind, []) if n_probes(t[0], t[1], t[2]) >= 10]
        hit_t = 0
        for c, s_, e_, t in tr:
            ok = any(CLASS.get(e["call"]) == t and e["chrom"] == c and ro((s_, e_), (e["start"], e["end"])) >= 0.5
                     for e in evs)  # fmt: skip
            cell = sens[t][bin_of(e_ - s_ + 1)]
            cell[0] += ok
            cell[1] += 1
            hit_t += ok
        hit_c = 0
        cnv_calls = [e for e in evs if e["call"] in CLASS]
        for e in cnv_calls:
            ok = any(t == CLASS[e["call"]] and c == e["chrom"] and ro((s_, e_), (e["start"], e["end"])) >= 0.5
                     for c, s_, e_, t in truth.get(ind, []))  # fmt: skip
            p = prec[CLASS[e["call"]]]
            p[0] += ok
            p[1] += 1
            hit_c += ok
        per_sample.append({"sample": sid, "individual": ind, "truth_detectable": len(tr), "truth_found": hit_t,
                           "cnv_calls": len(cnv_calls), "calls_in_truth": hit_c,
                           "loh_calls": sum(e["call_type"] == "LOH" for e in evs)})  # fmt: skip

    # replicate reproducibility of CNV and LOH calls
    by_ind: dict[str, list[str]] = defaultdict(list)
    for sid in calls:
        if sid in smap:
            by_ind[smap[sid]["individual"]].append(sid)
    rep = {"CNV": [0, 0], "LOH": [0, 0]}

    def kind_of(e: dict) -> str | None:
        if e["call_type"] == "LOH":
            return "LOH"
        return "CNV" if e["call"] in CLASS else None

    def matched(e: dict, others: list[dict], kind: str) -> bool:
        return any(
            o["chrom"] == e["chrom"]
            and kind_of(o) == kind
            and ro((o["start"], o["end"]), (e["start"], e["end"])) >= 0.5
            for o in others
        )

    for sids in by_ind.values():
        for i, s1 in enumerate(sids):
            for e in calls[s1]:
                kind = kind_of(e)
                if kind is None:
                    continue
                for j, s2 in enumerate(sids):
                    if j != i:
                        rep[kind][0] += matched(e, calls[s2], kind)
                        rep[kind][1] += 1

    def rate(v: list[int]) -> float | None:
        return round(v[0] / v[1], 4) if v[1] else None

    summary = {
        "samples": len(per_sample),
        "min_qual": a.min_qual,
        "sensitivity": {
            t: {b: {"found": v[0], "truth": v[1], "rate": rate(v)} for b, v in sorted(d.items())}
            for t, d in sens.items()
        },
        "precision": {t: {"in_truth": v[0], "calls": v[1], "rate": rate(v)} for t, v in prec.items()},
        "replicate_reproducibility": {
            k: {"matched": v[0], "comparisons": v[1], "rate": rate(v)} for k, v in rep.items()
        },
        "calls_per_sample_median": statistics.median([p["cnv_calls"] for p in per_sample]) if per_sample else None,
        "loh_calls_per_sample_median": statistics.median([p["loh_calls"] for p in per_sample]) if per_sample else None,
    }
    (a.out / "summary.json").write_text(json.dumps(summary, indent=1))
    with (a.out / "per_sample.tsv").open("w") as fh:
        w = csv.DictWriter(fh, fieldnames=list(per_sample[0]) if per_sample else ["sample"], delimiter="\t")
        w.writeheader()
        w.writerows(per_sample)
    print(json.dumps(summary, indent=1))
    return 0


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