"""
Unit tests for correlated measurement noise in GP posterior sampling.

Tests the three noise types supported by the Observations class:
1. Scalar noise: single variance applied to all observations
2. Diagonal noise: per-observation variances (heteroscedastic)
3. Full covariance noise: correlated measurement noise
"""

import numpy as np
import pytest

from gp4c import sample_posterior, Observations, SamplingSpec, PosteriorSamples


# =============================================================================
# Noise Validation Tests (types.py Observations class)
# =============================================================================


class TestScalarNoiseValidation:
    """Test validation of scalar noise parameters."""

    def test_positive_scalar_noise_f(self):
        """Scalar noise_f must be positive."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])

        # Valid: positive scalar
        obs = Observations(x_f=x, y_f=y, noise_f=0.01)
        assert obs.noise_f == 0.01

    def test_negative_scalar_noise_f_raises(self):
        """Negative scalar noise_f should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        with pytest.raises(ValueError, match="noise_f must be positive"):
            Observations(x_f=x, y_f=y, noise_f=-0.01)

    def test_zero_scalar_noise_f_raises(self):
        """Zero scalar noise_f should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        with pytest.raises(ValueError, match="noise_f must be positive"):
            Observations(x_f=x, y_f=y, noise_f=0.0)

    def test_positive_scalar_noise_g(self):
        """Scalar noise_g must be positive."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        obs = Observations(x_g=x, y_g=y, noise_g=0.02)
        assert obs.noise_g == 0.02

    def test_negative_scalar_noise_g_raises(self):
        """Negative scalar noise_g should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        with pytest.raises(ValueError, match="noise_g must be positive"):
            Observations(x_g=x, y_g=y, noise_g=-0.01)

    def test_positive_scalar_noise_h(self):
        """Scalar noise_h must be positive."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        obs = Observations(x_h=x, y_h=y, noise_h=0.015)
        assert obs.noise_h == 0.015

    def test_negative_scalar_noise_h_raises(self):
        """Negative scalar noise_h should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        with pytest.raises(ValueError, match="noise_h must be positive"):
            Observations(x_h=x, y_h=y, noise_h=-0.01)


class TestDiagonalNoiseValidation:
    """Test validation of diagonal (per-observation) noise."""

    def test_valid_diagonal_noise_f(self):
        """Valid diagonal noise_f (1D array) should be accepted."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        noise = np.array([0.01, 0.02, 0.015])

        obs = Observations(x_f=x, y_f=y, noise_f=noise)
        np.testing.assert_array_equal(obs.noise_f, noise)

    def test_diagonal_noise_f_wrong_length_raises(self):
        """Diagonal noise_f with wrong length should raise ValueError."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        noise = np.array([0.01, 0.02])  # Length 2, should be 3

        with pytest.raises(ValueError, match="noise_f diagonal must have length 3"):
            Observations(x_f=x, y_f=y, noise_f=noise)

    def test_diagonal_noise_f_negative_element_raises(self):
        """Diagonal noise_f with negative element should raise ValueError."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        noise = np.array([0.01, -0.02, 0.015])

        with pytest.raises(ValueError, match="noise_f diagonal elements must be positive"):
            Observations(x_f=x, y_f=y, noise_f=noise)

    def test_diagonal_noise_f_zero_element_raises(self):
        """Diagonal noise_f with zero element should raise ValueError."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        noise = np.array([0.01, 0.0, 0.015])

        with pytest.raises(ValueError, match="noise_f diagonal elements must be positive"):
            Observations(x_f=x, y_f=y, noise_f=noise)

    def test_valid_diagonal_noise_g(self):
        """Valid diagonal noise_g should be accepted."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        noise = np.array([0.01, 0.02])

        obs = Observations(x_g=x, y_g=y, noise_g=noise)
        np.testing.assert_array_equal(obs.noise_g, noise)

    def test_valid_diagonal_noise_h(self):
        """Valid diagonal noise_h should be accepted."""
        x = np.array([0.0, 1.0, 2.0, 3.0])
        y = np.array([0.5, 0.8, 0.3, 0.1])
        noise = np.array([0.01, 0.02, 0.015, 0.012])

        obs = Observations(x_h=x, y_h=y, noise_h=noise)
        np.testing.assert_array_equal(obs.noise_h, noise)


