SimurgArray/src/simurg/compare/records.py

"""GTC-vs-GTC and VCF-vs-VCF comparison on parsed data (spec CLI-038; plan §16.3)."""

from __future__ import annotations

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

import numpy as np
import numpy.typing as npt
from numpy.lib import recfunctions as rfn

from simurg.compare.core import (
    DISCREPANCY,
    EXACT,
    GTC_TOLERANCES,
    UNKNOWN,
    VCF_EXACT_FIELDS,
    VCF_FLOAT_TOL,
    FieldResult,
    Tolerance,
    categorize,
    compare_locus_field,
    compare_scalar,
    tol_text,
)

GTC_LOCUS_FIELDS = ("raw_x", "raw_y", "genotypes", "base_calls", "genotype_scores", "b_allele_freqs", "logr_ratios")
GTC_SCALARS = ("call_rate", "logr_dev", "gc10", "gc50", "gender", "num_calls", "num_no_calls", "num_intensity_only")
GTC_METADATA = ("sample_name", "sample_plate", "sample_well", "cluster_file", "snp_manifest", "imaging_date",
                "autocall_date", "autocall_version", "slide_identifier")  # fmt: skip


@dataclass
class Comparison:
    """Results of one pair (one sample)."""

    label: str
    kind: str  # "gtc" | "vcf"
    fields: list[FieldResult] = field(default_factory=list)
    metadata: dict[str, list[str]] = field(default_factory=dict)  # field -> [a, b] when they differ
    notes: list[str] = field(default_factory=list)


def compare_gtc(
    a: Any,
    b: Any,
    label: str,
    classes: Mapping[str, npt.NDArray[np.bool_]],
    names: Sequence[str],
    tolerances: Mapping[str, Tolerance] = GTC_TOLERANCES,
) -> Comparison:
    """Two parsed GTCs of the same manifest."""
    res = Comparison(label, "gtc")
    ctx = {"pair": label}
    if a.num_snps != b.num_snps:
        res.notes.append(f"locus count differs: {a.num_snps} vs {b.num_snps}")
        return res
    for f in GTC_LOCUS_FIELDS:
        va, vb = getattr(a, f, None), getattr(b, f, None)
        if (va is None or len(va) == 0) and (vb is None or len(vb) == 0):
            continue  # absent on both sides (e.g. BAF/LRR in GTC v3): not applicable
        res.fields.append(compare_locus_field(f, va, vb, tolerances[f], classes, names, ctx))
    # normalization transforms: one row per bin
    ta, tb = a.normalization_transforms, b.normalization_transforms
    fr = FieldResult("normalization_transforms", tol_text(tolerances["normalization_transforms"]))
    if len(ta) != len(tb) or not len(ta):
        fr.add("bins", np.full(max(len(ta), len(tb), 1), UNKNOWN, np.int8))
    else:
        ua = rfn.structured_to_unstructured(ta).astype(np.float32)
        ub = rfn.structured_to_unstructured(tb).astype(np.float32)
        cats = categorize(ua, ub, tolerances["normalization_transforms"]).max(axis=1)
        fr.add("bins", cats)
        for i in np.flatnonzero(cats == DISCREPANCY).tolist():
            fr.discrepancies.append({**ctx, "field": "normalization_transforms", "locus": f"bin{i}",
                                     "a": str(ua[i].tolist()), "b": str(ub[i].tolist())})  # fmt: skip
    res.fields.append(fr)
    # controls: one row per control bead
    for f in ("controls_x", "controls_y"):
        ca, cb = getattr(a, f), getattr(b, f)
        fr = FieldResult(f, tol_text(tolerances[f]))
        fr.add("controls", categorize(ca, cb, tolerances[f]) if len(ca) == len(cb) else np.array([UNKNOWN], np.int8))
        res.fields.append(fr)
    for f in GTC_SCALARS:
        res.fields.append(compare_scalar(f, _scalar(a, f), _scalar(b, f), tolerances[f], ctx))
    for f in ("percentiles_x", "percentiles_y"):
        pa, pb = getattr(a, f), getattr(b, f)
        fr = FieldResult(f, tol_text(tolerances[f]))
        fr.add("sample", categorize(pa, pb, tolerances[f]) if len(pa) == len(pb) else np.array([UNKNOWN], np.int8))
        res.fields.append(fr)
    for f in GTC_METADATA:
        va, vb = str(getattr(a, f, "")), str(getattr(b, f, ""))
        if va != vb:
            res.metadata[f] = [va, vb]
    if a.version != b.version:
        res.notes.append(f"GTC versions differ: {a.version} vs {b.version} (fields absent in one are UNKNOWN)")
    return res


