"""Parameter-group helpers for partial-gradient importance proxies."""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import nn


@dataclass(frozen=True)
class PartialGradientGroup:
    """A named, non-empty subset of action-model parameters."""

    name: str
    parameter_names: tuple[str, ...]
    parameter_count: int


def define_partial_gradient_groups(
    action_model: nn.Module,
) -> dict[str, PartialGradientGroup]:
    """Return the four parameter groups used for partial-gradient analysis."""
    named_parameters = dict(action_model.named_parameters())
    if not named_parameters:
        raise ValueError("action model has no parameters")

    transformer_blocks = getattr(
        getattr(action_model, "model", None),
        "transformer_blocks",
        None,
    )
    if transformer_blocks is None or len(transformer_blocks) == 0:
        raise ValueError("action model has no DiT transformer blocks")
    last_block_prefix = f"model.transformer_blocks.{len(transformer_blocks) - 1}."

    selectors = {
        "decoder": lambda name: name.startswith("action_decoder."),
        "output_head": lambda name: (
            name.startswith("action_decoder.")
            or name.startswith("model.proj_out_1.")
            or name.startswith("model.proj_out_2.")
        ),
        "last_block_and_head": lambda name: (
            name.startswith(last_block_prefix)
            or name.startswith("action_decoder.")
            or name.startswith("model.proj_out_1.")
            or name.startswith("model.proj_out_2.")
        ),
        "normalization": lambda name: "norm" in name.lower(),
    }

    groups = {}
    for group_name, selector in selectors.items():
        selected_names = tuple(
            name for name in named_parameters if selector(name)
        )
        if not selected_names:
            raise ValueError(
                f"partial-gradient group {group_name!r} has no parameters"
            )
        groups[group_name] = PartialGradientGroup(
            name=group_name,
            parameter_names=selected_names,
            parameter_count=sum(
                named_parameters[name].numel() for name in selected_names
            ),
        )
    return groups


def activate_partial_gradient_group(
    action_model: nn.Module,
    group: PartialGradientGroup,
) -> list[nn.Parameter]:
    """Freeze the action model except for one exact, validated group."""
    named_parameters = dict(action_model.named_parameters())
    requested = set(group.parameter_names)
    missing = requested - set(named_parameters)
    if missing:
        raise ValueError(
            f"group {group.name!r} contains unknown parameters: {sorted(missing)}"
        )

    action_model.requires_grad_(False)
    selected = []
    for name, parameter in named_parameters.items():
        if name in requested:
            parameter.requires_grad_(True)
            selected.append(parameter)
    if not selected:
        raise RuntimeError(f"group {group.name!r} activated no parameters")
    observed_count = sum(parameter.numel() for parameter in selected)
    if observed_count != group.parameter_count:
        raise RuntimeError(
            f"group {group.name!r} parameter count changed from "
            f"{group.parameter_count} to {observed_count}"
        )
    return selected


def finite_gradient_l2_norm(parameters: list[nn.Parameter]) -> float:
    """Compute an exact finite L2 norm over the gradients that were produced."""
    component_norms = []
    for parameter in parameters:
        if parameter.grad is None:
            continue
        gradient = parameter.grad.detach()
        if not torch.isfinite(gradient).all():
            raise FloatingPointError("partial-gradient proxy is non-finite")
        component_norms.append(torch.linalg.vector_norm(gradient.float()))
    if not component_norms:
        raise RuntimeError("partial-gradient group received no gradients")
    return float(torch.linalg.vector_norm(torch.stack(component_norms)).item())