model.py

"""
RheOFormer architecture for 2D viscoelastic flow on irregular meshes.

Overview
--------
                     [b, tw, n, 7]
                          │
                    IrregSTEncoder2D
                          │
                     [b, n, C]  ← latent node features
                          │
                    IrregSTDecoder2D
                          │
               [b, forward_steps, n, 5]

Attention
---------
Both encoder and decoder use linear-complexity Galerkin/Fourier attention
from "Choose a Transformer: Fourier or Galerkin" (Cao 2021), with 2-D
Rotary Position Embeddings (RoPE) for relative spatial encoding.

Encoder
-------
A temporal convolution collapses the tw-frame input window into one token
per mesh node.  A stack of Galerkin self-attention layers then propagates
information across nodes.

Decoder
-------
Mesh coordinates are mapped to latent space via random Fourier features and
cross-attended to the encoder latent.  A propagation MLP rolls out future
field states autoregressively.
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
from torch.nn.init import xavier_uniform_, orthogonal_


# ── Attention primitives ──────────────────────────────────────────────────────

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn   = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class GeGELU(nn.Module):
    """Gated GELU activation (https://paperswithcode.com/method/geglu)."""
    def __init__(self):
        super().__init__()
        self.fn = nn.GELU()

    def forward(self, x):
        c = x.shape[-1]
        return self.fn(x[..., :c // 2]) * x[..., c // 2:]


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim * 2),
            GeGELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class ReLUFeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


def masked_instance_norm(x, mask, eps=1e-5):
    """Instance normalisation that ignores padded (masked-out) positions.

    Args:
        x    : [N, L, C]
        mask : [N, L, 1]  1 = valid, 0 = padded
    """
    mask = mask.float()
    mean = (torch.sum(x * mask, 1) / torch.sum(mask, 1)).detach()
    var  = (torch.sum(((x - mean.unsqueeze(1)) * mask) ** 2, 1)
            / torch.sum(mask, 1)).detach()
    return (x - mean.unsqueeze(1)) / torch.sqrt(var.unsqueeze(1) + eps)


# ── Rotary Position Embedding (RoPE) ─────────────────────────────────────────

class RotaryEmbedding(nn.Module):
    """1-D rotary position embedding.

    Modified from https://github.com/lucidrains/x-transformers.
    """
    def __init__(self, dim, min_freq=1 / 64, scale=1.):
        super().__init__()
        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.min_freq = min_freq
        self.scale    = scale
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, coordinates, device):
        t     = coordinates.to(device).type_as(self.inv_freq)
        t     = t * (self.scale / self.min_freq)
        freqs = torch.einsum('... i , j -> ... i j', t, self.inv_freq)
        return torch.cat((freqs, freqs), dim=-1)


def rotate_half(x):
    x = rearrange(x, '... (j d) -> ... j d', j=2)
    x1, x2 = x.unbind(dim=-2)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(t, freqs):
    return (t * freqs.cos()) + (rotate_half(t) * freqs.sin())


def apply_2d_rotary_pos_emb(t, freqs_x, freqs_y):
    """Apply 2-D RoPE by splitting the head dimension in half."""
    d = t.shape[-1]
    t_x, t_y = t[..., :d // 2], t[..., d // 2:]
    return torch.cat((apply_rotary_pos_emb(t_x, freqs_x),
                      apply_rotary_pos_emb(t_y, freqs_y)), dim=-1)


# =============================================================================
# Attention Modules
# =============================================================================

# ── Standard (softmax) self-attention ────────────────────────────────────────

class StandardAttention(nn.Module):
    """Scaled dot-product attention."""
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim    = dim_head * heads
        project_out  = not (heads == 1 and dim_head == dim)
        self.heads   = heads
        self.scale   = dim_head ** -0.5
        self.attend  = nn.Softmax(dim=-1)
        self.to_qkv  = nn.Linear(dim, inner_dim * 3, bias=False)
        self.to_out  = (nn.Sequential(nn.Linear(inner_dim, dim),
                                      nn.Dropout(dropout))
                        if project_out else nn.Identity())

    def forward(self, x, mask=None):
        qkv      = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v  = map(lambda t: rearrange(t, 'b n (h d) -> b h n d',
                                           h=self.heads), qkv)
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
        if mask is not None:
            dots = dots.masked_fill(mask, -torch.finfo(dots.dtype).max)
        attn = self.attend(dots)
        out  = rearrange(torch.matmul(attn, v), 'b h n d -> b n (h d)')
        return self.to_out(out)


# ── Linear (Galerkin / Fourier) self-attention ────────────────────────────────

class LinearAttention(nn.Module):
    """Linear-complexity attention with optional 2-D RoPE.

    Supports both Galerkin type (normalise K, V) and
    Fourier type (normalise Q, K) as described in Cao (2021).
    """

    def __init__(self,
                 dim,
                 attn_type,           # 'galerkin' or 'fourier'
                 heads=8,
                 dim_head=64,
                 dropout=0.,
                 init_params=True,
                 relative_emb=False,
                 scale=1.,
                 init_method='orthogonal',
                 init_gain=None,
                 relative_emb_dim=2,
                 min_freq=1 / 64,
                 cat_pos=False,
                 pos_dim=2,
                 use_ln=False,
                 ):
        super().__init__()
        assert attn_type in ('galerkin', 'fourier')
        inner_dim    = dim_head * heads
        project_out  = not (heads == 1 and dim_head == dim)
        self.attn_type    = attn_type
        self.use_ln       = use_ln
        self.heads        = heads
        self.dim_head     = dim_head
        self.to_qkv       = nn.Linear(dim, inner_dim * 3, bias=False)

        if attn_type == 'galerkin':
            Norm = nn.LayerNorm(dim_head) if use_ln else nn.InstanceNorm1d(dim_head)
            self.k_norm = Norm
            self.v_norm = Norm.__class__(dim_head)
        else:  # fourier
            Norm = nn.LayerNorm(dim_head) if use_ln else nn.InstanceNorm1d(dim_head)
            self.q_norm = Norm
            self.k_norm = Norm.__class__(dim_head)

        if cat_pos:
            self.to_out = nn.Sequential(
                nn.Linear(inner_dim + pos_dim * heads, dim),
                nn.Dropout(dropout),
            )
        else:
            self.to_out = (nn.Sequential(nn.Linear(inner_dim, dim),
                                         nn.Dropout(dropout))
                           if project_out else nn.Identity())

        self.init_gain       = (1. / dim_head) if init_gain is None else init_gain
        self.diagonal_weight = self.init_gain
        self.init_method     = init_method
        if init_params:
            self._init_params()

        self.cat_pos          = cat_pos
        self.pos_dim          = pos_dim
        self.relative_emb     = relative_emb
        self.relative_emb_dim = relative_emb_dim
        if relative_emb:
            assert not cat_pos
            self.emb_module = RotaryEmbedding(
                dim_head // relative_emb_dim, min_freq=min_freq, scale=scale)

    def _init_params(self):
        init_fn = orthogonal_ if self.init_method == 'orthogonal' else xavier_uniform_
        for param in self.to_qkv.parameters():
            if param.ndim > 1:
                for h in range(self.heads):
                    if self.attn_type == 'fourier':
                        idx = (self.heads * 2 + h) * self.dim_head
                    else:
                        idx = h * self.dim_head
                    init_fn(param[idx:idx + self.dim_head, :],
                            gain=self.init_gain)
                    param.data[idx:idx + self.dim_head, :] += (
                        self.diagonal_weight
                        * torch.diag(torch.ones(param.size(-1)))
                    )

    def norm_wrt_domain(self, x, norm_fn):
        b = x.shape[0]
        return rearrange(
            norm_fn(rearrange(x, 'b h n d -> (b h) n d')),
            '(b h) n d -> b h n d', b=b,
        )

    def forward(self, x, pos=None, not_assoc=False, padding_mask=None):
        qkv     = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d',
                                          h=self.heads), qkv)

        if padding_mask is None:
            if self.attn_type == 'galerkin':
                k = self.norm_wrt_domain(k, self.k_norm)
                v = self.norm_wrt_domain(v, self.v_norm)
            else:
                q = self.norm_wrt_domain(q, self.q_norm)
                k = self.norm_wrt_domain(k, self.k_norm)
        else:
            grid_size = torch.sum(padding_mask, dim=[-1, -2]).view(-1, 1, 1, 1)
            pad_exp   = repeat(padding_mask, 'b n d -> (b h) n d', h=self.heads)
            if self.use_ln:
                if self.attn_type == 'galerkin':
                    k = self.k_norm(k); v = self.v_norm(v)
                else:
                    q = self.q_norm(q); k = self.k_norm(k)
            else:
                if self.attn_type == 'galerkin':
                    k = rearrange(k, 'b h n d -> (b h) n d')
                    v = rearrange(v, 'b h n d -> (b h) n d')
                    k = masked_instance_norm(k, pad_exp)
                    v = masked_instance_norm(v, pad_exp)
                    k = rearrange(k, '(b h) n d -> b h n d', h=self.heads)
                    v = rearrange(v, '(b h) n d -> b h n d', h=self.heads)
                else:
                    q = rearrange(q, 'b h n d -> (b h) n d')
                    k = rearrange(k, 'b h n d -> (b h) n d')
                    q = masked_instance_norm(q, pad_exp)
                    k = masked_instance_norm(k, pad_exp)
                    q = rearrange(q, '(b h) n d -> b h n d', h=self.heads)
                    k = rearrange(k, '(b h) n d -> b h n d', h=self.heads)
            padding_mask = rearrange(pad_exp, '(b h) n d -> b h n d',
                                     h=self.heads)

        if self.relative_emb:
            if self.relative_emb_dim == 2:
                fx = repeat(self.emb_module(pos[..., 0], x.device),
                            'b n d -> b h n d', h=q.shape[1])
                fy = repeat(self.emb_module(pos[..., 1], x.device),
                            'b n d -> b h n d', h=q.shape[1])
                q = apply_2d_rotary_pos_emb(q, fx, fy)
                k = apply_2d_rotary_pos_emb(k, fx, fy)
            elif self.relative_emb_dim == 1:
                f = repeat(self.emb_module(pos[..., 0], x.device),
                           'b n d -> b h n d', h=q.shape[1])
                q = apply_rotary_pos_emb(q, f)
                k = apply_rotary_pos_emb(k, f)
        elif self.cat_pos:
            pos_exp = pos.unsqueeze(1).repeat([1, self.heads, 1, 1])
            q, k, v = [torch.cat([pos_exp, t], dim=-1) for t in (q, k, v)]

        if not_assoc:
            score = torch.matmul(q, k.transpose(-1, -2))
            if padding_mask is not None:
                inv   = ~padding_mask
                score = score.masked_fill(
                    torch.matmul(inv, inv.transpose(-1, -2)), 0.)
                out   = torch.matmul(score, v) * (1. / grid_size)
            else:
                out = torch.matmul(score, v) * (1. / q.shape[2])
        else:
            if padding_mask is not None:
                q = q.masked_fill(~padding_mask, 0)
                k = k.masked_fill(~padding_mask, 0)
                v = v.masked_fill(~padding_mask, 0)
                dots = torch.matmul(k.transpose(-1, -2), v)
                out  = torch.matmul(q, dots) * (1. / grid_size)
            else:
                dots = torch.matmul(k.transpose(-1, -2), v)
                out  = torch.matmul(q, dots) * (1. / q.shape[2])

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)


