from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

import torch


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def check(name: str, tensor: torch.Tensor, expected_last_dim: int | None = None) -> None:
    if not torch.isfinite(tensor).all():
        raise RuntimeError(f"{name} produced non-finite values")
    if expected_last_dim is not None and tensor.shape[-1] != expected_last_dim:
        raise RuntimeError(
            f"{name} last dim expected {expected_last_dim}, got {tensor.shape[-1]}"
        )
    print(f"[OK] {name}: shape={tuple(tensor.shape)}")


@torch.no_grad()
def smoke_deeponet() -> None:
    from model.deeponet import DeepONetCartesianProd1D

    model = DeepONetCartesianProd1D(
        size=8,
        in_channel_branch=1,
        query_dim=1,
        out_channel=1,
        activation="relu",
        base_model="MLP",
    ).eval()
    x_func = torch.randn(2, 8, 1)
    x_loc = torch.linspace(0.0, 1.0, 8).unsqueeze(-1)
    out = model((x_func, x_loc))
    check("deeponet.DeepONetCartesianProd1D", out, expected_last_dim=1)


@torch.no_grad()
def smoke_fno() -> None:
    from model.fno import FNO2d

    model = FNO2d(num_channels=1, modes1=4, modes2=4, width=8, initial_step=1).eval()
    x = torch.randn(2, 16, 16, 1)
    coords = torch.linspace(0.0, 1.0, 16)
    xx, yy = torch.meshgrid(coords, coords, indexing="ij")
    grid = torch.stack((xx, yy), dim=-1).unsqueeze(0).repeat(2, 1, 1, 1)
    out = model(x, grid)
    check("fno.FNO2d", out, expected_last_dim=1)


@torch.no_grad()
def smoke_mpnn() -> None:
    from torch_geometric.data import Data

    from model.mpnn import MPNN

    pde = SimpleNamespace(
        name="smoke",
        tmin=0.0,
        tmax=1.0,
        resolution_t=10,
        spatial_domain=[(0.0, 1.0)],
        resolution=[4],
        spatial_dim=1,
        variables={},
    )
    model = MPNN(pde=pde, time_window=10, hidden_features=128, hidden_layers=1).eval()
    edge_index = torch.tensor(
        [[0, 1, 2, 3, 1, 2, 3, 0], [1, 2, 3, 0, 0, 1, 2, 3]],
        dtype=torch.long,
    )
    data = Data(
        x=torch.randn(4, 10, 1),
        x_pos=torch.linspace(0.0, 1.0, 4).unsqueeze(-1),
        t_pos=torch.zeros(4),
        edge_index=edge_index,
        batch=torch.zeros(4, dtype=torch.long),
        variables=torch.empty(4, 0),
    )
    out = model(data, v=0)
    check("mpnn.MPNN", out)


@torch.no_grad()
def smoke_pino_fno() -> None:
    from model.pino_fno import FNO1d

    model = FNO1d(
        modes=[4, 4],
        width=4,
        layers=[4, 4, 4],
        fc_dim=8,
        in_dim=2,
        out_dim=1,
        pad_ratio=[0.0, 0.0],
    ).eval()
    out = model(torch.randn(2, 16, 2))
    check("pino_fno.FNO1d", out, expected_last_dim=1)


@torch.no_grad()
def smoke_unet() -> None:
    from model.unet import UNet2d

    model = UNet2d(in_channels=1, out_channels=1, init_features=2).eval()
    out = model(torch.randn(1, 1, 16, 16))
    check("unet.UNet2d", out)


@torch.no_grad()
def smoke_uno() -> None:
    from model.uno import UNO1d

    model = UNO1d(num_channels=1, width=4, initial_step=1).eval()
    x = torch.randn(1, 64, 1)
    grid = torch.linspace(0.0, 1.0, 64).view(1, 64, 1)
    out = model(x, grid)
    check("uno.UNO1d", out, expected_last_dim=1)


def main() -> int:
    torch.manual_seed(0)
    smoke_deeponet()
    smoke_fno()
    smoke_mpnn()
    smoke_pino_fno()
    smoke_unet()
    smoke_uno()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())