SimurgArray/src/simurg/harmonize/combine.py

"""Genotypes of VCF sites from probe calls (spec OUT-018, OUT-031, OUT-032, OUT-037; plan §8.3).

Single-probe sites map the A/B call straight to allele indices. Sites with several probes use the exact candidate
algorithm of OUT-031: enumerate every genotype over {REF} ∪ ALT with the site's ploidy, predict each probe's call
from the alleles' signal classes (Infinium II: dye group of the extended base; Infinium I and indels: exact allele,
no signal otherwise) and keep the candidates consistent with every called probe.
[OSS: Illumina GTCtoVCF ``GenotypeFormat.RecordCombiner`` (Apache-2.0) — the consistency idea; the class model and
the ploidy-aware enumeration are ours.]
"""

from __future__ import annotations

from dataclasses import dataclass
from itertools import combinations_with_replacement

from simurg.harmonize.sites import GREEN, RED, ProbeSite, Site

NC, AA, AB, BB = 0, 1, 2, 3
CALL_TEXT = {NC: "NC", AA: "AA", AB: "AB", BB: "BB"}


def signal_class(p: ProbeSite, allele: str) -> str | None:
    """OUT-031 step 1."""
    a, b = p.alleles
    if p.kind == "snv":  # SNV probes read the single base at POS; alleles may be extended by --merge-snv-indel
        z, a0, b0 = allele[:1], a[:1], b[:1]
        if p.infinium == 2:
            grp = RED if z in RED else GREEN
            if a0 in grp:
                return "A"
            if b0 in grp:
                return "B"
            return None
        return "A" if z == a0 else "B" if z == b0 else None
    if allele == a:
        return "A"
    if allele == b:
        return "B"
    return None


def predicted_call(p: ProbeSite, genotype: tuple[str, ...]) -> int:
    classes = sorted(c for c in (signal_class(p, z) for z in genotype) if c is not None)
    if not classes:
        return NC
    u = set(classes)
    if u == {"A"}:
        return AA
    if u == {"B"}:
        return BB
    return AB


@dataclass(frozen=True)
class Outcome:
    alleles: tuple[int, ...] | None  # allele indices (sorted) or None for missing
    reason: str  # "" | "het_in_haploid" | W_MERGE_* | W_INF1_NC_CONFLICT


def single(site: Site, call: int, ploidy: int) -> Outcome:
    p = site.probes[0]
    if call == NC:
        return Outcome(None, "")
    ia, ib = site.allele_index(p.alleles[0]), site.allele_index(p.alleles[1])
    if ploidy == 1:
        if call == AB:
            return Outcome(None, "het_in_haploid")
        return Outcome((ia if call == AA else ib,), "")
    pair = {AA: (ia, ia), AB: (ia, ib), BB: (ib, ib)}[call]
    return Outcome(tuple(sorted(pair)), "")


def merged(site: Site, calls: tuple[int, ...], ploidy: int, inf1_nc_informative: bool = False) -> Outcome:
    """OUT-031 steps 3–6."""
    if all(c == NC for c in calls):
        return Outcome(None, "W_MERGE_ALL_NC")
    alleles = [site.ref, *site.alts]
    cands = list(combinations_with_replacement(range(len(alleles)), ploidy))
    keep = [
        g
        for g in cands
        if all(
            c == NC or predicted_call(p, tuple(alleles[i] for i in g)) == c
            for p, c in zip(site.probes, calls, strict=True)
        )
    ]
    reason = ""
    if inf1_nc_informative:
        for p, c in zip(site.probes, calls, strict=True):
            if c == NC and p.infinium == 1:
                sub = [g for g in keep if all(signal_class(p, alleles[i]) is None for i in g)]
                if sub:
                    keep = sub
                else:
                    reason = "W_INF1_NC_CONFLICT"
    if len(keep) == 1:
        return Outcome(keep[0], reason)
    return Outcome(None, "W_MERGE_INCONSISTENT" if not keep else "W_MERGE_AMBIGUOUS")


def genotype(site: Site, calls: tuple[int, ...], ploidy: int, inf1_nc_informative: bool = False) -> Outcome:
    if len(site.probes) == 1:
        return single(site, calls[0], ploidy)
    key = (calls, ploidy, inf1_nc_informative)
    hit = site.memo.get(key)
    if hit is None:
        hit = site.memo[key] = merged(site, calls, ploidy, inf1_nc_informative)
    assert isinstance(hit, Outcome)
    return hit


def gt_text(o: Outcome, ploidy: int) -> str:
    if o.alleles is None:
        return "." if ploidy == 1 else "./."
    return "/".join(str(i) for i in o.alleles)


def failure_message(site: Site, calls: tuple[int, ...], reason: str) -> str:
    what = {"W_MERGE_INCONSISTENT": "inconsistent", "W_MERGE_AMBIGUOUS": "ambiguous", "W_MERGE_ALL_NC": "all no-calls"}
    probes = ", ".join(
        f"{p.name} ({'InfiniumI' if p.infinium == 1 else 'InfiniumII'}): {CALL_TEXT[c]}"
        for p, c in zip(site.probes, calls, strict=True)
    )
    return f"Failed to combine genotypes ({what.get(reason, reason)}) - {probes}"