import pytest
import torch
from graph_weather.models.aurora.model import AuroraModel, EarthSystemLoss
from graph_weather.models.aurora.encoder import Swin3DEncoder
from graph_weather.models.aurora.decoder import Decoder3D
from graph_weather.models.aurora.processor import PerceiverProcessor, ProcessorConfig


# Fixtures
@pytest.fixture
def sample_3d_data():
    """Create small sample 3D data for testing"""
    batch_size = 2
    in_channels = 2  # Match swin3d_config
    depth = height = width = 8
    return torch.randn(batch_size, in_channels, depth, height, width)


@pytest.fixture
def swin3d_config():
    """Minimal configuration for Swin3D encoder"""
    return {"in_channels": 2, "embed_dim": 256}  # Match processor's input_dim


@pytest.fixture
def sample_unstructured_data():
    """Create sample unstructured point data for testing"""
    batch_size = 1
    num_points = 45  # Keeping under max_points=50
    num_features = 4

    points = torch.rand(batch_size, num_points, 2) * 360
    points[:, :, 0] = points[:, :, 0] - 180
    points[:, :, 1] = points[:, :, 1] - 90

    features = torch.randn(batch_size, num_points, num_features - 2)

    return points, features


@pytest.fixture
def model_config():
    """Basic model configuration for unstructured point data"""
    return {
        "input_features": 2,
        "output_features": 2,
        "latent_dim": 8,  # Adjusted
        "max_points": 50,  # Adjusted
        "max_seq_len": 128,  # Adjusted
    }


@pytest.fixture
def processor_config():
    """Basic configuration for PerceiverProcessor"""
    return ProcessorConfig(
        input_dim=256,  # Adjusted as per test needs
        latent_dim=64,  # Reduced from 512
        d_model=32,  # Adjusted
        max_seq_len=256,  # Adjusted
        num_self_attention_layers=2,  # Reduced
        num_cross_attention_layers=1,  # Reduced
        num_attention_heads=4,  # Reduced
        hidden_dropout=0.1,
        attention_dropout=0.1,
    )


@pytest.fixture
def test_gradient_checkpointing_config():
    """Test that gradient checkpointing can be configured"""
    # Test with checkpointing disabled
    config_no_checkpoint = {
        "input_features": 2,
        "output_features": 2,
        "latent_dim": 8,
        "max_points": 50,
        "max_seq_len": 128,
        "use_checkpointing": False,
    }
    model_no_checkpoint = AuroraModel(**config_no_checkpoint)
    assert not model_no_checkpoint.use_checkpointing

    # Test with checkpointing enabled
    config_with_checkpoint = {**config_no_checkpoint, "use_checkpointing": True}
    model_with_checkpoint = AuroraModel(**config_with_checkpoint)
    assert model_with_checkpoint.use_checkpointing


def get_test_configs():
    return {
        "swin3d_config": {
            "in_channels": 1,
            "embed_dim": 256,  # Consistent with processor input_dim
        },
        "processor_config": ProcessorConfig(
            input_dim=256,  # Match Swin3D output
            latent_dim=64,  # Reduced for testing
            d_model=256,  # Match input_dim
            max_seq_len=256,  # Reduced for testing
            num_self_attention_layers=2,
            num_attention_heads=8,
        ),
    }


def test_swin3d_encoder(sample_3d_data, swin3d_config):
    """Test Swin3D encoder with minimal 3D data"""
    encoder = Swin3DEncoder(**swin3d_config)

    # Check the input data size
    batch_size = 2
    depth = height = width = 8  # These should match sample_3d_data fixture

    # Create input with correct size
    input_data = torch.randn(
        batch_size,
        swin3d_config["in_channels"],  # This should be 2 based on config
        depth,
        height,
        width,
    )

    # Process through encoder
    output = encoder(input_data)

    # Expected sequence length after flattening spatial dimensions
    expected_seq_len = depth * height * width

    # Verify output shape
    assert output.shape == (
        batch_size,
        expected_seq_len,
        swin3d_config["embed_dim"],
    ), f"Expected shape {(batch_size, expected_seq_len, swin3d_config['embed_dim'])}, got {output.shape}"
    assert not torch.isnan(output).any()


