"""
Unit tests for the flexible GP sampler with derivative support.
"""

import numpy as np
import pytest

from gp4c import sample_prior, sample_posterior, SamplingSpec, GPSamples, Observations, PosteriorSamples


class TestSamplingSpec:
    """Tests for SamplingSpec validation."""

    def test_requires_at_least_one_array(self):
        """SamplingSpec must have at least one of x_f, x_g, x_h."""
        with pytest.raises(ValueError, match="At least one"):
            SamplingSpec()

    def test_flags_f_only(self):
        """Flags should be 0x1 when only f is specified."""
        spec = SamplingSpec(x_f=np.linspace(0, 1, 10))
        assert spec.flags == 0x1

    def test_flags_g_only(self):
        """Flags should be 0x2 when only g is specified."""
        spec = SamplingSpec(x_g=np.linspace(0, 1, 10))
        assert spec.flags == 0x2

    def test_flags_h_only(self):
        """Flags should be 0x4 when only h is specified."""
        spec = SamplingSpec(x_h=np.linspace(0, 1, 10))
        assert spec.flags == 0x4

    def test_flags_all(self):
        """Flags should be 0x7 when all are specified."""
        x = np.linspace(0, 1, 10)
        spec = SamplingSpec(x_f=x, x_g=x, x_h=x)
        assert spec.flags == 0x7

    def test_arrays_converted_to_float64(self):
        """Arrays should be converted to float64."""
        x_int = np.array([1, 2, 3])
        spec = SamplingSpec(x_f=x_int)
        assert spec.x_f.dtype == np.float64


