"""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"}


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 {"bp"} | _FEEDBACK_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 == "bp":
            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 {"bp"} | _FEEDBACK_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)

    def forward(self, tokens, targets: 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
        hidden = F.dropout(
            hidden, p=self.config.dropout, training=self.training)
        for block in self.blocks:
            hidden = block(hidden)
        logits = self.head(self.final_norm(hidden))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                targets.reshape(-1))
        return {"logits": logits, "loss": loss, "hidden": hidden}

    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):
        return sum(
            module.feedback.numel()
            for module in self.feedback_linears()
            if module.feedback is not None)

    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)