# ── Linear cross-attention ────────────────────────────────────────────────────

class CrossLinearAttention(nn.Module):
    """Linear cross-attention with 2-D RoPE support.

    Used in the decoder to attend from coordinate queries to the
    encoder latent field.
    """

    def __init__(self,
                 dim,
                 attn_type,
                 heads=8,
                 dim_head=64,
                 dropout=0.,
                 init_params=True,
                 relative_emb=False,
                 scale=1.,
                 init_method='orthogonal',
                 init_gain=None,
                 relative_emb_dim=2,
                 min_freq=1 / 64,
                 cat_pos=False,
                 pos_dim=2,
                 use_ln=False,
                 ):
        super().__init__()
        assert attn_type in ('galerkin', 'fourier')
        inner_dim    = dim_head * heads
        project_out  = not (heads == 1 and dim_head == dim)
        self.attn_type    = attn_type
        self.use_ln       = use_ln
        self.heads        = heads
        self.dim_head     = dim_head
        self.to_q         = nn.Linear(dim, inner_dim,     bias=False)
        self.to_kv        = nn.Linear(dim, inner_dim * 2, bias=False)

        if attn_type == 'galerkin':
            Norm = nn.LayerNorm(dim_head) if use_ln else nn.InstanceNorm1d(dim_head)
            self.k_norm = Norm
            self.v_norm = Norm.__class__(dim_head)
        else:
            Norm = nn.LayerNorm(dim_head) if use_ln else nn.InstanceNorm1d(dim_head)
            self.q_norm = Norm
            self.k_norm = Norm.__class__(dim_head)

        if cat_pos:
            self.to_out = nn.Sequential(
                nn.Linear(inner_dim + pos_dim * heads, dim),
                nn.Dropout(dropout),
            )
        else:
            self.to_out = (nn.Sequential(nn.Linear(inner_dim, dim),
                                         nn.Dropout(dropout))
                           if project_out else nn.Identity())

        self.init_gain       = (1. / dim_head) if init_gain is None else init_gain
        self.diagonal_weight = self.init_gain
        self.init_method     = init_method
        if init_params:
            self._init_params()

        self.cat_pos          = cat_pos
        self.pos_dim          = pos_dim
        self.relative_emb     = relative_emb
        self.relative_emb_dim = relative_emb_dim
        if relative_emb:
            self.emb_module = RotaryEmbedding(
                dim_head // relative_emb_dim, min_freq=min_freq, scale=scale)

    def _init_params(self):
        init_fn = orthogonal_ if self.init_method == 'orthogonal' else xavier_uniform_
        for param in self.to_kv.parameters():
            if param.ndim > 1:
                for h in range(self.heads):
                    for offset in (0, self.heads):
                        idx = (offset + h) * self.dim_head
                        init_fn(param[idx:idx + self.dim_head, :],
                                gain=self.init_gain)
                        param.data[idx:idx + self.dim_head, :] += (
                            self.diagonal_weight
                            * torch.diag(torch.ones(param.size(-1)))
                        )
        for param in self.to_q.parameters():
            if param.ndim > 1:
                for h in range(self.heads):
                    idx = h * self.dim_head
                    init_fn(param[idx:idx + self.dim_head, :],
                            gain=self.init_gain)
                    param.data[idx:idx + self.dim_head, :] += (
                        self.diagonal_weight
                        * torch.diag(torch.ones(param.size(-1)))
                    )

    def norm_wrt_domain(self, x, norm_fn):
        b = x.shape[0]
        return rearrange(
            norm_fn(rearrange(x, 'b h n d -> (b h) n d')),
            '(b h) n d -> b h n d', b=b,
        )

    def forward(self, x, z, x_pos=None, z_pos=None, padding_mask=None):
        # x : [b, n1, d]  query domain (output coordinates)
        # z : [b, n2, d]  key/value domain (encoder latent)
        n2 = z.shape[1]
        q  = rearrange(self.to_q(x), 'b n (h d) -> b h n d', h=self.heads)
        kv = self.to_kv(z).chunk(2, dim=-1)
        k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d',
                                       h=self.heads), kv)

        if padding_mask is None:
            grid_size = n2
            if self.attn_type == 'galerkin':
                k = self.norm_wrt_domain(k, self.k_norm)
                v = self.norm_wrt_domain(v, self.v_norm)
            else:
                q = self.norm_wrt_domain(q, self.q_norm)
                k = self.norm_wrt_domain(k, self.k_norm)
        else:
            grid_size = torch.sum(padding_mask, dim=1).view(-1, 1, 1, 1)
            pad_exp   = repeat(padding_mask, 'b n d -> (b h) n d', h=self.heads)
            if self.use_ln:
                if self.attn_type == 'galerkin':
                    k = self.k_norm(k); v = self.v_norm(v)
                else:
                    q = self.q_norm(q); k = self.k_norm(k)
            else:
                if self.attn_type == 'galerkin':
                    k = masked_instance_norm(
                        rearrange(k, 'b h n d -> (b h) n d'), pad_exp)
                    v = masked_instance_norm(
                        rearrange(v, 'b h n d -> (b h) n d'), pad_exp)
                    k = rearrange(k, '(b h) n d -> b h n d', h=self.heads)
                    v = rearrange(v, '(b h) n d -> b h n d', h=self.heads)
                else:
                    q = masked_instance_norm(
                        rearrange(q, 'b h n d -> (b h) n d'), pad_exp)
                    k = masked_instance_norm(
                        rearrange(k, 'b h n d -> (b h) n d'), pad_exp)
                    q = rearrange(q, '(b h) n d -> b h n d', h=self.heads)
                    k = rearrange(k, '(b h) n d -> b h n d', h=self.heads)
            padding_mask = rearrange(pad_exp, '(b h) n d -> b h n d',
                                     h=self.heads)

        if self.relative_emb:
            if self.relative_emb_dim == 2:
                xfx = repeat(self.emb_module(x_pos[..., 0], x.device),
                             'b n d -> b h n d', h=q.shape[1])
                xfy = repeat(self.emb_module(x_pos[..., 1], x.device),
                             'b n d -> b h n d', h=q.shape[1])
                zfx = repeat(self.emb_module(z_pos[..., 0], z.device),
                             'b n d -> b h n d', h=q.shape[1])
                zfy = repeat(self.emb_module(z_pos[..., 1], z.device),
                             'b n d -> b h n d', h=q.shape[1])
                q = apply_2d_rotary_pos_emb(q, xfx, xfy)
                k = apply_2d_rotary_pos_emb(k, zfx, zfy)
            elif self.relative_emb_dim == 1:
                xf = repeat(self.emb_module(x_pos[..., 0], x.device),
                            'b n d -> b h n d', h=q.shape[1])
                zf = repeat(self.emb_module(z_pos[..., 0], x.device),
                            'b n d -> b h n d', h=q.shape[1])
                q = apply_rotary_pos_emb(q, xf)
                k = apply_rotary_pos_emb(k, zf)
        elif self.cat_pos:
            xp = x_pos.unsqueeze(1).repeat([1, self.heads, 1, 1])
            zp = z_pos.unsqueeze(1).repeat([1, self.heads, 1, 1])
            q  = torch.cat([xp, q], dim=-1)
            k  = torch.cat([zp, k], dim=-1)
            v  = torch.cat([zp, v], dim=-1)

        if padding_mask is not None:
            q = q.masked_fill(~padding_mask, 0)
            k = k.masked_fill(~padding_mask, 0)
            v = v.masked_fill(~padding_mask, 0)
            out = torch.matmul(q, torch.matmul(k.transpose(-1, -2), v)) * (1. / grid_size)
        else:
            out = torch.matmul(q, torch.matmul(k.transpose(-1, -2), v)) * (1. / n2)

        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)


