"""
Tests for kernel functions, including Matérn 5/2 kernel support.

Tests cover:
1. Kernel function mathematical correctness
2. Symmetry properties (k_ff, k_gg, k_hh should be symmetric)
3. Anti-symmetry of k_fh
4. Derivative relationship tests (h correlates with numerical gradient)
5. Integral relationship tests (g correlates with numerical integral)
6. Numerical stability at various length scales
7. Comparison tests: Matérn 5/2 vs RBF behavior differences
"""

import numpy as np
import pytest
from scipy import integrate

# Import the sampler functions
from gp4c import sample_prior, sample_posterior, SamplingSpec, Observations


class TestMatern52KernelBasics:
    """Basic tests for Matérn 5/2 kernel functionality."""

    def test_sample_f_only(self):
        """Test sampling f with Matérn 5/2 kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, kernel='matern52', n_samples=5, seed=42)

        assert result.f is not None
        assert result.f.shape == (5, 50)
        assert result.g is None
        assert result.h is None

    def test_sample_f_and_h(self):
        """Test sampling f and h with Matérn 5/2 kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_h=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, kernel='matern52', n_samples=5, seed=42)

        assert result.f is not None
        assert result.h is not None
        assert result.f.shape == (5, 50)
        assert result.h.shape == (5, 50)

    def test_sample_all_three(self):
        """Test sampling f, g, h with Matérn 5/2 kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_g=x, x_h=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, kernel='matern52', n_samples=3, seed=42)

        assert result.f is not None
        assert result.g is not None
        assert result.h is not None
        assert result.f.shape == (3, 50)
        assert result.g.shape == (3, 50)
        assert result.h.shape == (3, 50)

    def test_invalid_kernel_raises(self):
        """Test that invalid kernel name raises ValueError."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        with pytest.raises(ValueError, match="Unknown kernel type"):
            sample_prior(spec, kernel='invalid_kernel')


class TestKernelSymmetry:
    """Test symmetry properties of covariance matrices."""

    def test_rbf_covariance_is_symmetric(self):
        """RBF covariance matrix should be symmetric."""
        x = np.linspace(0, 5, 30)
        spec = SamplingSpec(x_f=x, x_g=x, x_h=x)

        # Sample many times and check that covariance of samples is symmetric
        result = sample_prior(spec, ell=0.5, kernel='rbf', n_samples=1000, seed=42)

        # Covariance of f samples
        cov_f = np.cov(result.f.T)
        assert np.allclose(cov_f, cov_f.T, atol=0.1)

    def test_matern52_covariance_is_symmetric(self):
        """Matérn 5/2 covariance matrix should be symmetric."""
        x = np.linspace(0, 5, 30)
        spec = SamplingSpec(x_f=x, x_g=x, x_h=x)

        result = sample_prior(spec, ell=0.5, kernel='matern52', n_samples=1000, seed=42)

        # Covariance of f samples
        cov_f = np.cov(result.f.T)
        assert np.allclose(cov_f, cov_f.T, atol=0.1)


class TestDerivativeRelationship:
    """Test that h = f' relationship holds."""

    def test_rbf_h_approximates_numerical_derivative(self):
        """h should approximate the numerical derivative of f for RBF."""
        x = np.linspace(0.1, 4.9, 100)  # Avoid boundaries for numerical derivative
        spec = SamplingSpec(x_f=x, x_h=x)

        result = sample_prior(spec, ell=0.5, kernel='rbf', n_samples=10, seed=42)

        for i in range(result.f.shape[0]):
            f = result.f[i]
            h = result.h[i]

            # Numerical derivative using central differences
            dx = x[1] - x[0]
            h_numerical = np.gradient(f, dx)

            # Should be correlated (but not identical due to sampling)
            correlation = np.corrcoef(h, h_numerical)[0, 1]
            assert correlation > 0.8, f"Sample {i}: correlation {correlation} too low"

    def test_matern52_h_approximates_numerical_derivative(self):
        """h should approximate the numerical derivative of f for Matérn 5/2."""
        x = np.linspace(0.1, 4.9, 100)
        spec = SamplingSpec(x_f=x, x_h=x)

        result = sample_prior(spec, ell=0.5, kernel='matern52', n_samples=10, seed=42)

        for i in range(result.f.shape[0]):
            f = result.f[i]
            h = result.h[i]

            dx = x[1] - x[0]
            h_numerical = np.gradient(f, dx)

            correlation = np.corrcoef(h, h_numerical)[0, 1]
            assert correlation > 0.8, f"Sample {i}: correlation {correlation} too low"


