"""Targeted PGx copy number — spec PCN-001, PCN-009. Pure functions, no I/O.

Region value x = median over the region's unmasked probes of (sample-centred LRR − the probe's training reference),
minus the sample's scale (median of the same quantity over non-target autosomal probes).

CN mixture [OURS; array intensity compression as in cnvPartition/PennCNV, LIT: Wang et al. 2007 Genome Res 17:1665]:
component means µ_c = δ_r + γ_r · log2((c + 1) / 3) for CN c = 0 … 5 — a probe signal proportional to (c + β) with
a background β = 1 copy-equivalent, compressed by γ_r. δ_r (the region's CN-2 location), γ_r, σ_r and the weights
π_r are per region; γ_r has a Gaussian prior N(0.25, 0.08²). CN 0 may sit lower than the template (its own shift
ε₀ ≤ 0 and SD): GSTT1 probes drop by ~2 at CN 0, GSTM1 probes by ~0.3. The weights are constrained to Hardy–Weinberg
proportions of haplotype copy numbers (0 … 3 copies per haplotype, Dirichlet prior favouring one copy). Fitted by
EM from a grid of starts (each observed cluster labelled CN 0 … 4, CN 0 on the template or on the lowest values).
HWE and the template's spacing ratios identify the labels in multi-cluster regions (GSTM1 0/1/2 rather than 1/2/3,
which would imply unobserved CN-0 samples). On GSA v4 PGx CNV probes γ_r ≈ 0.19–0.56 (K-028).
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

F64 = npt.NDArray[np.float64]
CN_MAX = 5  # "5 or more"
CNS = np.arange(CN_MAX + 1)
TEMPLATE = np.log2((CNS + 1) / 3.0)
HAP_ALPHA = np.array([1.2, 3.0, 1.2, 1.05])  # Dirichlet prior on haplotype copy frequencies f_0..f_3 (f_1 favoured)
SIGMA_FLOOR = 0.012
GAMMA_PRIOR = (0.25, 0.08)  # mean, SD
GAMMA_RANGE = (0.15, 0.80)
PI_FLOOR = 0.002
CN0_PRIOR_N = 2.0


@dataclass(frozen=True)
class Region:
    chrom: str
    start0: int
    end: int
    name: str
    gene: str


def parse_regions(text: str) -> list[Region]:
    """PCN-001 BED: ``chrom start0 end region gene``; ``#`` comments."""
    out = []
    for ln in text.splitlines():
        if not ln.strip() or ln.startswith("#"):
            continue
        c, s, e, r, g = ln.split("\t")[:5]
        out.append(Region(c, int(s), int(e), r, g))
    return out


@dataclass
class RegionFit:
    delta: float
    gamma: float
    sigma: float
    weights: list[float]
    n: int
    cn0_shift: float = 0.0  # ≤ 0: extra drop of the CN-0 component below the template (probe-specific background)
    cn0_sigma: float = 0.0  # SD of the CN-0 component (≥ sigma; 0 = sigma)

    def means(self) -> F64:
        m = self.delta + self.gamma * TEMPLATE
        m[0] += self.cn0_shift
        return m

    def sds(self) -> F64:
        s = np.full(len(CNS), self.sigma)
        s[0] = max(self.sigma, self.cn0_sigma)
        return s


def hwe_weights(f: F64) -> F64:
    """CN weights under Hardy–Weinberg from haplotype copy-number frequencies f_k (k = 0 … 3 copies):
    π_c = Σ_{i+j=c} f_i f_j, CN ≥ 5 pooled into 5 [LIT: HWE; OURS: use as the CN mixture prior]."""
    pi = np.zeros(len(CNS))
    for i in range(len(f)):
        for j in range(len(f)):
            pi[min(i + j, CN_MAX)] += f[i] * f[j]
    return pi


