SimurgArray/src/simurg/qc/thresholds.py

"""QC thresholds: YAML schema, validation, profiles and precedence (spec QC-020 … QC-025).

Precedence, highest first: CLI > ``--config`` YAML > profile > built-in. Every effective value keeps its source.
"""

from __future__ import annotations

import math
from collections.abc import Mapping
from dataclasses import dataclass, field
from importlib import resources
from pathlib import Path
from typing import Any

import yaml

from simurg.core.errors import SimurgError
from simurg.qc.controls import DEFAULT_BG_OFFSET, METRICS

FUNCTIONAL_KEYS = ("callRate", "logRDev")
CONTROL_KEYS = (*(m.key for m in METRICS), "assayControl")
KNOWN_KEYS = (*FUNCTIONAL_KEYS, *CONTROL_KEYS, "bgOffset")
ALIASES = {"tgaControl": "assayControl", "call_rate_min": "callRate", "logr_dev_max": "logRDev"}
ABOVE_KEYS = {"logRDev"}  # flagged when value > threshold; every other key is flagged when value < threshold
BUILTIN: dict[str, float | None] = {"callRate": 0.98, "logRDev": 0.20}  # QC-023 [EMU]; every other key unset
PROFILES = ("genotyping", "genotyping_ex", "genotyping_example", "methylation_epic", "methylation_msa")
KEY_TO_COLUMN = {m.key: m.column for m in METRICS} | {"assayControl": "assay_control_norm_r"}


@dataclass(frozen=True)
class Threshold:
    value: float | None  # None: check disabled (QC-022)
    direction: str  # "below" | "above"
    source: str  # builtin | profile:<name> | config | cli

    @property
    def enabled(self) -> bool:
        return self.value is not None

    def fails(self, v: float) -> bool:
        """True when ``v`` is out of range. Equality passes (QC-020); NaN signal fails (QC-008)."""
        if self.value is None:
            return False
        if v is None or math.isnan(v):
            return True
        return v < self.value if self.direction == "below" else v > self.value


@dataclass
class Thresholds:
    items: dict[str, Threshold] = field(default_factory=dict)
    profile: str = ""

    def get(self, key: str) -> Threshold:
        return self.items.get(key, Threshold(None, "above" if key in ABOVE_KEYS else "below", "unset"))

    @property
    def bg_offset(self) -> float:
        t = self.items.get("bgOffset")
        return float(t.value) if t is not None and t.value is not None else DEFAULT_BG_OFFSET

    def as_records(self) -> list[dict[str, Any]]:
        return [
            {"key": k, "value": t.value, "direction": t.direction, "source": t.source}
            for k, t in sorted(self.items.items(), key=lambda kv: KNOWN_KEYS.index(kv[0]))
        ]


def _canon_key(key: str) -> str:
    key = ALIASES.get(key, key)
    if key not in KNOWN_KEYS:
        raise SimurgError(
            "E_QC_THRESHOLD_KEY", f"unknown QC threshold key {key!r}; valid keys: {', '.join(KNOWN_KEYS)}"
        )
    return key


def _parse_value(key: str, raw: Any, source: str) -> Threshold:
    direction = "above" if key in ABOVE_KEYS else "below"
    if isinstance(raw, Mapping):
        direction = str(raw.get("direction", direction))
        if direction not in ("below", "above"):
            raise SimurgError("E_QC_THRESHOLD_VALUE", f"{key}: direction must be 'below' or 'above', got {direction!r}")
        raw = raw.get("threshold")
    if raw is None or (isinstance(raw, str) and raw.strip().lower() in ("", "null", "none", "~")):
        return Threshold(None, direction, source)
    if isinstance(raw, bool):
        raise SimurgError("E_QC_THRESHOLD_VALUE", f"{key}: value must be numeric or null, got {raw!r}")
    try:
        v = float(raw)
    except (TypeError, ValueError):
        raise SimurgError("E_QC_THRESHOLD_VALUE", f"{key}: value must be numeric or null, got {raw!r}") from None
    if not math.isfinite(v):
        raise SimurgError("E_QC_THRESHOLD_VALUE", f"{key}: value must be finite, got {raw!r}")
    if key == "callRate" and not 0.0 <= v <= 1.0:
        raise SimurgError("E_QC_THRESHOLD_RANGE", f"callRate must be within [0, 1], got {v}")
    if v < 0:
        raise SimurgError("E_QC_THRESHOLD_RANGE", f"{key} must be >= 0, got {v}")
    return Threshold(v, direction, source)


