"""Minimum detectable effect and the price of an answer (docs/05-statistics.md).

This module exists so the tool can refuse to imply it sees effects it cannot
see. Its two jobs:

1. Before a run, say what effect size the design could resolve — so the user
   can decline to spend $28 learning "inconclusive."
2. After an INCONCLUSIVE verdict, say what a conclusive answer would cost, in
   tasks, trials, runs and dollars. That converts "the tool failed" into "the
   tool told me what an answer costs," which is a service.

For a paired design:

    SE(Δ̂) ≈ sqrt( (σ²_δ + E[p_A(1−p_A) + p_B(1−p_B)] / N) / T )

`T` divides the whole variance; `N` divides only the within-task component.
That asymmetry is why the guidance is always "more tasks before more trials."
"""

from __future__ import annotations

from dataclasses import dataclass
from statistics import NormalDist

_NORM = NormalDist()

#: Between-task SD of the *true* per-task effect. Tasks genuinely differ in how
#: much a config change helps them, and that heterogeneity does not shrink with
#: more trials. 0.05 is the planning figure from docs/05-statistics.md.
DEFAULT_EFFECT_SD = 0.05

#: Per-trial coefficient of variation for token counts, from docs/05-statistics.md.
DEFAULT_TOKEN_CV = 0.25


def _z_multiplier(alpha: float, power: float) -> float:
    """z_{1−α/2} + z_{power} — 2.80 at the conventional α=0.05, 80% power."""
    return _NORM.inv_cdf(1.0 - alpha / 2.0) + _NORM.inv_cdf(power)


def standard_error(
    n_tasks: int,
    n_trials: int,
    *,
    within_task_variance: float,
    effect_sd: float = DEFAULT_EFFECT_SD,
) -> float:
    """SE(Δ̂) for a balanced paired design."""
    if n_tasks <= 0 or n_trials <= 0:
        return float("inf")
    return ((effect_sd**2 + within_task_variance / n_trials) / n_tasks) ** 0.5


def mde_pass_rate(
    n_tasks: int,
    n_trials: int,
    *,
    p_a: float = 0.5,
    p_b: float = 0.5,
    effect_sd: float = DEFAULT_EFFECT_SD,
    alpha: float = 0.05,
    power: float = 0.80,
) -> float:
    """Smallest pass-rate difference this design could detect, as a fraction.

    Defaults are the worst case (p ≈ 0.5, where binomial variance peaks). Pass
    the observed rates for a design-specific figure; at p ≈ 0.8 the MDE improves
    by roughly 20% — not by an order of magnitude.
    """
    within = p_a * (1.0 - p_a) + p_b * (1.0 - p_b)
    se = standard_error(n_tasks, n_trials, within_task_variance=within, effect_sd=effect_sd)
    return _z_multiplier(alpha, power) * se


def mde_continuous(
    n_tasks: int,
    n_trials: int,
    *,
    cv: float = DEFAULT_TOKEN_CV,
    effect_sd: float = 0.0,
    alpha: float = 0.05,
    power: float = 0.80,
) -> float:
    """Smallest detectable *relative* change in a continuous metric.

    Expressed as a fraction of the mean, so 0.10 means "a 10% change in tokens
    would be detectable." `cv` is the per-trial coefficient of variation.
    """
    within = 2.0 * cv**2
    se = standard_error(n_tasks, n_trials, within_task_variance=within, effect_sd=effect_sd)
    return _z_multiplier(alpha, power) * se


@dataclass(frozen=True)
class DesignRequirement:
    """What it would take to resolve an effect of a given size."""

    n_tasks: int
    n_trials: int
    n_runs: int
    target: float
    achieved_mde: float
    feasible: bool

    def cost_usd(self, per_trial_usd: float) -> float:
        return self.n_runs * per_trial_usd


def required_design(
    target: float,
    *,
    n_trials: int = 10,
    p_a: float = 0.5,
    p_b: float = 0.5,
    effect_sd: float = DEFAULT_EFFECT_SD,
    alpha: float = 0.05,
    power: float = 0.80,
    max_tasks: int = 100_000,
    n_arms: int = 2,
) -> DesignRequirement:
    """Tasks needed to bring the pass-rate MDE down to `target`, at fixed N.

    Solved in closed form rather than searched: SE² is linear in 1/T, so

        T = (σ²_δ + V/N) / (target / z)²

    Holding N fixed and solving for T matches how the result is spoken about
    ("~10 trials × 40 tasks") and reflects the standing advice to buy tasks
    before trials.

    `feasible=False` means the between-task heterogeneity floor alone exceeds
    the target: no number of tasks at this N resolves it, because σ²_δ/T only
    vanishes as T → ∞ faster than the budget does. The honest answer there is
    "not at any price you will pay," and the report must say so.
    """
    z = _z_multiplier(alpha, power)
    if target <= 0:
        raise ValueError("target effect size must be positive")

    within = p_a * (1.0 - p_a) + p_b * (1.0 - p_b)
    variance_per_task = effect_sd**2 + within / n_trials
    needed = variance_per_task / (target / z) ** 2

    n_tasks = max(int(-(-needed // 1)), 1)  # ceil
    feasible = n_tasks <= max_tasks
    n_tasks = min(n_tasks, max_tasks)

    return DesignRequirement(
        n_tasks=n_tasks,
        n_trials=n_trials,
        n_runs=n_tasks * n_trials * n_arms,
        target=target,
        achieved_mde=mde_pass_rate(
            n_tasks, n_trials, p_a=p_a, p_b=p_b, effect_sd=effect_sd, alpha=alpha, power=power
        ),
        feasible=feasible,
    )