class TestIntegralRelationship:
    """Test that g = integral(f) relationship holds."""

    def test_rbf_g_approximates_numerical_integral(self):
        """g should approximate the numerical integral of f for RBF."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x, x_g=x)

        result = sample_prior(spec, ell=0.5, kernel='rbf', n_samples=10, seed=42)

        for i in range(result.f.shape[0]):
            f = result.f[i]
            g = result.g[i]

            # Numerical integral using cumulative trapezoid
            g_numerical = integrate.cumulative_trapezoid(f, x, initial=0)

            # Should be correlated
            correlation = np.corrcoef(g, g_numerical)[0, 1]
            assert correlation > 0.9, f"Sample {i}: correlation {correlation} too low"

    def test_matern52_g_approximates_numerical_integral(self):
        """g should approximate the numerical integral of f for Matérn 5/2."""
        x = np.linspace(0, 5, 100)
        spec = SamplingSpec(x_f=x, x_g=x)

        result = sample_prior(spec, ell=0.5, kernel='matern52', n_samples=10, seed=42)

        for i in range(result.f.shape[0]):
            f = result.f[i]
            g = result.g[i]

            g_numerical = integrate.cumulative_trapezoid(f, x, initial=0)

            correlation = np.corrcoef(g, g_numerical)[0, 1]
            assert correlation > 0.9, f"Sample {i}: correlation {correlation} too low"


class TestPosteriorSampling:
    """Test posterior sampling with both kernels."""

    def test_rbf_posterior_passes_through_observations(self):
        """Posterior should pass near observations for RBF."""
        x_train = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-6)

        spec = SamplingSpec(x_f=x_train)
        result = sample_posterior(obs, spec, ell=1.0, kernel='rbf', n_samples=10, seed=42)

        # Posterior mean should be very close to observations at training points
        assert np.allclose(result.f_mean, y_train, atol=0.1)

    def test_matern52_posterior_passes_through_observations(self):
        """Posterior should pass near observations for Matérn 5/2."""
        x_train = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-6)

        spec = SamplingSpec(x_f=x_train)
        result = sample_posterior(obs, spec, ell=1.0, kernel='matern52', n_samples=10, seed=42)

        # Posterior mean should be very close to observations at training points
        assert np.allclose(result.f_mean, y_train, atol=0.1)

    def test_posterior_predicts_derivative(self):
        """Posterior should correctly predict derivative from f observations."""
        x_train = np.array([0.0, 0.5, 1.0, 1.5, 2.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-6)

        x_test = np.linspace(0, 2, 50)
        spec = SamplingSpec(x_h=x_test)

        # Predict derivative using both kernels
        result_rbf = sample_posterior(obs, spec, ell=0.5, kernel='rbf', n_samples=1, seed=42)
        result_m52 = sample_posterior(obs, spec, ell=0.5, kernel='matern52', n_samples=1, seed=42)

        # True derivative is cos(x)
        h_true = np.cos(x_test)

        # Both should correlate with true derivative
        corr_rbf = np.corrcoef(result_rbf.h_mean, h_true)[0, 1]
        corr_m52 = np.corrcoef(result_m52.h_mean, h_true)[0, 1]

        assert corr_rbf > 0.9, f"RBF correlation {corr_rbf} too low"
        assert corr_m52 > 0.9, f"Matérn 5/2 correlation {corr_m52} too low"


class TestNumericalStability:
    """Test numerical stability at various parameter values."""

    @pytest.mark.parametrize("ell", [0.1, 0.5, 1.0, 2.0, 5.0])
    def test_rbf_various_length_scales(self, ell):
        """RBF should work at various length scales."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, ell=ell, kernel='rbf', n_samples=5, seed=42)
        assert result.f is not None
        assert not np.any(np.isnan(result.f))
        assert not np.any(np.isinf(result.f))

    @pytest.mark.parametrize("ell", [0.1, 0.5, 1.0, 2.0, 5.0])
    def test_matern52_various_length_scales(self, ell):
        """Matérn 5/2 should work at various length scales."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, ell=ell, kernel='matern52', n_samples=5, seed=42)
        assert result.f is not None
        assert not np.any(np.isnan(result.f))
        assert not np.any(np.isinf(result.f))

    @pytest.mark.parametrize("sigma2", [0.1, 1.0, 10.0])
    def test_various_variances(self, sigma2):
        """Both kernels should work at various variance levels."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)

        for kernel in ['rbf', 'matern52']:
            result = sample_prior(spec, sigma2=sigma2, ell=0.5, kernel=kernel, n_samples=5, seed=42)
            assert not np.any(np.isnan(result.f))
            assert not np.any(np.isinf(result.f))

            # Variance of samples should scale with sigma2
            empirical_var = result.f.var()
            assert 0.1 * sigma2 < empirical_var < 10 * sigma2


