# MIT License
#
# Copyright (c) 2025 Quandela
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Robustness and integration tests for Merlin.
"""

import torch
import torch.nn as nn

import merlin as ML


class TestRobustness:
    """Test suite for robustness and edge cases."""

    def test_large_batch_sizes(self):
        """Test handling of large batch sizes."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model = torch.nn.Sequential(layer, torch.nn.Linear(layer.output_size, 3))

        # Test with large batch
        large_batch_size = 1000
        x = torch.rand(large_batch_size, 2)

        output = model(x)

        assert output.shape == (large_batch_size, 3)
        assert torch.all(torch.isfinite(output))

    def test_extreme_input_values(self):
        """Test handling of extreme input values."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model = torch.nn.Sequential(layer, torch.nn.Linear(layer.output_size, 3))

        # Test boundary values
        boundary_inputs = torch.tensor([
            [0.0, 0.0],  # All zeros
            [1.0, 1.0],  # All ones
            [0.0, 1.0],  # Mixed
            [1.0, 0.0],  # Mixed reverse
        ])

        output = model(boundary_inputs)

        assert output.shape == (4, 3)
        assert torch.all(torch.isfinite(output))

    def test_numerical_stability(self):
        """Test numerical stability with repeated computations."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model = torch.nn.Sequential(layer, torch.nn.Linear(layer.output_size, 3))

        x = torch.rand(5, 2)

        # Run multiple times - should get identical results
        outputs = []
        for _ in range(10):
            with torch.no_grad():
                output = model(x)
                outputs.append(output)

        # All outputs should be identical (deterministic)
        for i in range(1, len(outputs)):
            assert torch.allclose(outputs[0], outputs[i], atol=1e-6)

    def test_gradient_accumulation(self):
        """Test gradient accumulation over multiple batches."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model = torch.nn.Sequential(layer, torch.nn.Linear(layer.output_size, 3))

        # Accumulate gradients over multiple batches
        total_loss = 0
        for _ in range(3):
            x = torch.rand(4, 2, requires_grad=True)
            output = model(x)
            loss = output.sum()
            loss.backward()
            total_loss += loss.item()

        # Check that gradients accumulated
        param_count = 0
        for param in layer.parameters():
            if param.requires_grad and param.grad is not None:
                param_count += 1
                assert not torch.allclose(param.grad, torch.zeros_like(param.grad))

        assert param_count > 0, "No parameters have gradients"

    def test_device_compatibility(self):
        """Test CPU compatibility (GPU testing would require CUDA)."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer_cpu = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
            device=torch.device("cpu"),
        )
        model_cpu = torch.nn.Sequential(
            layer_cpu,
            torch.nn.Linear(layer_cpu.output_size, 3, device=torch.device("cpu")),
        )

        x_cpu = torch.rand(3, 2, device="cpu")
        output_cpu = model_cpu(x_cpu)

        assert output_cpu.device.type == "cpu"
        assert output_cpu.shape == (3, 3)

    def test_different_dtypes(self):
        """Test different data types."""

        # Test float32
        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer_f32 = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
            dtype=torch.float32,
        )
        model_f32 = torch.nn.Sequential(
            layer_f32, torch.nn.Linear(layer_f32.output_size, 3, dtype=torch.float32)
        )
        x_f32 = torch.rand(2, 2, dtype=torch.float32)
        output_f32 = model_f32(x_f32)
        assert output_f32.dtype == torch.float32

        # Test float64
        layer_f64 = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
            dtype=torch.float64,
        )
        model_f64 = torch.nn.Sequential(
            layer_f64, torch.nn.Linear(layer_f64.output_size, 3, dtype=torch.float64)
        )
        x_f64 = torch.rand(2, 2, dtype=torch.float64)
        output_f64 = model_f64(x_f64)

        # The output dtype might be influenced by the underlying quantum simulation
        # So we'll be more flexible and just check that it's a valid float type
        assert output_f64.dtype in [
            torch.float32,
            torch.float64,
        ], f"Expected float type, got {output_f64.dtype}"

        # More importantly, check that the computation works correctly
        assert torch.all(torch.isfinite(output_f64))
        assert output_f64.shape == (2, 3)

    def test_parameter_initialization_consistency(self):
        """Test that parameter initialization is consistent."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        # Create multiple layers with same random seed
        torch.manual_seed(42)
        layer1 = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model1 = torch.nn.Sequential(layer1, torch.nn.Linear(layer1.output_size, 3))

        torch.manual_seed(42)
        layer2 = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
            measurement_strategy=ML.MeasurementStrategy.probs(),
        )

        model2 = torch.nn.Sequential(layer2, torch.nn.Linear(layer2.output_size, 3))

        # Parameters should be identical
        assert len(list(model1.parameters())) == len(list(model2.parameters())), (
            "Mismatch in number of parameters between model1 and model2"
        )
        for p1, p2 in zip(model1.parameters(), model2.parameters(), strict=True):
            assert torch.allclose(p1, p2, atol=1e-6)

    def test_memory_efficiency(self):
        """Test memory usage doesn't grow unexpectedly."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input", subset_combinations=True)
        builder.add_entangling_layer(trainable=True, name="U2")

        layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
        )

        model = torch.nn.Sequential(layer, torch.nn.Linear(layer.output_size, 3))

        # Run many forward passes
        for _ in range(100):
            x = torch.rand(10, 2)
            with torch.no_grad():
                output = model(x)
                del output, x  # Explicit cleanup

        # Should complete without memory issues


