"""MeshGraphNet (Pfaff et al., 2021) for cylinder_flow, laid out like a transformers LLM.

    MeshGraphNetForSimulation      backbone + decoder head + loss + rollout  (~ LlamaForCausalLM)
      model: MeshGraphNetModel     backbone                                   (~ LlamaModel)
        embeddings                 normalizers + node/edge encoder MLPs       (~ embed_tokens)
        layers[i]                  message passing: edge_block + node_block   (~ LlamaDecoderLayer)
"""

from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn
from transformers import PretrainedConfig, PreTrainedModel
from transformers.utils import ModelOutput


class MeshGraphNetConfig(PretrainedConfig):
    r"""Configuration for MeshGraphNet (Pfaff et al., 2021) on cylinder_flow.

    Defaults reproduce the DeepMind / medav cylinder_flow model.

    Args:
        hidden_size: Latent width of every node/edge representation.
        num_hidden_layers: Number of message-passing layers (processor steps).
        num_mlp_hidden_layers: Hidden layers inside every MLP (output layer excluded).
        velocity_size: Velocity components per node.
        num_node_types: Size of the node-type one-hot vector.
        edge_input_size: Raw edge features (relative mesh position + its norm).
        output_size: Predicted quantities per node (velocity change).
        split_first_linear: Hold the first Linear of each layer's edge/node MLP
            as one Linear per input (``W [h_s; h_r; e] = W_s h_s + W_r h_r + W_e e``).
            Sender/receiver projections are then computed once per node and
            gathered, instead of building the per-edge concatenation.
            Weights and outputs are identical to ``False``.
        layer_norm_eps: LayerNorm epsilon.
        normalizer_eps: Lower bound on the feature-normalizer std.
    """

    model_type = "meshgraphnet"

    def __init__(
        self,
        hidden_size: int = 128,
        num_hidden_layers: int = 15,
        num_mlp_hidden_layers: int = 2,
        velocity_size: int = 2,
        num_node_types: int = 9,
        edge_input_size: int = 3,
        output_size: int = 2,
        split_first_linear: bool = False,
        layer_norm_eps: float = 1e-5,
        normalizer_eps: float = 1e-8,
        **kwargs,
    ):
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_mlp_hidden_layers = num_mlp_hidden_layers
        self.velocity_size = velocity_size
        self.num_node_types = num_node_types
        self.edge_input_size = edge_input_size
        self.output_size = output_size
        self.split_first_linear = split_first_linear
        self.layer_norm_eps = layer_norm_eps
        self.normalizer_eps = normalizer_eps
        super().__init__(**kwargs)

    @property
    def node_input_size(self) -> int:
        return self.velocity_size + self.num_node_types


# Node types of the DeepMind datasets. Only NORMAL and OUTFLOW nodes are predicted.
NORMAL, OBSTACLE, AIRFOIL, HANDLE, INFLOW, OUTFLOW, WALL_BOUNDARY = range(7)


def prediction_mask(node_type: torch.Tensor) -> torch.Tensor:
    return (node_type == NORMAL) | (node_type == OUTFLOW)


def prepare_inputs(cells, mesh_pos, node_type, velocity) -> dict:
    """One mesh time step -> model inputs (the counterpart of a tokenizer).

    cells [M, 3] triangles -> edge_index [2, E] (senders, receivers), both directions.
    """
    cells = torch.as_tensor(cells).long()
    raw = torch.cat([cells[:, 0:2], cells[:, 1:3], cells[:, [2, 0]]])
    edges = torch.stack([raw.max(1).values, raw.min(1).values], 1).unique(dim=0)
    a, b = edges[:, 0], edges[:, 1]
    return {
        "velocity": torch.as_tensor(velocity).float(),
        "node_type": torch.as_tensor(node_type).long().reshape(-1),
        "mesh_pos": torch.as_tensor(mesh_pos).float(),
        "edge_index": torch.stack([torch.cat([a, b]), torch.cat([b, a])]),
    }


