"""Matched decoder-Transformer components for local-learning crossovers.

The forward graph is deliberately identical for BP, ordinary FA, clean KP,
and SDIL.  Only the vector transported through parameterized affine maps is
changed.  Parameter-free Jacobians (residual addition, LayerNorm, GELU, and
softmax attention) remain local and exact.
"""
from dataclasses import dataclass
import math
from typing import Dict, Iterable, Optional

import torch
from torch import nn
import torch.nn.functional as F


_FEEDBACK_METHODS = {"fa", "clean_kp", "sdil"}
_SUPPORTED_METHODS = {
    "bp", "dfa", "pepita", "ff", "dualprop", "ep"} | _FEEDBACK_METHODS


class _FeedbackLinearFunction(torch.autograd.Function):
    """Linear map with fixed or locally plastic feedback.

    ``feedback`` has the same orientation as ``weight``.  Consequently the
    transported vector is ``delta @ feedback``, while both plastic matrices
    can be updated from the locally available ``delta.T @ input`` correlation.
    The feedback correlation is recomputed rather than copied from the
    forward-weight gradient.
    """

    @staticmethod
    def forward(
            ctx, x, weight, feedback, bias, method_code, traffic_ratio,
            raw_rms, innovation_rms, traffic_rms):
        output = F.linear(x, weight, bias)
        ctx.save_for_backward(x, feedback, output)
        ctx.method_code = int(method_code)
        ctx.traffic_ratio = float(traffic_ratio)
        ctx.has_bias = bias is not None
        ctx.raw_rms = raw_rms
        ctx.innovation_rms = innovation_rms
        ctx.traffic_rms = traffic_rms
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, feedback, soma = ctx.saved_tensors
        task_instruction = grad_output
        traffic = torch.zeros_like(task_instruction)
        if ctx.method_code == 2 and ctx.traffic_ratio:
            # A paired neutral observation exposes the component predictable
            # from the local somatic response.  Match it to a frozen multiple
            # of task-instruction RMS without changing its somatic direction.
            task_rms = task_instruction.square().mean().sqrt()
            centered_soma = soma - soma.mean(dim=-1, keepdim=True)
            soma_rms = centered_soma.square().mean().sqrt().clamp_min(1e-30)
            traffic = (
                centered_soma * task_rms * ctx.traffic_ratio / soma_rms)
        raw_apical = task_instruction + traffic
        neutral_prediction = traffic
        innovation = raw_apical - neutral_prediction

        input_flat = x.reshape(-1, x.shape[-1])
        delta_flat = innovation.reshape(-1, innovation.shape[-1])
        grad_input = innovation @ feedback
        grad_weight = delta_flat.t() @ input_flat
        grad_feedback = None
        if ctx.method_code in (1, 2):
            # This is intentionally a second evaluation of the local
            # correlation, not an assignment from grad_weight.
            grad_feedback = delta_flat.t() @ input_flat
        grad_bias = None
        if ctx.has_bias:
            grad_bias = delta_flat.sum(dim=0)

        with torch.no_grad():
            ctx.raw_rms.copy_(raw_apical.square().mean().sqrt())
            ctx.innovation_rms.copy_(innovation.square().mean().sqrt())
            ctx.traffic_rms.copy_(traffic.square().mean().sqrt())
        return (
            grad_input, grad_weight, grad_feedback, grad_bias,
            None, None, None, None, None)


class FeedbackLinear(nn.Module):
    """A forward-matched affine map for BP, FA, clean KP, or SDIL."""

    def __init__(
            self, in_features: int, out_features: int, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator,
            bias: bool = False, init_std: float = 0.02,
            traffic_ratio: float = 4.0, dtype=torch.float32):
        super().__init__()
        if method not in _SUPPORTED_METHODS:
            raise ValueError(f"unsupported feedback-linear method: {method}")
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.method = method
        self.traffic_ratio = float(traffic_ratio if method == "sdil" else 0.0)
        self.weight = nn.Parameter(torch.empty(
            out_features, in_features, dtype=dtype))
        nn.init.normal_(
            self.weight, mean=0.0, std=init_std,
            generator=forward_generator)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=dtype))
        else:
            self.register_parameter("bias", None)

        if method in _FEEDBACK_METHODS:
            feedback = torch.empty(
                out_features, in_features, dtype=dtype)
            nn.init.normal_(
                feedback, mean=0.0, std=init_std,
                generator=feedback_generator)
            if method == "fa":
                self.register_buffer("feedback", feedback)
            else:
                self.feedback = nn.Parameter(feedback)
        else:
            self.register_buffer("feedback", None)
        self.register_buffer("last_raw_rms", torch.zeros((), dtype=dtype))
        self.register_buffer(
            "last_innovation_rms", torch.zeros((), dtype=dtype))
        self.register_buffer("last_traffic_rms", torch.zeros((), dtype=dtype))

    def forward(self, x):
        if self.method in {
                "bp", "dfa", "pepita", "ff", "dualprop", "ep"}:
            return F.linear(x, self.weight, self.bias)
        method_code = {"fa": 0, "clean_kp": 1, "sdil": 2}[self.method]
        return _FeedbackLinearFunction.apply(
            x, self.weight, self.feedback, self.bias, method_code,
            self.traffic_ratio, self.last_raw_rms,
            self.last_innovation_rms, self.last_traffic_rms)

    def extra_repr(self):
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, method={self.method}")