class TestShapes:
    """Test output shapes for all 7 combinations."""

    @pytest.fixture
    def x(self):
        return np.linspace(0, 5, 50)

    def test_f_only(self, x):
        """Sample only f."""
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_g_only(self, x):
        """Sample only g (integral)."""
        spec = SamplingSpec(x_g=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_h_only(self, x):
        """Sample only h (derivative)."""
        spec = SamplingSpec(x_h=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_f_and_g(self, x):
        """Sample f and g."""
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_f_and_h(self, x):
        """Sample f and h."""
        spec = SamplingSpec(x_f=x, x_h=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_g_and_h(self, x):
        """Sample g and h."""
        spec = SamplingSpec(x_g=x, x_h=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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

    def test_all_three(self, x):
        """Sample f, g, and h."""
        spec = SamplingSpec(x_f=x, x_g=x, x_h=x)
        result = sample_prior(spec, n_samples=3, seed=42)

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


class TestDerivativeRelationship:
    """Test that h ≈ f' (derivative relationship)."""

    def test_derivative_correlation(self):
        """h should correlate with numerical derivative of f."""
        x = np.linspace(0, 5, 200)
        spec = SamplingSpec(x_f=x, x_h=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, n_samples=5, seed=42)

        for i in range(5):
            # Numerical derivative
            h_numeric = np.gradient(result.f[i], x)

            # Should be highly correlated
            corr = np.corrcoef(result.h[i], h_numeric)[0, 1]
            assert corr > 0.9, f"Sample {i}: low derivative correlation: {corr}"

    def test_derivative_larger_ell(self):
        """With larger ell, derivative should be smoother and better correlated."""
        x = np.linspace(0, 5, 200)
        spec = SamplingSpec(x_f=x, x_h=x)

        # Larger length scale = smoother functions
        result = sample_prior(spec, sigma2=1.0, ell=1.0, n_samples=5, seed=42)

        for i in range(5):
            h_numeric = np.gradient(result.f[i], x)
            corr = np.corrcoef(result.h[i], h_numeric)[0, 1]
            assert corr > 0.95, f"Sample {i}: low correlation with ell=1.0: {corr}"


class TestIntegralRelationship:
    """Test that g ≈ ∫f (integral relationship)."""

    def test_integral_correlation(self):
        """g should correlate with numerical integral of f."""
        x = np.linspace(0, 5, 200)
        dx = x[1] - x[0]
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, n_samples=5, seed=42)

        for i in range(5):
            # Numerical integration (cumulative sum)
            g_numeric = np.cumsum(result.f[i]) * dx

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


class TestCrossValidation:
    """Cross-validate: derivative of integral should give back f."""

    def test_derivative_of_integral_gives_f(self):
        """np.gradient(g) should approximate f."""
        x = np.linspace(0.5, 4.5, 150)  # Avoid boundaries
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, sigma2=1.0, ell=0.5, n_samples=5, seed=42)

        for i in range(5):
            # Derivative of integral should give back f
            f_recovered = np.gradient(result.g[i], x)

            # Should be correlated with original f
            corr = np.corrcoef(result.f[i], f_recovered)[0, 1]
            assert corr > 0.85, f"Sample {i}: low f recovery correlation: {corr}"


class TestDifferentGrids:
    """Test sampling on different grids for f, g, h."""

    def test_different_grid_sizes(self):
        """f, g, h can be on different grids."""
        x_f = np.linspace(0, 5, 100)
        x_g = np.linspace(0, 5, 50)  # Coarser
        x_h = np.linspace(0.1, 4.9, 80)  # Avoid boundaries

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

        assert result.f.shape == (3, 100)
        assert result.g.shape == (3, 50)
        assert result.h.shape == (3, 80)

    def test_different_domains(self):
        """Grids can have different domains."""
        x_f = np.linspace(0, 5, 50)
        x_g = np.linspace(0, 3, 30)  # Subset
        x_h = np.linspace(1, 4, 40)  # Different range

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

        assert result.f.shape == (2, 50)
        assert result.g.shape == (2, 30)
        assert result.h.shape == (2, 40)


class TestNumericalStability:
    """Test numerical stability with different parameters."""

    def test_small_length_scale(self):
        """Small length scale should still work (with higher jitter)."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_h=x)

        # Small ell makes derivatives high-variance
        result = sample_prior(spec, sigma2=1.0, ell=0.2, n_samples=3, seed=42)

        assert result.f.shape == (3, 50)
        assert result.h.shape == (3, 50)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.h))

    def test_various_length_scales(self):
        """Test different length scales."""
        x = np.linspace(0, 5, 50)

        for ell in [0.1, 0.5, 1.0, 2.0]:
            spec = SamplingSpec(x_f=x, x_g=x, x_h=x)
            result = sample_prior(spec, ell=ell, n_samples=2, seed=42)

            assert np.all(np.isfinite(result.f)), f"NaN in f for ell={ell}"
            assert np.all(np.isfinite(result.g)), f"NaN in g for ell={ell}"
            assert np.all(np.isfinite(result.h)), f"NaN in h for ell={ell}"


class TestReproducibility:
    """Test that same seed gives same samples."""

    def test_same_seed_same_samples(self):
        """Same seed should give same samples."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_g=x, x_h=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)
        np.testing.assert_array_almost_equal(result1.h, result2.h)

    def test_different_seed_different_samples(self):
        """Different seeds should give different samples."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=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 TestInputValidation:
    """Test input validation and error handling."""

    def test_invalid_spec_type(self):
        """Should raise TypeError for non-SamplingSpec."""
        with pytest.raises(TypeError):
            sample_prior("not a spec")

    def test_negative_sigma2(self):
        """Should raise ValueError for negative sigma2."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

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

    def test_zero_ell(self):
        """Should raise ValueError for zero ell."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

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

    def test_zero_samples(self):
        """Should raise ValueError for zero n_samples."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)

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


class TestGPSamplesNamedTuple:
    """Test GPSamples named tuple behavior."""

    def test_is_named_tuple(self):
        """Result should be a GPSamples named tuple."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x)
        result = sample_prior(spec)

        assert isinstance(result, GPSamples)
        assert hasattr(result, "f")
        assert hasattr(result, "g")
        assert hasattr(result, "h")

    def test_unpacking(self):
        """Should support unpacking."""
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec)

        f, g, h, u = result
        assert f is not None
        assert g is not None
        assert h is None
        assert u is None


# =============================================================================
# Posterior Sampling Tests
# =============================================================================


class TestObservations:
    """Tests for Observations validation."""

    def test_requires_at_least_one_observation(self):
        """Observations must have at least one of f, g, h observations."""
        with pytest.raises(ValueError, match="At least one"):
            Observations()

    def test_requires_paired_x_y(self):
        """x and y must both be provided for each observation type."""
        with pytest.raises(ValueError, match="must both be provided"):
            Observations(x_f=np.array([1.0, 2.0]))

        with pytest.raises(ValueError, match="must both be provided"):
            Observations(y_f=np.array([1.0, 2.0]))

    def test_requires_same_length(self):
        """x and y must have the same length."""
        with pytest.raises(ValueError, match="must have same length"):
            Observations(x_f=np.array([1.0, 2.0]), y_f=np.array([1.0]))

    def test_valid_f_observations(self):
        """Valid f observations should be accepted."""
        obs = Observations(x_f=np.array([1.0, 2.0]), y_f=np.array([0.5, 0.8]))
        assert len(obs.x_f) == 2
        assert len(obs.y_f) == 2

    def test_noise_defaults(self):
        """Default noise should be 1e-6."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        assert obs.noise_f == 1e-6
        assert obs.noise_g == 1e-6
        assert obs.noise_h == 1e-6

    def test_arrays_converted_to_float64(self):
        """Arrays should be converted to float64."""
        obs = Observations(x_f=np.array([1, 2, 3]), y_f=np.array([0, 1, 2]))
        assert obs.x_f.dtype == np.float64
        assert obs.y_f.dtype == np.float64


class TestPosteriorShapes:
    """Test output shapes for posterior sampling."""

    def test_posterior_f_from_f(self):
        """Train on f, predict f."""
        x_train = np.array([0.0, 1.0, 2.0, 3.0])
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train)

        x_test = np.linspace(0, 4, 50)
        spec = SamplingSpec(x_f=x_test)

        result = sample_posterior(obs, spec, ell=1.0, n_samples=5)

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

    def test_posterior_h_from_f(self):
        """Train on f, predict derivative h."""
        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)

        x_test = np.linspace(0, 4, 30)
        spec = SamplingSpec(x_h=x_test)

        result = sample_posterior(obs, spec, ell=1.0, n_samples=3)

        assert result.h.shape == (3, 30)
        assert result.h_mean.shape == (30,)
        assert result.f is None
        assert result.g is None

    def test_posterior_f_from_h(self):
        """Train on derivative h, predict f."""
        x_train = np.array([0.5, 1.5, 2.5, 3.5])
        y_train = np.cos(x_train)  # Derivative of sin(x)
        obs = Observations(x_h=x_train, y_h=y_train)

        x_test = np.linspace(0, 4, 40)
        spec = SamplingSpec(x_f=x_test)

        result = sample_posterior(obs, spec, ell=1.0, n_samples=4)

        assert result.f.shape == (4, 40)
        assert result.f_mean.shape == (40,)

    def test_posterior_all_from_f(self):
        """Train on f, predict f, g, and h."""
        x_train = np.linspace(0, 4, 10)
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train)

        x_test = np.linspace(0, 4, 20)
        spec = SamplingSpec(x_f=x_test, x_g=x_test, x_h=x_test)

        result = sample_posterior(obs, spec, ell=1.0, n_samples=2)

        assert result.f.shape == (2, 20)
        assert result.g.shape == (2, 20)
        assert result.h.shape == (2, 20)
        assert result.f_mean.shape == (20,)
        assert result.g_mean.shape == (20,)
        assert result.h_mean.shape == (20,)