class TestKernelComparison:
    """Compare RBF and Matérn 5/2 kernel behaviors."""

    def test_same_seed_different_kernel_gives_different_samples(self):
        """Same seed but different kernel should give different samples."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)

        result_rbf = sample_prior(spec, ell=0.5, kernel='rbf', n_samples=1, seed=42)
        result_m52 = sample_prior(spec, ell=0.5, kernel='matern52', n_samples=1, seed=42)

        # Samples should be different
        assert not np.allclose(result_rbf.f, result_m52.f)

    def test_matern52_is_rougher_than_rbf(self):
        """Matérn 5/2 samples should be rougher (higher second derivative variance)."""
        x = np.linspace(0, 10, 200)
        spec = SamplingSpec(x_f=x)

        # Use same length scale
        result_rbf = sample_prior(spec, ell=1.0, kernel='rbf', n_samples=100, seed=42)
        result_m52 = sample_prior(spec, ell=1.0, kernel='matern52', n_samples=100, seed=42)

        # Compute roughness as variance of second differences
        def roughness(f):
            d2f = np.diff(f, n=2)
            return np.var(d2f)

        rbf_roughness = np.mean([roughness(result_rbf.f[i]) for i in range(100)])
        m52_roughness = np.mean([roughness(result_m52.f[i]) for i in range(100)])

        # Matérn 5/2 should be rougher than RBF
        assert m52_roughness > rbf_roughness, \
            f"Expected Matérn 5/2 roughness ({m52_roughness}) > RBF roughness ({rbf_roughness})"

    def test_reproducibility_with_seed(self):
        """Same seed should give same samples for each kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)

        for kernel in ['rbf', 'matern52']:
            result1 = sample_prior(spec, ell=0.5, kernel=kernel, n_samples=5, seed=123)
            result2 = sample_prior(spec, ell=0.5, kernel=kernel, n_samples=5, seed=123)
            assert np.allclose(result1.f, result2.f), f"Reproducibility failed for {kernel}"


class TestBackwardCompatibility:
    """Ensure backward compatibility - default kernel is RBF."""

    def test_default_kernel_is_rbf(self):
        """Calling without kernel parameter should use RBF (backward compatible)."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)

        # These should produce the same result
        result_default = sample_prior(spec, ell=0.5, n_samples=5, seed=42)
        result_rbf = sample_prior(spec, ell=0.5, kernel='rbf', n_samples=5, seed=42)

        assert np.allclose(result_default.f, result_rbf.f)

    def test_posterior_default_kernel_is_rbf(self):
        """Posterior sampling without kernel parameter should use RBF."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train)

        x_test = np.linspace(0, 3, 30)
        spec = SamplingSpec(x_f=x_test)

        result_default = sample_posterior(obs, spec, ell=1.0, n_samples=5, seed=42)
        result_rbf = sample_posterior(obs, spec, ell=1.0, kernel='rbf', n_samples=5, seed=42)

        assert np.allclose(result_default.f, result_rbf.f)
        assert np.allclose(result_default.f_mean, result_rbf.f_mean)


