"""Targeted PGx copy number (spec PCN-001 … PCN-017, CLI-031/032, IN-050)."""

from __future__ import annotations

import gzip
from pathlib import Path

import numpy as np
from typer.testing import CliRunner

from simurg.cli.main import app
from simurg.formats.gtc import Gtc, write_gtc
from simurg.pgx.cn import TEMPLATE, call_cn, fit_region, hwe_weights, parse_regions
from simurg.pipeline.pgx_cn import default_regions
from simurg.synth.files import SynthLocus, encode_bpm


def _cluster_values(rng: np.random.Generator, cns: np.ndarray, delta: float, gamma: float, sd: float,
                    cn0_drop: float = 0.0) -> np.ndarray:  # fmt: skip
    x = delta + gamma * TEMPLATE[cns] + rng.normal(0, sd, len(cns))
    return np.where(cns == 0, x + cn0_drop, x)


def _hwe_cns(rng: np.random.Generator, n: int, f: list[float]) -> np.ndarray:
    return rng.choice(len(f), n, p=f) + rng.choice(len(f), n, p=f)


def test_hwe_weights() -> None:  # PCN-009 [OURS]
    w = hwe_weights(np.array([0.3, 0.6, 0.1, 0.0]))
    assert abs(w.sum() - 1) < 1e-12
    assert np.allclose(w[:4], [0.09, 0.36, 0.06 + 0.36, 0.12])


def test_labels_identified() -> None:  # PCN-009: GSTM1-like, GSTT1-like and one-cluster regions
    rng = np.random.default_rng(1)
    cns = _hwe_cns(rng, 60, [0.7, 0.3])  # common deletion: mostly CN 0 and 1
    x = _cluster_values(rng, cns, 0.1, 0.22, 0.02)
    f = fit_region(x)
    assert np.mean([call_cn(v, f)[0] for v in x] == cns) > 0.95
    cns = _hwe_cns(rng, 60, [0.45, 0.55])
    x = _cluster_values(rng, cns, 0.3, 0.55, 0.02, cn0_drop=-1.2)  # CN 0 far below the template
    f = fit_region(x)
    assert f.cn0_shift < -0.5
    assert np.mean([call_cn(v, f)[0] for v in x] == cns) > 0.95
    cns = np.full(40, 2)
    cns[:2] = 1
    x = _cluster_values(rng, cns, -0.05, 0.25, 0.015)
    f = fit_region(x)
    assert [call_cn(v, f)[0] for v in x] == cns.tolist()


def test_packaged_regions() -> None:  # PCN-001
    name, text = default_regions()
    regs = parse_regions(text)
    assert name.endswith(".v1.bed") and len(regs) == 10
    assert {r.gene for r in regs} == {"CYP2D6", "CYP2A6", "GSTM1", "GSTT1", "UGT2B17", "SULT1A1", "CYP2E1"}
    assert next(r for r in regs if r.gene == "GSTT1").chrom == "chr22_KI270879v1_alt"


