"""Assemble the stress operator: encoder -> N pre-norm blocks -> decoder head.

Faithful reproduction of Transolver ``model/Transolver_Irregular_Mesh.py::Model`` (MIT). The
attention type is configurable (``physics`` for Stage 1, ``linearno`` for Stage 2) and is the
*only* thing that changes between gates — a controlled comparison.

Reproduction note: ``initialize_weights()`` runs after the blocks are built, so the global
``trunc_normal_(std=0.02)`` init is applied to every Linear, **overwriting** the orthogonal
init of each attention's ``in_project_slice``. This matches the upstream assembly order.
"""
from __future__ import annotations

import numpy as np
import torch
import torch.nn as nn

from .blocks import MLP, TransolverBlock
from .physics_attention import Physics_Attention_Irregular_Mesh


def make_attention(
    kind: str,
    dim,
    heads,
    dim_head,
    dropout,
    slice_num,
    linearno_variant: str = "shared_qk",
    linearno_project_out: bool = False,
    linearno_temperature: bool = False,
) -> nn.Module:
    if kind == "physics":
        return Physics_Attention_Irregular_Mesh(
            dim, heads=heads, dim_head=dim_head, dropout=dropout, slice_num=slice_num
        )
    if kind == "linearno":
        from .linear_no import LinearNO  # lazy: only needed at Stage 2

        return LinearNO(
            dim,
            heads=heads,
            dim_head=dim_head,
            slice_num=slice_num,
            dropout=dropout,
            variant=linearno_variant,
            project_out=linearno_project_out,
            temperature=linearno_temperature,
        )
    raise ValueError(f"unknown attention kind {kind!r} (expected 'physics' or 'linearno')")


class StressOperator(nn.Module):
    def __init__(
        self,
        attention: str = "physics",
        space_dim: int = 2,
        n_layers: int = 8,
        n_hidden: int = 128,
        dropout: float = 0.0,
        n_heads: int = 8,
        dim_head: int | None = None,
        mlp_ratio: int = 1,
        fun_dim: int = 0,
        out_dim: int = 1,
        slice_num: int = 64,
        unified_pos: bool = False,
        ref: int = 8,
        act: str = "gelu",
        linearno_variant: str = "shared_qk",
        linearno_project_out: bool = False,
        linearno_temperature: bool = False,
    ):
        super().__init__()
        if dim_head is None:
            dim_head = n_hidden // n_heads  # = 16 for the Elasticity config (repo value)
        self.attention_kind = attention
        self.unified_pos = unified_pos
        self.ref = ref
        self.n_hidden = n_hidden

        in_dim = (fun_dim + ref * ref) if unified_pos else (fun_dim + space_dim)
        self.preprocess = MLP(in_dim, n_hidden * 2, n_hidden, n_layers=0, res=False, act=act)

        self.blocks = nn.ModuleList(
            [
                TransolverBlock(
                    attention=make_attention(
                        attention, n_hidden, n_heads, dim_head, dropout, slice_num,
                        linearno_variant=linearno_variant,
                        linearno_project_out=linearno_project_out,
                        linearno_temperature=linearno_temperature,
                    ),
                    hidden_dim=n_hidden,
                    dropout=dropout,
                    act=act,
                    mlp_ratio=mlp_ratio,
                    last_layer=(i == n_layers - 1),
                    out_dim=out_dim,
                )
                for i in range(n_layers)
            ]
        )
        self.initialize_weights()
        self.placeholder = nn.Parameter((1 / n_hidden) * torch.rand(n_hidden, dtype=torch.float))

    def initialize_weights(self):
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_grid(self, x):
        """Unified positional grid (only used when ``unified_pos`` is True). Device-agnostic."""
        b = x.shape[0]
        device = x.device
        gx = torch.linspace(0, 1, self.ref, device=device).reshape(1, self.ref, 1, 1).repeat(b, 1, self.ref, 1)
        gy = torch.linspace(0, 1, self.ref, device=device).reshape(1, 1, self.ref, 1).repeat(b, self.ref, 1, 1)
        grid_ref = torch.cat((gx, gy), dim=-1).reshape(b, self.ref * self.ref, 2)
        pos = torch.sqrt(((x[:, :, None, :] - grid_ref[:, None, :, :]) ** 2).sum(-1))
        return pos.reshape(b, x.shape[1], self.ref * self.ref).contiguous()

    def forward(self, x, fx=None):
        # x: (B, N, space_dim) node coordinates; fx: optional extra input function
        if self.unified_pos:
            x = self.get_grid(x)
        fx = self.preprocess(x if fx is None else torch.cat((x, fx), dim=-1))
        fx = fx + self.placeholder[None, None, :]
        for block in self.blocks:
            fx = block(fx)
        return fx  # (B, N, out_dim)


def build_model(model_cfg: dict) -> StressOperator:
    """Instantiate :class:`StressOperator` from a config dict (configs/*.yaml ``model`` block)."""
    return StressOperator(
        attention=model_cfg.get("attention", "physics"),
        space_dim=model_cfg.get("space_dim", 2),
        n_layers=model_cfg.get("n_layers", 8),
        n_hidden=model_cfg.get("n_hidden", 128),
        dropout=model_cfg.get("dropout", 0.0),
        n_heads=model_cfg.get("n_heads", 8),
        dim_head=model_cfg.get("dim_head", None),
        mlp_ratio=model_cfg.get("mlp_ratio", 1),
        fun_dim=model_cfg.get("fun_dim", 0),
        out_dim=model_cfg.get("out_dim", 1),
        slice_num=model_cfg.get("slice_num", 64),
        unified_pos=model_cfg.get("unified_pos", False),
        ref=model_cfg.get("ref", 8),
        act=model_cfg.get("act", "gelu"),
        linearno_variant=model_cfg.get("linearno_variant", "shared_qk"),
        linearno_project_out=model_cfg.get("linearno_project_out", False),
        linearno_temperature=model_cfg.get("linearno_temperature", False),
    )


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)