class TestPosteriorInterpolation:
    """Test that posterior interpolates training data."""

    def test_posterior_mean_passes_through_training_points(self):
        """Posterior mean should pass through (or near) training points."""
        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-8)

        # Predict at training points
        spec = SamplingSpec(x_f=x_train)
        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=1)

        # Mean should match training data closely (small noise)
        np.testing.assert_allclose(result.f_mean, y_train, atol=0.01)

    def test_posterior_uncertainty_low_at_training_points(self):
        """Samples should have low variance at training points."""
        x_train = np.array([1.0, 2.0, 3.0])
        y_train = np.array([0.5, 1.0, 0.8])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-8)

        spec = SamplingSpec(x_f=x_train)
        result = sample_posterior(
            obs, spec, sigma2=1.0, ell=1.0, n_samples=50, seed=123
        )

        # Variance should be very low at training points
        sample_std = np.std(result.f, axis=0)
        assert np.all(sample_std < 0.1)


class TestPosteriorDerivativeConsistency:
    """Test consistency between posterior f and h."""

    def test_posterior_h_matches_numerical_derivative_of_f_mean(self):
        """Posterior h_mean should match numerical derivative of f_mean."""
        # Train on f
        x_train = np.linspace(0, 4, 15)
        y_train = np.sin(x_train)
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=1e-6)

        # Predict both f and h at same dense grid
        x_test = np.linspace(0.5, 3.5, 50)  # Avoid boundaries
        spec = SamplingSpec(x_f=x_test, x_h=x_test)

        result = sample_posterior(obs, spec, sigma2=1.0, ell=0.7, n_samples=1)

        # Numerical derivative of f_mean
        dx = x_test[1] - x_test[0]
        f_deriv_numerical = np.gradient(result.f_mean, dx)

        # Compare with h_mean (skip edges where gradient is less accurate)
        corr = np.corrcoef(result.h_mean[5:-5], f_deriv_numerical[5:-5])[0, 1]
        assert corr > 0.9, f"Correlation {corr} too low"