def _dataset(tmp: Path, n_samples: int = 40, n_outliers: int = 0) -> tuple[Path, Path, Path, Path]:
    rng = np.random.default_rng(5)
    loci = []
    for c in ("1", "2"):
        for k in range(1500):
            loci.append(SynthLocus(f"cnv{c}_{k}", c, 2000 * (k + 1), "[A/G]", len(loci) + 1, 0, 0))
    bpm = tmp / "m.bpm"
    bpm.write_bytes(encode_bpm(loci))
    fa = tmp / "g.fa"
    body = {c: "".join(rng.choice(list("ACGT"), 1500 * 2000 + 10_000)) for c in ("1", "2")}
    fa.write_text("".join(f">{c}\n" + "\n".join(s[i : i + 60] for i in range(0, len(s), 60)) + "\n"
                          for c, s in body.items()))  # fmt: skip
    regions = tmp / "regions.bed"
    regions.write_text("# test regions\n1\t100000\t160000\tGENEA\tGENEA\n2\t300000\t340000\tGENEB\tGENEB\n"
                       "3\t1000\t5000\tGENEC\tGENEC\n")  # fmt: skip
    n = len(loci)
    pos = np.array([lc.pos for lc in loci])
    chrom = np.array([lc.chrom for lc in loci])
    in_a = (chrom == "1") & (pos > 100_000) & (pos <= 160_000)
    in_b = (chrom == "2") & (pos > 300_000) & (pos <= 340_000)
    cn_a = _hwe_cns(rng, n_samples, [0.5, 0.5])
    gtcs = tmp / "gtcs"
    truth = {}
    for i in range(n_samples):
        sid = f"2000000000{i:02d}_R01C01"
        lrr = rng.normal(0, 0.08, n)
        lrr[in_a] += 0.24 * TEMPLATE[cn_a[i]] + (-0.2 if cn_a[i] == 0 else 0.0)
        cn_b = 3 if i == 0 else 2
        lrr[in_b] += 0.25 * TEMPLATE[cn_b]
        truth[sid] = (int(cn_a[i]), cn_b)
        g = Gtc(num_snps=n, genotypes=np.ones(n, np.uint8), genotype_scores=np.full(n, 0.9, np.float32),
                b_allele_freqs=np.zeros(n, np.float32), logr_ratios=lrr.astype(np.float32),
                raw_x=np.full(n, 1000, np.uint16), raw_y=np.full(n, 1000, np.uint16),
                base_calls=np.zeros((n, 2), np.uint8), gender=b"F", call_rate=0.995, logr_dev=0.12,
                cluster_file="c.egt")  # fmt: skip
        (gtcs / sid).mkdir(parents=True)
        write_gtc(g, gtcs / sid / f"{sid}.gtc")
    for j in range(n_outliers):  # EDGE-026: noisy, shifted arrays that still pass the LRR SD gate
        sid = f"2000000099{j:02d}_R01C01"
        lrr = rng.normal(0.3, 0.19, n)
        g = Gtc(num_snps=n, genotypes=np.ones(n, np.uint8), genotype_scores=np.full(n, 0.9, np.float32),
                b_allele_freqs=np.zeros(n, np.float32), logr_ratios=lrr.astype(np.float32),
                raw_x=np.full(n, 1000, np.uint16), raw_y=np.full(n, 1000, np.uint16),
                base_calls=np.zeros((n, 2), np.uint8), gender=b"F", call_rate=0.995, logr_dev=0.19,
                cluster_file="c.egt")  # fmt: skip
        (gtcs / sid).mkdir(parents=True)
        write_gtc(g, gtcs / sid / f"{sid}.gtc")
    (tmp / "truth.txt").write_text("\n".join(f"{k}\t{a}\t{b}" for k, (a, b) in truth.items()))
    return bpm, fa, gtcs, regions


