from __future__ import annotations

import unittest

from airfrans_frontier.runtime import remove_pythonpath_entries

remove_pythonpath_entries()

import torch

from airfrans_frontier.models import (
    DeepONetBranchTrunk,
    LocalPointTransformer,
    NeRFCFDMultiRes,
    PointContextPerceiver,
    PointwiseMLP,
    RasterFNOUNet,
    SirenConditionedINR,
)


class PointwiseMLPTests(unittest.TestCase):
    def test_forward_pass_returns_batch_by_target_dim(self) -> None:
        model = PointwiseMLP(input_dim=5, output_dim=4, hidden_width=16, depth=2, activation="gelu")
        batch = torch.randn(7, 5)

        output = model(batch)

        self.assertEqual(tuple(output.shape), (7, 4))

    def test_frontier_models_return_batch_by_target_dim(self) -> None:
        feature_names = ("x", "y", "sdf", "u_inf", "log_re", "aoa_deg")
        batch = torch.randn(8, len(feature_names))
        models = [
            NeRFCFDMultiRes(
                feature_names=feature_names,
                output_dim=4,
                coordinate_features=("x", "y", "sdf"),
                encoding_levels=2,
                hidden_width=16,
                depth=2,
                condition_width=12,
                condition_depth=2,
                activation="gelu",
            ),
            DeepONetBranchTrunk(
                feature_names=feature_names,
                output_dim=4,
                coordinate_features=("x", "y", "sdf"),
                fourier_scales=(1.0, 2.0),
                hidden_width=16,
                depth=2,
                condition_width=12,
                condition_depth=2,
                activation="gelu",
            ),
            PointContextPerceiver(
                input_dim=len(feature_names),
                output_dim=4,
                hidden_width=16,
                latent_width=12,
                context_points=6,
                attention_depth=2,
                activation="gelu",
            ),
            LocalPointTransformer(
                feature_names=feature_names,
                output_dim=4,
                coordinate_features=("x", "y", "sdf"),
                hidden_width=16,
                depth=2,
                neighbors=3,
                activation="gelu",
            ),
            RasterFNOUNet(
                feature_names=feature_names,
                output_dim=4,
                coordinate_features=("x", "y", "sdf"),
                grid_resolution=4,
                hidden_width=16,
                depth=2,
                condition_width=12,
                condition_depth=2,
                activation="gelu",
            ),
            SirenConditionedINR(
                feature_names=feature_names,
                output_dim=4,
                coordinate_features=("x", "y", "sdf"),
                hidden_width=16,
                depth=2,
                condition_width=12,
                condition_depth=2,
                omega0=10.0,
                activation="gelu",
            ),
        ]

        for model in models:
            with self.subTest(model=type(model).__name__):
                output = model(batch)
                self.assertEqual(tuple(output.shape), (8, 4))


if __name__ == "__main__":
    unittest.main()