class TestPosteriorMixedObservations:
    """Test posterior with mixed observation types."""

    def test_observe_f_and_h(self):
        """Train on both f and h observations."""
        # Observe f at some points
        x_f_train = np.array([0.0, 2.0, 4.0])
        y_f_train = np.sin(x_f_train)

        # Observe h at other points
        x_h_train = np.array([1.0, 3.0])
        y_h_train = np.cos(x_h_train)

        obs = Observations(x_f=x_f_train, y_f=y_f_train, x_h=x_h_train, y_h=y_h_train)

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

        result = sample_posterior(obs, spec, ell=1.0, n_samples=3)

        assert result.f.shape == (3, 30)
        assert result.f_mean.shape == (30,)


class TestPosteriorInputValidation:
    """Test input validation for posterior sampling."""

    def test_obs_must_be_observations(self):
        """obs must be an Observations instance."""
        spec = SamplingSpec(x_f=np.linspace(0, 1, 10))
        with pytest.raises(TypeError, match="Observations"):
            sample_posterior("not observations", spec)

    def test_spec_must_be_samplingspec(self):
        """spec must be a SamplingSpec instance."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        with pytest.raises(TypeError, match="SamplingSpec"):
            sample_posterior(obs, "not spec")

    def test_sigma2_must_be_positive(self):
        """sigma2 must be positive."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        spec = SamplingSpec(x_f=np.linspace(0, 2, 10))
        with pytest.raises(ValueError, match="sigma2"):
            sample_posterior(obs, spec, sigma2=-1.0)

    def test_ell_must_be_positive(self):
        """ell must be positive."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        spec = SamplingSpec(x_f=np.linspace(0, 2, 10))
        with pytest.raises(ValueError, match="ell"):
            sample_posterior(obs, spec, ell=0.0)

    def test_n_samples_must_be_positive(self):
        """n_samples must be positive."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        spec = SamplingSpec(x_f=np.linspace(0, 2, 10))
        with pytest.raises(ValueError, match="n_samples"):
            sample_posterior(obs, spec, n_samples=0)


class TestPosteriorReproducibility:
    """Test that posterior samples are reproducible."""

    def test_same_seed_same_samples(self):
        """Same seed should produce same samples."""
        obs = Observations(x_f=np.array([1.0, 2.0]), y_f=np.array([0.5, 0.8]))
        spec = SamplingSpec(x_f=np.linspace(0, 3, 20))

        result1 = sample_posterior(obs, spec, n_samples=5, seed=12345)
        result2 = sample_posterior(obs, spec, n_samples=5, seed=12345)

        np.testing.assert_array_equal(result1.f, result2.f)
        np.testing.assert_array_equal(result1.f_mean, result2.f_mean)

    def test_different_seed_different_samples(self):
        """Different seeds should produce different samples."""
        obs = Observations(x_f=np.array([1.0, 2.0]), y_f=np.array([0.5, 0.8]))
        spec = SamplingSpec(x_f=np.linspace(0, 3, 20))

        result1 = sample_posterior(obs, spec, n_samples=5, seed=100)
        result2 = sample_posterior(obs, spec, n_samples=5, seed=200)

        # Means should be the same
        np.testing.assert_array_equal(result1.f_mean, result2.f_mean)
        # Samples should differ
        assert not np.allclose(result1.f, result2.f)


class TestPosteriorSamplesNamedTuple:
    """Test PosteriorSamples named tuple."""

    def test_is_named_tuple(self):
        """PosteriorSamples should be a NamedTuple."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        spec = SamplingSpec(x_f=np.array([0.5, 1.0, 1.5]))
        result = sample_posterior(obs, spec, n_samples=1)

        assert isinstance(result, PosteriorSamples)
        assert hasattr(result, "f")
        assert hasattr(result, "g")
        assert hasattr(result, "h")
        assert hasattr(result, "f_mean")
        assert hasattr(result, "g_mean")
        assert hasattr(result, "h_mean")

    def test_can_unpack(self):
        """PosteriorSamples should be unpackable."""
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.5]))
        spec = SamplingSpec(x_f=np.array([0.5, 1.0, 1.5]))
        f, g, h, u, f_mean, g_mean, h_mean, u_mean, f_std, g_std, h_std, u_std = sample_posterior(obs, spec, n_samples=1)

        assert f is not None
        assert g is None
        assert h is None
        assert f_mean is not None


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