class TestFullCovarianceNoiseValidation:
    """Test validation of full covariance noise matrices."""

    def test_valid_full_covariance_f(self):
        """Valid full covariance noise_f should be accepted."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        # Symmetric positive definite 2x2 matrix
        noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])

        obs = Observations(x_f=x, y_f=y, noise_f=noise_cov)
        np.testing.assert_array_equal(obs.noise_f, noise_cov)

    def test_full_covariance_f_wrong_shape_raises(self):
        """Full covariance noise_f with wrong shape should raise ValueError."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        # 2x2 matrix but need 3x3
        noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])

        with pytest.raises(ValueError, match="noise_f covariance must be \\(3, 3\\)"):
            Observations(x_f=x, y_f=y, noise_f=noise_cov)

    def test_full_covariance_f_asymmetric_raises(self):
        """Asymmetric full covariance noise_f should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        # Asymmetric matrix
        noise_cov = np.array([[0.01, 0.002], [0.003, 0.015]])

        with pytest.raises(ValueError, match="noise_f covariance must be symmetric"):
            Observations(x_f=x, y_f=y, noise_f=noise_cov)

    def test_full_covariance_f_not_positive_definite_raises(self):
        """Non-positive-definite full covariance noise_f should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        # Symmetric but not positive definite (negative eigenvalue)
        noise_cov = np.array([[0.01, 0.02], [0.02, 0.01]])

        with pytest.raises(ValueError, match="noise_f covariance must be positive definite"):
            Observations(x_f=x, y_f=y, noise_f=noise_cov)

    def test_valid_full_covariance_g(self):
        """Valid full covariance noise_g should be accepted."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.5, 0.8, 0.3])
        noise_cov = np.array([
            [0.01, 0.002, 0.001],
            [0.002, 0.015, 0.002],
            [0.001, 0.002, 0.012]
        ])

        obs = Observations(x_g=x, y_g=y, noise_g=noise_cov)
        np.testing.assert_array_equal(obs.noise_g, noise_cov)

    def test_valid_full_covariance_h(self):
        """Valid full covariance noise_h should be accepted."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])

        obs = Observations(x_h=x, y_h=y, noise_h=noise_cov)
        np.testing.assert_array_equal(obs.noise_h, noise_cov)


class TestNoiseTypeValidation:
    """Test validation of noise parameter types."""

    def test_noise_must_be_scalar_or_array(self):
        """Noise must be float or numpy array."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        # List should raise (must be converted to array first)
        with pytest.raises(ValueError, match="noise_f must be float or numpy array"):
            Observations(x_f=x, y_f=y, noise_f=[0.01, 0.02])

    def test_noise_3d_array_raises(self):
        """3D noise array should raise ValueError."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])
        noise_3d = np.ones((2, 2, 2))

        with pytest.raises(ValueError, match="noise_f must be scalar, 1D, or 2D array"):
            Observations(x_f=x, y_f=y, noise_f=noise_3d)


class TestNoiseWithNoObservations:
    """Test that noise validation is skipped when no observations exist."""

    def test_noise_ignored_when_no_f_observations(self):
        """noise_f is ignored when x_f and y_f are None."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        # noise_f=-1.0 would normally raise, but should be ignored
        obs = Observations(x_g=x, y_g=y, noise_f=-1.0)
        assert obs.noise_f == -1.0  # Not validated

    def test_noise_ignored_when_no_g_observations(self):
        """noise_g is ignored when x_g and y_g are None."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        obs = Observations(x_f=x, y_f=y, noise_g=-1.0)
        assert obs.noise_g == -1.0  # Not validated

    def test_noise_ignored_when_no_h_observations(self):
        """noise_h is ignored when x_h and y_h are None."""
        x = np.array([0.0, 1.0])
        y = np.array([0.5, 0.8])

        obs = Observations(x_f=x, y_f=y, noise_h=-1.0)
        assert obs.noise_h == -1.0  # Not validated


# =============================================================================
# Functional Tests for sample_posterior with Different Noise Types
# =============================================================================