@dataclass(frozen=True)
class LocalTransformerConfig:
    vocab_size: int = 65
    context_length: int = 64
    depth: int = 4
    width: int = 128
    heads: int = 4
    mlp_ratio: int = 4
    dropout: float = 0.0
    bias: bool = False
    init_std: float = 0.02
    traffic_ratio: float = 4.0
    seed: int = 2027

    def __post_init__(self):
        if self.width % self.heads:
            raise ValueError("width must be divisible by heads")
        if self.depth < 1 or self.context_length < 1:
            raise ValueError("depth and context_length must be positive")


class LocalCausalSelfAttention(nn.Module):

    def __init__(
            self, config: LocalTransformerConfig, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator, dtype=torch.float32):
        super().__init__()
        self.heads = config.heads
        self.head_width = config.width // config.heads
        self.width = config.width
        common = dict(
            method=method, forward_generator=forward_generator,
            feedback_generator=feedback_generator, bias=config.bias,
            init_std=config.init_std, traffic_ratio=config.traffic_ratio,
            dtype=dtype)
        self.q = FeedbackLinear(config.width, config.width, **common)
        self.k = FeedbackLinear(config.width, config.width, **common)
        self.v = FeedbackLinear(config.width, config.width, **common)
        self.output = FeedbackLinear(config.width, config.width, **common)
        causal = torch.tril(torch.ones(
            config.context_length, config.context_length, dtype=torch.bool))
        self.register_buffer("causal_mask", causal, persistent=False)
        self.dropout = float(config.dropout)

    def forward(self, x):
        batch, time, width = x.shape

        def split_heads(value):
            return value.view(
                batch, time, self.heads, self.head_width).transpose(1, 2)

        query = split_heads(self.q(x))
        key = split_heads(self.k(x))
        value = split_heads(self.v(x))
        scores = query @ key.transpose(-2, -1)
        scores = scores * self.head_width ** -0.5
        mask = self.causal_mask[:time, :time]
        scores = scores.masked_fill(~mask, float("-inf"))
        attention = F.softmax(scores, dim=-1)
        attention = F.dropout(
            attention, p=self.dropout, training=self.training)
        mixed = attention @ value
        mixed = mixed.transpose(1, 2).contiguous().view(batch, time, width)
        return self.output(mixed)


class LocalTransformerBlock(nn.Module):

    def __init__(
            self, config: LocalTransformerConfig, method: str,
            forward_generator: torch.Generator,
            feedback_generator: torch.Generator, dtype=torch.float32):
        super().__init__()
        self.ln_attention = nn.LayerNorm(config.width, dtype=dtype)
        self.attention = LocalCausalSelfAttention(
            config, method, forward_generator, feedback_generator, dtype)
        self.ln_mlp = nn.LayerNorm(config.width, dtype=dtype)
        hidden = config.mlp_ratio * config.width
        common = dict(
            method=method, forward_generator=forward_generator,
            feedback_generator=feedback_generator, bias=config.bias,
            init_std=config.init_std, traffic_ratio=config.traffic_ratio,
            dtype=dtype)
        self.mlp_in = FeedbackLinear(config.width, hidden, **common)
        self.mlp_out = FeedbackLinear(hidden, config.width, **common)
        self.dropout = float(config.dropout)

    def forward(self, x):
        x = x + F.dropout(
            self.attention(self.ln_attention(x)),
            p=self.dropout, training=self.training)
        x = x + F.dropout(
            self.mlp_out(F.gelu(self.mlp_in(self.ln_mlp(x)))),
            p=self.dropout, training=self.training)
        return x


