"""Cyto model covariates — spec CYTO-038, IN-051 (plan §10.1). Pure functions, no I/O.

GC fraction of windows centred on each probe [LIT: Diskin et al. 2008 NAR 36:e126 (GC-wave), windows 1 kb / 100 kb /
1 Mb]: (G+C) / (A+C+G+T) over the window, N bases excluded; NaN when the window has no called base. Population
B-allele frequency per probe (PFB) [LIT: Wang et al. 2007 Genome Res 17:1665 (PennCNV)] from cluster-file counts:
(n_AB + 2·n_BB) / (2·N); 0.5 when unknown.
"""

from __future__ import annotations

import numpy as np
import numpy.typing as npt

GC_WINDOWS = (1_000, 100_000, 1_000_000)
_C, _G, _A, _T = ord("C"), ord("G"), ord("A"), ord("T")


def gc_fractions(
    seq: npt.NDArray[np.uint8], positions: npt.NDArray[np.int64], windows: tuple[int, ...] = GC_WINDOWS
) -> npt.NDArray[np.float32]:
    """GC fraction per (probe, window) for 1-based ``positions`` on one contig (upper-case ASCII ``seq``)."""
    n = len(seq)
    gc = np.zeros(n + 1, dtype=np.int32)
    called = np.zeros(n + 1, dtype=np.int32)
    step = 1 << 24
    run_gc = run_called = 0
    for s in range(0, n, step):  # chunked cumulative sums keep temporaries small
        chunk = np.asarray(seq[s : s + step])
        g = (chunk == _C) | (chunk == _G)
        c = g | (chunk == _A) | (chunk == _T)
        cg = np.cumsum(g, dtype=np.int32) + run_gc
        cc = np.cumsum(c, dtype=np.int32) + run_called
        gc[s + 1 : s + 1 + len(chunk)] = cg
        called[s + 1 : s + 1 + len(chunk)] = cc
        run_gc, run_called = int(cg[-1]), int(cc[-1])
    out = np.full((len(positions), len(windows)), np.nan, dtype=np.float32)
    centre = np.asarray(positions, dtype=np.int64) - 1  # 0-based
    for j, w in enumerate(windows):
        lo = np.clip(centre - w // 2, 0, n)
        hi = np.clip(centre - w // 2 + w, 0, n)
        num = gc[hi] - gc[lo]
        den = called[hi] - called[lo]
        ok = den > 0
        out[ok, j] = num[ok] / den[ok]
    return out


def population_bfreq(counts: npt.NDArray[np.int32] | None, n_loci: int) -> npt.NDArray[np.float32]:
    """PFB from cluster counts (n, 3) = (AA, AB, BB); 0.5 where there are no counts."""
    if counts is None:
        return np.full(n_loci, 0.5, dtype=np.float32)
    c = np.asarray(counts, dtype=np.float64)
    tot = c.sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        p = (c[:, 1] + 2 * c[:, 2]) / (2 * tot)
    return np.where(tot > 0, p, 0.5).astype(np.float32)


def reference_stats(
    lrr: npt.NDArray[np.float32], baf: npt.NDArray[np.float32]
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32], npt.NDArray[np.int32]]:
    """Batch reference (plan §10.1) from a samples × probes stack of normal samples: per-probe median LRR, per-probe
    median heterozygous BAF (samples with 0.2 < BAF < 0.8; 0.5 when none) and the number of such samples."""
    with np.errstate(all="ignore"):
        med_lrr = np.nanmedian(np.where(np.isfinite(lrr), lrr, np.nan), axis=0)
        het = np.where((baf > 0.2) & (baf < 0.8), baf, np.nan)
        n_het = np.isfinite(het).sum(axis=0).astype(np.int32)
        het_c = np.where(n_het > 0, np.nanmedian(het, axis=0), 0.5)
    return (
        np.nan_to_num(med_lrr, nan=0.0).astype(np.float32),
        np.nan_to_num(het_c, nan=0.5).astype(np.float32),
        n_het,
    )