class TestScalarNoisePosterior:
    """Test posterior sampling with scalar noise (baseline behavior)."""

    def test_scalar_noise_produces_valid_samples(self):
        """Scalar noise should produce valid posterior samples."""
        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=0.01)

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

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=10, seed=42)

        assert result.f.shape == (10, 30)
        assert result.f_mean.shape == (30,)
        assert result.f_std.shape == (30,)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.f_mean))
        assert np.all(np.isfinite(result.f_std))

    def test_scalar_noise_posterior_mean_near_observations(self):
        """Posterior mean should be close to observations with scalar noise."""
        x_train = np.array([1.0, 2.0, 3.0])
        y_train = np.array([0.5, 0.8, 0.3])
        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=1, seed=42)

        np.testing.assert_allclose(result.f_mean, y_train, atol=0.05)

    def test_scalar_noise_uncertainty_increases_away_from_data(self):
        """Posterior std should increase away from observations."""
        x_train = np.array([2.0])
        y_train = np.array([0.5])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=0.01)

        x_test = np.linspace(0, 4, 50)
        spec = SamplingSpec(x_f=x_test)
        result = sample_posterior(obs, spec, sigma2=1.0, ell=0.5, n_samples=1, seed=42)

        # Find index closest to training point
        idx_train = np.argmin(np.abs(x_test - 2.0))

        # Std should be lower near training point
        std_at_train = result.f_std[idx_train]
        std_far = result.f_std[0]  # At x=0, far from x=2

        assert std_at_train < std_far


class TestDiagonalNoisePosterior:
    """Test posterior sampling with diagonal noise (heteroscedastic)."""

    def test_diagonal_noise_produces_valid_samples(self):
        """Diagonal noise should produce valid posterior samples."""
        x_train = np.array([0.0, 1.0, 2.0, 3.0])
        y_train = np.sin(x_train)
        noise_variances = np.array([0.01, 0.02, 0.015, 0.01])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_variances)

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

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=10, seed=42)

        assert result.f.shape == (10, 20)
        assert result.f_mean.shape == (20,)
        assert result.f_std.shape == (20,)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.f_mean))
        assert np.all(np.isfinite(result.f_std))

    def test_diagonal_noise_higher_noise_gives_more_uncertainty(self):
        """Points with higher noise should have more posterior uncertainty."""
        # Two observations: one with low noise, one with high noise
        x_train = np.array([1.0, 3.0])
        y_train = np.array([0.5, 0.5])
        # First point has low noise, second has high noise
        noise_variances = np.array([1e-6, 1.0])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_variances)

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

        # Empirical std from samples
        sample_std = np.std(result.f, axis=0)

        # Point with higher noise should have higher uncertainty
        assert sample_std[1] > sample_std[0] * 2

    def test_diagonal_noise_for_g_observations(self):
        """Diagonal noise should work for g observations."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.array([0.0, 0.5, 1.2])
        noise_variances = np.array([0.01, 0.02, 0.015])
        obs = Observations(x_g=x_train, y_g=y_train, noise_g=noise_variances)

        x_test = np.linspace(0, 2, 15)
        spec = SamplingSpec(x_g=x_test)

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

        assert result.g.shape == (5, 15)
        assert result.g_mean.shape == (15,)
        assert np.all(np.isfinite(result.g))

    def test_diagonal_noise_for_h_observations(self):
        """Diagonal noise should work for h observations."""
        x_train = np.array([0.5, 1.5, 2.5])
        y_train = np.cos(x_train)  # Derivative of sin
        noise_variances = np.array([0.01, 0.02, 0.01])
        obs = Observations(x_h=x_train, y_h=y_train, noise_h=noise_variances)

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

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

        assert result.h.shape == (5, 20)
        assert result.h_mean.shape == (20,)
        assert np.all(np.isfinite(result.h))


class TestFullCovarianceNoisePosterior:
    """Test posterior sampling with full covariance noise (correlated)."""

    def test_full_covariance_noise_produces_valid_samples(self):
        """Full covariance noise should produce valid posterior samples."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.sin(x_train)
        # Correlated noise: observations 0 and 1 are correlated
        noise_cov = np.array([
            [0.01, 0.005, 0.001],
            [0.005, 0.015, 0.002],
            [0.001, 0.002, 0.012]
        ])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_cov)

        x_test = np.linspace(0, 2, 20)
        spec = SamplingSpec(x_f=x_test)

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=10, seed=42)

        assert result.f.shape == (10, 20)
        assert result.f_mean.shape == (20,)
        assert result.f_std.shape == (20,)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.f_mean))
        assert np.all(np.isfinite(result.f_std))

    def test_full_covariance_noise_for_two_observations(self):
        """Full covariance with 2x2 matrix should work."""
        x_train = np.array([1.0, 3.0])
        y_train = np.array([0.5, 0.8])
        # Positive correlation
        noise_cov = np.array([[0.01, 0.005], [0.005, 0.02]])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_cov)

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

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=10, seed=42)

        assert result.f.shape == (10, 30)
        assert np.all(np.isfinite(result.f))

    def test_full_covariance_noise_different_from_diagonal(self):
        """Correlated noise should give different results than diagonal."""
        x_train = np.array([1.0, 1.5])
        y_train = np.array([0.5, 0.6])

        # Diagonal noise
        noise_diag = np.array([0.01, 0.02])
        obs_diag = Observations(x_f=x_train, y_f=y_train, noise_f=noise_diag)

        # Full covariance with strong correlation
        noise_cov = np.array([[0.01, 0.009], [0.009, 0.02]])
        obs_cov = Observations(x_f=x_train, y_f=y_train, noise_f=noise_cov)

        x_test = np.linspace(0.5, 2.0, 20)
        spec = SamplingSpec(x_f=x_test)

        result_diag = sample_posterior(
            obs_diag, spec, sigma2=1.0, ell=0.5, n_samples=1, seed=42
        )
        result_cov = sample_posterior(
            obs_cov, spec, sigma2=1.0, ell=0.5, n_samples=1, seed=100
        )

        # Means should differ (different noise models)
        assert not np.allclose(result_diag.f_mean, result_cov.f_mean)

    def test_full_covariance_noise_for_g_observations(self):
        """Full covariance noise should work for g observations."""
        x_train = np.array([0.0, 1.0])
        y_train = np.array([0.0, 0.5])
        noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])
        obs = Observations(x_g=x_train, y_g=y_train, noise_g=noise_cov)

        x_test = np.linspace(0, 1, 10)
        spec = SamplingSpec(x_g=x_test)

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

        assert result.g.shape == (5, 10)
        assert np.all(np.isfinite(result.g))

    def test_full_covariance_noise_for_h_observations(self):
        """Full covariance noise should work for h observations."""
        x_train = np.array([0.5, 1.5, 2.5])
        y_train = np.cos(x_train)
        noise_cov = np.array([
            [0.01, 0.002, 0.001],
            [0.002, 0.015, 0.002],
            [0.001, 0.002, 0.012]
        ])
        obs = Observations(x_h=x_train, y_h=y_train, noise_h=noise_cov)

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

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

        assert result.h.shape == (5, 20)
        assert np.all(np.isfinite(result.h))