def _scalar(g: Any, f: str) -> Any:
    if f == "gender":
        return g.gender_char
    v = getattr(g, f, None)
    if isinstance(v, float) and math.isnan(v):
        return v
    return v


# ------------------------------------------------------------------------------------------------ VCF
@dataclass
class VcfRecord:
    chrom: str
    pos: int
    vid: str
    ref: str
    alt: str
    fields: dict[str, str]  # FORMAT key -> value of the compared sample


def vcf_classes(r: VcfRecord) -> list[str]:
    c = r.chrom[3:] if r.chrom.lower().startswith("chr") else r.chrom
    grp = "autosome" if c.isdigit() else {"X": "X", "Y": "Y", "MT": "MT", "M": "MT"}.get(c, "other")
    alts = r.alt.split(",")
    if all(len(x) == len(r.ref) == 1 for x in alts):
        vtype = "SNV"
    elif all(len(x) == len(r.ref) for x in alts):
        vtype = "MNV"
    else:
        vtype = "indel"
    out = ["all", f"chrom={grp}", f"type={vtype}"]
    if "," in r.vid or ";" in r.vid:
        out.append("merged_site")
    if len(alts) > 1:
        out.append("multiallelic")
    return out


def _gt_norm(gt: str) -> str:
    parts = gt.replace("|", "/").split("/")
    if any(p == "." for p in parts):
        return "." if len(parts) == 1 else "./."
    return "/".join(sorted(parts, key=int))


def _cat_scalar(va: str, vb: str, tol: Tolerance) -> int:
    if va == vb:
        return EXACT
    x, y = _num(va), _num(vb)
    if math.isnan(x) and math.isnan(y):
        return EXACT
    if math.isnan(x) or math.isnan(y):
        return DISCREPANCY
    if x == y:
        return EXACT
    if tol.abs is not None and abs(x - y) <= tol.abs:
        return 1  # NEAR
    return DISCREPANCY


def compare_vcf_records(
    a: Mapping[tuple[str, int, str, str, str], VcfRecord],
    b: Mapping[tuple[str, int, str, str, str], VcfRecord],
    label: str,
    float_tol: Tolerance = VCF_FLOAT_TOL,
) -> Comparison:
    """Records keyed by (CHROM, POS, REF, ALT, ID); FORMAT fields of one sample on each side."""
    res = Comparison(label, "vcf")
    ctx = {"pair": label}
    keys = list(a) + [k for k in b if k not in a]
    fields = sorted({f for r in a.values() for f in r.fields} | {f for r in b.values() for f in r.fields})
    counts: dict[str, dict[str, list[int]]] = {f: {} for f in ["record", *fields]}
    disc: dict[str, list[dict[str, Any]]] = {f: [] for f in ["record", *fields]}
    tols = {f: (Tolerance() if f in VCF_EXACT_FIELDS else float_tol) for f in fields}

    def bump(f: str, classes: list[str], cat: int) -> None:
        per = counts[f]
        for cls in classes:
            c = per.get(cls)
            if c is None:
                c = per[cls] = [0, 0, 0, 0]
            c[cat] += 1

    for k in keys:
        ra, rb = a.get(k), b.get(k)
        r = ra or rb
        assert r is not None
        classes = vcf_classes(r)
        where = f"{r.chrom}:{r.pos}:{r.ref}>{r.alt} {r.vid}"
        if ra is None or rb is None:
            bump("record", classes, UNKNOWN)
            disc["record"].append({**ctx, "field": "record", "locus": where,
                                   "a": "present" if ra else "absent", "b": "present" if rb else "absent"})  # fmt: skip
            continue
        bump("record", classes, EXACT)
        for f in fields:
            va, vb = ra.fields.get(f), rb.fields.get(f)
            if va is None and vb is None:
                continue  # the field does not apply to this record on either side (e.g. BAF of merged sites)
            if va is None or vb is None:
                cat = UNKNOWN
            elif f == "GT":
                cat = EXACT if _gt_norm(va) == _gt_norm(vb) else DISCREPANCY
            else:
                cat = _cat_scalar(va, vb, tols[f])
            bump(f, classes, cat)
            if cat == DISCREPANCY:
                disc[f].append({**ctx, "field": f, "locus": where, "a": va, "b": vb})
    for f in ["record", *fields]:
        fr = FieldResult(
            f, "present in both" if f == "record" else ("exact" if f in VCF_EXACT_FIELDS else tol_text(float_tol))
        )
        fr.counts = counts[f]
        fr.discrepancies = disc[f]
        res.fields.append(fr)
    return res


def _num(v: str) -> float:
    try:
        return float(v)
    except ValueError:
        return math.nan