SimurgArray/src/simurg/compare/core.py

"""Field-level comparison with tolerance categories (spec CLI-038; plan §16.3).

Every compared value falls in exactly one category:

* **EXACT** — bitwise equal (NaN equals NaN whatever its sign bit, K-014);
* **NEAR** — within the field's tolerance;
* **DISCREPANCY** — outside it (or NaN on one side only);
* **UNKNOWN** — one side lacks the field or the record.

Counts are kept per locus class (assay, variant type, chromosome group, intensity-only), and every discrepancy is
kept with its identifiers — nothing is hidden (plan §16.3).
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from dataclasses import field as dc_field
from typing import Any

import numpy as np
import numpy.typing as npt

EXACT, NEAR, DISCREPANCY, UNKNOWN = 0, 1, 2, 3
CATEGORY_NAMES = ("EXACT", "NEAR", "DISCREPANCY", "UNKNOWN")


@dataclass(frozen=True)
class Tolerance:
    abs: float | None = None  # absolute difference allowed
    ulp: int | None = None  # float32 units in the last place allowed

    @property
    def exact(self) -> bool:
        return self.abs is None and self.ulp is None


EXACT_TOL = Tolerance()
# plan §16.3 (bit-exact mode targets)
GTC_TOLERANCES: dict[str, Tolerance] = {
    "raw_x": EXACT_TOL,
    "raw_y": EXACT_TOL,
    "controls_x": EXACT_TOL,
    "controls_y": EXACT_TOL,
    "genotypes": EXACT_TOL,
    "base_calls": EXACT_TOL,
    "normalization_transforms": Tolerance(ulp=2),
    "genotype_scores": Tolerance(abs=1e-5),
    "b_allele_freqs": Tolerance(abs=1e-5),
    "logr_ratios": Tolerance(abs=1e-5),
    "call_rate": Tolerance(abs=1e-6),
    "logr_dev": Tolerance(abs=1e-6),
    "gc10": Tolerance(abs=1e-6),
    "gc50": Tolerance(abs=1e-6),
    "gender": EXACT_TOL,
    "num_calls": EXACT_TOL,
    "num_no_calls": EXACT_TOL,
    "num_intensity_only": EXACT_TOL,
    "percentiles_x": EXACT_TOL,
    "percentiles_y": EXACT_TOL,
}
VCF_FLOAT_TOL = Tolerance(abs=1e-5)  # GS, BAF, LRR, NORMX, NORMY, THETA, R
VCF_EXACT_FIELDS = {"GT", "X", "Y"}


def _float_ulps(a: npt.NDArray[np.float32], b: npt.NDArray[np.float32]) -> npt.NDArray[np.int64]:
    """Distance in float32 ULPs (monotone integer mapping of the bit patterns)."""

    def key(x: npt.NDArray[np.float32]) -> npt.NDArray[np.int64]:
        i = x.astype(np.float32).view(np.int32).astype(np.int64)
        return np.where(i < 0, -(i & 0x7FFFFFFF), i)

    return np.abs(key(a) - key(b))


def categorize(a: npt.ArrayLike, b: npt.ArrayLike, tol: Tolerance) -> npt.NDArray[np.int8]:
    """Element-wise category of two equally shaped arrays (numbers or strings)."""
    a, b = np.asarray(a), np.asarray(b)
    if a.shape != b.shape:
        raise ValueError(f"shape {a.shape} vs {b.shape}")
    out = np.full(a.shape, DISCREPANCY, np.int8)
    if a.dtype.kind in "fc" or b.dtype.kind in "fc":
        fa, fb = a.astype(np.float64), b.astype(np.float64)
        na, nb = np.isnan(fa), np.isnan(fb)
        same = (fa == fb) | (na & nb)
        if a.dtype == np.float32 and b.dtype == np.float32:
            same &= (a.view(np.uint32) == b.view(np.uint32)) | (na & nb) | ((fa == 0) & (fb == 0))
        out[same] = EXACT
        near = np.zeros(a.shape, np.bool_)
        valid = ~(na | nb) & ~same
        if tol.abs is not None:
            near |= valid & (np.abs(fa - fb) <= tol.abs)
        if tol.ulp is not None:
            near |= valid & (_float_ulps(a.astype(np.float32), b.astype(np.float32)) <= tol.ulp)
        out[near] = NEAR
        return out
    out[a == b] = EXACT
    return out


@dataclass
class FieldResult:
    """Category counts of one field: overall and per locus class, plus the discrepancies themselves."""

    field: str
    tolerance: str
    counts: dict[str, list[int]] = dc_field(default_factory=dict)  # class -> [EXACT, NEAR, DISCREPANCY, UNKNOWN]
    discrepancies: list[dict[str, Any]] = dc_field(default_factory=list)

    def add(self, cls: str, cats: npt.NDArray[np.int8]) -> None:
        c = self.counts.setdefault(cls, [0, 0, 0, 0])
        binc = np.bincount(cats.ravel().astype(np.int64), minlength=4)
        for k in range(4):
            c[k] += int(binc[k])

    @property
    def total(self) -> list[int]:
        return self.counts.get("all", [0, 0, 0, 0])


def tol_text(t: Tolerance) -> str:
    if t.exact:
        return "exact"
    parts = []
    if t.abs is not None:
        parts.append(f"abs<={t.abs:g}")
    if t.ulp is not None:
        parts.append(f"ulp<={t.ulp}")
    return ",".join(parts)


def locus_classes(
    chroms: Sequence[str], snps: Sequence[str], assay_types: npt.ArrayLike, intensity_only: npt.ArrayLike
) -> dict[str, npt.NDArray[np.bool_]]:
    """Locus class masks (plan §16.3): assay, variant type, chromosome group, intensity-only."""
    at = np.asarray(assay_types)
    io = np.asarray(intensity_only, dtype=np.bool_)

    def group(c: str) -> str:
        c = c[3:] if c.lower().startswith("chr") else c
        if c in ("0", ""):
            return "unmapped"
        if c == "XY":
            return "PAR"
        if c in ("X", "Y"):
            return c
        if c in ("MT", "M"):
            return "MT"
        return "autosome" if c.isdigit() else "other"

    groups = np.array([group(c) for c in chroms])
    indel = np.array([("D" in s or "I" in s) for s in snps])
    out = {"all": np.ones(len(groups), np.bool_), "assay=InfII": at == 0, "assay=InfI": at != 0,
           "type=SNV": ~indel, "type=indel": indel, "intensity_only": io}  # fmt: skip
    for gname in ("autosome", "X", "Y", "PAR", "MT", "unmapped", "other"):
        m = groups == gname
        if m.any():
            out[f"chrom={gname}"] = m
    return out


def compare_locus_field(
    name: str,
    a: npt.ArrayLike | None,
    b: npt.ArrayLike | None,
    tol: Tolerance,
    classes: Mapping[str, npt.NDArray[np.bool_]],
    ids: Sequence[str],
    context: Mapping[str, str],
) -> FieldResult:
    """Per-locus arrays (first axis = locus). Missing or empty on one side → UNKNOWN for every locus."""
    res = FieldResult(name, tol_text(tol))
    n = len(classes["all"])
    a_ok = a is not None and len(np.asarray(a)) == n
    b_ok = b is not None and len(np.asarray(b)) == n
    if not (a_ok and b_ok):
        cats = np.full(n, UNKNOWN, np.int8)
    else:
        el = categorize(np.asarray(a), np.asarray(b), tol)
        cats = el.max(axis=tuple(range(1, el.ndim))) if el.ndim > 1 else el  # worst element of the locus
        for i in np.flatnonzero(cats == DISCREPANCY).tolist():
            res.discrepancies.append(
                {**context, "field": name, "locus": ids[i], "a": _fmt(np.asarray(a)[i]), "b": _fmt(np.asarray(b)[i])}
            )
    for cls, mask in classes.items():
        res.add(cls, cats[mask])
    return res


def compare_scalar(name: str, a: Any, b: Any, tol: Tolerance, context: Mapping[str, str]) -> FieldResult:
    res = FieldResult(name, tol_text(tol))
    if a is None or b is None:
        cat = np.array([UNKNOWN], np.int8)
    else:
        aa, bb = np.asarray([a]), np.asarray([b])
        if aa.dtype.kind == "f" or bb.dtype.kind == "f":
            aa, bb = aa.astype(np.float64), bb.astype(np.float64)
        cat = categorize(aa, bb, tol)
    res.add("sample", cat)
    if cat[0] == DISCREPANCY:
        res.discrepancies.append({**context, "field": name, "locus": "", "a": _fmt(a), "b": _fmt(b)})
    return res


def _fmt(v: Any) -> str:
    arr = np.asarray(v)
    if arr.ndim == 0:
        x = arr.item()
        return repr(x) if isinstance(x, float) else str(x)
    return "[" + ",".join(_fmt(x) for x in arr.ravel().tolist()) + "]"


def merge_results(results: Sequence[FieldResult]) -> dict[str, FieldResult]:
    """Sum results of the same field (e.g. over samples)."""
    out: dict[str, FieldResult] = {}
    for r in results:
        m = out.setdefault(r.field, FieldResult(r.field, r.tolerance))
        for cls, c in r.counts.items():
            cur = m.counts.setdefault(cls, [0, 0, 0, 0])
            for k in range(4):
                cur[k] += c[k]
        m.discrepancies.extend(r.discrepancies)
    return out