class TestNoiseTypeComparison:
    """Compare behavior across different noise types."""

    def test_scalar_equals_diagonal_with_same_variance(self):
        """Scalar noise should equal diagonal noise with identical variances."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.sin(x_train)

        # Scalar noise
        obs_scalar = Observations(x_f=x_train, y_f=y_train, noise_f=0.01)

        # Diagonal noise with same variance for all observations
        obs_diagonal = Observations(
            x_f=x_train, y_f=y_train, noise_f=np.array([0.01, 0.01, 0.01])
        )

        x_test = np.linspace(0, 2, 15)
        spec = SamplingSpec(x_f=x_test)

        result_scalar = sample_posterior(
            obs_scalar, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=42
        )
        result_diagonal = sample_posterior(
            obs_diagonal, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=42
        )

        # Should give identical results
        np.testing.assert_allclose(result_scalar.f_mean, result_diagonal.f_mean)
        np.testing.assert_allclose(result_scalar.f_std, result_diagonal.f_std)
        np.testing.assert_allclose(result_scalar.f, result_diagonal.f)

    def test_diagonal_equals_full_covariance_when_uncorrelated(self):
        """Diagonal noise should equal full covariance with zero off-diagonals."""
        x_train = np.array([0.0, 1.0])
        y_train = np.array([0.5, 0.8])

        # Diagonal noise
        obs_diagonal = Observations(
            x_f=x_train, y_f=y_train, noise_f=np.array([0.01, 0.02])
        )

        # Full covariance with zero off-diagonals (uncorrelated)
        obs_full = Observations(
            x_f=x_train, y_f=y_train,
            noise_f=np.array([[0.01, 0.0], [0.0, 0.02]])
        )

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

        result_diagonal = sample_posterior(
            obs_diagonal, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=42
        )
        result_full = sample_posterior(
            obs_full, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=42
        )

        # Should give identical results
        np.testing.assert_allclose(result_diagonal.f_mean, result_full.f_mean)
        np.testing.assert_allclose(result_diagonal.f_std, result_full.f_std)
        np.testing.assert_allclose(result_diagonal.f, result_full.f)


# =============================================================================
# Edge Cases
# =============================================================================


class TestSingleObservationEdgeCases:
    """Test edge cases with single observations."""

    def test_single_observation_scalar_noise(self):
        """Single observation with scalar noise should work."""
        obs = Observations(
            x_f=np.array([1.5]),
            y_f=np.array([0.7]),
            noise_f=0.01
        )

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

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

        assert result.f.shape == (5, 20)
        assert result.f_mean.shape == (20,)
        assert np.all(np.isfinite(result.f))

    def test_single_observation_diagonal_noise(self):
        """Single observation with diagonal noise (1D array of length 1)."""
        obs = Observations(
            x_f=np.array([1.5]),
            y_f=np.array([0.7]),
            noise_f=np.array([0.01])
        )

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

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

        assert result.f.shape == (5, 20)
        assert np.all(np.isfinite(result.f))

    def test_single_observation_full_covariance_noise(self):
        """Single observation with full covariance (1x1 matrix)."""
        obs = Observations(
            x_f=np.array([1.5]),
            y_f=np.array([0.7]),
            noise_f=np.array([[0.01]])
        )

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

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

        assert result.f.shape == (5, 20)
        assert np.all(np.isfinite(result.f))


class TestMixedObservationTypes:
    """Test mixed observation types (f, g, h) with different noise types."""

    def test_f_and_g_with_different_noise_types(self):
        """f with scalar noise, g with diagonal noise."""
        x_f = np.array([0.0, 2.0])
        y_f = np.sin(x_f)
        x_g = np.array([0.0, 1.0, 2.0])
        y_g = np.array([0.0, 0.5, 1.0])

        obs = Observations(
            x_f=x_f, y_f=y_f, noise_f=0.01,
            x_g=x_g, y_g=y_g, noise_g=np.array([0.01, 0.02, 0.015])
        )

        x_test = np.linspace(0, 2, 15)
        spec = SamplingSpec(x_f=x_test, x_g=x_test)

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

        assert result.f.shape == (5, 15)
        assert result.g.shape == (5, 15)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.g))

    def test_f_g_h_with_all_noise_types(self):
        """f, g, h with scalar, diagonal, and full covariance noise."""
        x_f = np.array([0.0, 1.0, 2.0])
        y_f = np.sin(x_f)
        x_g = np.array([0.0, 1.0])
        y_g = np.array([0.0, 0.5])
        x_h = np.array([0.5, 1.5, 2.5, 3.5])
        y_h = np.cos(x_h)

        obs = Observations(
            x_f=x_f, y_f=y_f, noise_f=0.01,  # Scalar
            x_g=x_g, y_g=y_g, noise_g=np.array([0.01, 0.02]),  # Diagonal
            x_h=x_h, y_h=y_h, noise_h=np.array([  # Full covariance
                [0.01, 0.002, 0.001, 0.0],
                [0.002, 0.015, 0.002, 0.001],
                [0.001, 0.002, 0.012, 0.002],
                [0.0, 0.001, 0.002, 0.01]
            ])
        )

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

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

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

    def test_f_with_full_covariance_h_with_diagonal(self):
        """f with full covariance, h with diagonal noise."""
        x_f = np.array([0.0, 1.0])
        y_f = np.sin(x_f)
        x_h = np.array([0.5, 1.5, 2.5])
        y_h = np.cos(x_h)

        obs = Observations(
            x_f=x_f, y_f=y_f,
            noise_f=np.array([[0.01, 0.002], [0.002, 0.015]]),
            x_h=x_h, y_h=y_h,
            noise_h=np.array([0.01, 0.02, 0.01])
        )

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

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

        assert result.f.shape == (5, 25)
        assert result.h.shape == (5, 25)
        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.h))


class TestReproducibility:
    """Test reproducibility with different noise types."""

    def test_diagonal_noise_same_seed_same_result(self):
        """Same seed should produce same samples with diagonal noise."""
        x_train = np.array([0.0, 1.0, 2.0])
        y_train = np.array([0.5, 0.8, 0.3])
        noise = np.array([0.01, 0.02, 0.015])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise)

        x_test = np.linspace(0, 2, 15)
        spec = SamplingSpec(x_f=x_test)

        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)
        np.testing.assert_array_equal(result1.f_std, result2.f_std)

    def test_full_covariance_noise_same_seed_same_result(self):
        """Same seed should produce same samples with full covariance noise."""
        x_train = np.array([0.0, 1.0])
        y_train = np.array([0.5, 0.8])
        noise_cov = np.array([[0.01, 0.002], [0.002, 0.015]])
        obs = Observations(x_f=x_train, y_f=y_train, noise_f=noise_cov)

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

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

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


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