"""Control-probe QC metrics (spec QC-001 … QC-011; plan §9.1).

Controls are the non-polymorphic bead types listed in the manifest's control section. Each gives a raw green and
red intensity (IDAT means, no normalization). Derived metrics are signal/background ratios per assay step.
[DOC: Illumina BeadArray Controls Reporter software guide and Infinium controls training material — formulas only]
"""

from __future__ import annotations

import math
import re
from collections.abc import Sequence
from dataclasses import dataclass, field
from functools import cache
from importlib import resources
from typing import Any

import numpy as np
import numpy.typing as npt
import yaml

from simurg.formats.idat import Idat
from simurg.signal.raw_xy import AddressIndex

DEFAULT_BG_OFFSET = 3000.0


@dataclass(frozen=True)
class ControlBead:
    """One control bead type: a (manifest control line, distinct address) pair, in manifest order."""

    category: str
    color: str
    name: str
    address: int
    canonical: str  # canonical ID, "" when the name is not in control_names.yaml


@dataclass(frozen=True)
class Metric:
    column: str  # controls_qc.tsv column
    key: str  # threshold YAML key (QC-020)
    genotyping: bool
    methylation: bool


# QC-004, in table order.
METRICS: tuple[Metric, ...] = (
    Metric("restoration", "restoration", True, True),
    Metric("staining_green", "stainingGreen", True, True),
    Metric("staining_red", "stainingRed", True, True),
    Metric("extension_green", "extensionGreen", True, True),
    Metric("extension_red", "extensionRed", True, True),
    Metric("hyb_high_med", "hybGreenHighMed", True, True),
    Metric("hyb_med_low", "hybGreenMediumLow", True, True),
    Metric("target_removal_1", "targetRemovalI", True, True),
    Metric("target_removal_2", "targetRemovalII", False, True),
    Metric("stringency", "stringency", True, False),
    Metric("nsb_green", "nonSpecificBindingGreen", True, False),
    Metric("nsb_red", "nonSpecificBindingRed", True, False),
    Metric("np_green", "npGreen", True, True),
    Metric("np_red", "npRed", True, True),
    Metric("specificity_1_green", "specificityIGreen", False, True),
    Metric("specificity_1_red", "specificityIRed", False, True),
    Metric("specificity_2", "specificityII", False, True),
    Metric("specificity_2_bg", "specificityIIBg", False, True),
    Metric("bs_conv_1_green", "bisulfiteConversionIGreen", False, True),
    Metric("bs_conv_1_green_bg", "bisulfiteConversionIGreenBg", False, True),
    Metric("bs_conv_1_red", "bisulfiteConversionIRed", False, True),
    Metric("bs_conv_1_red_bg", "bisulfiteConversionIRedBg", False, True),
    Metric("bs_conv_2", "bisulfiteConversionII", False, True),
    Metric("bs_conv_2_bg", "bisulfiteConversionIIBg", False, True),
)
GTMM_PM = [f"gtmm_{i}_pm" for i in range(1, 7)]
GTMM_MM = [f"gtmm_{i}_mm" for i in range(1, 7)]
SPEC2 = ["spec2_1", "spec2_2", "spec2_3"]
BS2 = ["bs2_1", "bs2_2", "bs2_3", "bs2_4"]
ASSAY_CONTROL_COLUMN = "assay_control_norm_r"  # QC-005: copied from the sample summary
COMPAT_METRIC_NAMES = {
    "restoration": "RestorationQC",
    "staining_green": "StainingGreenQC",
    "staining_red": "StainingRedQC",
    "extension_green": "ExtensionGreenQC",
    "extension_red": "ExtensionRedQC",
    "hyb_high_med": "HybridizationHighMediumQC",
    "hyb_med_low": "HybridizationMediumLowQC",
    "target_removal_1": "TargetRemovalQC",
    "target_removal_2": "TargetRemoval2QC",
    "stringency": "StringencyQC",
    "nsb_green": "NonSpecificBindingGreenQC",
    "nsb_red": "NonSpecificBindingRedQC",
    "np_green": "NonPolymorphicGreenQC",
    "np_red": "NonPolymorphicRedQC",
    ASSAY_CONTROL_COLUMN: "AssayControlQC",
}


def metrics_for(array_type: str) -> list[Metric]:
    """QC-005 / QC-006 metric sets."""
    if array_type == "genotyping":
        return [m for m in METRICS if m.genotyping]
    return [m for m in METRICS if m.methylation]


# ------------------------------------------------------------------------------------------------ names
@cache
def name_config() -> dict[str, Any]:
    text = resources.files("simurg.qc").joinpath("control_names.yaml").read_text(encoding="utf-8")
    cfg: dict[str, Any] = yaml.safe_load(text)
    return cfg