def test_decoder3d():
    """Test 3D decoder with minimal dimensions"""
    batch_size = 2
    embed_dim = 32  # Reduced from 96 to 32
    target_shape = (8, 8, 8)  # Reduced from 32 to 8

    decoder = Decoder3D(output_channels=1, embed_dim=embed_dim, target_shape=target_shape)
    input_tensor = torch.randn(
        batch_size, target_shape[0] * target_shape[1] * target_shape[2], embed_dim
    )

    output = decoder(input_tensor)
    expected_shape = (batch_size, 1, *target_shape)
    assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}"
    assert not torch.isnan(output).any(), "Output contains NaN values"


def test_perceiver_processor_with_config(processor_config):
    """Test Perceiver processor with configuration"""
    batch_size = 2
    seq_len = 16

    # Match input dimension with processor's expected input_dim
    input_dim = processor_config.input_dim  # This should be 256 as per ProcessorConfig

    processor = PerceiverProcessor(processor_config)
    input_tensor = torch.randn(batch_size, seq_len, input_dim)

    output = processor(input_tensor)

    # Check output shape
    # Output should be (batch_size, latent_dim) after global average pooling
    assert output.shape == (
        batch_size,
        processor_config.latent_dim,
    ), f"Expected shape {(batch_size, processor_config.latent_dim)}, got {output.shape}"
    assert not torch.isnan(output).any()


def test_perceiver_processor_default_config():
    """Test Perceiver processor with default configuration"""
    batch_size = 2
    seq_len = 16

    processor = PerceiverProcessor()  # Uses default config
    input_dim = processor.config.input_dim  # Should be 256 from default config

    input_tensor = torch.randn(batch_size, seq_len, input_dim)
    output = processor(input_tensor)

    # Check output shape against default config
    assert output.shape == (batch_size, processor.config.latent_dim)
    assert not torch.isnan(output).any()


def test_perceiver_processor_4d_input(processor_config):
    """Test Perceiver processor with 4D input"""
    batch_size = 2
    seq_len = 8
    height = width = 4

    processor = PerceiverProcessor(processor_config)

    # Calculate input features to match processor's input_dim
    # Each position (h,w) needs to have features that sum to input_dim
    features_per_position = processor_config.input_dim // (height * width)

    # Create 4D input with correct dimensions for flattening
    input_tensor = torch.randn(batch_size, seq_len, features_per_position, height * width)

    # Reshape to (batch, seq, H*W*features)
    input_tensor = input_tensor.reshape(batch_size, seq_len, -1)

    output = processor(input_tensor)

    # Check output shape
    assert output.shape == (batch_size, processor_config.latent_dim)
    assert not torch.isnan(output).any()


def test_processor_config_validation():
    """Test validation of processor configuration parameters"""
    with pytest.raises(ValueError):
        # Test invalid input_dim
        ProcessorConfig(input_dim=-1)

    with pytest.raises(ValueError):
        # Test invalid max_seq_len
        ProcessorConfig(max_seq_len=0)

    with pytest.raises(ValueError):
        # Test invalid number of attention heads
        ProcessorConfig(num_attention_heads=0)

    with pytest.raises(ValueError):
        # Test invalid hidden_dropout
        ProcessorConfig(hidden_dropout=1.5)


