"""Two-stage ablation: screen cheap, confirm properly (docs/07-ablation.md).

Stage 1 ranks every section from a randomized subset design on a cheap model.
Stage 2 takes a shortlist and runs dedicated, pre-specified, full-N paired A/Bs
on the target model.

The split is not an optimization. Cheap-model screening assumes config
sensitivity transfers across model tiers, which is unverified — and separately,
ridge shrinkage flattens coefficients toward zero, which makes a section that
genuinely matters look deletable. Stage 2 is what contains both risks, and it
is why stage 1 output is labelled *screening — not confirmatory* everywhere it
appears.

Trials are obtained through an injected runner so the pipeline can be tested
end to end without an agent, an API key, or a dollar.
"""

from __future__ import annotations

from collections.abc import Callable, Sequence
from dataclasses import dataclass

import numpy as np

from ablate.ablation.design import ScreeningDesign, build_design
from ablate.ablation.regression import RidgeFit, SectionEffect, fit_sections
from ablate.config.parse import Section, render_variant
from ablate.config.variant import ConfigVariant
from ablate.errors import AblateError
from ablate.record import TrialRecord
from ablate.report.analysis import Analysis, analyze
from ablate.stats import DEFAULT_FDR_Q, Verdict, benjamini_hochberg

#: How many sections get a confirmatory A/B by default. Each one costs a full
#: paired run, so this is the main lever on stage-2 spend.
DEFAULT_CONFIRMATIONS = 3

#: Trials per (variant, task) during screening. One is the point: task fixed
#: effects plus R ≈ 2K+8 variants already give ~R observations per task.
SCREENING_TRIALS = 1

RunnerFn = Callable[[Sequence[ConfigVariant], int, str], list[TrialRecord]]
"""Run every (variant × task × replicate) cell and return the records.

Takes the variants, the replicate count, and a stage label for progress
output. The pipeline never touches worktrees or agents itself.
"""


@dataclass(frozen=True)
class ScreenRow:
    """One section's screening result, ready to render."""

    section: Section
    effect: SectionEffect
    q_value: float
    significant: bool

    @property
    def tokens(self) -> int:
        return self.section.tokens

    @property
    def is_deletion_candidate(self) -> bool:
        """Consistent with no effect — so its token cost is unpaid-for weight.

        These are the commercially interesting rows: an EQUIVALENT verdict on a
        large section is what licenses a deletion.
        """
        return self.effect.spans_zero


@dataclass(frozen=True)
class ScreeningResult:
    design: ScreeningDesign
    rows: tuple[ScreenRow, ...]
    fit: RidgeFit
    model: str
    n_records: int
    reference_pass_rate: float
    floor_pass_rate: float

    @property
    def ranked(self) -> list[ScreenRow]:
        return sorted(self.rows, key=lambda r: -abs(r.effect.beta))

    @property
    def floor_delta(self) -> float:
        """Full config minus empty config — "what if I delete the whole thing?"."""
        return self.reference_pass_rate - self.floor_pass_rate


@dataclass(frozen=True)
class Confirmation:
    """A pre-specified A/B: the whole config, versus the config minus one section."""

    section: Section
    analysis: Analysis

    @property
    def pass_rate_verdict(self) -> Verdict:
        metric = self.analysis.metric("pass_rate")
        return metric.verdict.verdict if metric else Verdict.INCONCLUSIVE

    @property
    def safe_to_delete(self) -> bool:
        """Only EQUIVALENT licenses a deletion.

        INCONCLUSIVE must never reach here: "we didn't look hard enough" is not
        permission to delete, and collapsing it with "we looked and found
        nothing" is the specific error the four-verdict taxonomy exists to
        prevent.
        """
        return self.pass_rate_verdict is Verdict.EQUIVALENT