# =============================================================================
# Transformer Blocks
# =============================================================================


class TransformerCatNoCls(nn.Module):
    """Stack of self-attention + FFN layers with optional layer-norm.

    Supports standard (softmax) attention and linear Galerkin / Fourier
    attention.  Position encoding is handled by the attention modules via
    2-D rotary embeddings (relative_emb=True).
    """

    def __init__(self,
                 dim,
                 depth,
                 heads,
                 dim_head,
                 mlp_dim,
                 attn_type,           # 'standard' | 'galerkin' | 'fourier'
                 use_ln=False,
                 scale=16,            # int or list of ints, one per layer
                 dropout=0.,
                 relative_emb_dim=2,
                 min_freq=1 / 64,
                 attention_init='orthogonal',
                 init_gain=None,
                 use_relu=False,
                 cat_pos=False,
                 ):
        super().__init__()
        assert attn_type in ('standard', 'galerkin', 'fourier')

        if isinstance(scale, int):
            scale = [scale] * depth
        assert len(scale) == depth

        self.use_ln    = use_ln
        self.attn_type = attn_type
        self.layers    = nn.ModuleList()

        FFN = ReLUFeedForward if use_relu else FeedForward

        if attn_type == 'standard':
            for _ in range(depth):
                self.layers.append(nn.ModuleList([
                    PreNorm(dim, StandardAttention(
                        dim, heads=heads, dim_head=dim_head, dropout=dropout)),
                    PreNorm(dim, FFN(dim, mlp_dim, dropout=dropout)),
                ]))
        else:
            for d in range(depth):
                if scale[d] != -1 or not cat_pos:
                    attn = LinearAttention(
                        dim, attn_type,
                        heads=heads, dim_head=dim_head, dropout=dropout,
                        relative_emb=True, scale=scale[d],
                        relative_emb_dim=relative_emb_dim,
                        min_freq=min_freq,
                        init_method=attention_init,
                        init_gain=init_gain,
                        use_ln=False,
                    )
                else:
                    attn = LinearAttention(
                        dim, attn_type,
                        heads=heads, dim_head=dim_head, dropout=dropout,
                        cat_pos=True, pos_dim=relative_emb_dim,
                        relative_emb=False,
                        init_method=attention_init,
                        init_gain=init_gain,
                    )

                if use_ln:
                    self.layers.append(nn.ModuleList([
                        nn.LayerNorm(dim), attn,
                        nn.LayerNorm(dim), FFN(dim, mlp_dim, dropout=dropout),
                    ]))
                else:
                    self.layers.append(nn.ModuleList([
                        attn, FFN(dim, mlp_dim, dropout=dropout),
                    ]))

    def forward(self, x, pos_embedding):
        """
        Args:
            x             : [b, n, c]
            pos_embedding : [b, n, 2]
        """
        for layer in self.layers:
            if not self.use_ln:
                attn, ffn = layer
                x = attn(x, pos_embedding) + x
                x = ffn(x) + x
            else:
                ln1, attn, ln2, ffn = layer
                x = attn(ln1(x), pos_embedding) + x
                x = ffn(ln2(x)) + x
        return x


