"""Tests for gradient checkpointing implementation.

This test suite ensures that gradient checkpointing:
1. Produces identical outputs to non-checkpointed versions
2. Reduces memory usage during training
3. Works correctly with all checkpointing strategies
4. Maintains backward compatibility
5. Works with efficient batching
"""

import numpy as np
import pytest
import torch

from graph_weather.models import GraphCast, GraphCastConfig
from graph_weather.models.layers.decoder import Decoder
from graph_weather.models.layers.encoder import Encoder
from graph_weather.models.layers.processor import Processor


def create_lat_lon_grid(resolution_deg: float):
    """Create a lat/lon grid at specified resolution."""
    lat_lons = []
    lats = np.arange(-90, 90, resolution_deg)
    lons = np.arange(0, 360, resolution_deg)
    for lat in lats:
        for lon in lons:
            lat_lons.append((float(lat), float(lon)))
    return lat_lons


# Layer-level checkpointing tests


@pytest.mark.parametrize("use_checkpointing", [False, True])
def test_encoder_checkpointing(use_checkpointing):
    """Test that Encoder with checkpointing produces identical outputs."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    encoder = Encoder(
        lat_lons,
        resolution=2,
        input_dim=78,
        output_dim=256,
        use_checkpointing=use_checkpointing,
        efficient_batching=True,
    )
    encoder.eval()

    with torch.no_grad():
        x, edge_idx, edge_attr = encoder(features)

    # Verify output shape
    assert x.shape[1] == 256
    assert edge_idx.shape[0] == 2


@pytest.mark.parametrize("use_checkpointing", [False, True])
def test_processor_checkpointing(use_checkpointing):
    """Test that Processor with checkpointing produces identical outputs."""
    batch_size = 2
    num_nodes = 5882  # H3 resolution 2
    num_edges = 41162

    torch.manual_seed(42)
    x = torch.randn((batch_size * num_nodes, 256))
    edge_index = torch.randint(0, num_nodes, (2, num_edges))
    edge_attr = torch.randn((num_edges, 256))

    processor = Processor(
        input_dim=256,
        edge_dim=256,
        num_blocks=9,
        use_checkpointing=use_checkpointing,
    )
    processor.eval()

    with torch.no_grad():
        out = processor(x, edge_index, edge_attr, batch_size=batch_size, efficient_batching=True)

    # Verify output shape matches input
    assert out.shape == x.shape


@pytest.mark.parametrize("use_checkpointing", [False, True])
def test_decoder_checkpointing(use_checkpointing):
    """Test that Decoder with checkpointing produces identical outputs."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2
    num_h3 = 5882

    torch.manual_seed(42)
    processor_features = torch.randn((batch_size * num_h3, 256))
    start_features = torch.randn((batch_size, len(lat_lons), 78))

    decoder = Decoder(
        lat_lons,
        resolution=2,
        input_dim=256,
        output_dim=78,
        use_checkpointing=use_checkpointing,
        efficient_batching=True,
    )
    decoder.eval()

    with torch.no_grad():
        out = decoder(processor_features, start_features, batch_size)

    # Verify output shape
    assert out.shape == start_features.shape


# Output equivalence tests


