SimurgArray/src/simurg/signal/raw_xy.py

"""Raw X/Y assembly from a Red/Grn IDAT pair — plan §7.1, spec IN-027.

For every manifest locus (BPM names order, i.e. GTC order):

* Infinium II (assay 0):      X = Red.Mean[addrA], Y = Grn.Mean[addrA]
* Infinium I A/T (assay 1):   X = Red.Mean[addrA], Y = Red.Mean[addrB]
* Infinium I C/G (assay 2):   X = Grn.Mean[addrA], Y = Grn.Mean[addrB]

A locus is (0, 0) — "missing" — unless every address it uses is present and has NBeads >= 2 in the channel(s)
read. Control probes (BPM control config, flattened in order) take X = Red, Y = Grn with no bead-count filter.
[OSS: gtc2vcf ``fill_array`` / ``fill_controls_array`` (MIT); LIT: Steemers 2006 for the chemistry]
"""

from __future__ import annotations

import hashlib
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

from simurg.core.errors import SimurgError
from simurg.formats.bpm import ASSAY_INF_I_AT, ASSAY_INF_I_GC, ASSAY_INF_II, Manifest
from simurg.formats.idat import Idat

MIN_BEADS = 2
COVERAGE_ERROR = 0.99


@dataclass
class RawIntensities:
    """Per-sample raw intensities in manifest order."""

    raw_x: npt.NDArray[np.uint16]
    raw_y: npt.NDArray[np.uint16]
    missing: npt.NDArray[np.bool_]  # address absent or NBeads < MIN_BEADS
    controls_x: npt.NDArray[np.uint16]
    controls_y: npt.NDArray[np.uint16]
    control_addresses: npt.NDArray[np.int32]

    @property
    def n_missing(self) -> int:
        return int(self.missing.sum())


class AddressIndex:
    """Vectorised address → row lookup for one IDAT (sorted-array search, no per-address Python loop)."""

    def __init__(self, idat: Idat) -> None:
        ids = np.asarray(idat.address_ids, dtype=np.int64)
        self.order = np.argsort(ids, kind="stable")
        self.sorted = ids[self.order]

    def lookup(self, addresses: npt.ArrayLike) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.bool_]]:
        """Row index of each address (0 where absent) and a presence mask."""
        a = np.asarray(addresses, dtype=np.int64)
        if len(self.sorted) == 0:
            return np.zeros(len(a), dtype=np.intp), np.zeros(len(a), dtype=np.bool_)
        pos = np.minimum(np.searchsorted(self.sorted, a), len(self.sorted) - 1)
        found = self.sorted[pos] == a
        return np.where(found, self.order[pos], 0), found


def control_addresses(manifest: Manifest) -> npt.NDArray[np.int32]:
    """Flattened control addresses from the BPM control config (``a:b:c:d,Category,Color,Name`` per line)."""
    out: list[int] = []
    for line in manifest.control_config.splitlines():
        head = line.split(",", 1)[0].strip()
        if not head:
            continue
        for part in head.split(":"):
            part = part.strip()
            if part:
                out.append(int(part))
    return np.asarray(out, dtype=np.int32)


def manifest_addresses(manifest: Manifest) -> npt.NDArray[np.int32]:
    """Every bead address the manifest's loci read (A for all loci, B for Infinium I)."""
    b = manifest.address_b[manifest.address_b != 0]
    return np.unique(np.concatenate([manifest.address_a, b]))


def address_coverage(idat: Idat, manifest: Manifest) -> float:
    """Fraction of manifest bead addresses present in the IDAT (IN-027)."""
    wanted = manifest_addresses(manifest)
    if len(wanted) == 0:
        return 1.0
    _, found = AddressIndex(idat).lookup(wanted)
    return float(found.mean())


def check_coverage(
    idat: Idat, manifest: Manifest, cache: LocusMapCache | None = None
) -> tuple[float, list[tuple[str, str]]]:
    """IN-027: < 0.99 → E_IDAT_MANIFEST_MISMATCH (raised); < 1 → W_IDAT_LOW_ADDRESS_COVERAGE (returned)."""
    cov = cache.coverage(idat) if cache is not None and cache.manifest is manifest else address_coverage(idat, manifest)
    notes: list[tuple[str, str]] = []
    if cov < COVERAGE_ERROR:
        raise SimurgError(
            "E_IDAT_MANIFEST_MISMATCH",
            f"{idat.path}: only {cov:.2%} of the addresses of manifest {manifest.manifest_name!r} are present "
            f"(chip type {idat.chip_type!r}); the IDATs and the manifest belong to different products",
            path=idat.path,
        )
    if cov < 1.0:
        notes.append(
            (
                "W_IDAT_LOW_ADDRESS_COVERAGE",
                f"{idat.path}: {1 - cov:.4%} of manifest addresses are absent from the IDAT",
            )
        )
    if (
        idat.manifest
        and manifest.manifest_name
        and idat.manifest.split(".")[0].lower() not in manifest.manifest_name.lower()
    ):
        notes.append(
            (
                "W_IDAT_MANIFEST_NAME",
                f"{idat.path}: IDAT names manifest {idat.manifest!r}, BPM is {manifest.manifest_name!r}",
            )
        )
    return cov, notes


