"""Physics-Attention for irregular meshes / point clouds.

This module is reused **verbatim** from the Transolver project and is licensed under the MIT
License. Attribution preserved per the license terms.

    Transolver: A Fast Transformer Solver for PDEs on General Geometries
    Wu, Luo, Wang, Wang, Long (THUML / Tsinghua University), ICML 2024 Spotlight.
    Source: https://github.com/thuml/Transolver  (MIT License)
    Paper:  arXiv:2402.02366

Key constants preserved exactly: temperature init 0.5, orthogonal init of the slice
projection, slice-norm epsilon 1e-5, attention scale dim_head**-0.5.

Note: in the assembled model (see transolver.py), the operator's global ``trunc_normal_``
weight initialization runs *after* this block is constructed and overwrites the orthogonal
init of ``in_project_slice`` — this matches the upstream assembly order and is what produced
the published 0.0064. We keep the orthogonal init here to stay verbatim with upstream.
"""
import torch
import torch.nn as nn
from einops import rearrange


class Physics_Attention_Irregular_Mesh(nn.Module):
    ## for irregular meshes in 1D, 2D or 3D space
    # NOTE: the dim_head=64 default is upstream's generic default. For the Elasticity config the
    # assembler (transolver.py) always passes dim_head = n_hidden // n_heads = 16; the default here
    # is never used in this repo. See docs/RECONCILIATION.md §3.
    def __init__(self, dim, heads=8, dim_head=64, dropout=0., slice_num=64):
        super().__init__()
        inner_dim = dim_head * heads
        self.dim_head = dim_head
        self.heads = heads
        self.scale = dim_head ** -0.5
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)

        self.in_project_x = nn.Linear(dim, inner_dim)
        self.in_project_fx = nn.Linear(dim, inner_dim)
        self.in_project_slice = nn.Linear(dim_head, slice_num)
        for l in [self.in_project_slice]:
            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # B N C
        B, N, C = x.shape

        ### (1) Slice
        fx_mid = self.in_project_fx(x).reshape(B, N, self.heads, self.dim_head) \
            .permute(0, 2, 1, 3).contiguous()  # B H N C
        x_mid = self.in_project_x(x).reshape(B, N, self.heads, self.dim_head) \
            .permute(0, 2, 1, 3).contiguous()  # B H N C
        slice_weights = self.softmax(self.in_project_slice(x_mid) / self.temperature)  # B H N G
        slice_norm = slice_weights.sum(2)  # B H G
        slice_token = torch.einsum("bhnc,bhng->bhgc", fx_mid, slice_weights)
        slice_token = slice_token / ((slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head))

        ### (2) Attention among slice tokens
        q_slice_token = self.to_q(slice_token)
        k_slice_token = self.to_k(slice_token)
        v_slice_token = self.to_v(slice_token)
        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
        attn = self.softmax(dots)
        attn = self.dropout(attn)
        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D

        ### (3) Deslice
        out_x = torch.einsum("bhgc,bhng->bhnc", out_slice_token, slice_weights)
        out_x = rearrange(out_x, 'b h n d -> b n (h d)')
        return self.to_out(out_x)