def _hap_update(r: F64, f: F64) -> F64:
    """M-step for f: split each sample's CN responsibilities over the haplotype pairs (i, j) with i + j = c."""
    counts = HAP_ALPHA - 1.0
    counts = counts.astype(np.float64).copy()
    pi = hwe_weights(f)
    for i in range(len(f)):
        for j in range(len(f)):
            c = min(i + j, CN_MAX)
            share = f[i] * f[j] / max(pi[c], 1e-300)
            w = float(r[:, c].sum()) * share
            counts[i] += w
            counts[j] += w
    counts = np.maximum(counts, 1e-9)
    out: F64 = counts / counts.sum()
    return out


def _norm_pdf(x: F64, mu: F64, sd: float | F64) -> F64:
    sd_ = np.broadcast_to(np.asarray(sd, dtype=np.float64), mu.shape)[None, :]
    return np.exp(-0.5 * ((x[:, None] - mu[None, :]) / sd_) ** 2) / (sd_ * math.sqrt(2 * math.pi))


def _em_region(x: F64, gamma: float, delta: float, iters: int, eps0: float = 0.0) -> tuple[RegionFit, float]:
    """EM for one region; returns the fit and its penalised log-likelihood (HWE weights from haplotype copy
    frequencies with a Dirichlet prior, Gaussian prior on γ).
    CN 0 has its own shift ε₀ ≤ 0 below the template and its own SD, both shrunk towards the template with a prior
    weight of ``CN0_PRIOR_N`` samples (on some probe sets CN 0 is far lower than the compressed template, GSTT1)."""
    n = len(x)
    f = np.array([0.1, 0.8, 0.07, 0.03])
    pi = hwe_weights(f)
    sd = max(SIGMA_FLOOR, float(np.std(x)) if n > 1 else 0.05)
    sd0 = sd
    if eps0 < 0:  # a start with a far CN-0 cluster: the other components get the spread of the rest
        rest = x[x > delta + gamma * TEMPLATE[0] + eps0 / 2]
        sd = max(SIGMA_FLOOR, float(np.std(rest)) if len(rest) > 1 else 0.05)
        sd0 = sd
    g0, gs = GAMMA_PRIOR
    ll = -math.inf
    for _ in range(iters):
        mu = delta + gamma * TEMPLATE
        mu[0] += eps0
        sds = np.full(len(CNS), sd)
        sds[0] = sd0
        dens = pi[None, :] * _norm_pdf(x, mu, sds) + 1e-300
        tot = dens.sum(axis=1, keepdims=True)
        ll = float(np.log(tot).sum() + ((HAP_ALPHA - 1) * np.log(np.maximum(f, 1e-12))).sum()
                   - 0.5 * ((gamma - g0) / gs) ** 2)  # fmt: skip
        r = dens / tot
        f = _hap_update(r, f)
        pi = hwe_weights(f)
        # δ and γ: weighted least squares of x on (1, t_c) over CN ≥ 1, γ shrunk towards its prior
        rr = r.copy()
        rr[:, 0] = 0.0
        w = rr.sum(axis=0)
        sw, st, stt = float(w.sum()), float((w * TEMPLATE).sum()), float((w * TEMPLATE**2).sum())
        sx, stx = float((rr * x[:, None]).sum()), float((rr * x[:, None] * TEMPLATE[None, :]).sum())
        lam = sd**2 / gs**2
        det = sw * (stt + lam) - st * st
        if det > 0:
            delta = (sx * (stt + lam) - st * (stx + lam * g0)) / det
            gamma = float(np.clip((sw * (stx + lam * g0) - st * sx) / det, *GAMMA_RANGE))
        mu = delta + gamma * TEMPLATE
        n1 = float(w.sum())
        if n1 > 0:
            sd = max(SIGMA_FLOOR, math.sqrt(float((rr * (x[:, None] - mu[None, :]) ** 2).sum() / n1)))
        r0 = r[:, 0]
        n0 = float(r0.sum())
        eps0 = min(0.0, float((r0 * (x - mu[0])).sum()) / (n0 + CN0_PRIOR_N))
        v0 = float((r0 * (x - mu[0] - eps0) ** 2).sum())
        sd0 = max(sd, math.sqrt((v0 + CN0_PRIOR_N * sd**2) / (n0 + CN0_PRIOR_N)))
    return RegionFit(float(delta), float(gamma), sd, pi.tolist(), n, float(eps0), float(sd0)), ll


