"""The screening design: which sections are present in which variant.

`R = 2K + 8` variants. Two are fixed and carry meaning on their own:

- **variant 0, all kept** — the reference, i.e. the user's actual config.
- **variant 1, all dropped** — the empty-config floor, which answers the
  question most users actually have ("what happens if I delete the whole
  thing?") and comes out of this design for free.

The rest keep each section independently with p=0.5. That is a standard
fractional-factorial screening design, and the reason it beats exhaustive
leave-one-out at equal budget: every trial varies many sections at once, so
every trial informs every β. Leave-one-out spends a whole task-set learning
about one section.

**Conditioning is checked, not assumed** (ADR-0007). A design where a section
never varies, or where two sections are strongly collinear, produces
confidently wrong coefficients — the exact failure mode this project exists to
avoid. Designs that fail are resampled rather than fitted and caveated.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ablate.errors import AblateError

#: Beyond 2K, extra variants buy precision on every β at once. Eight is the
#: constant from docs/07-ablation.md; it also keeps very small K from producing
#: a design too short to condition.
EXTRA_VARIANTS = 8

#: Above this, two sections' effects cannot be told apart and the coefficients
#: become an artefact of which one the fit happened to favour.
MAX_PAIRWISE_CORR = 0.5

#: Resampling is cheap (microseconds) and failure is loud, so this can be
#: generous without hiding a real problem.
MAX_ATTEMPTS = 500


@dataclass(frozen=True)
class DesignConditioning:
    """Why a candidate design was accepted or rejected."""

    ok: bool
    constant_columns: tuple[int, ...]
    max_abs_corr: float
    worst_pair: tuple[int, int] | None

    def describe(self, section_ids: tuple[str, ...]) -> str:
        if self.ok:
            return f"well-conditioned, max pairwise |corr| = {self.max_abs_corr:.2f}"
        if self.constant_columns:
            names = ", ".join(section_ids[i] for i in self.constant_columns)
            return f"section(s) {names} never varied across variants"
        if self.worst_pair is not None:
            a, b = self.worst_pair
            return (
                f"{section_ids[a]} and {section_ids[b]} are collinear "
                f"(|corr| = {self.max_abs_corr:.2f})"
            )
        return "ill-conditioned"


@dataclass(frozen=True)
class ScreeningDesign:
    """R variants over K ablatable sections.

    `masks[v][k]` is True when section k is present in variant v. Pinned
    sections are not represented here at all — they are present in every
    variant by definition, so they have no coefficient to estimate and would
    be a constant column if included.
    """

    section_ids: tuple[str, ...]
    masks: tuple[tuple[bool, ...], ...]
    conditioning: DesignConditioning
    attempts: int
    seed: int | None

    @property
    def n_variants(self) -> int:
        return len(self.masks)

    @property
    def n_sections(self) -> int:
        return len(self.section_ids)

    @property
    def reference_index(self) -> int:
        """The all-kept variant: the config exactly as the user wrote it."""
        return 0

    @property
    def floor_index(self) -> int:
        """The all-dropped variant: the empty-config floor."""
        return 1

    def matrix(self) -> np.ndarray:
        return np.array(self.masks, dtype=float)

    def n_trials(self, n_tasks: int) -> int:
        """Screening runs one trial per (variant, task)."""
        return self.n_variants * n_tasks


def _conditioning(matrix: np.ndarray) -> DesignConditioning:
    """Reject constant columns first, then collinear pairs.

    Order matters: a constant column has zero variance, so its correlation is
    undefined, and reporting "collinear" for it would misdescribe the problem.
    """
    constant = tuple(int(i) for i in np.flatnonzero(matrix.std(axis=0) == 0))
    if constant:
        return DesignConditioning(False, constant, float("nan"), None)

    if matrix.shape[1] < 2:
        return DesignConditioning(True, (), 0.0, None)

    corr = np.corrcoef(matrix, rowvar=False)
    np.fill_diagonal(corr, 0.0)
    flat = np.abs(np.nan_to_num(corr))
    worst = int(np.argmax(flat))
    a, b = divmod(worst, flat.shape[1])
    max_corr = float(flat[a, b])

    return DesignConditioning(
        ok=max_corr < MAX_PAIRWISE_CORR,
        constant_columns=(),
        max_abs_corr=max_corr,
        worst_pair=(a, b) if max_corr >= MAX_PAIRWISE_CORR else None,
    )


def build_design(
    section_ids: list[str] | tuple[str, ...],
    *,
    seed: int | None = None,
    n_variants: int | None = None,
    max_attempts: int = MAX_ATTEMPTS,
) -> ScreeningDesign:
    """Sample a well-conditioned screening design over `section_ids`.

    Raises if no conditioned design is found. That is deliberate: silently
    proceeding on a collinear design would produce a confident ranking that
    means nothing, and there is no honest caveat that repairs it.
    """
    ids = tuple(section_ids)
    k = len(ids)
    if k == 0:
        raise AblateError(
            "nothing to ablate — every section is pinned with <!-- ablate:keep -->, "
            "or the config has no sections"
        )

    total = n_variants if n_variants is not None else 2 * k + EXTRA_VARIANTS
    if total < 3:
        raise AblateError(f"a screening design needs at least 3 variants, got {total}")

    rng = np.random.default_rng(seed)
    reference = np.ones(k, dtype=bool)
    floor = np.zeros(k, dtype=bool)

    conditioning = DesignConditioning(False, (), float("nan"), None)
    for attempt in range(1, max_attempts + 1):
        random_block = rng.random((total - 2, k)) < 0.5
        matrix = np.vstack([reference, floor, random_block])

        conditioning = _conditioning(matrix.astype(float))
        if conditioning.ok:
            return ScreeningDesign(
                section_ids=ids,
                masks=tuple(tuple(bool(v) for v in row) for row in matrix),
                conditioning=conditioning,
                attempts=attempt,
                seed=seed,
            )

    raise AblateError(
        f"could not find a well-conditioned design for {k} sections in "
        f"{max_attempts} attempts ({conditioning.describe(ids)}). With very few "
        f"sections the all-kept and all-dropped variants dominate the design; "
        f"try --granularity bullet to split the config further, or compare "
        f"whole configs with `ablate compare` instead."
    )
