#!/usr/bin/env python3
"""
Kuramoto block for transformer hybrids — PyTorch port of the kolm.py core.

AKOrN-style (Miyato et al., ICLR 2025): instead of one oscillator step per
sequence position (kolm.py), the block runs K settle steps per token, in
parallel over all positions. The transformer hidden state is the
conditioning stimulus c; H unit vectors on S^{N-1} relax under

    y = J x + c
    x <- normalize(x + dt (A x + (y - (x.y) x)))     A antisymmetric

and the readout returns flat phases + G trained order parameters
r_g = ||sum_i M_gi x_i|| through a zero-initialized projection, added
residually — at init the block is an exact identity, so it can be inserted
into a frozen pretrained model without degrading it.

`steps` can be raised at inference ("more thinking = more settling").

    python3 kuramoto_torch.py    # parity vs kolm.py NumPy step + identity test
"""

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


class KuramotoBlock(nn.Module):
    def __init__(self, d_model: int, H: int = 256, N: int = 4,
                 groups: int = 32, steps: int = 4, dt: float = 0.5,
                 frustrated: bool = False, grad_steps: int = 0):
        super().__init__()
        assert N % 2 == 0
        self.H, self.N, self.G = H, N, groups
        self.steps, self.dt = steps, dt
        self.grad_steps = grad_steps        # 0 = backprop through all steps
        self.cond = nn.Linear(d_model, H * N)
        self.x0 = nn.Parameter(torch.randn(H, N))
        self.J = nn.Parameter(torch.randn(H, H) / H ** 0.5)
        # frustration (FSN, arXiv 2606.18694): quadrature coupling Jq acts on
        # 90-deg-rotated states, so (J, Jq) span learned per-pair phase
        # offsets. Zero init = pure consensus dynamics at start.
        self.Jq = nn.Parameter(torch.zeros(H, H)) if frustrated else None
        self.Om = nn.Parameter(0.1 * torch.randn(H, N, N))
        self.M = nn.Parameter(torch.randn(groups, H) / H ** 0.5)
        self.out = nn.Linear(H * N + groups, d_model)
        nn.init.zeros_(self.out.weight)     # identity at init (residual)
        nn.init.zeros_(self.out.bias)

    @staticmethod
    def rot90(x):
        """Rotate each coordinate pair 90 deg: (a,b) -> (-b,a)."""
        a, b = x[..., 0::2], x[..., 1::2]
        return torch.stack((-b, a), dim=-1).flatten(-2)

    def dynamics(self, c, steps):
        """c: (B*, H, N) drive. Returns settled x: (B*, H, N) unit vectors."""
        B, H, N = c.shape
        A = self.Om - self.Om.transpose(1, 2)
        x = F.normalize(self.x0, dim=-1).expand_as(c).contiguous()
        ng = steps - self.grad_steps if self.grad_steps else 0
        for i in range(steps):
            if i < ng and torch.is_grad_enabled():
                with torch.no_grad():
                    x = self._step(x, c, A, B, H, N)
                x = x.detach()
            else:
                x = self._step(x, c, A, B, H, N)
        return x

    def _step(self, x, c, A, B, H, N):
        # one (H,H)@(H,B*N) GEMM instead of B broadcast matmuls
        xr = x.transpose(0, 1).reshape(H, B * N)
        y = (self.J @ xr).view(H, B, N).transpose(0, 1) + c
        if self.Jq is not None:
            qr = self.rot90(x).transpose(0, 1).reshape(H, B * N)
            y = y + (self.Jq @ qr).view(H, B, N).transpose(0, 1)
        d = (x * y).sum(-1, keepdim=True)
        # rotation as one bmm over H units: (H,N,N)@(H,N,B)
        R = torch.bmm(A, x.permute(1, 2, 0)).permute(2, 0, 1)
        x = x + self.dt * (R + y - d * x)
        return F.normalize(x, dim=-1)

    def forward(self, h, steps: int = None):
        """h: (B, T, D) hidden states -> (B, T, D), residual."""
        B, T, D = h.shape
        hf = h.reshape(B * T, D).to(self.cond.weight.dtype)
        c = self.cond(hf).view(B * T, self.H, self.N)
        x = self.dynamics(c, steps or self.steps)
        m = (self.M @ x.transpose(0, 1).reshape(self.H, B * T * self.N)) \
            .view(self.G, B * T, self.N).transpose(0, 1)
        r = torch.sqrt((m * m).sum(-1) + 1e-8)
        f = torch.cat([x.reshape(B * T, self.H * self.N), r], dim=1)
        return h + self.out(f).view(B, T, D).to(h.dtype)


def _parity_test():
    """One dynamics step must match the kolm.py NumPy update exactly."""
    import numpy as np
    import kolm
    torch.manual_seed(0)
    B, H, N, dt = 3, 8, 4, 0.5
    blk = KuramotoBlock(d_model=16, H=H, N=N, groups=4, steps=1, dt=dt).double()
    c = torch.randn(B, H, N, dtype=torch.float64)
    xt = blk.dynamics(c, steps=1)

    # same update in kolm's NumPy: x0 = normalized blk.x0, drive = c
    p = {"J": blk.J.detach().numpy(), "Om": blk.Om.detach().numpy()}
    x0 = (blk.x0 / torch.linalg.norm(blk.x0, dim=-1, keepdim=True))
    x0 = x0.detach().numpy()[None].repeat(B, 0)
    A = p["Om"] - p["Om"].transpose(0, 2, 1)
    y = kolm.couple(p["J"], x0) + c.numpy()
    d = (x0 * y).sum(-1, keepdims=True)
    R = np.einsum("inm,bim->bin", A, x0)
    z = x0 + dt * (R + y - d * x0)
    z /= np.sqrt((z * z).sum(-1, keepdims=True))

    err = float(np.abs(xt.detach().numpy() - z).max())
    print(f"parity vs kolm.py NumPy step: max abs err {err:.2e} "
          f"({'PASS' if err < 1e-12 else 'FAIL'})")

    # identity at init: inserting the block must not change hidden states
    h = torch.randn(2, 5, 16, dtype=torch.float64)
    delta = float((blk(h) - h).abs().max())
    print(f"identity at init: max |delta| {delta:.2e} "
          f"({'PASS' if delta == 0.0 else 'FAIL'})")

    # gradients flow to every parameter
    blk(h).sum().backward()
    missing = [n for n, q in blk.named_parameters() if q.grad is None]
    print(f"grad flow: {'PASS' if not missing else 'FAIL ' + str(missing)}")
    return err < 1e-12 and delta == 0.0 and not missing


if __name__ == "__main__":
    import sys
    sys.exit(0 if _parity_test() else 1)