#!/usr/bin/env python3
"""Sensitivity of `simurg cn call` on implanted events (spec CYTO-022, plan §16.4).

Real GTCs of normal samples are copied and events are written into their LRR/BAF:
  constitutional DEL (CN1: LRR − 0.45, heterozygous BAF → 0/1), DUP (CN3: LRR + 0.3, het → 1/3 or 2/3), LOH (het → 0/1);
  mosaic loss / gain / CN-LOH with cell fraction f (LRR shifted by f × the constitutional shift, heterozygous BAF
  moved to 0.5 ± d(f), sign random per probe as in unphased data), keeping each probe's original noise.
Every synthetic sample gets several events on different chromosomes. `simurg cn call` then runs on them and an
implanted event counts as detected when a called event of the same class (loss / gain / LOH) overlaps ≥ 50 % of it.

    python scripts/v2/cn_synthetic_eval.py --bpm M.bpm --model MODEL --gtcs BASE_DIR --out OUT [--threads N]
"""

from __future__ import annotations

import argparse
import copy
import gzip
import json
import subprocess
import sys
from pathlib import Path

import numpy as np

from simurg.formats.bpm import read_bpm
from simurg.formats.discovery import discover_gtcs
from simurg.formats.gtc import read_gtc, write_gtc

KINDS = [("DEL", 1.0), ("DUP", 1.0), ("LOH", 1.0)] + [
    (k, f) for k in ("MDEL", "MDUP", "MLOH") for f in (0.05, 0.10, 0.15, 0.20, 0.30)
]
SIZES = (1_000_000, 5_000_000, 20_000_000)
CLASS = {"DEL": "loss", "MDEL": "loss", "DUP": "gain", "MDUP": "gain", "LOH": "loh", "MLOH": "loh", "GAINLOH": "gain"}
CHROMS = [str(c) for c in range(1, 21)]  # leave 21, 22, X, Y untouched