def normalise_name(name: str) -> str:
    """QC-002 normalisation: case-fold, collapse whitespace, no spaces before or inside parentheses, bgnd = bkg."""
    s = " ".join(name.casefold().split())
    s = re.sub(r"\s*\(\s*", "(", s)
    s = re.sub(r"\s*\)", ")", s)
    return re.sub(r"\bbgnd\b|\bbkgd\b", "bkg", s)


def canonical_id(name: str) -> str:
    return str(name_config()["canonical"].get(normalise_name(name), ""))


def excluded_from_report(bead: ControlBead) -> bool:
    """QC-010."""
    ex = name_config()["report_exclude"]
    return bead.category.casefold() in ex["categories"] or normalise_name(bead.name) in ex["names"]


def is_methylation_controls(beads: Sequence[ControlBead]) -> bool:
    cats = set(name_config()["methylation_categories"])
    return any(b.category.casefold() in cats for b in beads)


def beads_from_definitions(defs: Sequence[dict[str, str]]) -> list[ControlBead]:
    """Manifest control lines (``address`` may be ``a:b:…``) → bead types in manifest order (QC-001)."""
    out: list[ControlBead] = []
    for d in defs:
        seen: set[int] = set()
        for part in str(d["address"]).split(":"):
            part = part.strip()
            if not part:
                continue
            addr = int(part)
            if addr in seen:
                continue
            seen.add(addr)
            out.append(ControlBead(d["category"], d["color"], d["name"], addr, canonical_id(d["name"])))
    return out


def compare_definitions(a: Sequence[ControlBead], b: Sequence[ControlBead]) -> list[str]:
    """Differences between two control lists (BPM vs CSV, QC-001); empty when they agree."""
    ka = {(x.address, normalise_name(x.name)) for x in a}
    kb = {(x.address, normalise_name(x.name)) for x in b}
    diffs = [f"only in BPM: {addr} {nm}" for addr, nm in sorted(ka - kb)[:5]]
    diffs += [f"only in CSV: {addr} {nm}" for addr, nm in sorted(kb - ka)[:5]]
    return diffs