@dataclass(frozen=True)
class AblationResult:
    sections: tuple[Section, ...]
    screening: ScreeningResult
    confirmations: tuple[Confirmation, ...]
    config_path: str
    total_tokens: int

    @property
    def deletions(self) -> list[Confirmation]:
        return [c for c in self.confirmations if c.safe_to_delete]

    @property
    def keeps(self) -> list[Confirmation]:
        return [
            c for c in self.confirmations
            if c.pass_rate_verdict in (Verdict.REGRESSION, Verdict.IMPROVEMENT)
        ]

    @property
    def unresolved(self) -> list[Confirmation]:
        return [c for c in self.confirmations if c.pass_rate_verdict is Verdict.INCONCLUSIVE]

    @property
    def tokens_saved(self) -> int:
        return sum(c.section.tokens for c in self.deletions)

    @property
    def unscreened(self) -> list[ScreenRow]:
        """Sections that never reached a confirmatory run.

        Named and priced in the report rather than dropped. "Six sections we
        couldn't resolve, here's what resolving them costs" is a service;
        omitting them is a lie by selection.
        """
        confirmed = {c.section.id for c in self.confirmations}
        return [r for r in self.screening.rows if r.section.id not in confirmed]


def ablatable(sections: Sequence[Section]) -> list[Section]:
    """Sections eligible for a coefficient.

    Pinned sections are excluded: they appear in every variant, so they have no
    estimable effect and would be a constant column in the design.
    """
    return [s for s in sections if not s.required]


def screen(
    sections: Sequence[Section],
    runner: RunnerFn,
    *,
    model: str,
    seed: int | None = None,
    alpha: float = 0.05,
    q: float = DEFAULT_FDR_Q,
    n_resamples: int | None = None,
) -> ScreeningResult:
    """Stage 1: rank every ablatable section from a randomized subset design."""
    candidates = ablatable(sections)
    design = build_design([s.id for s in candidates], seed=seed)

    variants = [
        ConfigVariant(
            id=f"v{index:02d}",
            text=render_variant(list(sections), _full_mask(sections, candidates, mask)),
            source=_variant_label(index, design),
        )
        for index, mask in enumerate(design.masks)
    ]

    records = [r for r in runner(variants, SCREENING_TRIALS, "screening") if r.error is None]
    if not records:
        raise AblateError("screening produced no usable trials")

    by_variant = {v.id: index for index, v in enumerate(variants)}
    rows_mask, rows_task, rows_y = [], [], []
    for record in records:
        index = by_variant.get(record.variant_id)
        if index is None:
            continue
        rows_mask.append(design.masks[index])
        rows_task.append(record.task_id)
        rows_y.append(float(record.passed))

    fit = fit_sections(
        design.section_ids,
        np.array(rows_mask, dtype=float),
        np.array(rows_task),
        np.array(rows_y, dtype=float),
        alpha=alpha,
        seed=seed,
        **({"n_resamples": n_resamples} if n_resamples else {}),
    )

    fdr = benjamini_hochberg([e.p_value for e in fit.effects], q=q)
    by_id = {s.id: s for s in candidates}
    rows = tuple(
        ScreenRow(
            section=by_id[effect.section_id],
            effect=effect,
            q_value=fdr.adjusted[index],
            significant=fdr.rejected[index],
        )
        for index, effect in enumerate(fit.effects)
    )

    return ScreeningResult(
        design=design,
        rows=rows,
        fit=fit,
        model=model,
        n_records=len(records),
        reference_pass_rate=_variant_pass_rate(records, variants[design.reference_index].id),
        floor_pass_rate=_variant_pass_rate(records, variants[design.floor_index].id),
    )


def _full_mask(
    sections: Sequence[Section], candidates: Sequence[Section], mask: Sequence[bool]
) -> list[bool]:
    """Widen a design mask over ablatable sections to every section.

    Pinned sections are marked True here for clarity; `render_variant` forces
    them in regardless, so the two agree.
    """
    chosen = dict(zip((s.id for s in candidates), mask, strict=True))
    return [chosen.get(s.id, True) for s in sections]