class TestIntegrationScenarios:
    """Integration tests for realistic usage scenarios."""

    def test_training_loop_simulation(self):
        """Simulate a realistic training loop."""
        # Create a simple dataset
        n_samples = 200
        X = torch.rand(n_samples, 3)
        y = (X.sum(dim=1) > 1.5).long()  # Simple binary classification

        # Create model
        class SimpleQuantumModel(nn.Module):
            def __init__(self):
                super().__init__()
                builder = ML.CircuitBuilder(n_modes=4)
                builder.add_entangling_layer(trainable=True, name="U1")
                builder.add_angle_encoding(
                    modes=[0, 1, 2], name="input", subset_combinations=True
                )
                builder.add_entangling_layer(trainable=True, name="U2")

                self.quantum = ML.QuantumLayer(
                    input_size=3,
                    input_state=[1, 0, 1, 0],
                    builder=builder,
                )
                self.classifier = nn.Linear(self.quantum.output_size, 2)

            def forward(self, x):
                x = torch.sigmoid(x)  # Normalize for quantum layer
                x = self.quantum(x)
                return self.classifier(x)

        model = SimpleQuantumModel()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        criterion = nn.CrossEntropyLoss()

        # Training loop
        model.train()
        initial_loss = None
        final_loss = None

        for epoch in range(10):
            epoch_loss = 0
            for i in range(0, len(X), 32):  # Batch size 32
                batch_X = X[i : i + 32]
                batch_y = y[i : i + 32]

                optimizer.zero_grad()
                outputs = model(batch_X)
                loss = criterion(outputs, batch_y)
                loss.backward()
                optimizer.step()

                epoch_loss += loss.item()

            if epoch == 0:
                initial_loss = epoch_loss
            elif epoch == 9:
                final_loss = epoch_loss

        # Loss should decrease (learning is happening)
        assert final_loss < initial_loss, "Model should learn and reduce loss"

    def test_hybrid_architecture(self):
        """Test complex hybrid classical-quantum architecture."""

        class ComplexHybridModel(nn.Module):
            def __init__(self):
                super().__init__()

                # Classical preprocessing
                self.pre_classical = nn.Sequential(
                    nn.Linear(8, 6), nn.ReLU(), nn.Linear(6, 3)
                )

                # First quantum layer
                builder = ML.CircuitBuilder(n_modes=4)
                builder.add_entangling_layer(trainable=True, name="U1")
                builder.add_angle_encoding(
                    modes=[0, 1, 2], name="input", subset_combinations=True
                )
                builder.add_entangling_layer(trainable=True, name="U2")

                self.quantum1 = ML.QuantumLayer(
                    input_size=3,
                    input_state=[1, 0, 1, 0],
                    builder=builder,
                )

                # Middle classical processing
                self.mid_classical = nn.Sequential(
                    nn.Linear(self.quantum1.output_size, 5), nn.ReLU()
                )

                # Second quantum layer (reservoir)
                builder = ML.CircuitBuilder(n_modes=5)
                builder.add_entangling_layer(trainable=True, name="U1")
                builder.add_angle_encoding(
                    modes=[0, 1, 2, 3], name="input", subset_combinations=True
                )
                builder.add_entangling_layer(trainable=True, name="U2")
                self.quantum2 = ML.QuantumLayer(
                    input_size=4,
                    input_state=[1, 0, 1, 0, 0],
                    builder=builder,
                )
                self.quantum2.requires_grad_(False)  # Freeze reservoir layer
                # Final classical layer
                self.final_classical = nn.Linear(self.quantum2.output_size, 2)

            def forward(self, x):
                x = self.pre_classical(x)
                x = torch.sigmoid(x)  # Normalize for quantum
                x = self.quantum1(x)
                x = self.mid_classical(x)
                x = torch.sigmoid(x)  # Normalize for quantum
                x = self.quantum2(x)
                x = self.final_classical(x)
                return x

        model = ComplexHybridModel()

        # Test forward pass
        x = torch.rand(16, 8)
        output = model(x)

        assert output.shape == (16, 2)
        assert torch.all(torch.isfinite(output))

        # Test backward pass
        loss = output.sum()
        loss.backward()

        # Check gradients exist for trainable parameters
        trainable_params = 0
        for _name, param in model.named_parameters():
            if param.requires_grad and param.grad is not None:
                trainable_params += 1

        assert trainable_params > 0, "Should have trainable parameters with gradients"

    def test_ensemble_quantum_models(self):
        """Test ensemble of quantum models."""

        class QuantumEnsemble(nn.Module):
            def __init__(self, n_models=3):
                super().__init__()

                self.models = nn.ModuleList()

                for _i in range(n_models):
                    builder = ML.CircuitBuilder(n_modes=4)
                    builder.add_entangling_layer(trainable=True, name="U1")
                    builder.add_angle_encoding(modes=[0, 1], name="input")
                    builder.add_entangling_layer(trainable=True, name="U2")
                    layer = ML.QuantumLayer(
                        input_size=2,
                        input_state=[1, 0, 1, 0],
                        builder=builder,
                    )

                    model = torch.nn.Sequential(
                        layer, torch.nn.Linear(layer.output_size, 3)
                    )
                    self.models.append(model)

            def forward(self, x):
                outputs = []
                for model in self.models:
                    normalized_x = torch.sigmoid(x)
                    output = model(normalized_x)
                    outputs.append(output)

                # Average ensemble predictions
                return torch.stack(outputs).mean(dim=0)

        ensemble = QuantumEnsemble(n_models=3)

        x = torch.rand(5, 2)
        output = ensemble(x)

        assert output.shape == (5, 3)
        assert torch.all(torch.isfinite(output))

        # Test that individual models produce different outputs
        individual_outputs = []
        normalized_x = torch.sigmoid(x)
        for model in ensemble.models:
            with torch.no_grad():
                individual_output = model(normalized_x)
                individual_outputs.append(individual_output)

        # Outputs should be different (different random initializations)
        for i in range(len(individual_outputs)):
            for j in range(i + 1, len(individual_outputs)):
                assert not torch.allclose(
                    individual_outputs[i], individual_outputs[j], atol=1e-3
                ), f"Models {i} and {j} produced identical outputs"

    def test_saving_and_loading(self):
        """Test model saving and loading."""

        builder = ML.CircuitBuilder(n_modes=4)
        builder.add_entangling_layer(trainable=True, name="U1")
        builder.add_angle_encoding(modes=[0, 1], name="input")
        builder.add_entangling_layer(trainable=True, name="U2")
        original_layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
        )

        model = torch.nn.Sequential(
            original_layer, torch.nn.Linear(original_layer.output_size, 3)
        )

        x = torch.rand(3, 2)
        original_output = model(x)

        # Save model state
        state_dict = model.state_dict()

        # Create new model and load state
        new_layer = ML.QuantumLayer(
            input_size=2,
            input_state=[1, 0, 1, 0],
            builder=builder,
        )

        new_model = torch.nn.Sequential(
            new_layer, torch.nn.Linear(new_layer.output_size, 3)
        )
        new_model.load_state_dict(state_dict)

        # Test that outputs are identical
        new_output = new_model(x)
        assert torch.allclose(original_output, new_output, atol=1e-6)