class GaussianFourierFeatureTransform(nn.Module):
    """Random Fourier feature mapping for coordinate inputs.

    Maps [b, n, d_in] → [b, n, 2*mapping_size] via:
        x → [sin(2π B x), cos(2π B x)]

    Reference: "Fourier Features Let Networks Learn High Frequency
    Functions in Low Dimensional Domains" (Tancik et al., NeurIPS 2020)
    https://arxiv.org/abs/2006.10739
    """

    def __init__(self, num_input_channels, mapping_size=256, scale=10):
        super().__init__()
        self._B = nn.Parameter(
            torch.randn(num_input_channels, mapping_size) * scale,
            requires_grad=False,
        )

    def forward(self, x):
        b, n, _ = x.shape
        x_flat  = rearrange(x, 'b n c -> (b n) c')
        proj    = rearrange(x_flat @ self._B.to(x.device),
                            '(b n) c -> b n c', b=b)
        proj    = 2 * np.pi * proj
        return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)


class CrossFormer(nn.Module):
    """One cross-attention layer with an optional residual FFN.

    Queries come from the output coordinate space; keys/values come from
    the encoder latent field.
    """

    def __init__(self,
                 dim,
                 attn_type,
                 heads,
                 dim_head,
                 mlp_dim,
                 residual=True,
                 use_ffn=True,
                 use_ln=False,
                 relative_emb=False,
                 scale=1.,
                 relative_emb_dim=2,
                 min_freq=1 / 64,
                 dropout=0.,
                 cat_pos=False,
                 ):
        super().__init__()
        self.residual = residual
        self.use_ffn  = use_ffn
        self.use_ln   = use_ln

        self.cross_attn = CrossLinearAttention(
            dim, attn_type,
            heads=heads, dim_head=dim_head, dropout=dropout,
            relative_emb=relative_emb, scale=scale,
            relative_emb_dim=relative_emb_dim, min_freq=min_freq,
            init_method='orthogonal',
            cat_pos=cat_pos, pos_dim=relative_emb_dim,
            use_ln=False,
        )

        if use_ln:
            self.ln1 = nn.LayerNorm(dim)
            self.ln2 = nn.LayerNorm(dim)
        if use_ffn:
            self.ffn = FeedForward(dim, mlp_dim, dropout)

    def forward(self, x, z, x_pos=None, z_pos=None):
        """
        Args:
            x     : [b, n1, c]  query (output coordinates)
            z     : [b, n2, c]  key/value (encoder latent)
            x_pos : [b, n1, 2]
            z_pos : [b, n2, 2]
        """
        if self.use_ln:
            z = self.ln1(z)
        attn_out = self.cross_attn(x, z, x_pos, z_pos)
        x = (attn_out + x) if self.residual else attn_out
        if self.use_ln:
            x = self.ln2(x)
        if self.use_ffn:
            x = self.ffn(x) + x
        return x