def test_pgx_cn_train_and_call(tmp_path: Path) -> None:  # CLI-031/032, PCN-002/004/011/012/013, IN-050
    bpm, fa, gtcs, regions = _dataset(tmp_path)
    run = CliRunner()
    r = run.invoke(app, ["pgx", "copy-number", "train", "--gtcs", str(gtcs), "--bpm", str(bpm), "--fasta", str(fa),
                         "--out", str(tmp_path / "model"), "--regions", str(regions), "--platform", "EX",
                         "--quiet"])  # fmt: skip
    assert r.exit_code == 2 and "E_PGX_TRAIN_TOO_SMALL" in r.output  # 40 < 96 (PCN-005)
    r = run.invoke(app, ["pgx", "copy-number", "train", "--gtcs", str(gtcs), "--bpm", str(bpm), "--fasta", str(fa),
                         "--out", str(tmp_path / "model"), "--regions", str(regions), "--min-samples", "30", "--platform", "EX",
                         "--min-passing", "30", "--quiet"])  # fmt: skip
    assert r.exit_code == 0, r.output
    model = tmp_path / "model" / "m.simurg-pgxcn"
    assert model.exists() and (tmp_path / "model" / "train_exclusions.tsv").exists()
    out = tmp_path / "call"
    r = run.invoke(app, ["pgx", "copy-number", "call", "--gtcs", str(gtcs), "--model", str(model), "--fasta", str(fa),
                         "--out", str(out), "--min-batch", "50", "--quiet"])  # fmt: skip
    assert r.exit_code == 2 and "E_PGX_BATCH_TOO_SMALL" in r.output
    r = run.invoke(app, ["pgx", "copy-number", "call", "--gtcs", str(gtcs), "--model", str(model), "--fasta", str(fa),
                         "--out", str(out), "--quiet"])  # fmt: skip
    assert r.exit_code == 0, r.output
    truth = {ln.split("\t")[0]: (int(ln.split("\t")[1]), int(ln.split("\t")[2]))
             for ln in (tmp_path / "truth.txt").read_text().splitlines()}  # fmt: skip
    ok_a = ok_b = 0
    for sid, (a, b) in truth.items():
        with gzip.open(out / f"{sid}.pgx_cn.vcf.gz", "rt") as fh:
            text = fh.read()
        assert (out / f"{sid}.pgx_cn.vcf.gz.tbi").exists()
        recs = [ln.split("\t") for ln in text.splitlines() if not ln.startswith("#")]
        assert [x[7].split("REGION=")[1].split(";")[0] for x in recs] == ["GENEA", "GENEB", "GENEC"]  # PCN-002
        for x in recs:  # PCN-011: POS = S − 1, ID coordinates = POS + 1 … END
            end = int(x[7].split("END=")[1].split(";")[0])
            s_, e_ = x[2].split(":")[-1].split("-")
            assert int(s_) == int(x[1]) + 1 and int(e_) == end and x[8] == "CN:NR"
        genec = recs[2]
        assert genec[6] == "NoCall" and genec[9].startswith(".:")
        ok_a += int(recs[0][9].split(":")[0]) == a
        ok_b += int(recs[1][9].split(":")[0]) == b
    assert ok_a >= 0.95 * len(truth) and ok_b >= 0.9 * len(truth), (ok_a, ok_b)
    rows = [ln.split("\t") for ln in (out / "pgx_cn_batch_summary.tsv").read_text().splitlines()]
    assert rows[0] == ["region", "gene", "n_gain", "n_loss", "n_neutral", "n_nocall", "n_filtered"]
    assert [r_[0] for r_ in rows[1:]] == ["GENEA", "GENEB", "GENEC", "ALL"]
    for r_ in rows[1:-1]:  # PCN-012: the four counts add up to N
        assert sum(int(v) for v in r_[2:6]) == len(truth)
    head = (out / "pgx_cn_sample_summary.tsv").read_text().splitlines()[0].split("\t")
    assert head[:6] == ["sample_id", "sample_name", "sample_folder", "logr_dev", "passed_batch_qc",
                        "used_in_reference"] and head[-1] == "flags"  # fmt: skip


def _calls(out: Path) -> dict[str, str]:
    rows = [ln.split("\t") for ln in (out / "pgx_cn_sample_summary.tsv").read_text().splitlines()]
    head = rows[0]
    return {r[0]: ":".join(r[head.index("cn_GENEA")] + r[head.index("cn_GENEB")]) for r in rows[1:]}


def test_outlier_robustness(tmp_path: Path) -> None:  # PCN-017 / EDGE-026: 10 % outliers leave other calls alone
    run = CliRunner()
    res = {}
    for tag, n_out in (("clean", 0), ("outliers", 4)):
        d = tmp_path / tag
        d.mkdir()
        bpm, fa, gtcs, regions = _dataset(d, 40, n_out)
        r = run.invoke(app, ["pgx", "copy-number", "train", "--gtcs", str(gtcs), "--bpm", str(bpm), "--fasta", str(fa),
                             "--out", str(d / "model"), "--regions", str(regions), "--min-samples", "30",
                             "--min-passing", "30", "--platform", "EX", "--quiet"])  # fmt: skip
        assert r.exit_code == 0, r.output
        r = run.invoke(app, ["pgx", "copy-number", "call", "--gtcs", str(gtcs), "--model", str(d / "model" / "m.simurg-pgxcn"),
                             "--fasta", str(fa), "--out", str(d / "call"), "--quiet"])  # fmt: skip
        assert r.exit_code == 0, r.output
        res[tag] = _calls(d / "call")
        assert (d / "call" / "pgx_cn_baf_evidence.tsv").exists()  # PCN-010 output
    same = sum(res["clean"][k] == res["outliers"][k] for k in res["clean"])
    assert same >= 0.99 * len(res["clean"])


def test_platform_required(tmp_path: Path) -> None:  # CLI-032: --platform auto needs a CSV manifest
    bpm, fa, gtcs, regions = _dataset(tmp_path, 30)
    r = CliRunner().invoke(app, ["pgx", "copy-number", "train", "--gtcs", str(gtcs), "--bpm", str(bpm), "--fasta",
                                 str(fa), "--out", str(tmp_path / "m"), "--regions", str(regions), "--quiet"])  # fmt: skip
    assert r.exit_code == 2 and "E_PGX_PLATFORM_UNKNOWN" in r.output