def fit_region(x: F64, iters: int = 200) -> RegionFit:
    """PCN-009 step 4 for one region: the best of EM runs started with each observed cluster labelled CN 0 … 4, each
    with CN 0 on the template or on the lowest values."""
    g0 = GAMMA_PRIOR[0]
    cands = sorted({round(float(q) - g0 * float(t), 4) for q in np.quantile(x, [0.1, 0.3, 0.5, 0.7, 0.9])
                    for t in TEMPLATE[:5]})  # fmt: skip
    lo = float(np.quantile(x, 0.02))
    best: tuple[float, float, float] | None = None
    for d0 in cands:
        for e0 in sorted({0.0, round(min(0.0, lo - (d0 + g0 * float(TEMPLATE[0]))), 4)}):
            _, ll = _em_region(x, g0, d0, 80, e0)
            if best is None or ll > best[0] + 1e-9:
                best = (ll, d0, e0)
    assert best is not None
    return _em_region(x, g0, best[1], iters, best[2])[0]


def fit_mixture(values: dict[str, F64], iters: int = 200) -> dict[str, RegionFit]:
    """Per-region fits for every region with ≥ 3 finite values."""
    out = {}
    for name, x in values.items():
        x = x[np.isfinite(x)]
        if len(x) >= 3:
            out[name] = fit_region(x, iters)
    return out


def posterior(x: float, fit: RegionFit) -> F64:
    """P(CN = c | x) for c = 0 … 5, with the training weights (floored) as prior."""
    pi = np.maximum(np.asarray(fit.weights), PI_FLOOR)
    pi = pi / pi.sum()
    d = pi * _norm_pdf(np.array([x]), fit.means(), fit.sds())[0]
    s = float(d.sum())
    out: F64 = d / s if s > 0 else pi
    return out


def baf_loglik(baf: F64, pfb: F64 | None = None, sd: float = 0.05) -> F64:
    """PCN-010: per-CN-state log-likelihood of a region's SNP BAFs: for CN c the BAF is k/c with binomial weights
    C(c, k) p^k (1 − p)^(c − k) from the probe's population B-allele frequency p (0.5 when unknown); CN 0 → uniform
    [LIT: Wang et al. 2007 (PennCNV BAF emissions); OURS: region-level use]."""
    ok = np.isfinite(baf)
    b = baf[ok]
    p = np.clip(np.nan_to_num(pfb[ok] if pfb is not None else np.full(len(b), 0.5), nan=0.5), 0.01, 0.99)
    out = np.zeros(len(CNS))
    if len(b) == 0:
        return out
    for c in CNS:
        if c == 0:
            out[c] = 0.0  # uniform on [0, 1]
            continue
        k = np.arange(c + 1)
        means = k / c
        w = np.array([math.comb(c, int(j)) for j in k])[None, :] * p[:, None] ** k[None, :] * (1 - p[:, None]) ** (
            c - k[None, :])  # fmt: skip
        dens = np.exp(-0.5 * ((b[:, None] - means[None, :]) / sd) ** 2) / (sd * math.sqrt(2 * math.pi))
        out[c] = float(np.log(0.98 * (w * dens).sum(axis=1) + 0.02).sum())
    return out


def call_cn(x: float, fit: RegionFit, baf_ll: F64 | None = None) -> tuple[int, int, F64]:
    """PCN-009 steps 5–6: (CN, QUAL, posterior); ``baf_ll`` (PCN-010, experimental ``--bafs-in-cn``) is added to the
    log posterior."""
    p = posterior(x, fit)
    if baf_ll is not None:
        lp = np.log(np.maximum(p, 1e-300)) + baf_ll
        p = np.exp(lp - lp.max())
        p = p / p.sum()
    cn = int(np.argmax(p))
    pmax = float(p[cn])
    qual = 60 if pmax >= 1.0 else int(min(60, max(0, round(-10 * math.log10(max(1e-300, 1 - pmax))))))
    return cn, qual, p