class LocalDecoderTransformer(nn.Module):
    """Depth-scaled, forward-matched character decoder."""

    def __init__(
            self, config: LocalTransformerConfig, method: str = "bp",
            dtype=torch.float32):
        super().__init__()
        if method not in _SUPPORTED_METHODS:
            raise ValueError(f"unsupported Transformer method: {method}")
        self.config = config
        self.method = method
        forward_generator = torch.Generator().manual_seed(config.seed)
        feedback_generator = torch.Generator().manual_seed(config.seed + 1)
        self.token_embedding = nn.Embedding(
            config.vocab_size, config.width, dtype=dtype)
        self.position_embedding = nn.Parameter(torch.empty(
            config.context_length, config.width, dtype=dtype))
        nn.init.normal_(
            self.token_embedding.weight, mean=0.0, std=config.init_std,
            generator=forward_generator)
        nn.init.normal_(
            self.position_embedding, mean=0.0, std=config.init_std,
            generator=forward_generator)
        self.blocks = nn.ModuleList([
            LocalTransformerBlock(
                config, method, forward_generator, feedback_generator, dtype)
            for _ in range(config.depth)])
        self.final_norm = nn.LayerNorm(config.width, dtype=dtype)
        self.head = FeedbackLinear(
            config.width, config.vocab_size, method, forward_generator,
            feedback_generator, bias=False, init_std=config.init_std,
            traffic_ratio=config.traffic_ratio, dtype=dtype)
        if method == "dfa":
            dfa_generator = torch.Generator().manual_seed(config.seed + 2)
            feedback_scale = config.init_std
            self.register_buffer("dfa_block_feedback", torch.empty(
                config.depth, config.vocab_size, config.width, dtype=dtype))
            self.register_buffer("dfa_embedding_feedback", torch.empty(
                config.vocab_size, config.width, dtype=dtype))
            self.register_buffer("dfa_final_norm_feedback", torch.empty(
                config.vocab_size, config.width, dtype=dtype))
            nn.init.normal_(
                self.dfa_block_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
            nn.init.normal_(
                self.dfa_embedding_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
            nn.init.normal_(
                self.dfa_final_norm_feedback, mean=0.0, std=feedback_scale,
                generator=dfa_generator)
        else:
            self.register_buffer("dfa_block_feedback", None)
            self.register_buffer("dfa_embedding_feedback", None)
            self.register_buffer("dfa_final_norm_feedback", None)
        if method == "pepita":
            pepita_generator = torch.Generator().manual_seed(
                config.seed + 3)
            limit = math.sqrt(6.0 / config.width) * 0.05
            projection = (
                2.0 * torch.rand(
                    config.vocab_size, config.width,
                    generator=pepita_generator, dtype=dtype) - 1.0
            ) * limit
            self.register_buffer("pepita_input_feedback", projection)
        else:
            self.register_buffer("pepita_input_feedback", None)

    def forward(
            self, tokens, targets: Optional[torch.Tensor] = None,
            return_cache: bool = False,
            embedding_offset: Optional[torch.Tensor] = None):
        if tokens.ndim != 2:
            raise ValueError("tokens must have shape (batch, time)")
        if tokens.shape[1] > self.config.context_length:
            raise ValueError("sequence exceeds configured context length")
        positions = self.position_embedding[:tokens.shape[1]]
        hidden = self.token_embedding(tokens) + positions
        if embedding_offset is not None:
            if embedding_offset.shape != hidden.shape:
                raise ValueError("embedding_offset shape does not match tokens")
            hidden = hidden + embedding_offset
        embedded = hidden
        hidden = F.dropout(
            hidden, p=self.config.dropout, training=self.training)
        block_inputs = []
        block_outputs = []
        for block in self.blocks:
            if return_cache:
                block_inputs.append(hidden.detach())
            hidden = block(hidden)
            if return_cache:
                block_outputs.append(hidden.detach())
        final_input = hidden
        normalized = self.final_norm(hidden)
        logits = self.head(normalized)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                targets.reshape(-1))
        result = {"logits": logits, "loss": loss, "hidden": hidden}
        if return_cache:
            result["block_inputs"] = block_inputs
            result["block_outputs"] = block_outputs
            result["embedded"] = embedded.detach()
            result["final_input"] = final_input.detach()
            result["normalized"] = normalized.detach()
        return result

    @staticmethod
    def _assign_gradients(parameters, gradients):
        for parameter, gradient in zip(parameters, gradients):
            if parameter.grad is None:
                parameter.grad = gradient.detach().clone()
            else:
                parameter.grad.copy_(gradient.detach())

    def dfa_gradients(
            self, tokens, targets: Optional[torch.Tensor] = None,
            cache: Optional[Dict[str, torch.Tensor]] = None,
            output_error: Optional[torch.Tensor] = None):
        """Populate strict block-DFA gradients without a task-loss backward.

        Every decoder block receives a separate fixed projection of the
        analytical output error.  Inputs are detached at block boundaries;
        autograd is used only for each explicitly local block objective.
        Embeddings and the final normalization receive their own fixed direct
        projections, while the vocabulary head uses its exact local delta.
        """
        if self.method != "dfa":
            raise ValueError("dfa_gradients is only valid for method='dfa'")
        if cache is None:
            with torch.no_grad():
                cache = self.forward(tokens, return_cache=True)
        if output_error is None:
            if targets is None:
                raise ValueError("targets or output_error must be provided")
            probabilities = torch.softmax(cache["logits"], dim=-1)
            one_hot = F.one_hot(
                targets, self.config.vocab_size).to(probabilities.dtype)
            output_error = (
                probabilities - one_hot) / targets.numel()
        output_error = output_error.detach()
        for parameter in self.parameters():
            parameter.grad = None

        error_flat = output_error.reshape(-1, self.config.vocab_size)
        normalized_flat = cache["normalized"].reshape(
            -1, self.config.width)
        self.head.weight.grad = error_flat.t() @ normalized_flat
        if self.head.bias is not None:
            self.head.bias.grad = error_flat.sum(dim=0)

        final_input = cache["final_input"].detach()
        normalized = self.final_norm(final_input)
        final_field = output_error @ self.dfa_final_norm_feedback
        final_parameters = tuple(self.final_norm.parameters())
        final_objective = torch.sum(normalized * final_field)
        final_gradients = torch.autograd.grad(
            final_objective, final_parameters)
        self._assign_gradients(final_parameters, final_gradients)

        for index, (block, block_input) in enumerate(zip(
                self.blocks, cache["block_inputs"])):
            local_input = block_input.detach()
            local_output = block(local_input)
            local_field = output_error @ self.dfa_block_feedback[index]
            local_parameters = tuple(block.parameters())
            local_objective = torch.sum(local_output * local_field)
            local_gradients = torch.autograd.grad(
                local_objective, local_parameters)
            self._assign_gradients(local_parameters, local_gradients)

        embedding_field = output_error @ self.dfa_embedding_feedback
        token_activity = self.token_embedding(tokens)
        position_activity = self.position_embedding[:tokens.shape[1]]
        embedding_objective = (
            torch.sum(token_activity * embedding_field)
            + torch.sum(position_activity * embedding_field.sum(dim=0)))
        embedding_parameters = (
            self.token_embedding.weight, self.position_embedding)
        embedding_gradients = torch.autograd.grad(
            embedding_objective, embedding_parameters)
        self._assign_gradients(
            embedding_parameters, embedding_gradients)
        return {
            "output_error_rms": float(
                output_error.square().mean().sqrt()),
            "block_field_rms": [
                float((output_error @ feedback).square().mean().sqrt())
                for feedback in self.dfa_block_feedback],
            "uses_task_loss_backward": False,
            "detached_block_boundaries": len(self.blocks),
        }

    def pepita_gradients(self, tokens, targets):
        """Populate two-presentation PEPITA/ERIN local gradients.

        The analytical output error is projected into the continuous token
        embedding stream.  Each block then uses its first-minus-second output
        difference and the second-presentation input.  Block boundaries are
        detached.  Since discrete token IDs cannot themselves be perturbed,
        the embedding table and positional code use the directly observable
        embedding difference as their local field.
        """
        if self.method != "pepita":
            raise ValueError(
                "pepita_gradients is only valid for method='pepita'")
        with torch.no_grad():
            clean = self.forward(tokens, return_cache=True)
            one_hot = F.one_hot(
                targets, self.config.vocab_size).to(clean["logits"].dtype)
            clean_error = torch.softmax(clean["logits"], dim=-1) - one_hot
            embedding_offset = clean_error @ self.pepita_input_feedback
            modulated = self.forward(
                tokens, return_cache=True,
                embedding_offset=embedding_offset)
            modulated_error = (
                torch.softmax(modulated["logits"], dim=-1) - one_hot)
        for parameter in self.parameters():
            parameter.grad = None
        observations = targets.numel()

        error_flat = modulated_error.reshape(
            -1, self.config.vocab_size) / observations
        normalized_flat = modulated["normalized"].reshape(
            -1, self.config.width)
        self.head.weight.grad = error_flat.t() @ normalized_flat
        if self.head.bias is not None:
            self.head.bias.grad = error_flat.sum(dim=0)

        final_input = modulated["final_input"].detach()
        normalized = self.final_norm(final_input)
        final_field = (
            clean["normalized"] - modulated["normalized"]).detach()
        final_parameters = tuple(self.final_norm.parameters())
        final_objective = torch.sum(
            normalized * final_field) / observations
        final_gradients = torch.autograd.grad(
            final_objective, final_parameters)
        self._assign_gradients(final_parameters, final_gradients)

        for block, local_input, clean_output, modulated_output in zip(
                self.blocks, modulated["block_inputs"],
                clean["block_outputs"], modulated["block_outputs"]):
            local_output = block(local_input.detach())
            local_field = (clean_output - modulated_output).detach()
            local_parameters = tuple(block.parameters())
            local_objective = torch.sum(
                local_output * local_field) / observations
            local_gradients = torch.autograd.grad(
                local_objective, local_parameters)
            self._assign_gradients(local_parameters, local_gradients)

        base_embedding = (
            self.token_embedding(tokens)
            + self.position_embedding[:tokens.shape[1]])
        embedding_field = (
            clean["embedded"] - modulated["embedded"]).detach()
        embedding_objective = torch.sum(
            base_embedding * embedding_field) / observations
        embedding_parameters = (
            self.token_embedding.weight, self.position_embedding)
        embedding_gradients = torch.autograd.grad(
            embedding_objective, embedding_parameters)
        self._assign_gradients(
            embedding_parameters, embedding_gradients)
        return {
            "clean_loss": float(F.cross_entropy(
                clean["logits"].reshape(-1, self.config.vocab_size),
                targets.reshape(-1))),
            "embedding_offset_rms": float(
                embedding_offset.square().mean().sqrt()),
            "mean_block_field_rms": float(torch.stack([
                (first - second).square().mean().sqrt()
                for first, second in zip(
                    clean["block_outputs"],
                    modulated["block_outputs"])]).mean()),
            "training_presentations": 2,
            "uses_task_loss_backward": False,
            "detached_block_boundaries": len(self.blocks),
        }

    def feedback_linears(self) -> Iterable[FeedbackLinear]:
        return (
            module for module in self.modules()
            if isinstance(module, FeedbackLinear))

    def forward_parameters(self) -> Iterable[nn.Parameter]:
        feedback_ids = {
            id(module.feedback)
            for module in self.feedback_linears()
            if isinstance(module.feedback, nn.Parameter)}
        return (
            parameter for parameter in self.parameters()
            if id(parameter) not in feedback_ids)

    @property
    def n_forward_parameters(self):
        return sum(parameter.numel() for parameter in self.forward_parameters())

    @property
    def n_feedback_parameters(self):
        affine_feedback = sum(
            module.feedback.numel()
            for module in self.feedback_linears()
            if module.feedback is not None)
        dfa_feedback = sum(
            tensor.numel() for tensor in (
                self.dfa_block_feedback, self.dfa_embedding_feedback,
                self.dfa_final_norm_feedback)
            if tensor is not None)
        pepita_feedback = (
            self.pepita_input_feedback.numel()
            if self.pepita_input_feedback is not None else 0)
        return affine_feedback + dfa_feedback + pepita_feedback

    def teaching_statistics(self) -> Dict[str, float]:
        modules = list(self.feedback_linears())
        if not modules:
            return {
                "raw_rms": 0.0, "innovation_rms": 0.0,
                "traffic_rms": 0.0}
        return {
            name: float(torch.stack([
                getattr(module, f"last_{name}")
                for module in modules]).mean())
            for name in ("raw_rms", "innovation_rms", "traffic_rms")}

    @torch.no_grad()
    def set_feedback_equal_to_forward(self):
        for module in self.feedback_linears():
            if module.feedback is None:
                raise ValueError("BP modules do not contain feedback tensors")
            module.feedback.copy_(module.weight)

    @staticmethod
    def _ff_normalize(value):
        return value / (
            value.square().sum(dim=-1, keepdim=True).sqrt() + 1e-8)

    def _ff_overlay(self, embedded, candidate_labels):
        if candidate_labels.shape != embedded.shape[:2]:
            raise ValueError("candidate labels must match token positions")
        if self.config.width < self.config.vocab_size:
            raise ValueError(
                "Forward-Forward overlay requires width >= vocabulary")
        overlaid = embedded.clone()
        overlaid[..., :self.config.vocab_size] = 0.0
        magnitude = embedded.detach().abs().amax(
            dim=-1, keepdim=True).clamp_min(1e-4)
        overlaid.scatter_(
            -1, candidate_labels.unsqueeze(-1), magnitude)
        return overlaid

    def ff_forward(self, tokens, candidate_labels):
        """Run candidate-labelled, normalized data through the FF graph."""
        if self.method != "ff":
            raise ValueError("ff_forward is only valid for method='ff'")
        positions = self.position_embedding[:tokens.shape[1]]
        base = self.token_embedding(tokens) + positions
        embedded = self._ff_overlay(base, candidate_labels)
        hidden = embedded
        block_inputs = []
        block_outputs = []
        for block in self.blocks:
            local_input = self._ff_normalize(hidden)
            block_inputs.append(local_input)
            hidden = block(local_input)
            block_outputs.append(hidden)
        final_input = self._ff_normalize(hidden)
        logits = self.head(self.final_norm(final_input))
        return {
            "embedded": embedded,
            "block_inputs": block_inputs,
            "block_outputs": block_outputs,
            "final_input": final_input,
            "logits": logits,
        }

    @property
    def ff_num_layers(self):
        return self.config.depth + 2

    def ff_layer_parameters(self, layer_index):
        if layer_index == 0:
            return (self.token_embedding.weight, self.position_embedding)
        if 1 <= layer_index <= self.config.depth:
            return tuple(self.blocks[layer_index - 1].parameters())
        if layer_index == self.config.depth + 1:
            return (
                *tuple(self.final_norm.parameters()),
                *tuple(self.head.parameters()))
        raise ValueError("invalid Forward-Forward layer index")

    def _ff_local_outputs(self, layer_index, tokens, labels, cached):
        if layer_index == 0:
            base = (
                self.token_embedding(tokens)
                + self.position_embedding[:tokens.shape[1]])
            return self._ff_overlay(base, labels)
        if 1 <= layer_index <= self.config.depth:
            local_input = cached["block_inputs"][
                layer_index - 1].detach()
            return self.blocks[layer_index - 1](local_input)
        if layer_index == self.config.depth + 1:
            local_input = cached["final_input"].detach()
            return self.head(self.final_norm(local_input))
        raise ValueError("invalid Forward-Forward layer index")

    def ff_local_loss(
            self, layer_index, tokens, positive_labels, negative_labels,
            threshold=2.0):
        """Return one greedy FF objective with every prefix detached."""
        if self.method != "ff":
            raise ValueError("ff_local_loss is only valid for method='ff'")
        if not 0 <= layer_index < self.ff_num_layers:
            raise ValueError("invalid Forward-Forward layer index")
        with torch.no_grad():
            positive_cache = self.ff_forward(tokens, positive_labels)
            negative_cache = self.ff_forward(tokens, negative_labels)
        positive_output = self._ff_local_outputs(
            layer_index, tokens, positive_labels, positive_cache)
        negative_output = self._ff_local_outputs(
            layer_index, tokens, negative_labels, negative_cache)
        positive_goodness = positive_output.square().mean(dim=-1)
        negative_goodness = negative_output.square().mean(dim=-1)
        loss = (
            F.softplus(-positive_goodness + threshold)
            + F.softplus(negative_goodness - threshold)
        ).mean()
        return loss, {
            "positive_goodness": float(
                positive_goodness.mean().detach()),
            "negative_goodness": float(
                negative_goodness.mean().detach()),
            "pair_accuracy": float(
                (positive_goodness > negative_goodness).float().mean()),
        }

    @torch.no_grad()
    def ff_candidate_scores(self, tokens, score_from_layer=1):
        """Score every next-token candidate; cost is charged as V passes."""
        if not 0 <= score_from_layer < self.ff_num_layers:
            raise ValueError("invalid Forward-Forward score range")
        scores = []
        for candidate in range(self.config.vocab_size):
            labels = torch.full_like(tokens, candidate)
            candidate_forward = self.ff_forward(tokens, labels)
            layer_outputs = (
                [candidate_forward["embedded"]]
                + candidate_forward["block_outputs"]
                + [candidate_forward["logits"]])
            goodness = [
                output.square().mean(dim=-1)
                for output in layer_outputs]
            scores.append(sum(goodness[score_from_layer:]))
        return torch.stack(scores, dim=-1)

    def _embedding_prediction(self, tokens):
        return (
            self.token_embedding(tokens)
            + self.position_embedding[:tokens.shape[1]])

    @staticmethod
    def _local_vjp(function, state, field):
        """Evaluate one explicitly local state VJP at a detached boundary."""
        with torch.enable_grad():
            local_state = state.detach().requires_grad_(True)
            prediction = function(local_state)
            objective = torch.sum(prediction * field.detach())
            result, = torch.autograd.grad(objective, local_state)
        return result.detach()

    def _head_prediction(self, state):
        return self.head(self.final_norm(state))

    def dualprop_states(
            self, tokens, targets, alpha=0.0, beta=0.1,
            inference_passes=16, clean=None):
        """Infer author-style plus/minus states over Transformer block edges."""
        if self.method != "dualprop":
            raise ValueError(
                "dualprop_states is only valid for method='dualprop'")
        if not 0.0 <= alpha <= 1.0 or beta <= 0 or inference_passes < 1:
            raise ValueError("invalid Dual Propagation settings")
        if clean is None:
            with torch.no_grad():
                clean = self.forward(tokens, return_cache=True)
        hidden_clean = (
            [clean["embedded"]] + list(clean["block_outputs"]))
        plus = [state.detach().clone() for state in hidden_clean]
        minus = [state.detach().clone() for state in hidden_clean]
        plus.append(clean["logits"].detach().clone())
        minus.append(clean["logits"].detach().clone())
        one_hot = F.one_hot(
            targets, self.config.vocab_size).to(clean["logits"].dtype)
        fixed_error = torch.softmax(
            clean["logits"].detach(), dim=-1) - one_hot

        for _ in range(inference_passes):
            for index in range(self.config.depth + 1):
                states = [
                    alpha * positive + (1.0 - alpha) * negative
                    for positive, negative in zip(plus[:-1], minus[:-1])]
                if index == 0:
                    prediction = self._embedding_prediction(tokens)
                else:
                    prediction = self.blocks[index - 1](
                        states[index - 1].detach())
                if index < self.config.depth:
                    feedback = self._local_vjp(
                        self.blocks[index], states[index],
                        plus[index + 1] - minus[index + 1])
                else:
                    feedback = self._local_vjp(
                        self._head_prediction, states[index],
                        plus[-1] - minus[-1])
                plus[index] = (
                    prediction + (1.0 - alpha) * feedback).detach()
                minus[index] = (
                    prediction - alpha * feedback).detach()

            states = [
                alpha * positive + (1.0 - alpha) * negative
                for positive, negative in zip(plus[:-1], minus[:-1])]
            prediction = self._head_prediction(states[-1].detach())
            output_field = beta * fixed_error
            plus[-1] = (
                prediction - (1.0 - alpha) * output_field).detach()
            minus[-1] = (
                prediction + alpha * output_field).detach()
        return plus, minus

    def dualprop_gradients(
            self, tokens, targets, alpha=0.0, beta=0.1,
            inference_passes=16):
        """Populate local Dual-Propagation contrastive gradients."""
        if beta <= 0:
            raise ValueError("Dual Propagation beta must be positive")
        with torch.no_grad():
            clean = self.forward(tokens, return_cache=True)
        plus, minus = self.dualprop_states(
            tokens, targets, alpha=alpha, beta=beta,
            inference_passes=inference_passes, clean=clean)
        states = [
            (alpha * positive + (1.0 - alpha) * negative).detach()
            for positive, negative in zip(plus[:-1], minus[:-1])]
        deltas = [
            ((positive - negative) / beta).detach()
            for positive, negative in zip(plus, minus)]
        for parameter in self.parameters():
            parameter.grad = None
        observations = targets.numel()

        embedding_parameters = (
            self.token_embedding.weight, self.position_embedding)
        embedding_objective = -torch.sum(
            self._embedding_prediction(tokens) * deltas[0]) / observations
        embedding_gradients = torch.autograd.grad(
            embedding_objective, embedding_parameters)
        self._assign_gradients(
            embedding_parameters, embedding_gradients)

        for index, block in enumerate(self.blocks):
            prediction = block(states[index])
            parameters = tuple(block.parameters())
            objective = -torch.sum(
                prediction * deltas[index + 1]) / observations
            gradients = torch.autograd.grad(objective, parameters)
            self._assign_gradients(parameters, gradients)

        output_prediction = self._head_prediction(states[-1])
        output_parameters = (
            *tuple(self.final_norm.parameters()),
            *tuple(self.head.parameters()))
        output_objective = -torch.sum(
            output_prediction * deltas[-1]) / observations
        output_gradients = torch.autograd.grad(
            output_objective, output_parameters)
        self._assign_gradients(output_parameters, output_gradients)
        return {
            "clean_loss": float(F.cross_entropy(
                clean["logits"].reshape(-1, self.config.vocab_size),
                targets.reshape(-1))),
            "state_difference_rms": float(torch.stack([
                (positive - negative).square().mean().sqrt()
                for positive, negative in zip(plus, minus)]).mean()),
            "inference_passes": int(inference_passes),
            "local_vjp_evaluations":
                int((self.config.depth + 1) * inference_passes),
            "uses_task_loss_backward": False,
        }

    @staticmethod
    def ep_rho(state):
        return state.clamp(0.0, 1.0)

    @staticmethod
    def ep_rhop(state):
        return ((state >= 0.0) & (state <= 1.0)).to(state.dtype)

    def ep_settle(
            self, tokens, targets, beta=0.0, steps=20, dt=0.5,
            initial_states=None):
        """Settle synchronous hard-sigmoid states in a symmetric block graph."""
        if self.method != "ep":
            raise ValueError("ep_settle is only valid for method='ep'")
        if steps < 1 or not 0.0 < dt <= 1.0:
            raise ValueError("invalid Equilibrium Propagation dynamics")
        with torch.no_grad():
            clean = self.forward(tokens, return_cache=True)
        if initial_states is None:
            hidden_clean = (
                [clean["embedded"]] + list(clean["block_outputs"]))
            states = [torch.zeros_like(state) for state in hidden_clean]
            states.append(torch.zeros_like(clean["logits"]))
        else:
            states = [state.detach().clone() for state in initial_states]
        one_hot = F.one_hot(
            targets, self.config.vocab_size).to(clean["logits"].dtype)

        for _ in range(steps):
            rho_hidden = [
                self.ep_rho(state).detach() for state in states[:-1]]
            rho_output = self.ep_rho(states[-1]).detach()
            updated = []
            for index, state in enumerate(states[:-1]):
                if index == 0:
                    prediction = self._embedding_prediction(tokens).detach()
                else:
                    prediction = self.blocks[index - 1](
                        rho_hidden[index - 1]).detach()
                if index < self.config.depth:
                    feedback = self._local_vjp(
                        self.blocks[index], rho_hidden[index],
                        rho_hidden[index + 1])
                else:
                    feedback = self._local_vjp(
                        self._head_prediction, rho_hidden[index],
                        rho_output)
                drive = -self.ep_rho(state) + prediction + feedback
                new_state = (
                    state + dt * self.ep_rhop(state) * drive)
                updated.append(self.ep_rho(new_state).detach())

            output_prediction = self._head_prediction(
                rho_hidden[-1]).detach()
            output_state = states[-1]
            output_drive = (
                -self.ep_rho(output_state) + output_prediction)
            if beta:
                output_drive = output_drive + 2.0 * beta * (
                    one_hot - self.ep_rho(output_state))
            new_output = (
                output_state
                + dt * self.ep_rhop(output_state) * output_drive)
            updated.append(self.ep_rho(new_output).detach())
            states = updated
        return states

    def _ep_phase_gradients(self, tokens, states):
        rho_hidden = [
            self.ep_rho(state).detach() for state in states[:-1]]
        rho_output = self.ep_rho(states[-1]).detach()
        observations = tokens.numel()
        groups = []

        embedding_parameters = (
            self.token_embedding.weight, self.position_embedding)
        embedding_correlation = torch.sum(
            self._embedding_prediction(tokens)
            * rho_hidden[0]) / observations
        groups.append((
            embedding_parameters,
            torch.autograd.grad(
                embedding_correlation, embedding_parameters)))

        for index, block in enumerate(self.blocks):
            parameters = tuple(block.parameters())
            correlation = torch.sum(
                block(rho_hidden[index]) * rho_hidden[index + 1]
            ) / observations
            groups.append((
                parameters, torch.autograd.grad(correlation, parameters)))

        output_parameters = (
            *tuple(self.final_norm.parameters()),
            *tuple(self.head.parameters()))
        output_correlation = torch.sum(
            self._head_prediction(rho_hidden[-1])
            * rho_output) / observations
        groups.append((
            output_parameters,
            torch.autograd.grad(output_correlation, output_parameters)))
        return groups

    def ep_gradients(
            self, tokens, targets, ep_beta=0.5, dt=0.5,
            free_steps=20, nudge_steps=4, beta_sign=1.0):
        """Populate the free-versus-nudged EP contrastive gradients."""
        signed_beta = float(beta_sign) * float(ep_beta)
        if signed_beta == 0:
            raise ValueError("EP contrast requires a nonzero nudge")
        free = self.ep_settle(
            tokens, targets, beta=0.0, steps=free_steps, dt=dt)
        nudged = self.ep_settle(
            tokens, targets, beta=signed_beta, steps=nudge_steps, dt=dt,
            initial_states=free)
        free_groups = self._ep_phase_gradients(tokens, free)
        nudged_groups = self._ep_phase_gradients(tokens, nudged)
        for parameter in self.parameters():
            parameter.grad = None
        for (parameters, free_gradients), (
                nudged_parameters, nudged_gradients) in zip(
                    free_groups, nudged_groups):
            if tuple(map(id, parameters)) != tuple(
                    map(id, nudged_parameters)):
                raise AssertionError("EP phase parameter groups differ")
            # The optimizer descends, so negate the EP ascent direction.
            gradients = [
                -(nudged_gradient - free_gradient) / signed_beta
                for nudged_gradient, free_gradient in zip(
                    nudged_gradients, free_gradients)]
            self._assign_gradients(parameters, gradients)
        one_hot = F.one_hot(
            targets, self.config.vocab_size).to(free[-1].dtype)
        return {
            "free_mse": float(F.mse_loss(free[-1], one_hot)),
            "free_steps": int(free_steps),
            "nudge_steps": int(nudge_steps),
            "signed_beta": signed_beta,
            "local_vjp_evaluations": int(
                (self.config.depth + 1)
                * (free_steps + nudge_steps)),
            "uses_task_loss_backward": False,
        }, free, nudged