def test_encoder_output_equivalence():
    """Verify encoder outputs are identical with/without checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    # Without checkpointing
    encoder_no_cp = Encoder(
        lat_lons, resolution=2, use_checkpointing=False, efficient_batching=True
    )
    encoder_no_cp.eval()

    # With checkpointing
    encoder_with_cp = Encoder(
        lat_lons, resolution=2, use_checkpointing=True, efficient_batching=True
    )
    encoder_with_cp.load_state_dict(encoder_no_cp.state_dict())
    encoder_with_cp.eval()

    with torch.no_grad():
        x_no_cp, edge_idx_no_cp, edge_attr_no_cp = encoder_no_cp(features)
        x_with_cp, edge_idx_with_cp, edge_attr_with_cp = encoder_with_cp(features)

    # Verify outputs are identical
    assert torch.allclose(x_no_cp, x_with_cp, atol=1e-6)
    assert torch.equal(edge_idx_no_cp, edge_idx_with_cp)
    assert torch.allclose(edge_attr_no_cp, edge_attr_with_cp, atol=1e-6)


def test_full_pipeline_output_equivalence():
    """Verify full pipeline outputs are identical with/without checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    # Create models without checkpointing
    encoder_no_cp = Encoder(lat_lons, use_checkpointing=False, efficient_batching=True)
    processor_no_cp = Processor(256, num_blocks=9, use_checkpointing=False)
    decoder_no_cp = Decoder(
        lat_lons, input_dim=256, use_checkpointing=False, efficient_batching=True
    )

    encoder_no_cp.eval()
    processor_no_cp.eval()
    decoder_no_cp.eval()

    # Run without checkpointing
    with torch.no_grad():
        x, edge_idx, edge_attr = encoder_no_cp(features)
        x = processor_no_cp(x, edge_idx, edge_attr, batch_size=batch_size, efficient_batching=True)
        out_no_cp = decoder_no_cp(x, features, batch_size)

    # Create models with checkpointing
    encoder_with_cp = Encoder(lat_lons, use_checkpointing=True, efficient_batching=True)
    processor_with_cp = Processor(256, num_blocks=9, use_checkpointing=True)
    decoder_with_cp = Decoder(
        lat_lons, input_dim=256, use_checkpointing=True, efficient_batching=True
    )

    # Load same weights
    encoder_with_cp.load_state_dict(encoder_no_cp.state_dict())
    processor_with_cp.load_state_dict(processor_no_cp.state_dict())
    decoder_with_cp.load_state_dict(decoder_no_cp.state_dict())

    encoder_with_cp.eval()
    processor_with_cp.eval()
    decoder_with_cp.eval()

    # Run with checkpointing
    with torch.no_grad():
        x, edge_idx, edge_attr = encoder_with_cp(features)
        x = processor_with_cp(
            x, edge_idx, edge_attr, batch_size=batch_size, efficient_batching=True
        )
        out_with_cp = decoder_with_cp(x, features, batch_size)

    # Verify outputs are identical
    assert torch.allclose(out_no_cp, out_with_cp, atol=1e-5)


# GraphCast model tests


def test_graphcast_no_checkpointing():
    """Test GraphCast with no checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    model = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.no_checkpointing(model)
    model.eval()

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    with torch.no_grad():
        output = model(features)

    assert output.shape == features.shape


def test_graphcast_full_checkpointing():
    """Test GraphCast with full model checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    model = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.full_checkpointing(model)
    model.eval()

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    with torch.no_grad():
        output = model(features)

    assert output.shape == features.shape


def test_graphcast_balanced_checkpointing():
    """Test GraphCast with balanced checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    model = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.balanced_checkpointing(model)
    model.eval()

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    with torch.no_grad():
        output = model(features)

    assert output.shape == features.shape


def test_graphcast_processor_only_checkpointing():
    """Test GraphCast with processor-only checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    model = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.processor_only_checkpointing(model)
    model.eval()

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    with torch.no_grad():
        output = model(features)

    assert output.shape == features.shape