# ------------------------------------------------------------------------------------------------ intensities
def control_intensities(
    grn: Idat, red: Idat, beads: Sequence[ControlBead]
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Raw green/red means per bead type; NaN where the IDAT lacks the address."""
    addrs = np.array([b.address for b in beads], dtype=np.int64)
    out = []
    for idat in (grn, red):
        rows, found = AddressIndex(idat).lookup(addrs)
        v = idat.mean[rows].astype(np.float64)
        v[~found] = math.nan
        out.append(v)
    return out[0], out[1]


# ------------------------------------------------------------------------------------------------ metrics
@dataclass
class ControlQC:
    values: dict[str, float]  # column -> value (NaN when not evaluable or degenerate)
    missing: dict[str, list[str]] = field(default_factory=dict)  # column -> canonical IDs absent from the manifest
    zero_divisor: list[str] = field(default_factory=list)  # columns with a zero denominator (QC-008)

    @property
    def na_metrics(self) -> list[str]:
        return list(self.missing)


class _Probes:
    def __init__(self, beads: Sequence[ControlBead], green: npt.NDArray[np.float64], red: npt.NDArray[np.float64]):
        self.idx: dict[str, list[int]] = {}
        for i, b in enumerate(beads):
            if b.canonical:
                self.idx.setdefault(b.canonical, []).append(i)
        self.ch = {"G": green, "R": red}
        self.absent: set[str] = set()

    def vals(self, ids: Sequence[str], ch: str) -> npt.NDArray[np.float64]:
        """Intensities of every bead of ``ids`` in channel ``ch``. Bead addresses absent from the IDAT (NaN; e.g.
        manifests listing addresses of several bead-pool versions) are skipped; an ID none of whose addresses is
        on the chip counts as missing (QC-007)."""
        out: list[float] = []
        for cid in ids:
            v = self.ch[ch][self.idx.get(cid, [])]
            v = v[~np.isnan(v)]
            if not len(v):
                self.absent.add(cid)
            out.extend(v.tolist())
        return np.asarray(out, dtype=np.float64)

    def mean(self, cid: str, ch: str) -> float:
        v = self.vals([cid], ch)
        return float(v.mean()) if len(v) else math.nan

    def min(self, ids: Sequence[str], ch: str) -> float:
        v = self.vals(ids, ch)
        return float(v.min()) if len(v) else math.nan

    def max(self, ids: Sequence[str], ch: str) -> float:
        v = self.vals(ids, ch)
        return float(v.max()) if len(v) else math.nan


def compute_control_metrics(
    beads: Sequence[ControlBead],
    green: npt.NDArray[np.float64],
    red: npt.NDArray[np.float64],
    array_type: str = "genotyping",
    bg_offset: float = DEFAULT_BG_OFFSET,
    bs_conv_red_bg_set: str = "U3-U5",
) -> ControlQC:
    """QC-003/QC-004 metrics of one sample. Single-probe values are bead means; min/max run over all beads (QC-002)."""
    p = _Probes(beads, green, red)
    x = bg_offset
    bs_red_u = ["bs1_U3", "bs1_U4", "bs1_U5"] if bs_conv_red_bg_set == "U3-U5" else ["bs1_U4", "bs1_U5", "bs1_U6"]
    formulas: dict[str, tuple[Any, Any]] = {
        # column: (numerator thunk, denominator thunk); each thunk records the probes it needs
        "restoration": (lambda: p.mean("restoration", "G"), lambda: p.max(["ext_A", "ext_T"], "G") + x),
        "staining_green": (lambda: p.mean("biotin_high", "G"), lambda: p.mean("biotin_bkg", "G")),
        "staining_red": (lambda: p.mean("dnp_high", "R"), lambda: p.mean("dnp_bkg", "R")),
        "extension_green": (lambda: p.min(["ext_C", "ext_G"], "G"), lambda: p.max(["ext_A", "ext_T"], "G")),
        "extension_red": (lambda: p.min(["ext_A", "ext_T"], "R"), lambda: p.max(["ext_C", "ext_G"], "R")),
        "hyb_high_med": (lambda: p.mean("hyb_high", "G"), lambda: p.mean("hyb_medium", "G")),
        "hyb_med_low": (lambda: p.mean("hyb_medium", "G"), lambda: p.mean("hyb_low", "G")),
        "target_removal_1": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.mean("target_removal_1", "G")),
        "stringency": (lambda: p.mean("string_pm", "R"), lambda: p.mean("string_mm", "R")),
        "nsb_green": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.max(["nsb_bkg"], "G")),
        "nsb_red": (lambda: p.max(["ext_C", "ext_G"], "R") + x, lambda: p.max(["nsb_bkg"], "R")),
        "np_green": (lambda: p.min(["np_C", "np_G"], "G"), lambda: p.max(["np_A", "np_T"], "G")),
        "np_red": (lambda: p.min(["np_A", "np_T"], "R"), lambda: p.max(["np_C", "np_G"], "R")),
        # methylation (QC-004 rows 9, 15-24; MET-006)
        "target_removal_2": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.mean("target_removal_2", "G")),
        "specificity_1_green": (lambda: p.min(GTMM_PM[:3], "G"), lambda: p.max(GTMM_MM[:3], "G")),
        "specificity_1_red": (lambda: p.min(GTMM_PM[3:], "R"), lambda: p.max(GTMM_MM[3:], "R")),
        "specificity_2": (lambda: p.min(SPEC2, "R"), lambda: p.max(SPEC2, "G")),
        "specificity_2_bg": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.max(SPEC2, "G")),
        "bs_conv_1_green": (lambda: p.min(["bs1_C1", "bs1_C2"], "G"), lambda: p.max(["bs1_U1", "bs1_U2"], "G")),
        "bs_conv_1_green_bg": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.max(["bs1_U1", "bs1_U2"], "G")),
        "bs_conv_1_red": (lambda: p.min(["bs1_C3", "bs1_C4", "bs1_C5"], "R"), lambda: p.max(bs_red_u, "R")),
        "bs_conv_1_red_bg": (lambda: p.max(["ext_C", "ext_G"], "R") + x, lambda: p.max(bs_red_u, "R")),
        "bs_conv_2": (lambda: p.min(BS2, "R"), lambda: p.max(BS2, "G")),
        "bs_conv_2_bg": (lambda: p.max(["ext_A", "ext_T"], "G") + x, lambda: p.max(BS2, "G")),
    }
    res = ControlQC(values={})
    for m in metrics_for(array_type):
        p.absent = set()
        num, den = formulas[m.column][0](), formulas[m.column][1]()
        if p.absent:
            res.missing[m.column] = sorted(p.absent)
            res.values[m.column] = math.nan
        elif den == 0:  # QC-008: signal over a zero background is +Inf (passes); 0/0 is NaN (flagged)
            res.zero_divisor.append(m.column)
            res.values[m.column] = math.inf if num > 0 else math.nan
        else:
            res.values[m.column] = num / den
    return res