"""
Test mean function support in GP sampler.
"""

import numpy as np
import pytest
from gp4c import (
    sample_prior,
    sample_posterior,
    MeanSpec,
    SamplingSpec,
    Observations,
)


class TestMeanFunctions:
    """Test suite for mean function support."""

    def test_zero_mean_default(self):
        """Default behavior should be zero mean."""
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_g=x)
        result = sample_prior(spec, n_samples=100, seed=42)

        # With 100 samples, empirical mean should be close to 0
        assert abs(result.f.mean()) < 0.3, "Zero mean: f should have mean ≈ 0"
        assert abs(result.g.mean()) < 0.5, "Zero mean: g should have mean ≈ 0"

    def test_constant_mean_basic(self):
        """Test constant mean function."""
        x = np.linspace(0, 5, 50)
        mean_value = 2.5
        mean = MeanSpec(type="constant", value=mean_value)

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

        # Check f has correct mean
        f_mean_empirical = result.f.mean()
        assert abs(f_mean_empirical - mean_value) < 0.3, (
            f"Constant mean: f should have mean ≈ {mean_value}, got {f_mean_empirical}"
        )

        # Check g has correct mean: m_g(x) = c*x, so E[g] = c * E[x]
        expected_g_mean = mean_value * x.mean()
        g_mean_empirical = result.g.mean()
        assert abs(g_mean_empirical - expected_g_mean) < 0.5, (
            f"Constant mean: g should have mean ≈ {expected_g_mean}, got {g_mean_empirical}"
        )

    def test_mean_variance_unchanged(self):
        """Mean function should not change the variance/covariance structure."""
        x = np.linspace(0, 5, 50)
        n_samples = 500
        seed = 42

        # Sample with zero mean
        spec_zero = SamplingSpec(x_f=x, x_g=x)
        result_zero = sample_prior(spec_zero, n_samples=n_samples, seed=seed)

        # Sample with constant mean
        mean = MeanSpec(type="constant", value=3.0)
        spec_const = SamplingSpec(x_f=x, x_g=x, mean=mean)
        result_const = sample_prior(spec_const, n_samples=n_samples, seed=seed)

        # Center the constant mean samples
        f_centered = result_const.f - 3.0
        g_centered = result_const.g - 3.0 * x  # Subtract m_g(x) = c*x

        # Variances should be approximately equal
        f_zero_std = result_zero.f.std()
        f_centered_std = f_centered.std()
        assert abs(f_zero_std - f_centered_std) < 0.1, (
            "Variance should be same after centering"
        )

        g_zero_std = result_zero.g.std()
        g_centered_std = g_centered.std()
        assert abs(g_zero_std - g_centered_std) < 0.2, (
            "Variance of g should be same after centering"
        )

    def test_mean_with_flexible_sampling(self):
        """Test mean function with flexible sampling API."""
        x = np.linspace(0, 5, 50)
        mean = MeanSpec(type="constant", value=1.5)

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

        # Check means
        assert result.f is not None
        assert result.g is not None
        assert abs(result.f.mean() - 1.5) < 0.3
        assert abs(result.g.mean() - 1.5 * x.mean()) < 0.5

    def test_mean_with_posterior(self):
        """Test mean function with posterior sampling."""
        # Generate synthetic observations with known mean
        x_obs = np.array([0.5, 1.5, 2.5, 3.5, 4.5])
        true_mean = 2.0
        # Observations should be close to the mean with small noise
        np.random.seed(42)
        y_obs = true_mean + 0.1 * np.random.randn(len(x_obs))

        mean = MeanSpec(type="constant", value=true_mean)
        obs = Observations(x_f=x_obs, y_f=y_obs, noise_f=0.01, mean=mean)

        # Predict at new points
        x_pred = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x_pred, mean=mean)

        result = sample_posterior(obs, spec, n_samples=50, seed=42)

        # Posterior mean should be close to true mean
        assert result.f_mean is not None
        assert abs(result.f_mean.mean() - true_mean) < 0.5

    def test_invalid_mean_type(self):
        """Test that invalid mean types raise errors."""
        # Invalid type should raise an error (caught by dataclass validation or at runtime)
        with pytest.raises((ValueError, TypeError)):
            mean = MeanSpec(type="invalid", value=1.0)  # type: ignore

    def test_negative_mean_value(self):
        """Test that negative mean values work correctly."""
        x = np.linspace(0, 5, 50)
        mean = MeanSpec(type="constant", value=-1.5)

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

        # Check negative mean is applied
        assert result.f.mean() < 0, "Mean should be negative"
        assert abs(result.f.mean() - (-1.5)) < 0.3


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