def parse_report_mapping(doc: Any, where: str) -> tuple[dict[str, Any], str | None]:
    """``qc: {profile, report: {...}}`` (or ``QC:``/``Report:``) → (report mapping, profile)."""
    if not isinstance(doc, Mapping):
        raise SimurgError("E_QC_CONFIG_STRUCTURE", f"{where}: expected a 'qc:' mapping with a 'report:' section")
    top = doc.get("qc", doc.get("QC"))
    if not isinstance(top, Mapping):
        raise SimurgError("E_QC_CONFIG_STRUCTURE", f"{where}: missing top-level 'qc:' (or 'QC:') mapping")
    report = top.get("report", top.get("Report"))
    if report is None:
        report = {}
    if not isinstance(report, Mapping):
        raise SimurgError("E_QC_CONFIG_STRUCTURE", f"{where}: 'report:' (or 'Report:') must be a mapping")
    profile = top.get("profile")
    return dict(report), (str(profile) if profile else None)


def load_config(path: Path) -> tuple[dict[str, Any], str | None]:
    """Read and structurally validate a threshold YAML (QC-021)."""
    p = Path(path)
    if not p.exists():
        raise SimurgError("E_QC_CONFIG_MISSING", f"QC config file was provided but does not exist: {p}", path=str(p))
    if p.suffix.lower() not in (".yaml", ".yml"):
        raise SimurgError("E_QC_CONFIG_EXT", f"QC config file must have a .yaml or .yml extension: {p}", path=str(p))
    try:
        doc = yaml.safe_load(p.read_text(encoding="utf-8-sig"))
    except yaml.YAMLError as e:
        raise SimurgError("E_QC_CONFIG_PARSE", f"QC config file is not valid YAML: {p}: {e}", path=str(p)) from None
    return parse_report_mapping(doc, str(p))


def profile_values(name: str) -> dict[str, Any]:
    """Shipped profile (QC-024); ``+ffpe`` suffix enables restoration: 1."""
    base, _, suffix = name.partition("+")
    if base not in PROFILES or suffix not in ("", "ffpe"):
        raise SimurgError("E_QC_PROFILE", f"unknown QC profile {name!r}; choose from {', '.join(PROFILES)} or auto")
    text = resources.files("simurg.qc").joinpath("profiles", f"{base}.yaml").read_text(encoding="utf-8")
    report, _ = parse_report_mapping(yaml.safe_load(text), f"profile {base}")
    if suffix == "ffpe":
        report["restoration"] = 1
    return report


def auto_profile(array_type: str, assay_format: str = "", manifest_name: str = "") -> str:
    """QC-025."""
    if array_type == "methylation":
        return "methylation_msa" if "msa" in manifest_name.lower() else "methylation_epic"
    return "genotyping_ex" if assay_format.strip().upper() == "EX" else "genotyping"


def parse_cli_pairs(pairs: list[str]) -> dict[str, Any]:
    out: dict[str, Any] = {}
    for item in pairs:
        key, sep, val = item.partition("=")
        if not sep:
            raise SimurgError("E_QC_THRESHOLD_VALUE", f"--threshold expects KEY=VALUE, got {item!r}")
        out[key.strip()] = val.strip()
    return out


def resolve_thresholds(
    *,
    config: Path | None = None,
    profile: str | None = None,
    cli: Mapping[str, Any] | None = None,
    array_type: str = "genotyping",
    assay_format: str = "",
    manifest_name: str = "",
) -> Thresholds:
    """Merge built-in < profile < config < CLI (QC-023). A ``profile:`` in the YAML is used when no CLI profile."""
    layers: list[tuple[str, Mapping[str, Any]]] = [("builtin", BUILTIN)]
    cfg_report: dict[str, Any] = {}
    cfg_profile = None
    if config is not None:
        cfg_report, cfg_profile = load_config(config)
    chosen = profile or cfg_profile
    if chosen == "auto":
        chosen = auto_profile(array_type, assay_format, manifest_name)
    if chosen:
        layers.append((f"profile:{chosen}", profile_values(chosen)))
    if cfg_report:
        layers.append(("config", cfg_report))
    if cli:
        layers.append(("cli", cli))
    th = Thresholds(profile=chosen or "")
    for source, mapping in layers:
        for raw_key, raw in mapping.items():
            key = _canon_key(str(raw_key))
            th.items[key] = _parse_value(key, raw, source)
    return th