class TestMatern32Kernel:
    """Tests for Matérn 3/2 kernel (f and g only, no derivatives)."""

    def test_sample_f_only(self):
        """Test sampling f with Matérn 3/2 kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, kernel='matern32', n_samples=5, seed=42)

        assert result.f is not None
        assert result.f.shape == (5, 50)
        assert result.g is None
        assert result.h is None

    def test_sample_f_and_g(self):
        """Test sampling f and g with Matérn 3/2 kernel."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, kernel='matern32', n_samples=5, seed=42)

        assert result.f is not None
        assert result.g is not None
        assert result.f.shape == (5, 50)
        assert result.g.shape == (5, 50)

    def test_derivative_sampling_works(self):
        """Test that requesting h with Matérn 3/2 works (once mean-square differentiable)."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x, x_h=x)
        result = sample_prior(spec, kernel='matern32', n_samples=3, seed=42)
        assert result.f is not None
        assert result.h is not None
        assert result.f.shape == (3, 10)
        assert result.h.shape == (3, 10)

    def test_posterior_f_from_f(self):
        """Test posterior sampling of f from f observations with Matérn 3/2."""
        x_train = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-6)

        spec = SamplingSpec(x_f=x_train)
        result = sample_posterior(obs, spec, ell=1.0, kernel='matern32', n_samples=10, seed=42)

        # Posterior mean should be close to observations at training points
        assert np.allclose(result.f_mean, y_train, atol=0.1)

    def test_posterior_derivative_prediction_works(self):
        """Test that posterior with h prediction works for Matérn 3/2."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train)

        x_test = np.linspace(0, 3, 10)
        spec = SamplingSpec(x_h=x_test)

        result = sample_posterior(obs, spec, kernel='matern32', n_samples=3, seed=42)
        assert result.h is not None
        assert result.h.shape == (3, 10)

    def test_posterior_h_observation_works(self):
        """Test that posterior with h observations works for Matérn 3/2."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.cos(x_train)  # derivative of sin
        obs = Observations(x_h=x_train, y_h=y_train)

        x_test = np.linspace(0, 3, 10)
        spec = SamplingSpec(x_f=x_test)

        result = sample_posterior(obs, spec, kernel='matern32', n_samples=3, seed=42)
        assert result.f is not None
        assert result.f.shape == (3, 10)

    def test_integral_relationship(self):
        """Test that g approximates the numerical integral of f for Matérn 3/2."""
        x = np.linspace(0, 5, 50)  # Use smaller grid due to K_gg quadrature limitations
        spec = SamplingSpec(x_f=x, x_g=x)

        result = sample_prior(spec, ell=0.5, kernel='matern32', n_samples=10, seed=42)

        for i in range(result.f.shape[0]):
            f = result.f[i]
            g = result.g[i]

            g_numerical = integrate.cumulative_trapezoid(f, x, initial=0)

            correlation = np.corrcoef(g, g_numerical)[0, 1]
            assert correlation > 0.9, f"Sample {i}: correlation {correlation} too low"

    def test_matern32_is_rougher_than_matern52(self):
        """Matérn 3/2 samples should be rougher than Matérn 5/2."""
        x = np.linspace(0, 10, 200)
        spec = SamplingSpec(x_f=x)

        result_m32 = sample_prior(spec, ell=1.0, kernel='matern32', n_samples=100, seed=42)
        result_m52 = sample_prior(spec, ell=1.0, kernel='matern52', n_samples=100, seed=42)

        def roughness(f):
            d2f = np.diff(f, n=2)
            return np.var(d2f)

        m32_roughness = np.mean([roughness(result_m32.f[i]) for i in range(100)])
        m52_roughness = np.mean([roughness(result_m52.f[i]) for i in range(100)])

        # Matérn 3/2 should be rougher than Matérn 5/2
        assert m32_roughness > m52_roughness, \
            f"Expected Matérn 3/2 roughness ({m32_roughness}) > Matérn 5/2 roughness ({m52_roughness})"

    def test_reproducibility_with_seed(self):
        """Same seed should give same samples for Matérn 3/2."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)

        result1 = sample_prior(spec, ell=0.5, kernel='matern32', n_samples=5, seed=123)
        result2 = sample_prior(spec, ell=0.5, kernel='matern32', n_samples=5, seed=123)
        assert np.allclose(result1.f, result2.f)

    @pytest.mark.parametrize("ell", [0.1, 0.5, 1.0, 2.0, 5.0])
    def test_various_length_scales(self, ell):
        """Matérn 3/2 should work at various length scales."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, ell=ell, kernel='matern32', n_samples=5, seed=42)
        assert result.f is not None
        assert not np.any(np.isnan(result.f))
        assert not np.any(np.isinf(result.f))
