"""Methylation signal processing — spec MET-005, MET-010 … MET-012 (plan §7, v2). Pure functions, no I/O.

* Raw signals (MET-005): Infinium I — M = channel[AddressB], U = channel[AddressA] in the probe's colour channel,
  out-of-band = the other channel; Infinium II — M = Grn[AddressA], U = Red[AddressA]. Addresses missing from the
  IDAT or read from < 3 beads are masked.
* Detection p-value, pOOBAH [LIT: Zhou, Laird & Shen 2017 NAR 45:e22; OSS: sesame `pOOBAH`, MIT]: empirical CDF of
  the out-of-band Infinium I signal plus the negative controls per channel; p = 1 − max(F(M), F(U)) with each signal
  read in its own channel (Infinium II: M green, U red).
* Background correction, noob [LIT: Triche et al. 2013 NAR 41:e90; OSS: sesame `noob`, MIT; limma `normexp.signal`,
  GPL — formula only]: normal-exponential convolution per channel, background = out-of-band + negative controls,
  µ and σ from Huber's M-estimator (MASS `huber`, k = 1.5, MAD scale), α = max(Huber(in-band) − µ, 10); offset 15.
* Dye bias (MET-010 step 3) [OURS]: linear — scale red and green so the means of the normalization controls
  (NORM_A/T red, NORM_C/G green) meet at their average; nonlinear — quantile matching of the Infinium I red and green
  in-band signals to their average quantile function.
* β = M / (M + U + α) with signals floored at 1; M-value = log2(β' / (1 − β')), β' clipped to [1e-6, 1 − 1e-6].
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

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

F64 = npt.NDArray[np.float64]
Mask = npt.NDArray[np.bool_]
MIN_BEADS = 3
NOOB_OFFSET = 15.0


@dataclass
class MethylSignals:
    """In-band M/U per probe (float64, NaN where masked), out-of-band signal of Infinium I probes, control
    intensities (manifest control order) and the probe layout needed to split them by channel."""

    m: F64
    u: F64
    oob_m: F64  # Infinium I: the other channel at the M address; NaN for Infinium II
    oob_u: F64
    masked: Mask
    col: npt.NDArray[np.str_]  # "R" / "G" / "" (Infinium II)
    ctrl_green: F64
    ctrl_red: F64
    ctrl_type: list[str]

    @property
    def type_i(self) -> Mask:
        return np.asarray(self.col != "", dtype=np.bool_)


def assemble(grn: Idat, red: Idat, man: MethylManifest, *, min_beads: int = MIN_BEADS) -> MethylSignals:
    """MET-005 raw signal assembly."""
    gi, ri = AddressIndex(grn), AddressIndex(red)

    def read(addrs: npt.NDArray[np.int64]) -> tuple[F64, F64, Mask]:
        rg, fg = gi.lookup(addrs)
        rr, fr = ri.lookup(addrs)
        g = grn.mean[rg].astype(np.float64)
        r = red.mean[rr].astype(np.float64)
        ok = fg & fr & (grn.n_beads[rg] >= min_beads) & (red.n_beads[rr] >= min_beads) & (addrs != 0)
        return g, r, ok

    ga, ra, oka = read(man.u)
    gb, rb, okb = read(np.where(man.m != 0, man.m, man.u))
    t1 = man.m != 0
    is_r = man.col == "R"
    # Infinium II: M green, U red at address A; Infinium I: both addresses in the probe's channel
    m = np.where(t1, np.where(is_r, rb, gb), ga)
    u = np.where(t1, np.where(is_r, ra, ga), ra)
    oob_m = np.where(t1, np.where(is_r, gb, rb), np.nan)
    oob_u = np.where(t1, np.where(is_r, ga, ra), np.nan)
    masked = ~(oka & np.where(t1, okb, True)) | (t1 & ~np.isin(man.col, ("R", "G")))
    for a in (m, u, oob_m, oob_u):
        a[masked] = np.nan
    caddr = np.array([c.address for c in man.controls], dtype=np.int64)
    cg, cr, cok = read(caddr)
    cg[~cok] = np.nan
    cr[~cok] = np.nan
    return MethylSignals(m, u, oob_m, oob_u, masked, man.col, cg, cr, [c.type for c in man.controls])


def _finite(a: F64) -> F64:
    return a[np.isfinite(a)]


def background(s: MethylSignals, combine_neg: bool = True) -> tuple[F64, F64]:
    """Out-of-band background per channel (green, red): Infinium I OOB signal (+ negative controls)."""
    neg = np.array([combine_neg and t.casefold() == "negative" for t in s.ctrl_type], dtype=np.bool_)
    red_i, grn_i = s.col == "R", s.col == "G"
    oob_g = np.concatenate([s.oob_m[red_i], s.oob_u[red_i], s.ctrl_green[neg]])
    oob_r = np.concatenate([s.oob_m[grn_i], s.oob_u[grn_i], s.ctrl_red[neg]])
    return _finite(oob_g), _finite(oob_r)


def in_band(s: MethylSignals) -> tuple[F64, F64]:
    """In-band foreground per channel (green, red): Infinium I M/U in their channel, Infinium II M green / U red."""
    red_i, grn_i, t2 = s.col == "R", s.col == "G", s.col == ""
    g = np.concatenate([s.m[grn_i], s.u[grn_i], s.m[t2]])
    r = np.concatenate([s.m[red_i], s.u[red_i], s.u[t2]])
    return _finite(g), _finite(r)


# ------------------------------------------------------------------------------------------------ detection
def _ecdf(ref: F64):  # type: ignore[no-untyped-def]
    srt = np.sort(ref)
    n = len(srt)

    def f(x: F64) -> F64:
        if n == 0:
            return np.zeros(len(x))
        return np.searchsorted(srt, x, side="right") / n  # R ecdf: P(X <= x)

    return f


def detection_p(s: MethylSignals, combine_neg: bool = True) -> F64:
    """pOOBAH detection p-value per probe; 1 where masked."""
    bg_g, bg_r = background(s, combine_neg)
    fg, fr = _ecdf(bg_g), _ecdf(bg_r)
    red_i, t2 = s.col == "R", s.col == ""
    m, u = np.nan_to_num(s.m, nan=0.0), np.nan_to_num(s.u, nan=0.0)
    fm = np.where(red_i, fr(m), fg(m))  # Infinium II M is green, Infinium I green M is green
    fu = np.where(red_i | t2, fr(u), fg(u))  # Infinium II U is red
    p = 1.0 - np.maximum(fm, fu)
    p[s.masked] = 1.0
    return np.asarray(p, dtype=np.float64)


# ------------------------------------------------------------------------------------------------ noob
def huber(y: F64, k: float = 1.5, tol: float = 1e-6) -> tuple[float, float]:
    """MASS::huber: Huber M-estimator of location with MAD scale (returns the location before the last update)."""
    y = _finite(np.asarray(y, dtype=np.float64))
    n = len(y)
    if n == 0:
        raise ValueError("huber: no finite values")
    mu = float(np.median(y))
    s = float(1.4826 * np.median(np.abs(y - mu)))
    if s == 0:
        raise ValueError("huber: cannot estimate scale, MAD is zero")
    for _ in range(1000):
        yy = np.minimum(np.maximum(mu - k * s, y), mu + k * s)
        mu1 = float(yy.sum() / n)
        if abs(mu - mu1) < tol * s:
            break
        mu = mu1
    return mu, s


_erfc = np.frompyfunc(math.erfc, 1, 1)


def _mills(z: F64) -> F64:
    """φ(z)/Φ(z) (inverse Mills ratio), stable far into the lower tail."""
    z = np.asarray(z, dtype=np.float64)
    out = np.empty_like(z)
    lo = z < -30.0
    zz = z[~lo]
    phi = np.exp(-0.5 * zz * zz) / math.sqrt(2.0 * math.pi)
    big_phi = 0.5 * np.asarray(_erfc(-zz / math.sqrt(2.0)), dtype=np.float64)
    out[~lo] = phi / big_phi
    zl = z[lo]
    r = 1.0 / (zl * zl)
    out[lo] = -zl / (1.0 - r + 3.0 * r * r - 15.0 * r**3 + 105.0 * r**4)
    return out


def normexp_signal(mu: float, sigma: float, alpha: float, x: F64) -> F64:
    """limma ``normexp.signal`` (formula): E[signal | observed x] under x = N(µ, σ²) + Exp(α)."""
    if alpha <= 0 or sigma <= 0:
        raise ValueError("normexp: alpha and sigma must be positive")
    mu_sf = x - mu - sigma * sigma / alpha
    sig = mu_sf + sigma * _mills(mu_sf / sigma)
    ok = np.isfinite(sig)
    if (sig[ok] < 0).any():
        sig[ok] = np.maximum(sig[ok], 1e-6)
    return sig


@dataclass
class NoobFit:
    mu: float
    sigma: float
    alpha: float


def noob_fit(ib: F64, oob: F64) -> NoobFit:
    ib = np.where(ib < 1, 1.0, ib)
    mu, s = huber(oob)
    return NoobFit(mu, s, max(huber(ib)[0] - mu, 10.0))


def noob(
    s: MethylSignals, offset: float = NOOB_OFFSET, combine_neg: bool = True
) -> tuple[MethylSignals, NoobFit, NoobFit]:
    """Background-corrected copy of ``s`` (probe and control signals); returns the green and red fits."""
    bg_g, bg_r = background(s, combine_neg)
    ib_g, ib_r = in_band(s)
    fit_g, fit_r = noob_fit(ib_g, bg_g), noob_fit(ib_r, bg_r)

    def corr(x: F64, f: NoobFit) -> F64:
        y = np.where(np.isfinite(x), np.maximum(x, 1.0), np.nan)
        out = np.full_like(y, np.nan)
        ok = np.isfinite(y)
        out[ok] = normexp_signal(f.mu, f.sigma, f.alpha, y[ok]) + offset
        return out

    red_i, t2 = s.col == "R", s.col == ""
    fm = np.where(red_i, 1, 0)  # channel of M: red for Infinium I red, else green
    fu = np.where(red_i | t2, 1, 0)
    m = np.where(fm == 1, corr(s.m, fit_r), corr(s.m, fit_g))
    u = np.where(fu == 1, corr(s.u, fit_r), corr(s.u, fit_g))
    out = MethylSignals(
        m, u, s.oob_m, s.oob_u, s.masked, s.col, corr(s.ctrl_green, fit_g), corr(s.ctrl_red, fit_r), s.ctrl_type
    )
    return out, fit_g, fit_r


# ------------------------------------------------------------------------------------------------ dye bias
def _scale(s: MethylSignals, fg: float, fr: float) -> MethylSignals:
    red_i, t2 = s.col == "R", s.col == ""
    m = s.m * np.where(red_i, fr, fg)
    u = s.u * np.where(red_i | t2, fr, fg)
    return MethylSignals(m, u, s.oob_m, s.oob_u, s.masked, s.col, s.ctrl_green * fg, s.ctrl_red * fr, s.ctrl_type)


def dye_bias_linear(s: MethylSignals) -> MethylSignals:
    """Scale the channels so the normalization-control means (NORM_A/T red, NORM_C/G green) meet at their average."""
    t = np.array([x.upper() for x in s.ctrl_type])
    red = _finite(s.ctrl_red[np.isin(t, ("NORM_A", "NORM_T"))])
    grn = _finite(s.ctrl_green[np.isin(t, ("NORM_C", "NORM_G"))])
    if not len(red) or not len(grn):
        raise ValueError("dye bias: normalization controls missing")
    mr, mg = float(red.mean()), float(grn.mean())
    ref = (mr + mg) / 2.0
    return _scale(s, ref / mg, ref / mr)


def dye_bias_nonlinear(s: MethylSignals) -> MethylSignals:
    """Quantile-match the Infinium I red and green in-band signals to their average quantile function; every signal
    of a channel is mapped through that channel's curve (linear extrapolation beyond the Infinium I range)."""
    red_i, grn_i, t2 = s.col == "R", s.col == "G", s.col == ""
    r_ref = np.sort(_finite(np.concatenate([s.m[red_i], s.u[red_i]])))
    g_ref = np.sort(_finite(np.concatenate([s.m[grn_i], s.u[grn_i]])))
    if len(r_ref) < 2 or len(g_ref) < 2:
        raise ValueError("dye bias: too few Infinium I probes")
    q = np.linspace(0, 1, 1001)
    rq, gq = np.quantile(r_ref, q), np.quantile(g_ref, q)
    target = (rq + gq) / 2.0

    def mapper(src_q: F64):  # type: ignore[no-untyped-def]
        lo_slope = target[0] / src_q[0] if src_q[0] > 0 else 1.0
        hi_slope = target[-1] / src_q[-1] if src_q[-1] > 0 else 1.0

        def f(x: F64) -> F64:
            y = np.interp(x, src_q, target)
            y = np.where(x < src_q[0], x * lo_slope, y)
            return np.where(x > src_q[-1], x * hi_slope, y)

        return f

    fr, fg = mapper(rq), mapper(gq)
    m = np.where(red_i, fr(s.m), fg(s.m))
    u = np.where(red_i | t2, fr(s.u), fg(s.u))
    return MethylSignals(m, u, s.oob_m, s.oob_u, s.masked, s.col, fg(s.ctrl_green), fr(s.ctrl_red), s.ctrl_type)


# ------------------------------------------------------------------------------------------------ beta
def beta_values(s: MethylSignals, offset: float = 0.0) -> F64:
    m = np.maximum(s.m, 1.0)
    u = np.maximum(s.u, 1.0)
    return m / (m + u + offset)


def m_values(beta: F64) -> F64:
    b = np.clip(beta, 1e-6, 1 - 1e-6)
    return np.log2(b / (1 - b))


def average_signals(s: MethylSignals, analytical: Mask) -> tuple[float, float]:
    """MET-012: mean green and red in-band signal over analytical probes."""
    red_i, grn_i, t2 = s.col == "R", s.col == "G", s.col == ""
    a = analytical
    g = _finite(np.concatenate([s.m[grn_i & a], s.u[grn_i & a], s.m[t2 & a]]))
    r = _finite(np.concatenate([s.m[red_i & a], s.u[red_i & a], s.u[t2 & a]]))
    return (float(g.mean()) if len(g) else math.nan, float(r.mean()) if len(r) else math.nan)