@dataclass
class LocusMap:
    """Row indices of every manifest address in one IDAT address layout (reusable across samples of a chip)."""

    a: npt.NDArray[np.intp]
    a_ok: npt.NDArray[np.bool_]
    b: npt.NDArray[np.intp]
    b_ok: npt.NDArray[np.bool_]
    c: npt.NDArray[np.intp]
    c_ok: npt.NDArray[np.bool_]


def build_locus_map(idat: Idat, manifest: Manifest, caddr: npt.NDArray[np.int32]) -> LocusMap:
    ix = AddressIndex(idat)
    a, a_ok = ix.lookup(manifest.address_a)
    b, b_ok = ix.lookup(manifest.address_b)
    c, c_ok = ix.lookup(caddr)
    return LocusMap(a, a_ok, b, b_ok, c, c_ok)


class LocusMapCache:
    """Caches :class:`LocusMap` per (IDAT address layout, manifest). Samples of one chip type share a layout."""

    def __init__(self, manifest: Manifest, max_entries: int = 8) -> None:
        self.manifest = manifest
        self.caddr = control_addresses(manifest)
        self._maps: dict[bytes, LocusMap] = {}
        self._coverage: dict[bytes, float] = {}
        self.max_entries = max_entries

    @staticmethod
    def _key(idat: Idat) -> bytes:
        return hashlib.blake2b(np.ascontiguousarray(idat.address_ids).tobytes(), digest_size=16).digest()

    def coverage(self, idat: Idat) -> float:
        """``address_coverage`` memoised per address layout."""
        key = self._key(idat)
        if (cov := self._coverage.get(key)) is None:
            if len(self._coverage) >= self.max_entries:
                self._coverage.pop(next(iter(self._coverage)))
            cov = self._coverage[key] = address_coverage(idat, self.manifest)
        return cov

    def get(self, idat: Idat) -> LocusMap:
        key = self._key(idat)
        lm = self._maps.get(key)
        if lm is None:
            if len(self._maps) >= self.max_entries:
                self._maps.pop(next(iter(self._maps)))
            lm = build_locus_map(idat, self.manifest, self.caddr)
            self._maps[key] = lm
        return lm


def assemble_raw_xy(
    grn: Idat, red: Idat, manifest: Manifest, *, min_beads: int = MIN_BEADS, cache: LocusMapCache | None = None
) -> RawIntensities:
    """Build per-locus raw X/Y and control intensities (plan §7.1)."""
    n = manifest.num_loci
    at = manifest.assay_types
    if (at > ASSAY_INF_I_GC).any():
        bad = int(np.flatnonzero(at > ASSAY_INF_I_GC)[0])
        raise SimurgError("E_BPM_FORMAT", f"locus {manifest.names[bad]!r} has unknown assay type {int(at[bad])}")
    cache = cache or LocusMapCache(manifest)
    gm, rm = cache.get(grn), cache.get(red)

    red_mean, grn_mean = red.mean, grn.mean
    red_nb, grn_nb = red.n_beads, grn.n_beads
    x = np.zeros(n, dtype=np.uint16)
    y = np.zeros(n, dtype=np.uint16)
    ok = np.zeros(n, dtype=np.bool_)

    # (assay, X source, Y source): (mean, n_beads, row index, present)
    plans = (
        (ASSAY_INF_II, (red_mean, red_nb, rm.a, rm.a_ok), (grn_mean, grn_nb, gm.a, gm.a_ok)),
        (ASSAY_INF_I_AT, (red_mean, red_nb, rm.a, rm.a_ok), (red_mean, red_nb, rm.b, rm.b_ok)),
        (ASSAY_INF_I_GC, (grn_mean, grn_nb, gm.a, gm.a_ok), (grn_mean, grn_nb, gm.b, gm.b_ok)),
    )
    for assay, (xm, xnb, xi, xok), (ym, ynb, yi, yok) in plans:
        sel = (at == assay) & xok & yok
        sel[sel] &= (xnb[xi[sel]] >= min_beads) & (ynb[yi[sel]] >= min_beads)
        x[sel] = xm[xi[sel]]
        y[sel] = ym[yi[sel]]
        ok |= sel

    both = rm.c_ok & gm.c_ok
    cx = np.zeros(len(cache.caddr), dtype=np.uint16)
    cy = np.zeros(len(cache.caddr), dtype=np.uint16)
    cx[both] = red_mean[rm.c[both]]
    cy[both] = grn_mean[gm.c[both]]
    return RawIntensities(raw_x=x, raw_y=y, missing=~ok, controls_x=cx, controls_y=cy, control_addresses=cache.caddr)