"""
Unit tests for GP prior sampler.
"""

import numpy as np
from gp4c import sample_prior, SamplingSpec
import pytest


class TestOutputShapes:
    """Test output shapes for GP prior sampler."""

    def test_output_shapes_correct(self):
        """Test that output shapes are correct."""
        x_f = np.linspace(0, 5, 50)
        x_g = np.linspace(0, 5, 30)
        n_samples = 10

        spec = SamplingSpec(x_f=x_f, x_g=x_g)
        result = sample_prior(spec, n_samples=n_samples, seed=42)

        assert result.f.shape == (n_samples, len(x_f))
        assert result.g.shape == (n_samples, len(x_g))

    def test_single_sample(self):
        """Test that single sample works correctly."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, n_samples=1, seed=42)

        assert result.f.shape == (1, 50)
        assert result.g.shape == (1, 50)

    def test_different_grids(self):
        """Test that different grids for f and g work correctly."""
        x_f = np.linspace(0, 5, 100)
        x_g = np.linspace(0, 3, 60)

        spec = SamplingSpec(x_f=x_f, x_g=x_g)
        result = sample_prior(spec, n_samples=3, seed=42)

        assert result.f.shape == (3, 100)
        assert result.g.shape == (3, 60)


class TestReproducibility:
    """Test reproducibility with seeds."""

    def test_same_seed_same_samples(self):
        """Test that same seed gives same samples."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x, x_g=x)

        result1 = sample_prior(spec, seed=123, n_samples=3)
        result2 = sample_prior(spec, seed=123, n_samples=3)

        np.testing.assert_array_almost_equal(result1.f, result2.f)
        np.testing.assert_array_almost_equal(result1.g, result2.g)

    def test_different_seed_different_samples(self):
        """Test that different seeds give different samples."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x, x_g=x)

        result1 = sample_prior(spec, seed=123, n_samples=3)
        result2 = sample_prior(spec, seed=456, n_samples=3)

        assert not np.allclose(result1.f, result2.f)


class TestIntegralRelationship:
    """Test the integral relationship between f and g."""

    def test_g_approximates_integral_of_f(self):
        """Test that g is approximately integral of f."""
        x = np.linspace(0, 5, 200)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, n_samples=5, seed=99)

        for i in range(5):
            # Numerical integration
            g_numeric = np.cumsum(result.f[i]) * (x[1] - x[0])

            # Should be highly correlated
            corr = np.corrcoef(result.g[i], g_numeric)[0, 1]
            assert corr > 0.95, f"Low correlation: {corr}"


class TestStatisticalProperties:
    """Test statistical properties of samples."""

    def test_zero_mean(self):
        """Test that samples have approximately zero mean."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, n_samples=1000, seed=42)

        f_mean = np.mean(result.f, axis=0)
        g_mean = np.mean(result.g, axis=0)

        # Mean should be close to zero
        assert np.abs(np.mean(f_mean)) < 0.1
        assert np.abs(np.mean(g_mean)) < 0.5  # g can drift

    def test_variance_parameter_affects_output(self):
        """Test that variance parameter affects output variance."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x)

        # Low variance
        result_low = sample_prior(spec, sigma2=0.1, ell=0.5, n_samples=100, seed=42)

        # High variance
        result_high = sample_prior(spec, sigma2=10.0, ell=0.5, n_samples=100, seed=42)

        var_low = np.var(result_low.f)
        var_high = np.var(result_high.f)

        # Higher sigma2 should give higher variance
        assert var_high > var_low * 10


class TestInputValidation:
    """Test error handling for invalid inputs."""

    def test_negative_variance_raises(self):
        """Test that negative variance raises ValueError."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

        with pytest.raises(ValueError):
            sample_prior(spec, sigma2=-1.0)

    def test_zero_length_scale_raises(self):
        """Test that zero length scale raises ValueError."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

        with pytest.raises(ValueError):
            sample_prior(spec, ell=0.0)

    def test_zero_samples_raises(self):
        """Test that zero samples raises ValueError."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

        with pytest.raises(ValueError):
            sample_prior(spec, n_samples=0)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
