"""Two-axis sample QC status and informational flags (spec QC-030, QC-031, QC-033, QC-034).

* functional status (PASS/FAIL/NA) — call rate and LogR deviation only;
* control status (CLEAR/FLAGGED) — control-probe metrics including the assay control;
* flags — informational, never change either status.
"""

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

from simurg.qc.thresholds import CONTROL_KEYS, KEY_TO_COLUMN, Thresholds

CONTAMINATION_FLAG = 0.02  # plan §9.3 [OURS]
HET_Z = 3.0
ASSAY_CONTROL_MIN = 1.0  # OUT-062


@dataclass
class SampleStatus:
    functional_status: str = "NA"
    control_status: str = ""
    control_flag_codes: list[str] = field(default_factory=list)
    na_metrics: list[str] = field(default_factory=list)
    flags: list[str] = field(default_factory=list)

    def columns(self) -> dict[str, str]:
        return {
            "functional_status": self.functional_status,
            "control_status": self.control_status,
            "control_flag_codes": ";".join(self.control_flag_codes),
            "flags": ";".join(self.flags),
        }


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


def functional_status(row: Mapping[str, Any] | None, th: Thresholds) -> str:
    """QC-030. ``row`` is a sample-summary row; None means no genotyping summary (NA)."""
    if row is None:
        return "NA"
    if row.get("status", "OK") != "OK":
        return "FAIL"
    cr, ld = th.get("callRate"), th.get("logRDev")
    if cr.fails(_num(row.get("autosomal_call_rate"))) or ld.fails(_num(row.get("logr_dev"))):
        return "FAIL"
    return "PASS"


def control_status(
    values: Mapping[str, float] | None, na_metrics: Sequence[str], th: Thresholds
) -> tuple[str, list[str]]:
    """QC-031: FLAGGED iff an enabled, evaluable control metric is out of range. Codes in QC-004 order."""
    if values is None:
        return "", []
    codes = []
    for key in CONTROL_KEYS:
        col = KEY_TO_COLUMN[key]
        t = th.get(key)
        if not t.enabled or col not in values or col in na_metrics:
            continue
        if t.fails(_num(values[col])):
            codes.append(key)
    return ("FLAGGED" if codes else "CLEAR"), codes


def het_zscores(het: Sequence[float]) -> list[float]:
    """z-scores of heterozygosity within the batch (NaN where undefined or with < 3 finite values)."""
    a = np.array([_num(h) for h in het], dtype=np.float64)
    ok = np.isfinite(a)
    if ok.sum() < 3:
        return [math.nan] * len(a)
    mu, sd = float(a[ok].mean()), float(a[ok].std(ddof=1))
    if sd == 0:
        return [0.0 if o else math.nan for o in ok]
    return [float((x - mu) / sd) if o else math.nan for x, o in zip(a, ok, strict=True)]


def sample_flags(
    row: Mapping[str, Any],
    th: Thresholds,
    *,
    het_z: float = math.nan,
    probe_columns: Sequence[str] = (),
    duplicates: Sequence[str] = (),
    unexpected_relatedness: bool = False,
) -> list[str]:
    """QC-033 flags in a fixed order."""
    flags: list[str] = []
    ok = row.get("status", "OK") == "OK"
    if ok and th.get("callRate").fails(_num(row.get("autosomal_call_rate"))):
        flags.append("low_call_rate")
    if ok and th.get("logRDev").fails(_num(row.get("logr_dev"))):
        flags.append("high_lrr_sd")
    if row.get("sex_check") == "mismatch":
        flags.append("sex_mismatch")
    if math.isfinite(het_z) and abs(het_z) > HET_Z:
        flags.append("het_outlier")
    c = _num(row.get("contamination"))
    if math.isfinite(c) and c > CONTAMINATION_FLAG:
        flags.append("contaminated")
    if ok and any(_num(row.get(p)) < ASSAY_CONTROL_MIN for p in probe_columns):
        flags.append("assay_control_fail")
    flags.extend(f"duplicate_of:{d}" for d in duplicates)
    if unexpected_relatedness:
        flags.append("unexpected_relatedness")
    return flags