def test_processor_attention_mask():
    """Test processor with attention mask"""
    config = ProcessorConfig(input_dim=32, latent_dim=64, max_seq_len=128)

    processor = PerceiverProcessor(config)
    batch_size = 2
    seq_len = 16

    input_tensor = torch.randn(batch_size, seq_len, config.input_dim)
    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
    attention_mask[:, seq_len // 2 :] = False  # Mask out second half of sequence

    output_masked = processor(input_tensor, attention_mask=attention_mask)
    output_unmasked = processor(input_tensor)

    # Outputs should be different when mask is applied
    assert not torch.allclose(output_masked, output_unmasked)
    assert not torch.isnan(output_masked).any()


def test_processor_dropout():
    """Test processor dropout behavior"""
    config = ProcessorConfig(input_dim=32, latent_dim=64, hidden_dropout=0.5, attention_dropout=0.5)

    processor = PerceiverProcessor(config)
    batch_size = 2
    seq_len = 16

    input_tensor = torch.randn(batch_size, seq_len, config.input_dim)

    processor.train()
    output1 = processor(input_tensor)
    output2 = processor(input_tensor)

    # Outputs should be different in training mode due to dropout
    assert not torch.allclose(output1, output2)

    processor.eval()
    output1 = processor(input_tensor)
    output2 = processor(input_tensor)

    # Outputs should be identical in eval mode
    assert torch.allclose(output1, output2)


def test_full_pipeline_integration(processor_config):
    """Test minimal pipeline integration with processor config"""
    batch_size = 2
    channels = 1
    size = 8

    # Create input data
    input_3d_data = torch.randn(batch_size, channels, size, size, size)

    # Initialize components with matching dimensions
    encoder = Swin3DEncoder(
        in_channels=channels, embed_dim=processor_config.input_dim  # Match processor's input_dim
    )

    processor = PerceiverProcessor(processor_config)

    # Process through encoder
    encoded = encoder(input_3d_data)  # Shape: (batch, seq, embed_dim)

    # Process through processor
    processed = processor(encoded)  # Shape: (batch, latent_dim)

    # Check shapes at each step
    assert encoded.shape == (
        batch_size,
        size * size * size,
        processor_config.input_dim,
    ), f"Encoder output shape mismatch: {encoded.shape}"
    assert processed.shape == (
        batch_size,
        processor_config.latent_dim,
    ), f"Processor output shape mismatch: {processed.shape}"
    assert not torch.isnan(processed).any()


def test_aurora_model_with_3d(sample_3d_data, model_config):
    """Test AuroraModel with minimal 3D data"""
    model = AuroraModel(**model_config)

    # Reduce dimensions to stay under max_points=50
    batch_size = 2
    width = height = depth = 3  # 3x3x3 = 27 points, well under max_points
    channels = 1

    # Create smaller sample data
    sample_3d_data = torch.randn(batch_size, channels, depth, height, width)

    x, y, z = torch.meshgrid(
        torch.linspace(-180, 180, width),
        torch.linspace(-90, 90, height),
        torch.linspace(-1, 1, depth),
        indexing="ij",
    )

    points = torch.stack([x, y, z], dim=-1).reshape(-1, 3)
    points = points.unsqueeze(0).repeat(batch_size, 1, 1)

    features = sample_3d_data.reshape(batch_size, channels, -1)
    features = features.transpose(1, 2)

    if features.shape[-1] != model_config["input_features"]:
        features = torch.cat(
            [
                features,
                torch.zeros(
                    batch_size,
                    features.shape[1],
                    model_config["input_features"] - features.shape[-1],
                ),
            ],
            dim=-1,
        )

    output = model(points, features)
    expected_output_shape = (batch_size, points.shape[1], model_config["output_features"])
    assert output.shape == expected_output_shape
    assert not torch.isnan(output).any()


def test_aurora_point_processing(sample_unstructured_data, model_config):
    """Test model's ability to process unstructured point data"""
    points, features = sample_unstructured_data

    model = AuroraModel(**model_config)
    output = model(points, features)

    assert (
        output.shape[0] == points.shape[0]
    ), f"Expected batch size {points.shape[0]}, got {output.shape[0]}"
    assert (
        output.shape[1] == points.shape[1]
    ), f"Expected sequence length {points.shape[1]}, got {output.shape[1]}"
    assert (
        output.shape[2] == model_config["output_features"]
    ), f"Expected features {model_config['output_features']}, got {output.shape[2]}"
    assert not torch.isnan(output).any(), "Output contains NaN values"


def test_variable_point_counts(model_config):
    """Test model with varying numbers of points per batch"""
    model = AuroraModel(**model_config)
    point_counts = [20, 35, 45]  # Reduced from [100, 200, 300]

    for num_points in point_counts:
        points = torch.rand(1, num_points, 2) * 360 - 180
        features = torch.randn(1, num_points, 2)

        output = model(points, features)
        expected_shape = (1, num_points, model_config["output_features"])
        assert output.shape == expected_shape
        assert not torch.isnan(output).any()


def test_point_ordering_invariance(sample_unstructured_data, model_config):
    """Test point ordering invariance with carefully controlled transformations"""
    model = AuroraModel(**model_config)
    model.eval()

    with torch.no_grad():
        points = torch.tensor(
            [
                [
                    [0.0, 0.0],
                    [1.0, 0.0],
                    [0.0, 1.0],
                    [1.0, 1.0],
                ]
            ],
            dtype=torch.float32,
        )

        points[..., 0] = points[..., 0] * 360 - 180
        points[..., 1] = points[..., 1] * 180 - 90

        features = torch.ones((1, 4, 2), dtype=torch.float32)
        baseline_output = model(points, features)

        transformations = [
            ([0, 1, 2, 3], "Identity"),
            ([1, 0, 2, 3], "Swap adjacent horizontal"),
            ([2, 3, 0, 1], "Swap top and bottom halves"),
            ([3, 2, 1, 0], "Complete reversal"),
        ]

        for perm, name in transformations:
            perm_points = points[:, perm, :]
            perm_features = features[:, perm, :]
            perm_output = model(perm_points, perm_features)

            unperm_output = torch.zeros_like(perm_output)
            unperm_output[:, perm, :] = perm_output

            diff = (baseline_output - unperm_output).abs()
            final_diff = diff.max().item()

            assert final_diff < 0.2, (
                f"Point ordering invariance test failed for {name} transformation.\n"
                f"Maximum difference: {final_diff:.6f}\n"
                f"This suggests the model may be:\n"
                f"1. Using point indices in its computation\n"
                f"2. Not properly normalizing attention weights\n"
                f"3. Having numerical stability issues in its transformations"
            )


def test_earth_system_loss_point_data():
    """Test the custom loss function with point data"""
    loss_fn = EarthSystemLoss(alpha=0.5, beta=0.3, gamma=0.2)

    batch_size = 1
    num_points = 100
    num_features = 2

    points = torch.zeros(batch_size, num_points, 2)
    points[:, :, 0] = torch.linspace(-180, 180, num_points)
    points[:, :, 1] = torch.linspace(-90, 90, num_points)
    expected_points_shape = (batch_size, num_points, 2)
    assert (
        points.shape == expected_points_shape
    ), f"Expected points shape {expected_points_shape}, got {points.shape}"

    latitude_factor = 1.0 - torch.abs(points[:, :, 1]) / 90.0
    base_temp = 273.15 + latitude_factor * 30.0

    pred = (
        base_temp.unsqueeze(-1).repeat(1, 1, num_features)
        + torch.randn(batch_size, num_points, num_features) * 2
    )
    target = pred + torch.randn(batch_size, num_points, num_features)

    expected_shape = (batch_size, num_points, num_features)
    assert (
        pred.shape == expected_shape
    ), f"Expected predictions shape {expected_shape}, got {pred.shape}"
    assert (
        target.shape == expected_shape
    ), f"Expected target shape {expected_shape}, got {target.shape}"

    loss_dict = loss_fn(pred, target, points)

    assert loss_dict["mse_loss"] >= 0, "MSE loss should be non-negative"
    assert (
        loss_dict["spatial_correlation_loss"] >= 0
    ), "Spatial correlation loss should be non-negative"
    assert loss_dict["physical_loss"] >= 0, "Physical loss should be non-negative"
    assert loss_dict["total_loss"] >= 0, "Total loss should be non-negative"


def test_sparse_dense_distributions(model_config):
    """Test model with varying point densities"""
    model = AuroraModel(**model_config)

    sparse_points = torch.linspace(-180, 180, 5).view(1, 5, 1).repeat(1, 1, 2)
    dense_points = torch.linspace(-180, 180, 45).view(1, 45, 1).repeat(1, 1, 2)  # Reduced from 100
    sparse_features = torch.randn(1, 5, 2)
    dense_features = torch.randn(1, 45, 2)  # Adjusted to match points

    sparse_output = model(sparse_points, sparse_features)
    dense_output = model(dense_points, dense_features)

    assert sparse_output.shape == (1, 5, model_config["output_features"])
    assert dense_output.shape == (1, 45, model_config["output_features"])
    assert not torch.isnan(sparse_output).any()
    assert not torch.isnan(dense_output).any()


def test_pole_handling(model_config):
    """Test model behavior near the poles"""
    model = AuroraModel(**model_config)
    north_pole_points = torch.zeros(1, 50, 2)
    north_pole_points[:, :, 1] = 89.9 + torch.rand(50) * 0.1
    south_pole_points = torch.zeros(1, 50, 2)
    south_pole_points[:, :, 1] = -90 + torch.rand(50) * 0.1
    features = torch.randn(1, 50, 2)
    north_output = model(north_pole_points, features)
    south_output = model(south_pole_points, features)
    assert not torch.isnan(north_output).any()
    assert not torch.isnan(south_output).any()


def test_temporal_sequence(model_config):
    """Test processing of temporal sequences"""
    model = AuroraModel(**model_config)
    model.eval()  # Ensure model is in evaluation mode

    num_points = 45
    num_timesteps = 3

    # Generate consistent points across timesteps
    points = torch.rand(1, num_points, 2)
    points[..., 0] = points[..., 0] * 360 - 180  # longitude
    points[..., 1] = points[..., 1] * 180 - 90  # latitude

    # Generate temporally coherent features with even smaller variations
    base_features = torch.randn(1, num_points, model_config["input_features"])
    features = []
    for t in range(num_timesteps):
        # Further reduce the temporal noise
        temporal_noise = torch.randn_like(base_features) * 0.01  # Reduced from 0.05
        features.append(base_features + temporal_noise * t)

    # Process sequence with gradient tracking disabled
    outputs = []
    with torch.no_grad():  # Disable gradient computation for consistency
        for feat in features:
            output = model(points, feat)
            assert output.shape == (1, num_points, model_config["output_features"])
            assert not torch.isnan(output).any()
            outputs.append(output)

    # Stack outputs and check temporal consistency
    outputs = torch.stack(outputs, dim=1)
    temporal_diff = torch.diff(outputs, dim=1)

    # Check if the maximum difference exceeds the threshold
    max_diff = torch.abs(temporal_diff).max().item()
    max_allowed_diff = 2.0

    assert max_diff < max_allowed_diff, (
        f"Temporal differences exceed {max_allowed_diff}, got {max_diff}. "
        f"This suggests the model is too sensitive to small input changes."
    )


def test_missing_data(model_config):
    """Test handling of missing data points"""
    model = AuroraModel(**model_config)
    points = torch.rand(1, 50, 2) * 360 - 180
    features = torch.randn(1, 50, 2)
    mask = torch.ones(1, 50, dtype=bool)
    mask[:, ::4] = False
    output = model(points, features, mask=mask)
    assert not torch.isnan(output).any()


def test_model_save_load(tmp_path, model_config):
    """Test model serialization and deserialization"""
    model = AuroraModel(**model_config)
    points = torch.rand(1, 50, 2) * 360 - 180
    features = torch.randn(1, 50, 2)
    output_before = model(points, features)
    save_path = tmp_path / "aurora_model.pt"
    torch.save(model.state_dict(), save_path)
    loaded_model = AuroraModel(**model_config)
    loaded_model.load_state_dict(torch.load(save_path))
    output_after = loaded_model(points, features)
    assert torch.allclose(output_before, output_after)


def test_gradient_flow(model_config):
    model = AuroraModel(**model_config)
    model.train()
    points = torch.rand(1, 50, 2) * 360 - 180
    features = torch.randn(1, 50, 2)
    output = model(points, features)
    output.sum().backward()  # Ensure gradients flow
    for param in model.parameters():
        assert param.grad is not None, "Gradient not flowing through the model."