def _variant_label(index: int, design: ScreeningDesign) -> str:
    if index == design.reference_index:
        return "all sections (reference)"
    if index == design.floor_index:
        return "no sections (empty-config floor)"
    return f"screening variant {index}"


def _variant_pass_rate(records: Sequence[TrialRecord], variant_id: str) -> float:
    subset = [r for r in records if r.variant_id == variant_id]
    return sum(r.passed for r in subset) / len(subset) if subset else 0.0


def select_candidates(
    screening: ScreeningResult, *, limit: int = DEFAULT_CONFIRMATIONS
) -> list[Section]:
    """Choose which sections earn a confirmatory A/B.

    Two kinds, both wanted for different reasons:

    - the strongest mover by |β|, because if screening thinks a section
      carries the config, that is the claim most worth checking before anyone
      deletes around it;
    - the largest **deletion candidates** — big sections whose effect is
      consistent with zero — because those are the ones a user would actually
      act on, and because ridge shrinkage can push a section that matters into
      this group.
    """
    if limit <= 0:
        return []

    ranked = screening.ranked
    chosen: list[ScreenRow] = []

    movers = [r for r in ranked if not r.is_deletion_candidate]
    if movers:
        chosen.append(movers[0])

    for row in sorted(
        (r for r in ranked if r.is_deletion_candidate), key=lambda r: -r.tokens
    ):
        if len(chosen) >= limit:
            break
        chosen.append(row)

    for row in ranked:  # backfill if there were too few of either kind
        if len(chosen) >= limit:
            break
        if row not in chosen:
            chosen.append(row)

    return [r.section for r in chosen[:limit]]


def confirm(
    sections: Sequence[Section],
    candidates: Sequence[Section],
    runner: RunnerFn,
    *,
    trials: int,
    margin_overrides: dict[str, float] | None = None,
    alpha: float = 0.05,
    seed: int | None = None,
    agent: str = "",
    agent_version: str | None = None,
) -> list[Confirmation]:
    """Stage 2: a dedicated paired A/B per candidate, at full N on the target model.

    **Arm A is the config *without* the section; arm B is the full config.**
    A is the proposed change everywhere else in this tool — in CI it is the
    PR's config and B is the base branch — and the proposed change here is the
    deletion. Orienting it that way makes every verdict read directly:
    REGRESSION means dropping the section hurts, EQUIVALENT means it is safe to
    drop, and `--fail-on regression` keeps its meaning. Putting the full config
    in A would invert every row under a "dropping…" heading.

    The full-config arm is rendered once and shared across candidates rather
    than re-run per comparison. It is byte-identical every time, so re-running
    it would buy nothing but the same trials at three times the price. Each
    comparison is still pre-specified and reported on its own, with no
    multiplicity adjustment claimed across them.
    """
    if not candidates:
        return []

    all_kept = [True] * len(sections)
    variants = [
        ConfigVariant(id="B", text=render_variant(list(sections), all_kept), source="full config")
    ]
    for section in candidates:
        mask = [s.id != section.id for s in sections]
        variants.append(
            ConfigVariant(
                id=f"A-{section.id}",
                text=render_variant(list(sections), mask),
                source=f"without {section.id}",
            )
        )

    records = runner(variants, trials, "confirmatory")
    baseline = [r for r in records if r.variant_id == "B"]

    confirmations = []
    for section in candidates:
        without = [r for r in records if r.variant_id == f"A-{section.id}"]
        if not without or not baseline:
            continue
        paired = baseline + [_relabel(r, "A") for r in without]
        confirmations.append(
            Confirmation(
                section=section,
                analysis=analyze(
                    paired,
                    ("A", "B"),
                    margin_overrides=margin_overrides,
                    alpha=alpha,
                    seed=seed,
                    agent=agent,
                    agent_version=agent_version,
                ),
            )
        )
    return confirmations


def _relabel(record: TrialRecord, variant_id: str) -> TrialRecord:
    return record.model_copy(update={"variant_id": variant_id})