# =============================================================================
# Encoder
# =============================================================================

class Encoder1D(nn.Module):
    """
    Encoder: embeds input function values with coordinates, applies self-attention,
    and projects to latent space.

    Args:
        input_channels: Number of input feature channels (including coordinate).
        in_emb_dim: Internal embedding dimension.
        out_seq_emb_dim: Output latent dimension.
        depth: Number of self-attention layers.
        res: Temporal resolution for frequency scaling.
    """

    def __init__(
        self,
        input_channels: int,
        in_emb_dim: int,
        out_seq_emb_dim: int,
        depth: int,
        res: int = 2048,
    ):
        super().__init__()
        self.to_embedding = nn.Sequential(
            nn.Linear(input_channels, in_emb_dim - 1, bias=False),
        )
        self.transformer = TransformerCatNoCls(
            in_emb_dim,
            depth,
            heads=1,
            dim_head=in_emb_dim,
            mlp_dim=in_emb_dim,
            attn_type="fourier",
            scale=[8.0] + [4.0] * 2 + [1.0] * (depth - 3),
            relative_emb_dim=1,
            min_freq=1 / res,
            attention_init="orthogonal",
        )
        self.project_to_latent = nn.Linear(in_emb_dim, out_seq_emb_dim, bias=False)

    def forward(self, x: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
        x = self.to_embedding(x)
        x = torch.cat((x, input_pos), dim=-1)
        x = self.transformer(x, input_pos)
        return self.project_to_latent(x)


class IrregSTEncoder2D(nn.Module):
    """Spatiotemporal encoder for irregular 2-D meshes.

    A temporal convolution collapses the input time window into a single
    token per node; a stack of Galerkin-attention layers then exchanges
    information across nodes.

    Args:
        input_channels : Number of input field channels (e.g. 7 for
                         v_x, v_y, σ_xx, σ_yy, σ_xy, pos_x, pos_y).
        time_window    : Length of the input time window (tw).
        in_emb_dim     : Internal embedding dimension.
        out_chanels    : Output (latent) channel dimension.
        heads          : Number of attention heads.
        depth          : Number of transformer layers.
        res            : Characteristic mesh resolution (used to set RoPE
                         scale sequence; 200 for the triangle dataset).
        use_ln         : Use LayerNorm inside transformer layers.
        emb_dropout    : Dropout applied after initial embedding.
    """

    def __init__(self,
                 input_channels: int,
                 time_window: int,
                 in_emb_dim: int,
                 out_chanels: int,
                 heads: int,
                 depth: int,
                 res: int,
                 use_ln: bool = True,
                 emb_dropout: float = 0.0,
                 ):
        super().__init__()
        self.tw = time_window

        # Temporal convolution: compress [tw, n, c] → [n, emb]
        self.to_embedding = nn.Sequential(
            Rearrange('b t n c -> b c t n'),
            nn.Conv2d(input_channels, in_emb_dim,
                      kernel_size=(time_window, 1),
                      stride=(time_window, 1), bias=False),
            nn.GELU(),
            nn.Conv2d(in_emb_dim, in_emb_dim,
                      kernel_size=(1, 1), bias=False),
            Rearrange('b c 1 n -> b n c'),
        )

        self.dropout = nn.Dropout(emb_dropout)

        # Scale schedule: broad→local across the attention stack
        if depth > 4:
            scales = [32, 16, 8, 8] + [1] * (depth - 4)
        else:
            scales = [32] + [16] * (depth - 2) + [1]

        self.s_transformer = TransformerCatNoCls(
            in_emb_dim, depth, heads, in_emb_dim, in_emb_dim,
            'galerkin', use_ln,
            scale=scales,
            min_freq=1 / res,
            attention_init='orthogonal',
        )

        self.ln     = nn.LayerNorm(in_emb_dim)
        self.to_out = nn.Sequential(
            nn.Linear(in_emb_dim, in_emb_dim, bias=False),
            nn.ReLU(),
            nn.Linear(in_emb_dim, out_chanels, bias=False),
        )

    def forward(self, x, pos):
        """
        Args:
            x   : Tensor [b, tw, n, c_in]
            pos : Tensor [b, n,  2]        mesh node coordinates

        Returns:
            Tensor [b, n, out_chanels]  latent node features
        """
        z     = self.to_embedding(x)
        z_res = z
        z     = self.dropout(z)
        z     = self.s_transformer(z, pos)
        z     = self.ln(z + z_res)
        return self.to_out(z)


# =============================================================================
# Decoder
# =============================================================================


class PointWiseDecoder1D(nn.Module):
    """
    Decoder with cross-attention and latent-space propagator.

    Receives the encoded latent representation, attends to query positions via
    cross-attention, propagates dynamics through residual feed-forward blocks,
    and decodes to output field values.

    Args:
        latent_channels: Latent dimension from encoder.
        out_channels: Number of output channels.
        decoding_depth: Number of residual propagator blocks.
        scale: Scale for Gaussian Fourier features.
        res: Temporal resolution.
    """

    def __init__(
        self,
        latent_channels: int,
        out_channels: int,
        decoding_depth: int,
        scale: float = 8.0,
        res: int = 2048,
    ):
        super().__init__()
        self.latent_channels = latent_channels
        self.out_channels = out_channels

        # Coordinate projection via random Fourier features
        self.coordinate_projection = nn.Sequential(
            GaussianFourierFeatureTransform(1, latent_channels, scale=scale),
            nn.GELU(),
            nn.Linear(latent_channels * 2, latent_channels, bias=False),
        )

        # Cross-attention from query positions to encoder output
        self.decoding_transformer = CrossFormer(
            dim=latent_channels,
            attn_type="fourier",
            heads=8,
            dim_head=latent_channels,
            mlp_dim=latent_channels,
            relative_emb=True,
            scale=1,
            relative_emb_dim=1,
            min_freq=1 / res,
        )

        # Latent-space propagator (residual feed-forward blocks)
        self.propagator = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(latent_channels, latent_channels, bias=False),
                    nn.GELU(),
                    nn.Linear(latent_channels, latent_channels, bias=False),
                    nn.GELU(),
                    nn.Linear(latent_channels, latent_channels, bias=False),
                )
                for _ in range(decoding_depth)
            ]
        )
        self._init_propagator_params()

        # Output projection
        self.to_out = nn.Sequential(
            nn.Linear(latent_channels, latent_channels // 2, bias=False),
            nn.GELU(),
            nn.Linear(latent_channels // 2, out_channels, bias=True),
        )

    def _init_propagator_params(self):
        for block in self.propagator:
            for layer in block:
                for param in layer.parameters():
                    if param.ndim > 1:
                        in_c = param.size(-1)
                        orthogonal_(param[:in_c], gain=1 / in_c)
                        param.data[:in_c] += (1 / in_c) * torch.diag(
                            torch.ones(param.size(-1), dtype=torch.float32)
                        )
                        if param.size(-2) != param.size(-1):
                            orthogonal_(param[in_c:], gain=1 / in_c)
                            param.data[in_c:] += (1 / in_c) * torch.diag(
                                torch.ones(param.size(-1), dtype=torch.float32)
                            )

    def forward(
        self,
        z: torch.Tensor,
        propagate_pos: torch.Tensor,
        input_pos: torch.Tensor = None,
    ) -> torch.Tensor:
        x = self.coordinate_projection(propagate_pos)
        z = self.decoding_transformer(x, z, propagate_pos, input_pos)
        for layer in self.propagator:
            z = z + layer(z)
        return self.to_out(z)


class IrregSTDecoder2D(nn.Module):
    """Autoregressive decoder for 2-D irregular meshes.

    Pipeline per forward call
    -------------------------
    1. Project mesh coordinates to latent space via Fourier features
       + cross-attention (decode from encoder latent).
    2. Mix the decoded features with a self-attention layer.
    3. Expand channels for the propagation MLP.
    4. Roll out ``forward_steps`` future states by iterating the
       propagation MLP and decoding to output channels each step.

    Args:
        latent_channels : Encoder output / decoder hidden dimension.
        out_channels    : Number of predicted field channels (5 for
                          v_x, v_y, σ_xx, σ_yy, σ_xy).
        res             : Characteristic mesh resolution for RoPE scale.
        scale           : Fourier feature scale factor.
        dropout         : Dropout on the encoder latent before cross-attn.
    """

    def __init__(self,
                 latent_channels: int,
                 out_channels: int,
                 res: int = 200,
                 scale: int = 8,
                 dropout: float = 0.1,
                 **kwargs,
                 ):
        super().__init__()
        self.out_channels    = out_channels
        self.latent_channels = latent_channels
        C = latent_channels

        # Map mesh coordinates to latent space
        self.coordinate_projection = nn.Sequential(
            GaussianFourierFeatureTransform(2, C // 2, scale=scale),
            nn.Linear(C, C, bias=False),
            nn.GELU(),
            nn.Linear(C, C, bias=False),
        )

        self.input_dropout = nn.Dropout(dropout)

        # Cross-attention: coordinate queries ← encoder latent
        self.decoding_transformer = CrossFormer(
            C, 'galerkin', heads=4, dim_head=C, mlp_dim=C,
            relative_emb=True, scale=32., relative_emb_dim=2,
            min_freq=1 / res,
        )

        # Self-attention mix layer
        self.mix_layer = LinearAttention(
            C, 'galerkin', heads=1, dim_head=C,
            relative_emb=True, scale=32, relative_emb_dim=2,
            min_freq=1 / res, use_ln=False,
        )

        # Expand to 2C for the propagation MLP
        self.expand_layer = nn.Linear(C, C * 2, bias=False)

        # Single propagation step (MLP residual block)
        self.propagator = nn.ModuleList([
            nn.ModuleList([
                nn.LayerNorm(C * 2),
                nn.Sequential(
                    nn.Linear(C * 2 + 2, C * 2, bias=False), nn.GELU(),
                    nn.Linear(C * 2,     C * 2, bias=False), nn.GELU(),
                    nn.Linear(C * 2,     C * 2, bias=False), nn.GELU(),
                    nn.Linear(C * 2,     C * 2, bias=False),
                ),
            ])
        ])

        # Output projection
        self.out_norm = nn.LayerNorm(C * 2)
        self.to_out   = nn.Sequential(
            nn.Linear(C * 2, C * 2, bias=False), nn.ReLU(inplace=True),
            nn.Linear(C * 2, C,     bias=False), nn.ReLU(inplace=True),
            nn.Linear(C,     out_channels, bias=True),
        )

    def _propagate(self, z, pos):
        """One step of the latent propagation MLP."""
        for norm_fn, ffn in self.propagator:
            z = ffn(torch.cat((norm_fn(z), pos), dim=-1)) + z
        return z

    def _decode(self, z):
        return self.to_out(self.out_norm(z))

    def forward(self, z, propagate_pos, forward_steps, input_pos):
        """Roll out ``forward_steps`` future field states.

        Args:
            z              : [b, n, C]   encoder latent node features
            propagate_pos  : [b, n, 2]   mesh node coordinates
            forward_steps  : int         number of autoregressive steps
            input_pos      : [b, n, 2]   encoder input positions
                             (same as propagate_pos for non-superresolution)

        Returns:
            Tensor [b, forward_steps, n, out_channels]
        """
        x = self.coordinate_projection(propagate_pos)
        z = self.input_dropout(z)
        z = self.decoding_transformer(x, z, propagate_pos, input_pos)
        z = self.mix_layer(z, propagate_pos) + z
        z = self.expand_layer(z)

        history = []
        for _ in range(forward_steps):
            z = self._propagate(z, propagate_pos)
            history.append(self._decode(z))

        return torch.stack(history, dim=1)   # [b, forward_steps, n, out_channels]