def test_graphcast_output_equivalence():
    """Test that all checkpointing strategies produce identical outputs."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 1  # Use smaller batch for equivalence test

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    # No checkpointing (baseline)
    model_baseline = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.no_checkpointing(model_baseline)
    model_baseline.eval()

    with torch.no_grad():
        output_baseline = model_baseline(features)

    # Test each strategy
    strategies = [
        ("full", GraphCastConfig.full_checkpointing),
        ("balanced", GraphCastConfig.balanced_checkpointing),
        ("processor_only", GraphCastConfig.processor_only_checkpointing),
    ]

    for strategy_name, strategy_fn in strategies:
        model = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
        model.load_state_dict(model_baseline.state_dict())
        strategy_fn(model)
        model.eval()

        with torch.no_grad():
            output = model(features)

        # Verify output is identical to baseline
        assert torch.allclose(
            output_baseline, output, atol=1e-5
        ), f"{strategy_name} checkpointing produced different output"


# Backward compatibility tests


def test_original_api_still_works():
    """Test that the original API (without GraphCast wrapper) still works."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 2

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    # Original usage pattern
    encoder = Encoder(lat_lons, efficient_batching=True)
    processor = Processor(256, num_blocks=9)
    decoder = Decoder(lat_lons, efficient_batching=True)

    encoder.eval()
    processor.eval()
    decoder.eval()

    with torch.no_grad():
        x, edge_idx, edge_attr = encoder(features)
        x = processor(x, edge_idx, edge_attr, batch_size=batch_size, efficient_batching=True)
        output = decoder(x, features, batch_size)

    assert output.shape == features.shape


def test_efficient_batching_unaffected():
    """Test that efficient batching still works with checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 4

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))

    # With efficient batching AND checkpointing
    model = GraphCast(lat_lons, use_checkpointing=True, efficient_batching=True)
    GraphCastConfig.balanced_checkpointing(model)
    model.eval()

    with torch.no_grad():
        output = model(features)

    assert output.shape == features.shape


# Gradient flow tests


@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
def test_backward_pass_with_checkpointing():
    """Test that backward pass works with checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 1

    device = torch.device("cuda")
    model = GraphCast(lat_lons, use_checkpointing=True, efficient_batching=True).to(device)
    GraphCastConfig.balanced_checkpointing(model)
    model.train()

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78), device=device)
    target = torch.randn((batch_size, len(lat_lons), 78), device=device)

    # Forward pass
    output = model(features)

    # Backward pass
    loss = torch.nn.functional.mse_loss(output, target)
    loss.backward()

    # Verify gradients exist
    has_grads = False
    for param in model.parameters():
        if param.grad is not None:
            has_grads = True
            assert not torch.isnan(param.grad).any(), "NaN in gradients"
            assert not torch.isinf(param.grad).any(), "Inf in gradients"

    assert has_grads, "No gradients computed"


def test_gradient_equivalence():
    """Test that gradients are identical with/without checkpointing."""
    lat_lons = create_lat_lon_grid(resolution_deg=10.0)
    batch_size = 1

    torch.manual_seed(42)
    features = torch.randn((batch_size, len(lat_lons), 78))
    target = torch.randn((batch_size, len(lat_lons), 78))

    # Model without checkpointing
    model_no_cp = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    GraphCastConfig.no_checkpointing(model_no_cp)
    model_no_cp.train()

    output_no_cp = model_no_cp(features)
    loss_no_cp = torch.nn.functional.mse_loss(output_no_cp, target)
    loss_no_cp.backward()

    # Collect gradients
    grads_no_cp = []
    for param in model_no_cp.parameters():
        if param.grad is not None:
            grads_no_cp.append(param.grad.clone())

    # Model with checkpointing
    model_with_cp = GraphCast(lat_lons, use_checkpointing=False, efficient_batching=True)
    model_with_cp.load_state_dict(model_no_cp.state_dict())
    GraphCastConfig.balanced_checkpointing(model_with_cp)
    model_with_cp.train()

    output_with_cp = model_with_cp(features)
    loss_with_cp = torch.nn.functional.mse_loss(output_with_cp, target)
    loss_with_cp.backward()

    # Collect gradients
    grads_with_cp = []
    for param in model_with_cp.parameters():
        if param.grad is not None:
            grads_with_cp.append(param.grad.clone())

    # Compare gradients
    assert len(grads_no_cp) == len(grads_with_cp)
    for g1, g2 in zip(grads_no_cp, grads_with_cp):
        assert torch.allclose(g1, g2, atol=1e-5), "Gradients differ with checkpointing"