class MeshGraphNetNormalizer(nn.Module):
    """Feature normalizer that accumulates mean/std in training mode (DeepMind normalization.py).

    Statistics are accumulated for the first ``max_accumulations`` training calls, then frozen.
    """

    def __init__(self, size: int, eps: float = 1e-8, max_accumulations: int = 10**6):
        super().__init__()
        self.eps, self.max_accumulations = eps, max_accumulations
        self.register_buffer("acc_count", torch.zeros(()))
        self.register_buffer("num_accumulations", torch.zeros(()))
        self.register_buffer("acc_sum", torch.zeros(size))
        self.register_buffer("acc_sum_squared", torch.zeros(size))

    def _stats(self):
        count = self.acc_count.clamp(min=1)
        mean = self.acc_sum / count
        std = (self.acc_sum_squared / count - mean**2).clamp(min=0).sqrt().clamp(min=self.eps)
        return mean, std

    def forward(self, x):
        if self.training:  # no host sync: the update is masked once accumulation is done
            keep = (self.num_accumulations < self.max_accumulations).to(x.dtype)
            self.acc_count += keep * x.shape[0]
            self.num_accumulations += keep
            self.acc_sum += keep * x.detach().sum(0)
            self.acc_sum_squared += keep * (x.detach() ** 2).sum(0)
        mean, std = self._stats()
        return (x - mean) / std

    def inverse(self, x):
        mean, std = self._stats()
        return x * std + mean


class MeshGraphNetMLP(nn.Module):
    """``sizes=[in, h, h, out]`` -> Linear ReLU Linear ReLU Linear (+ LayerNorm)."""

    def __init__(self, sizes: list[int], layer_norm: bool, eps: float = 1e-5):
        super().__init__()
        self.layers = nn.ModuleList(nn.Linear(i, o) for i, o in zip(sizes[:-1], sizes[1:]))
        self.layer_norm = nn.LayerNorm(sizes[-1], eps=eps) if layer_norm else None

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(torch.relu(x) if i > 0 else x)
        return self.layer_norm(x) if self.layer_norm is not None else x


class MeshGraphNetEmbeddings(nn.Module):
    def __init__(self, config: MeshGraphNetConfig):
        super().__init__()
        d, hidden, eps = config.hidden_size, [config.hidden_size] * config.num_mlp_hidden_layers, config.layer_norm_eps
        self.num_node_types = config.num_node_types
        self.node_normalizer = MeshGraphNetNormalizer(config.node_input_size, config.normalizer_eps)
        self.edge_normalizer = MeshGraphNetNormalizer(config.edge_input_size, config.normalizer_eps)
        self.node_encoder = MeshGraphNetMLP([config.node_input_size, *hidden, d], True, eps)
        self.edge_encoder = MeshGraphNetMLP([config.edge_input_size, *hidden, d], True, eps)

    def forward(self, velocity, node_type, mesh_pos, edge_index):
        senders, receivers = edge_index
        one_hot = nn.functional.one_hot(node_type, self.num_node_types).to(velocity.dtype)
        rel = mesh_pos[senders] - mesh_pos[receivers]
        node_features = torch.cat([velocity, one_hot], -1)
        edge_features = torch.cat([rel, rel.norm(dim=-1, keepdim=True)], -1)
        return (self.node_encoder(self.node_normalizer(node_features)),
                self.edge_encoder(self.edge_normalizer(edge_features)))


class MeshGraphNetEdgeBlock(nn.Module):
    """New edge message = MLP([h_sender, h_receiver, e]).

    split_first_linear=True holds the first Linear as ``sender``/``receiver``/``edge``:
    sender/receiver projections run once per node and are gathered (no [E, 3d] concat).
    """

    def __init__(self, config: MeshGraphNetConfig):
        super().__init__()
        d = config.hidden_size
        self.split = config.split_first_linear
        if self.split:
            self.sender = nn.Linear(d, d, bias=False)
            self.receiver = nn.Linear(d, d, bias=False)
            self.edge = nn.Linear(d, d)
            for m in (self.sender, self.receiver, self.edge):
                m.concat_in_features = 3 * d  # init as one Linear(3d, d)
        else:
            self.input = nn.Linear(3 * d, d)
        self.mlp = MeshGraphNetMLP([d] * (config.num_mlp_hidden_layers + 1), True, config.layer_norm_eps)

    def forward(self, h, e, edge_index):
        s, r = edge_index
        if self.split:
            x = self.sender(h)[s] + self.receiver(h)[r] + self.edge(e)
        else:
            x = self.input(torch.cat([h[s], h[r], e], -1))
        return self.mlp(torch.relu(x))


