"""Wrapper around the OFFICIAL Transolver++ model (thuml/Transolver_plus, ICML 2025)
so it trains/evaluates in our harness on our PyG graphs, apples to apples with GeoReNet
and the Transolver baseline (same data, split, loss, eval). We load the authors'
models/Transolver_plus.py verbatim and only adapt the I/O plus two environment shims:

  1. timm: the official file does `from timm.layers import trunc_normal_`; we stub that
     one symbol with torch's native trunc_normal_ (avoids the timm dependency that
     conflicts with our pinned torch 2.11+cu128).
  2. torch.distributed.nn.all_reduce: the eidetic attention sums slice statistics across
     GPUs via all_reduce (the paper's input-invariant parallelism). On a single GPU,
     all_reduce over one process is the identity, so we replace it with a passthrough.
     This is the faithful single-GPU behavior of the authors' code, not a model change.

The Model.forward consumes (x, pos, condition): x = [coords + node features], pos = coords
(used only if unified_pos), condition = optional global token. We feed coords + node
features + broadcast (Re, yaw) globals so Transolver++ gets the SAME conditioning signal
GeoReNet's FiLM uses (fair), and condition=None.
"""
from __future__ import annotations
import os, sys, types, importlib.util
import torch
import torch.nn as nn

_OFFICIAL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                         "Transolver_plus-main", "models", "Transolver_plus.py")


def _stub_timm():
    if "timm.layers" in sys.modules:
        return
    timm = types.ModuleType("timm")
    layers = types.ModuleType("timm.layers"); layers.trunc_normal_ = torch.nn.init.trunc_normal_
    models = types.ModuleType("timm.models")
    mlayers = types.ModuleType("timm.models.layers"); mlayers.trunc_normal_ = torch.nn.init.trunc_normal_
    timm.layers = layers; timm.models = models; models.layers = mlayers
    sys.modules.update({"timm": timm, "timm.layers": layers,
                        "timm.models": models, "timm.models.layers": mlayers})


def _load_official():
    _stub_timm()
    spec = importlib.util.spec_from_file_location("transolverpp_official", _OFFICIAL)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    # single-GPU shims: ensure ReduceOp exists, then make all_reduce the identity
    # (all_reduce over one process is a passthrough = faithful single-GPU behavior).
    import torch.distributed as _td
    if not hasattr(mod.dist_nn, "ReduceOp"):
        mod.dist_nn.ReduceOp = getattr(_td, "ReduceOp", types.SimpleNamespace(SUM=0))
    mod.dist_nn.all_reduce = lambda t, op=None: t
    return mod


class _NS:
    pass


class TransolverPlusWrapper(nn.Module):
    """Drop-in for our training loop: forward(data) -> [N_nodes, out_dim].
    Config sized to ~GeoReNet/Transolver scale for a fair comparison; align to the
    authors' main_airplane.py settings if exact faithfulness to their aero run is wanted."""
    def __init__(self, node_in, space_dim=3, global_dim=2, out_dim=7, n_hidden=256,
                 n_layers=8, n_head=8, slice_num=64, mlp_ratio=2):
        super().__init__()
        mod = _load_official()
        self.space_dim = space_dim
        self.global_dim = global_dim
        # Model preprocess input = fun_dim + space_dim; we put coords in space_dim and
        # node features + broadcast globals in fun_dim.
        self.net = mod.Model(space_dim=space_dim, fun_dim=node_in + global_dim, out_dim=out_dim,
                             n_hidden=n_hidden, n_layers=n_layers, n_head=n_head,
                             slice_num=slice_num, mlp_ratio=mlp_ratio, unified_pos=False)

    def forward(self, data):
        # PER GRAPH: a PyG DataLoader concatenates graphs into one node tensor with a
        # `batch` vector. Transolver++ attention treats its input as ONE point cloud, so
        # we must run each graph separately (else nodes leak across graphs and only graph
        # 0's globals are used). Output is reassembled in the original node order.
        N = data.x.shape[0]
        batch = getattr(data, "batch", None)
        if batch is None:
            batch = torch.zeros(N, dtype=torch.long, device=data.x.device)
        B = int(batch.max().item()) + 1
        pos_all = data.pos
        if pos_all.shape[1] < self.space_dim:
            pos_all = torch.cat([pos_all, pos_all.new_zeros(pos_all.shape[0],
                                 self.space_dim - pos_all.shape[1])], 1)
        gf_all = data.global_feat.reshape(B, -1)[:, :self.global_dim] if self.global_dim else None
        outs = []
        for b in range(B):
            m = batch == b
            pos = pos_all[m]
            parts = [pos[:, :self.space_dim], data.x[m]]
            if self.global_dim:
                parts.append(gf_all[b:b + 1].expand(pos.shape[0], -1))
            x = torch.cat(parts, dim=-1).unsqueeze(0)        # [1, n_b, space+fun]
            outs.append(self.net((x, pos.unsqueeze(0), None))[0])   # [n_b, out_dim]
        return torch.cat(outs, dim=0)                        # [N, out_dim], original order


def count_params(m):
    return sum(p.numel() for p in m.parameters())