"""Programmatic parameter accounting and budget validation (PLAN section 6).

Section 6 asks for two things that the previous implementation did not do: a
*concrete* allocation across buckets, and programmatic validation of it.  Only
the model total was checked before, so an allocation that put 46% of the budget
in PAT (against a 15-25% target) and 1.4% in the heads (against 10-15%) passed.

:func:`validate_budget` therefore checks every bucket fraction as well as the
total, and reports all violations at once rather than failing on the first.
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Protocol, runtime_checkable

from torch import nn

#: Section 6 budget ranges, as fractions of total trainable parameters.
BUDGET: Mapping[str, tuple[float, float]] = {
    "graph_encoder": (0.10, 0.20),
    "pat": (0.15, 0.25),
    "reasoning_transformer": (0.45, 0.60),
    "heads": (0.10, 0.15),
    "embeddings": (0.05, 0.10),
}


@runtime_checkable
class BucketedModule(Protocol):
    """A module that can partition its own parameters into budget buckets."""

    def parameter_buckets(self) -> Mapping[str, int]: ...


@dataclass(frozen=True)
class ParameterReport:
    total: int
    trainable: int
    buckets: Mapping[str, int]

    @property
    def fractions(self) -> Mapping[str, float]:
        return {
            name: (count / self.total if self.total else 0.0)
            for name, count in self.buckets.items()
        }

    def to_dict(self) -> dict[str, object]:
        return {
            "total_parameters": self.total,
            "trainable_parameters": self.trainable,
            **dict(self.buckets),
            "fractions": {
                name: round(value, 5) for name, value in self.fractions.items()
            },
        }


def count_parameters(model: nn.Module) -> ParameterReport:
    total = sum(parameter.numel() for parameter in model.parameters())
    trainable = sum(
        parameter.numel() for parameter in model.parameters() if parameter.requires_grad
    )
    buckets = (
        dict(model.parameter_buckets()) if isinstance(model, BucketedModule) else {}
    )
    return ParameterReport(total, trainable, buckets)


def validate_target(
    report: ParameterReport, target: int, tolerance: float = 0.20
) -> None:
    """Check the trainable total lies within ``tolerance`` of ``target``."""
    if target <= 0:
        raise ValueError("target must be positive")
    lower, upper = int(target * (1 - tolerance)), int(target * (1 + tolerance))
    if not lower <= report.trainable <= upper:
        raise ValueError(
            f"trainable parameter count {report.trainable:,} is outside "
            f"[{lower:,}, {upper:,}] for target {target:,}"
        )


def budget_violations(
    report: ParameterReport, budget: Mapping[str, tuple[float, float]] | None = None
) -> tuple[str, ...]:
    """Return a human-readable violation per bucket outside its budget range."""
    ranges = budget or BUDGET
    fractions = report.fractions
    violations: list[str] = []
    for name, (low, high) in ranges.items():
        value = fractions.get(name)
        if value is None:
            violations.append(f"{name}: bucket missing from parameter report")
        elif not low <= value <= high:
            violations.append(
                f"{name}: {value * 100:.1f}% outside target {low * 100:.0f}-{high * 100:.0f}%"
            )
    return tuple(violations)


def validate_budget(
    report: ParameterReport, budget: Mapping[str, tuple[float, float]] | None = None
) -> None:
    """Fail with every budget violation listed, not just the first."""
    violations = budget_violations(report, budget)
    if violations:
        raise ValueError(
            "parameter allocation violates the PLAN section 6 budget:\n  "
            + "\n  ".join(violations)
        )