class MeshGraphNetNodeBlock(nn.Module):
    """Node update = MLP([h, sum of incoming messages]).

    split_first_linear=True holds the first Linear as ``node``/``aggregated``.
    """

    def __init__(self, config: MeshGraphNetConfig):
        super().__init__()
        d = config.hidden_size
        self.split = config.split_first_linear
        if self.split:
            self.node = nn.Linear(d, d)
            self.aggregated = nn.Linear(d, d, bias=False)
            for m in (self.node, self.aggregated):
                m.concat_in_features = 2 * d  # init as one Linear(2d, d)
        else:
            self.input = nn.Linear(2 * d, d)
        self.mlp = MeshGraphNetMLP([d] * (config.num_mlp_hidden_layers + 1), True, config.layer_norm_eps)

    def forward(self, h, messages, edge_index):
        agg = torch.zeros_like(h).index_add_(0, edge_index[1], messages)
        x = self.node(h) + self.aggregated(agg) if self.split else self.input(torch.cat([h, agg], -1))
        return self.mlp(torch.relu(x))


class MeshGraphNetLayer(nn.Module):
    def __init__(self, config: MeshGraphNetConfig):
        super().__init__()
        self.edge_block = MeshGraphNetEdgeBlock(config)
        self.node_block = MeshGraphNetNodeBlock(config)

    def forward(self, h, e, edge_index):
        messages = self.edge_block(h, e, edge_index)
        return h + self.node_block(h, messages, edge_index), e + messages


@dataclass
class MeshGraphNetOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    delta_velocity: Optional[torch.FloatTensor] = None  # [N, 2], de-normalized
    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None  # node states per layer


class MeshGraphNetPreTrainedModel(PreTrainedModel):
    config_class = MeshGraphNetConfig
    base_model_prefix = "model"
    _no_split_modules = ["MeshGraphNetLayer"]

    def _init_weights(self, module):
        # Sonnet (DeepMind) default: truncated normal, std 1/sqrt(fan_in), cut at 2 std; zero bias.
        if isinstance(module, nn.Linear):
            std = getattr(module, "concat_in_features", module.in_features) ** -0.5
            nn.init.trunc_normal_(module.weight, std=std, a=-2 * std, b=2 * std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)


class MeshGraphNetModel(MeshGraphNetPreTrainedModel):
    def __init__(self, config: MeshGraphNetConfig):
        super().__init__(config)
        self.embeddings = MeshGraphNetEmbeddings(config)
        self.layers = nn.ModuleList(MeshGraphNetLayer(config) for _ in range(config.num_hidden_layers))
        self.post_init()

    def forward(self, velocity, node_type, mesh_pos, edge_index, output_hidden_states=False):
        h, e = self.embeddings(velocity, node_type, mesh_pos, edge_index)
        hidden_states = (h,) if output_hidden_states else None
        for layer in self.layers:
            h, e = layer(h, e, edge_index)
            if output_hidden_states:
                hidden_states += (h,)
        return h, hidden_states


class MeshGraphNetForSimulation(MeshGraphNetPreTrainedModel):
    """Predicts the per-node velocity change for one time step (dt = 0.01)."""

    prepare_inputs = staticmethod(prepare_inputs)

    def __init__(self, config: MeshGraphNetConfig):
        super().__init__(config)
        hidden = [config.hidden_size] * config.num_mlp_hidden_layers
        self.model = MeshGraphNetModel(config)
        self.decoder = MeshGraphNetMLP([config.hidden_size, *hidden, config.output_size], False)
        self.output_normalizer = MeshGraphNetNormalizer(config.output_size, config.normalizer_eps)
        self.post_init()

    def forward(self, velocity, node_type, mesh_pos, edge_index, target_velocity=None,
                output_hidden_states=False):
        h, hidden_states = self.model(velocity, node_type, mesh_pos, edge_index, output_hidden_states)
        out = self.decoder(h)
        loss = None
        if target_velocity is not None:
            # DeepMind: MSE of the normalized velocity change on NORMAL/OUTFLOW nodes.
            target = self.output_normalizer(target_velocity - velocity)
            mask = prediction_mask(node_type).to(out.dtype)
            loss = (((target - out) ** 2).sum(-1) * mask).sum() / mask.sum()
        return MeshGraphNetOutput(loss, self.output_normalizer.inverse(out), hidden_states)

    @torch.no_grad()
    def rollout(self, velocity, node_type, mesh_pos, edge_index, num_steps: int):
        """Autoregressive prediction (~ generate). Boundary nodes keep their initial
        velocity. Returns [num_steps + 1, N, 2]."""
        mask = prediction_mask(node_type).unsqueeze(-1)
        trajectory = [velocity]
        for _ in range(num_steps):
            delta = self(velocity, node_type, mesh_pos, edge_index).delta_velocity
            velocity = torch.where(mask, velocity + delta, velocity)
            trajectory.append(velocity)
        return torch.stack(trajectory)