"""Ridge with task fixed effects, and cluster-bootstrap intervals on β.

The model (docs/07-ablation.md):

    y_{v,t} = α_t + Σ_k β_k · mask_{v,k} + ε

`β_k` is the marginal effect of *keeping* section k, averaged over the contexts
in which it appeared. The task fixed effects `α_t` absorb between-task
difficulty, which is the largest variance component and the reason N=1 per
(variant, task) is affordable at all: pooling across R ≈ 2K+8 variants gives
~R observations per task, and the fit borrows strength across tasks.

`α_t` is absorbed by the within transformation — demeaning y and every mask
column inside each task — rather than by adding T dummy columns. Identical
estimates, and it keeps the penalty off the fixed effects, which matters:
shrinking task difficulty toward zero would be actively wrong.

**Ridge, not LASSO** (ADR-0007). The interesting output is the sections whose
coefficient is reliably *near zero* — the deletion candidates — and LASSO zeroes
coefficients as a modelling choice, destroying the distinction between
"measured to be zero" and "shrunk to zero." That distinction is the product.

**What these intervals are and are not.** Two things make stage 1 output
screening-grade rather than confirmatory, and both are measured rather than
assumed (`tests/test_regression.py`):

- Ridge is a biased estimator. Coefficients land at roughly 75% of their true
  magnitude in simulation, so `|β|` is closer to a lower bound than to an
  estimate, and the interval covers the *penalized* estimand rather than the
  true effect. The ranking survives this — ~80% top-3 recall — which is why the
  ranking is the usable output and the numbers are not.
- Coverage on true-null sections is ~88% against a nominal 95%, in a world
  where other sections carry real effects; in a purely null world it is ~97%.
  λ is re-tuned inside every bootstrap resample to get there. Holding it fixed
  — the obvious implementation, and 25x cheaper — measures 82% and 89%
  respectively, because λ is chosen *from the data* and holding it constant
  treats a fitted quantity as if it were known in advance.

  How much that costs depends on the world, and both figures are worth stating:
  in a null world the fixed-λ interval is roughly **4x too narrow**; where real
  effects are present the gap falls to ~1.1x. The null world is the one that
  matters most for the decision this stage feeds — a deletion candidate *is* a
  section whose effect looks like zero — which is why the expensive path is the
  default.

ADR-0007 originally described the shrinkage as "conservative in the right
direction" for the deletion question. That is backwards: shrinking β toward
zero makes a section that *does* matter look more deletable, not less. What
actually contains the risk is stage 2 — a shrunk-looking section is exactly
what gets promoted to a confirmatory A/B — which is why that stage is not
optional. See the amendment in the ADR.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ablate.errors import AblateError
from ablate.stats.bootstrap import (
    MIN_TASKS_FOR_CI,
    bca_bounds,
    bca_corrections,
)

#: Penalties tried by leave-one-task-out CV. Spans "essentially unpenalized" to
#: "almost fully shrunk"; the data picks.
LAMBDA_GRID = np.logspace(-3, 3, 25)

#: Fewer resamples than the paired comparison's 10,000, because each one here
#: re-runs leave-one-task-out CV to pick its own λ — roughly 25x the work per
#: resample. 2,000 keeps a full screening analysis to a few seconds, which is
#: noise beside the minutes of agent runs that produced the data.
SCREENING_RESAMPLES = 2_000


@dataclass(frozen=True)
class SectionEffect:
    """One section's screening result.

    `p_value` is a bootstrap two-sided p, used only to order sections for
    Benjamini-Hochberg. It is a screening statistic and the report never
    presents it as a test.
    """

    section_id: str
    beta: float
    lo: float | None
    hi: float | None
    p_value: float

    @property
    def has_ci(self) -> bool:
        return self.lo is not None and self.hi is not None

    @property
    def spans_zero(self) -> bool:
        """True when the interval is consistent with no effect — which, for a
        large section, is what makes it a deletion candidate."""
        return not self.has_ci or (self.lo <= 0.0 <= self.hi)


@dataclass(frozen=True)
class RidgeFit:
    effects: tuple[SectionEffect, ...]
    lam: float
    n_tasks: int
    n_observations: int
    alpha: float
    n_resamples: int
    ci_withheld: str | None = None

    def by_id(self, section_id: str) -> SectionEffect | None:
        return next((e for e in self.effects if e.section_id == section_id), None)


@dataclass(frozen=True)
class _Blocks:
    """Per-task sufficient statistics, stacked, after the within transformation.

    Storing these rather than the raw rows is what makes the cluster bootstrap
    affordable: the normal equations decompose as sums over tasks, so a
    resample is a sum of K×K matrices rather than a refit over every
    observation.

    They are held as stacked arrays rather than a list of per-task objects for
    one specific reason — leave-one-task-out CV runs *inside* every resample,
    and rebuilding each fold's Gram sum from a Python list made a single
    screening analysis take sixteen minutes. With the sums stacked, a fold is
    `total - grams[i]`, which is arithmetic instead of iteration.
    """

    grams: np.ndarray  # (T, K, K)
    moments: np.ndarray  # (T, K)
    sum_squares: np.ndarray  # (T,)

    def __len__(self) -> int:
        return len(self.grams)

    @property
    def k(self) -> int:
        return self.grams.shape[1]

    def totals(self, indices: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
        if indices is None:
            return self.grams.sum(axis=0), self.moments.sum(axis=0)
        return self.grams[indices].sum(axis=0), self.moments[indices].sum(axis=0)


def _blocks(masks: np.ndarray, tasks: np.ndarray, y: np.ndarray) -> _Blocks:
    grams, moments, sums = [], [], []
    for task in np.unique(tasks):
        rows = tasks == task
        x = masks[rows].astype(float)
        target = y[rows].astype(float)
        # The within transformation: demeaning inside the task absorbs α_t
        # exactly, without spending a column on it or letting the ridge
        # penalty touch it.
        x = x - x.mean(axis=0)
        target = target - target.mean()
        grams.append(x.T @ x)
        moments.append(x.T @ target)
        sums.append(float(target @ target))

    k = masks.shape[1]
    return _Blocks(
        grams=np.array(grams).reshape(-1, k, k),
        moments=np.array(moments).reshape(-1, k),
        sum_squares=np.array(sums, dtype=float),
    )


def _solve(gram: np.ndarray, moment: np.ndarray, lam: float) -> np.ndarray:
    k = gram.shape[0]
    try:
        return np.linalg.solve(gram + lam * np.eye(k), moment)
    except np.linalg.LinAlgError:  # pragma: no cover — ridge makes this unreachable
        return np.linalg.lstsq(gram + lam * np.eye(k), moment, rcond=None)[0]


def _fit(blocks: _Blocks, lam: float, indices: np.ndarray | None = None) -> np.ndarray:
    gram, moment = blocks.totals(indices)
    return _solve(gram, moment, lam)


def _choose_lambda(blocks: _Blocks, indices: np.ndarray | None = None) -> float:
    """Leave-one-task-out CV (ADR-0007).

    Held out at *task* level, not observation level, because tasks are the
    resampling unit everywhere else, and a penalty tuned on held-out
    observations from a task the model has already seen would be tuned on
    information the estimator does not really have.

    Prediction error is evaluated on the held-out task's *demeaned* outcomes,
    so no estimate of that task's α is needed. Each fold's design is obtained
    by subtracting the held-out task from the total rather than re-summing the
    rest, which is what keeps this affordable inside a bootstrap loop.
    """
    grams = blocks.grams if indices is None else blocks.grams[indices]
    moments = blocks.moments if indices is None else blocks.moments[indices]
    squares = blocks.sum_squares if indices is None else blocks.sum_squares[indices]

    n = len(grams)
    if n < 2:
        return float(LAMBDA_GRID[len(LAMBDA_GRID) // 2])

    total_gram = grams.sum(axis=0)
    total_moment = moments.sum(axis=0)
    eye = np.eye(blocks.k)

    best_lam, best_error = float(LAMBDA_GRID[0]), np.inf
    for lam in LAMBDA_GRID:
        error = 0.0
        for held_out in range(n):
            beta = np.linalg.solve(
                total_gram - grams[held_out] + lam * eye,
                total_moment - moments[held_out],
            )
            # ||ỹ - X̃β||² expanded into the stored sufficient statistics.
            error += (
                squares[held_out]
                - 2.0 * float(beta @ moments[held_out])
                + float(beta @ grams[held_out] @ beta)
            )
        if error < best_error:
            best_lam, best_error = float(lam), error

    return best_lam


def fit_sections(
    section_ids: list[str] | tuple[str, ...],
    masks: np.ndarray,
    tasks: np.ndarray,
    y: np.ndarray,
    *,
    lam: float | None = None,
    alpha: float = 0.05,
    n_resamples: int = SCREENING_RESAMPLES,
    seed: int | None = None,
    retune_lambda: bool = True,
) -> RidgeFit:
    """Estimate each section's marginal effect, with BCa intervals.

    `masks` is (n_observations, K); `tasks` labels each observation's task;
    `y` is the outcome. One row per (variant, task, trial).

    Uncertainty comes from the same cluster bootstrap as every other interval
    in this tool — resample tasks, refit, take BCa percentiles — rather than
    from model-based standard errors. One inferential story throughout.
    """
    ids = tuple(section_ids)
    masks = np.asarray(masks, dtype=float)
    tasks = np.asarray(tasks)
    y = np.asarray(y, dtype=float)

    if masks.ndim != 2 or masks.shape[1] != len(ids):
        raise AblateError(
            f"design matrix is {masks.shape}, expected (n_observations, {len(ids)})"
        )
    if not (len(masks) == len(tasks) == len(y)):
        raise AblateError("masks, tasks and y must describe the same observations")

    blocks = _blocks(masks, tasks, y)
    n_tasks = len(blocks)
    if n_tasks == 0:
        raise AblateError("no observations to fit")

    penalty = lam if lam is not None else _choose_lambda(blocks)
    beta = _fit(blocks, penalty)

    if n_tasks < MIN_TASKS_FOR_CI:
        withheld = (
            f"{n_tasks} tasks — screening intervals need at least "
            f"{MIN_TASKS_FOR_CI}. Coefficients are shown without intervals."
        )
        return RidgeFit(
            effects=tuple(
                SectionEffect(sid, float(b), None, None, float("nan"))
                for sid, b in zip(ids, beta, strict=True)
            ),
            lam=penalty,
            n_tasks=n_tasks,
            n_observations=len(y),
            alpha=alpha,
            n_resamples=0,
            ci_withheld=withheld,
        )

    rng = np.random.default_rng(seed)
    draws = rng.integers(0, n_tasks, size=(n_resamples, n_tasks))

    # λ is re-selected inside every resample. It is chosen *from the data*, so
    # holding it fixed treats a fitted quantity as if it were known in advance,
    # and the resulting interval is narrower than the evidence supports —
    # measured at 89% coverage in a null world (nominal 95%), where it runs ~4x
    # too narrow. Re-tuning restores ~97% there, and lifts 82% to 88% in a world
    # with real effects. See the module docstring for both figures.
    def refit(indices: np.ndarray) -> np.ndarray:
        lam_i = _choose_lambda(blocks, indices) if retune_lambda else penalty
        return _fit(blocks, lam_i, indices)

    resampled = np.array([refit(idx) for idx in draws])
    jackknife = np.array(
        [refit(np.delete(np.arange(n_tasks), i)) for i in range(n_tasks)]
    )

    effects = []
    for k, sid in enumerate(ids):
        column = resampled[:, k]
        corrections = bca_corrections(column, float(beta[k]), jackknife[:, k])
        lo, hi = bca_bounds(column, corrections, alpha)
        effects.append(
            SectionEffect(
                section_id=sid,
                beta=float(beta[k]),
                lo=lo,
                hi=hi,
                p_value=_bootstrap_p(column),
            )
        )

    return RidgeFit(
        effects=tuple(effects),
        lam=penalty,
        n_tasks=n_tasks,
        n_observations=len(y),
        alpha=alpha,
        n_resamples=n_resamples,
    )


def _bootstrap_p(resampled: np.ndarray) -> float:
    """Two-sided bootstrap p for the null β = 0.

    Used only to order sections for Benjamini-Hochberg. The report ranks by
    effect size and never prints this as a test result.
    """
    if resampled.size == 0:
        return float("nan")
    below = float((resampled <= 0.0).mean())
    above = float((resampled >= 0.0).mean())
    return float(min(1.0, 2.0 * min(below, above)))