def implant(lrr: np.ndarray, baf: np.ndarray, idx: np.ndarray, kind: str, f: float, rng: np.random.Generator) -> None:
    het = (baf[idx] > 0.15) & (baf[idx] < 0.85)
    h = idx[het]
    noise = baf[h] - 0.5
    sign = rng.choice([-1.0, 1.0], len(h))
    if kind in ("DEL", "LOH"):
        if kind == "DEL":
            lrr[idx] -= 0.45
        baf[h] = np.clip(np.where(sign > 0, 1.0, 0.0) + noise * 0.2, 0, 1)
    elif kind == "DUP":
        lrr[idx] += 0.30
        baf[h] = np.clip(np.where(sign > 0, 2 / 3, 1 / 3) + noise, 0, 1)
    else:
        d = {"MDEL": f / (2 * (2 - f)), "MDUP": f / (2 * (2 + f)), "MLOH": f / 2}[kind]
        lrr[idx] += {"MDEL": -0.45 * f, "MDUP": 0.30 * f, "MLOH": 0.0}[kind]
        baf[h] = np.clip(0.5 + sign * d + noise, 0, 1)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--bpm", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--gtcs", required=True, type=Path, help="normal base samples (not in the model's reference)")
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--threads", type=int, default=0)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--reps", type=int, default=4, help="replicates of every (kind, fraction, size) cell")
    a = ap.parse_args()
    rng = np.random.default_rng(a.seed)
    man = read_bpm(a.bpm)
    chroms = np.array(man.chroms)
    pos = man.map_infos.astype(np.int64)
    base = sorted(discover_gtcs(a.gtcs).values())
    syn = a.out / "gtcs"
    syn.mkdir(parents=True, exist_ok=True)
    events = [(k, f, s) for k, f in KINDS for s in SIZES for _ in range(a.reps)]
    rng.shuffle(events)
    truth = []
    per_sample = len(CHROMS) // 2  # one event per chromosome pair slot, spaced out
    for n, start in enumerate(range(0, len(events), per_sample)):
        chunk = events[start : start + per_sample]
        src = base[n % len(base)]
        g = copy.deepcopy(read_gtc(src))
        lrr = g.logr_ratios.astype(np.float64)
        baf = g.b_allele_freqs.astype(np.float64)
        chosen = rng.permutation(CHROMS)[: len(chunk)]
        sid = f"syn{n:03d}"
        for (kind, f, size), c in zip(chunk, chosen, strict=True):
            sel = np.flatnonzero((chroms == c) & (pos > 0))
            lo, hi = int(pos[sel].min()), int(pos[sel].max())
            s0 = int(rng.integers(lo + 2_000_000, max(lo + 2_000_001, hi - size - 2_000_000)))
            idx = sel[(pos[sel] >= s0) & (pos[sel] < s0 + size)]
            implant(lrr, baf, idx, kind, f, rng)
            truth.append({"sample": sid, "chrom": c, "start": s0, "end": s0 + size - 1, "kind": kind, "fraction": f,
                          "size": size, "probes": len(idx), "base": src.stem})  # fmt: skip
        g.logr_ratios = lrr.astype(np.float32)
        g.b_allele_freqs = baf.astype(np.float32)
        (syn / sid).mkdir(exist_ok=True)
        write_gtc(g, syn / sid / f"{sid}.gtc")
    (a.out / "truth.json").write_text(json.dumps(truth, indent=1))
    simurg = str(Path(sys.executable).with_name("simurg"))
    cmd = [simurg, "cn", "call", "--bpm", a.bpm, "--model", a.model, "--gtcs", str(syn), "--out", str(a.out / "call"),
           "--threads", str(a.threads), "--unknown-sex-policy", "female", "--force", "--quiet"]  # fmt: skip
    subprocess.run(cmd, check=True)
    calls: dict[str, list[tuple[str, int, int, str, str]]] = {}
    for vcf in (a.out / "call").glob("*.cyto.cnv.vcf.gz"):
        sid = vcf.name.split(".")[0]
        with gzip.open(vcf, "rt") as fh:
            lines = [ln for ln in fh if not ln.startswith("#")]
        for line in lines:
            c, p, vid, _, _alt, _q, filt, info = line.split("\t")[:8]
            end = int(next(x for x in info.split(";") if x.startswith("END="))[4:])
            call = vid.split(":")[0]
            cls = {"DEL": "loss", "MosaicDEL": "loss", "DUP": "gain", "MosaicDUP": "gain", "AOH": "loh",
                   "CNLOH": "loh", "GAINLOH": "gain"}[call]  # fmt: skip
            calls.setdefault(sid, []).append((c, int(p) + 1, end, cls, filt))
    for t in truth:
        cover = 0
        seen = []
        for c, s, e, cls, filt in calls.get(t["sample"], []):
            ov = max(0, min(e, t["end"]) - max(s, t["start"]) + 1) if c == t["chrom"] else 0
            if ov:
                seen.append(f"{cls}:{filt}:{ov / (t['end'] - t['start'] + 1):.2f}")
                if cls == CLASS[t["kind"]]:
                    cover += ov  # fragments of one event add up
        frac = cover / (t["end"] - t["start"] + 1)
        t["detected"] = frac >= 0.5
        t["overlap"] = round(frac, 3)
        t["overlapping_calls"] = seen
    cells: dict[str, dict[str, list[bool]]] = {}
    for t in truth:
        key = t["kind"] if t["fraction"] == 1.0 else f"{t['kind']} f={t['fraction']:.2f}"
        cells.setdefault(key, {}).setdefault(f"{t['size'] // 1_000_000} Mb", []).append(t["detected"])
    table = {k: {s: f"{sum(v)}/{len(v)}" for s, v in d.items()} for k, d in cells.items()}
    n_calls = sum(len(v) for v in calls.values())
    implanted_hits = sum(t["detected"] for t in truth)
    summary = {"events": len(truth), "detected": implanted_hits, "table": table, "calls_total": n_calls}
    (a.out / "summary.json").write_text(json.dumps(summary, indent=1))
    (a.out / "truth_detected.json").write_text(json.dumps(truth, indent=1))
    for k, v in table.items():
        print(f"{k:14s} " + "  ".join(f"{s}: {v.get(s, '-')}" for s in ("1 Mb", "5 Mb", "20 Mb")))
    print(f"detected {implanted_hits}/{len(truth)}; {n_calls} calls in {len(calls)} samples")
    for t in truth:
        if not t["detected"] and t["kind"] in ("DUP", "DEL", "LOH") and t["size"] >= 5_000_000:
            print("  missed", t["kind"], t["size"], t["probes"], "probes:", t["overlapping_calls"])
    return 0


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