"""Explicit toroidal field dynamics and RK4 integration.

The dynamics layer deliberately contains no attention/Transformer block and no
learned feature-mixing MLP.  The state evolves through shared physical-style
rules: toroidal neighbour coupling, phase rotation, input driving, and energy
decay.  RK4 is only the numerical integrator for this continuous system.
"""

from __future__ import annotations

import torch
import torch.nn as nn


class ToroidalDynamics(nn.Module):
    """Shared-rule dynamics for the spectral toroidal field.

    ``alpha[m]`` is the amplitude/state carried by toroidal mode ``m``.
    Mode phase is fixed by its position on the torus; the incoming atom is
    injected as an external drive.  All trainable dynamics parameters are
    shared globally rather than being per-mode feature transforms.
    """

    def __init__(self, d_model: int = 256, n_modes: int = 256) -> None:
        super().__init__()
        self.d_model = d_model
        self.n_modes = n_modes
        self.coupling_scale = nn.Parameter(torch.tensor(1.0))
        self.energy_decay = nn.Parameter(torch.tensor(0.99))
        self.phase_sync = nn.Parameter(torch.tensor(0.5))
        self.register_buffer(
            "mode_phase",
            torch.arange(n_modes, dtype=torch.float32) * (2.0 * torch.pi / n_modes),
        )

    def forward(
        self,
        alpha: torch.Tensor,
        input_token: torch.Tensor | None = None,
        context: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if alpha.shape[-2:] != (self.n_modes, self.d_model):
            raise ValueError(
                f"expected alpha [..., {self.n_modes}, {self.d_model}], got {tuple(alpha.shape)}"
            )

        # Toroidal nearest-neighbour coupling.  Periodic roll closes the ring.
        left = torch.roll(alpha, shifts=1, dims=-2)
        right = torch.roll(alpha, shifts=-1, dims=-2)
        laplacian = left + right - 2.0 * alpha
        coupling = self.phase_sync * laplacian

        # Continuous phase rotation of each mode.  This is the field analogue
        # of oscillator evolution; it is not token attention.
        phase = self.mode_phase.to(device=alpha.device, dtype=alpha.dtype)
        phase = phase.view(*([1] * (alpha.ndim - 2)), self.n_modes, 1)
        rotated = torch.sin(phase) * alpha

        # External token/context drive.  Mode-modulated (choice A): each mode
        # gets a content×phase Hadamard of the token — not the same broadcast
        # vector on every mode (mean-broadcast made distinct texts collinear).
        drive = torch.zeros_like(alpha)
        dim_phase = torch.arange(
            self.d_model, device=alpha.device, dtype=alpha.dtype
        ).view(*([1] * (alpha.ndim - 1)), self.d_model)
        if input_token is not None:
            token = input_token
            while token.ndim > 1:
                token = token.mean(dim=0)
            tok = token.reshape(*([1] * (alpha.ndim - 2)), 1, self.d_model)
            # drive_m,d = token_d * cos(mode_phase_m + 0.5 * token_d * d)
            drive = tok * torch.cos(phase + 0.5 * tok * dim_phase)
            drive = drive.expand(*alpha.shape[:-2], self.n_modes, self.d_model)
            drive = self.coupling_scale * drive

        if context is not None:
            ctx = context
            while ctx.ndim > 1:
                ctx = ctx.mean(dim=0)
            ctx = ctx.reshape(*([1] * (alpha.ndim - 2)), 1, self.d_model)
            ctx = ctx * torch.cos(phase + 0.5 * ctx * dim_phase)
            ctx = ctx.expand(*alpha.shape[:-2], self.n_modes, self.d_model)
            drive = drive + 0.5 * self.coupling_scale * ctx

        decay = (self.energy_decay - 1.0) * alpha
        return coupling + rotated + drive + decay

    def _mode_freqs(self, device: torch.device) -> torch.Tensor:
        return torch.arange(self.n_modes, device=device, dtype=torch.float32) / self.n_modes * 2 * torch.pi

    def _compute_phase_laplacian(self, alpha: torch.Tensor) -> torch.Tensor:
        left = torch.roll(alpha, shifts=1, dims=-2)
        right = torch.roll(alpha, shifts=-1, dims=-2)
        return torch.sin(left - alpha) + torch.sin(right - alpha)

    def get_shared_params(self) -> dict:
        return {
            "coupling_scale": self.coupling_scale.detach().cpu(),
            "energy_decay": self.energy_decay.detach().cpu(),
            "phase_sync": self.phase_sync.detach().cpu(),
        }

    def load_shared_params(self, params: dict) -> None:
        for key in ("coupling_scale", "energy_decay", "phase_sync"):
            if key in params:
                getattr(self, key).data.copy_(params[key].to(getattr(self, key).device))


class RK4DynamicsEngine(nn.Module):
    """RK4 numerical integration of the toroidal dynamics."""

    def __init__(self, dynamics: ToroidalDynamics, dt: float = 0.01, n_steps: int = 4) -> None:
        super().__init__()
        self.dynamics = dynamics
        self.dt = dt
        self.n_steps = n_steps

    def evolve(
        self,
        alpha: torch.Tensor,
        input_token: torch.Tensor | None = None,
        context: torch.Tensor | None = None,
    ) -> torch.Tensor:
        def f(state: torch.Tensor) -> torch.Tensor:
            return self.dynamics(state, input_token, context)

        k1 = f(alpha)
        k2 = f(alpha + self.dt * 0.5 * k1)
        k3 = f(alpha + self.dt * 0.5 * k2)
        k4 = f(alpha + self.dt * k3)
        return alpha + (self.dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)

    def evolve_multi(
        self,
        alpha: torch.Tensor,
        input_token: torch.Tensor | None = None,
        context: torch.Tensor | None = None,
        n_steps: int | None = None,
    ) -> torch.Tensor:
        current = alpha
        for _ in range(n_steps or self.n_steps):
            current = self.evolve(current, input_token, context)
        return current