"""
Tests for u(x) = f''(x) second-derivative GP support.

Covers:
1. Basic prior/posterior sampling shapes and types
2. Kernel cross-covariance formulas validated against finite differences
3. u ≈ numerical second derivative of f (relationship test)
4. Posterior conditioning on u observations
5. Kernel restriction enforcement (non-RBF raises ValueError)
6. Numerical stability across length scales and variances

The test fixtures and parametrisation are designed so that adding a new
kernel only requires updating U_SUPPORTED_KERNELS.
"""

import numpy as np
import pytest

from gp4c import Observations, SamplingSpec, sample_posterior, sample_prior

# ── Kernels that support u=f'' ────────────────────────────────────────────────
# Extend this list when additional kernels gain u support.
U_SUPPORTED_KERNELS = ["rbf", "matern52", "periodic", "locally_periodic"]

# Kernels that must raise ValueError when u is requested
U_UNSUPPORTED_KERNELS = ["matern32"]

# Kernels that don't support g (integral) observations or predictions
KERNELS_WITHOUT_G = {"periodic", "locally_periodic"}

# Extra kwargs for sample_prior / sample_posterior per kernel.
# For periodic/locally_periodic: use period=2.5 so test points like [1, 2, 3]
# are not integer multiples of the period (avoiding rank-deficient K_uu).
KERNEL_KWARGS: dict = {k: {} for k in U_SUPPORTED_KERNELS}
KERNEL_KWARGS["periodic"] = {"period": 2.5}
KERNEL_KWARGS["locally_periodic"] = {"period": 2.5}

# Extra kwargs for physics tests that match a specific function's period.
# sin(x) has period 2π, so periodic kernels must be configured accordingly.
KERNEL_PHYSICS_KWARGS: dict = {k: {} for k in U_SUPPORTED_KERNELS}
KERNEL_PHYSICS_KWARGS["periodic"] = {"period": 2 * np.pi}
KERNEL_PHYSICS_KWARGS["locally_periodic"] = {"period": 2 * np.pi}


# =============================================================================
# Helpers
# =============================================================================


def empirical_covariance(samples):
    """Return the (n_pts, n_pts) empirical covariance from (n_samples, n_pts) array."""
    return np.cov(samples.T)


def rbf_kuu_diagonal(sigma2, ell):
    """Analytical K_uu(x, x): σ²·3/ℓ⁴  (r=0 → 3 - 0 + 0 = 3)."""
    return 3.0 * sigma2 / ell**4


def matern52_kuu_diagonal(sigma2, ell):
    """Analytical K_uu(x, x) for Matérn 5/2: (25/3)·σ²·3ℓ²/ℓ⁶ = 25σ²/ℓ⁴."""
    return 25.0 * sigma2 / ell**4


def periodic_kuu_diagonal(sigma2, ell, period=1.0):
    """Analytical K_uu(x, x) for periodic: 16π⁴σ²(ℓ²+3)/(ℓ⁴p⁴)."""
    return 16.0 * np.pi**4 * sigma2 * (ell**2 + 3.0) / (ell**4 * period**4)


def locally_periodic_kuu_diagonal(sigma2, ell, ell_p=1.0, period=1.0):
    """Analytical K_uu(x,x) for locally-periodic at t=0:
    σ²(16π⁴ℓ⁴(ℓ_p²+3) + 24π²ℓ²ℓ_p²p² + 3ℓ_p⁴p⁴) / (ℓ⁴ℓ_p⁴p⁴)."""
    return sigma2 * (
        16.0 * np.pi**4 * ell**4 * (ell_p**2 + 3.0)
        + 24.0 * np.pi**2 * ell**2 * ell_p**2 * period**2
        + 3.0 * ell_p**4 * period**4
    ) / (ell**4 * ell_p**4 * period**4)


KUU_DIAGONAL = {
    "rbf": rbf_kuu_diagonal,
    "matern52": matern52_kuu_diagonal,
    "periodic": periodic_kuu_diagonal,
    "locally_periodic": locally_periodic_kuu_diagonal,
}


def rbf_kfu_at_zero(sigma2, ell):
    """Analytical K_fu(x, x): (σ²/ℓ²)(0 - 1) = -σ²/ℓ²  (r=0)."""
    return -sigma2 / ell**2


def rbf_khu_at_zero(sigma2, ell):
    """Analytical K_hu(x, x): 0  (r=0)."""
    return 0.0


def matern52_kfu_at_zero(sigma2, ell):
    """Analytical K_fu(x, x) for Matérn 5/2: (5σ²/3ℓ⁴)(-ℓ²) = -5σ²/(3ℓ²)  (r=0)."""
    return -5.0 * sigma2 / (3.0 * ell**2)


def periodic_kfu_at_zero(sigma2, ell, period=1.0):
    """Analytical K_fu(x, x) for periodic: -4π²σ²/(ℓ²p²)  (t=0).
    Here ell is the periodic length scale (ell_p in C)."""
    return -4.0 * np.pi**2 * sigma2 / (ell**2 * period**2)


def locally_periodic_kfu_at_zero(sigma2, ell, ell_p=1.0, period=1.0):
    """Analytical K_fu(x, x) for locally-periodic: -σ²(4π²/(ell_p²p²) + 1/ell²)  (t=0)."""
    return -sigma2 * (4.0 * np.pi**2 / (ell_p**2 * period**2) + 1.0 / ell**2)


KFU_DIAGONAL = {
    "rbf": rbf_kfu_at_zero,
    "matern52": matern52_kfu_at_zero,
    "periodic": periodic_kfu_at_zero,
    "locally_periodic": locally_periodic_kfu_at_zero,
}


# =============================================================================
# 1. Basic prior sampling
# =============================================================================


class TestSecondDerivativePriorBasics:
    """Output shapes and None/non-None fields for prior sampling with u."""

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_sample_u_only(self, kernel):
        x = np.linspace(0, 5, 40)
        spec = SamplingSpec(x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=5, seed=0)

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

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_sample_f_and_u(self, kernel):
        x = np.linspace(0, 5, 40)
        spec = SamplingSpec(x_f=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=5, seed=0)

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

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_sample_f_h_u(self, kernel):
        x = np.linspace(0, 5, 40)
        spec = SamplingSpec(x_f=x, x_h=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=5, seed=0)

        assert result.f is not None
        assert result.h is not None
        assert result.u is not None

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_sample_all_four(self, kernel):
        x = np.linspace(0, 5, 30)
        x_g = None if kernel in KERNELS_WITHOUT_G else x
        spec = SamplingSpec(x_f=x, x_g=x_g, x_h=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, n_samples=3, seed=0,
                              kernel=kernel, **KERNEL_KWARGS[kernel])

        for arr, name in [(result.f, "f"), (result.h, "h"), (result.u, "u")]:
            assert arr is not None, f"{name} should not be None"
            assert arr.shape == (3, 30), f"{name}.shape mismatch"
        if kernel not in KERNELS_WITHOUT_G:
            assert result.g is not None and result.g.shape == (3, 30)
        else:
            assert result.g is None

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_no_nan_inf(self, kernel):
        x = np.linspace(0, 5, 50)
        spec = SamplingSpec(x_f=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=10, seed=0)

        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.u))

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_reproducibility(self, kernel):
        x = np.linspace(0, 5, 30)
        spec = SamplingSpec(x_f=x, x_u=x)
        r1 = sample_prior(spec, ell=1.0, kernel=kernel, n_samples=5, seed=99)
        r2 = sample_prior(spec, ell=1.0, kernel=kernel, n_samples=5, seed=99)

        assert np.allclose(r1.u, r2.u)
        assert np.allclose(r1.f, r2.f)


# =============================================================================
# 2. Kernel formula verification via finite differences
# =============================================================================


class TestKernelFormulas:
    """Verify K_uu, K_fu, K_hu against analytical formulae for all supported kernels.

    Approach: draw a large number of joint samples and compute the empirical
    covariance, then compare to the analytical value.  Tolerances are generous
    (Monte Carlo error at N=20000 is O(1/sqrt(N)) ≈ 0.7 %).
    """

    N_SAMPLES = 20000
    SIGMA2 = 1.0
    ELL = 1.0
    ATOL_DIAG = 0.08  # ~8 % relative tolerance on diagonal

    def _prior(self, spec, kernel="rbf", seed=0):
        return sample_prior(
            spec,
            sigma2=self.SIGMA2,
            ell=self.ELL,
            kernel=kernel,
            n_samples=self.N_SAMPLES,
            seed=seed,
            **KERNEL_KWARGS[kernel],
        )

    def _kuu_expected(self, kernel):
        extra = KERNEL_KWARGS[kernel]
        period = extra.get("period", 1.0)
        if kernel == "locally_periodic":
            return KUU_DIAGONAL[kernel](self.SIGMA2, self.ELL, ell_p=self.ELL, period=period)
        elif kernel == "periodic":
            return KUU_DIAGONAL[kernel](self.SIGMA2, self.ELL, period=period)
        return KUU_DIAGONAL[kernel](self.SIGMA2, self.ELL)

    def _kfu_expected(self, kernel):
        extra = KERNEL_KWARGS[kernel]
        period = extra.get("period", 1.0)
        if kernel == "locally_periodic":
            return KFU_DIAGONAL[kernel](self.SIGMA2, self.ELL, ell_p=self.ELL, period=period)
        elif kernel == "periodic":
            return KFU_DIAGONAL[kernel](self.SIGMA2, self.ELL, period=period)
        return KFU_DIAGONAL[kernel](self.SIGMA2, self.ELL)

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_kuu_diagonal_matches_formula(self, kernel):
        """Empirical Var[u(x)] matches analytical K_uu(x, x)."""
        x = np.array([0.3, 1.1, 2.4])
        spec = SamplingSpec(x_u=x)
        result = self._prior(spec, kernel=kernel)

        expected = self._kuu_expected(kernel)
        empirical_var = result.u.var(axis=0)

        np.testing.assert_allclose(empirical_var, expected, atol=self.ATOL_DIAG * expected)

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_kfu_at_same_point_matches_formula(self, kernel):
        """Empirical Cov[f(x), u(x)] matches analytical K_fu(x, x)."""
        x = np.array([0.7, 1.4, 2.1])
        spec = SamplingSpec(x_f=x, x_u=x)
        result = self._prior(spec, kernel=kernel)

        expected = self._kfu_expected(kernel)
        atol = max(self.ATOL_DIAG * abs(expected), 0.08)
        for i in range(len(x)):
            cov = np.cov(result.f[:, i], result.u[:, i])[0, 1]
            assert abs(cov - expected) < atol, (
                f"{kernel} K_fu diagonal [{i}]: got {cov:.4f}, expected {expected:.4f}"
            )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_khu_at_same_point_is_zero(self, kernel):
        """Corr[h(x), u(x)] ≈ 0  (K_hu is anti-symmetric → diagonal=0).

        Uses the correlation coefficient (bounded in [-1,1]) so the test is
        robust to kernels with large marginal variances (e.g. periodic).
        SE[ρ] ≈ 1/sqrt(N) ≈ 0.007 at N=20000.
        """
        x = np.array([0.7, 1.4, 2.1])
        spec = SamplingSpec(x_h=x, x_u=x)
        result = self._prior(spec, kernel=kernel)

        for i in range(len(x)):
            rho = np.corrcoef(result.h[:, i], result.u[:, i])[0, 1]
            assert abs(rho) < 0.05, (
                f"{kernel} K_hu diagonal [{i}]: correlation={rho:.4f}, expected ≈0"
            )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_kfu_is_symmetric(self, kernel):
        """K_fu(x, x') = K_fu(x', x)  (symmetric cross-covariance).

        Uses joint sampling of (f(a), f(b), u(a), u(b)) from the same draw so
        both estimates share the same noise, giving better cancellation.
        """
        xa, xb = 0.5, 1.5
        spec = SamplingSpec(x_f=np.array([xa, xb]), x_u=np.array([xa, xb]))
        result = self._prior(spec, kernel=kernel, seed=1)

        cov = np.cov(np.hstack([result.f, result.u]).T)
        # K_fu(a,b) = Cov[f(a), u(b)], K_fu(b,a) = Cov[f(b), u(a)]
        k_fu_ab = cov[0, 3]
        k_fu_ba = cov[1, 2]
        tol = max(self.ATOL_DIAG, 0.15 * abs(k_fu_ab))
        assert abs(k_fu_ab - k_fu_ba) < tol, (
            f"{kernel} K_fu not symmetric: K_fu(a,b)={k_fu_ab:.4f}, K_fu(b,a)={k_fu_ba:.4f}"
        )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_khu_is_antisymmetric(self, kernel):
        """K_hu(x, x') = -K_hu(x', x)  (anti-symmetric).

        Uses joint sampling of (h(a), h(b), u(a), u(b)) from the same draw so
        both estimates share the same noise, giving better cancellation.
        """
        xa, xb = 0.5, 1.5
        spec = SamplingSpec(x_h=np.array([xa, xb]), x_u=np.array([xa, xb]))
        result = self._prior(spec, kernel=kernel, seed=2)

        cov = np.cov(np.hstack([result.h, result.u]).T)
        # K_hu(a,b) = Cov[h(a), u(b)], K_hu(b,a) = Cov[h(b), u(a)]
        k_hu_ab = cov[0, 3]
        k_hu_ba = cov[1, 2]
        tol = max(self.ATOL_DIAG, 0.05 * abs(k_hu_ab))
        assert abs(k_hu_ab + k_hu_ba) < tol, (
            f"{kernel} K_hu not anti-symmetric: K_hu(a,b)={k_hu_ab:.4f}, K_hu(b,a)={k_hu_ba:.4f}"
        )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_kuu_is_symmetric(self, kernel):
        """K_uu(x, x') = K_uu(x', x)."""
        spec = SamplingSpec(x_u=np.array([0.5, 1.5]))
        result = self._prior(spec, kernel=kernel)

        cov_mat = empirical_covariance(result.u)
        assert abs(cov_mat[0, 1] - cov_mat[1, 0]) < self.ATOL_DIAG


# =============================================================================
# 3. u ≈ numerical second derivative of f
# =============================================================================


class TestSecondDerivativeRelationship:
    """Joint (f, u) samples must have u correlated with f''."""

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_u_correlates_with_numerical_second_derivative_of_f(self, kernel):
        """u should have high correlation with the numerical second derivative of f."""
        # Dense grid so finite-difference approximation is accurate
        x = np.linspace(0.2, 4.8, 120)
        spec = SamplingSpec(x_f=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=20, seed=7)

        dx = x[1] - x[0]
        for i in range(result.f.shape[0]):
            f = result.f[i]
            u = result.u[i]

            # Central second difference, drop 1-point boundary
            f_dd = np.gradient(np.gradient(f, dx), dx)[1:-1]
            u_inner = u[1:-1]

            corr = np.corrcoef(u_inner, f_dd)[0, 1]
            assert corr > 0.8, (
                f"Kernel {kernel}, sample {i}: Corr[u, f''] = {corr:.3f} too low"
            )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_u_correlates_with_numerical_derivative_of_h(self, kernel):
        """u = h', so u should correlate with the numerical derivative of h."""
        x = np.linspace(0.2, 4.8, 120)
        spec = SamplingSpec(x_h=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=20, seed=8)

        dx = x[1] - x[0]
        for i in range(result.h.shape[0]):
            h = result.h[i]
            u = result.u[i]

            h_d = np.gradient(h, dx)[1:-1]
            u_inner = u[1:-1]

            corr = np.corrcoef(u_inner, h_d)[0, 1]
            assert corr > 0.8, (
                f"Kernel {kernel}, sample {i}: Corr[u, h'] = {corr:.3f} too low"
            )


# =============================================================================
# 4. Posterior sampling
# =============================================================================


class TestSecondDerivativePosteriorBasics:
    """Output shapes and types for posterior sampling with u."""

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_predict_u_from_f_observations(self, kernel):
        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)

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

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

        assert result.u is not None
        assert result.u.shape == (5, 30)
        assert result.u_mean is not None
        assert result.u_mean.shape == (30,)
        assert result.u_std is not None
        assert result.u_std.shape == (30,)
        assert result.f is None
        assert np.all(np.isfinite(result.u))
        assert np.all(np.isfinite(result.u_mean))
        assert np.all(result.u_std >= 0)

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_predict_f_from_u_observations(self, kernel):
        x_train = np.array([0.5, 1.0, 1.5, 2.0, 2.5])
        y_train = -np.sin(x_train)  # f'' of sin(x)
        obs = Observations(x_u=x_train, y_u=y_train, noise_u=1e-4)

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

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

        assert result.f is not None
        assert result.f.shape == (5, 30)
        assert result.f_mean is not None
        assert np.all(np.isfinite(result.f))

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_predict_all_four_from_u_observations(self, kernel):
        x_train = np.array([1.0, 2.0, 3.0])
        y_train = -np.sin(x_train)
        obs = Observations(x_u=x_train, y_u=y_train, noise_u=1e-4)

        x_test = np.linspace(0.5, 3.5, 20)
        x_g = None if kernel in KERNELS_WITHOUT_G else x_test
        spec = SamplingSpec(x_f=x_test, x_g=x_g, x_h=x_test, x_u=x_test)

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=3, seed=0,
                                  kernel=kernel, **KERNEL_KWARGS[kernel])

        for attr in ["f", "h", "u"]:
            arr = getattr(result, attr)
            assert arr is not None, f"{attr} should not be None"
            assert arr.shape == (3, 20), f"{attr}.shape mismatch"
            assert np.all(np.isfinite(arr)), f"{attr} contains non-finite values"
        if kernel not in KERNELS_WITHOUT_G:
            assert result.g is not None and result.g.shape == (3, 20)

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_u_observations_condition_u_mean(self, kernel):
        """Posterior u_mean should pass near the u observations at training locations."""
        x_train = np.array([1.0, 2.0, 3.0])
        y_train = -np.sin(x_train)  # exact f'' values for sin(x) GP with the right ell
        obs = Observations(x_u=x_train, y_u=y_train, noise_u=1e-6)

        spec = SamplingSpec(x_u=x_train)
        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=0,
                                  kernel=kernel, **KERNEL_KWARGS[kernel])

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

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_posterior_u_std_shrinks_near_observations(self, kernel):
        """Posterior u_std should be smaller near the u observations."""
        x_train = np.array([2.0, 3.0])
        y_train = np.zeros(2)
        obs = Observations(x_u=x_train, y_u=y_train, noise_u=1e-6)

        x_near = np.array([2.0, 3.0])
        x_far = np.array([0.0, 5.0])
        spec = SamplingSpec(x_u=np.concatenate([x_near, x_far]))

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, n_samples=1, seed=0,
                                  kernel=kernel, **KERNEL_KWARGS[kernel])

        std_near = result.u_std[:2].mean()
        std_far = result.u_std[2:].mean()
        assert std_near < std_far, (
            f"Expected std near obs ({std_near:.4f}) < std far ({std_far:.4f})"
        )


class TestSecondDerivativePosteriorPhysics:
    """Physics-based checks: posterior u mean should approximate the true f'' of known functions."""

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_predicting_u_from_f_sin_function(self, kernel):
        """From f=sin(x) observations, posterior u_mean ≈ -sin(x)."""
        # For periodic kernel match the period of sin(x) = 2π.
        # For other kernels a short ell captures the local curvature.
        extra = KERNEL_PHYSICS_KWARGS[kernel]
        ell = 0.5 if kernel != "periodic" else 1.0
        x_train = np.linspace(0, 2 * np.pi, 30)
        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.5, 2 * np.pi - 0.5, 40)
        spec = SamplingSpec(x_u=x_test)

        result = sample_posterior(obs, spec, sigma2=1.0, ell=ell, n_samples=1, seed=0,
                                  kernel=kernel, **extra)

        u_true = -np.sin(x_test)
        corr = np.corrcoef(result.u_mean, u_true)[0, 1]
        assert corr > 0.9, f"Kernel {kernel}: Corr[u_mean, -sin] = {corr:.3f} too low"

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_conditioning_on_u_informs_f_curvature(self, kernel):
        """When we observe u=0 everywhere (flat curvature), f posterior should be near-linear."""
        x_train = np.linspace(0.5, 4.5, 8)
        y_u = np.zeros_like(x_train)  # zero second derivative → linear function
        obs = Observations(x_u=x_train, y_u=y_u, noise_u=1e-4)

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

        result = sample_posterior(obs, spec, sigma2=1.0, ell=1.0, kernel=kernel, n_samples=1, seed=0)

        # f_mean should be approximately linear → its second difference should be small
        f2 = np.diff(result.f_mean, n=2)
        assert np.all(np.abs(f2) < 0.5), "f_mean should have low curvature when u=0"


# =============================================================================
# 5. Kernel restriction enforcement
# =============================================================================


class TestSecondDerivativeKernelRestrictions:
    """Non-RBF kernels must raise ValueError when u is requested."""

    @pytest.mark.parametrize("kernel", U_UNSUPPORTED_KERNELS)
    def test_prior_raises_for_unsupported_kernel(self, kernel):
        x = np.linspace(0, 5, 10)
        spec = SamplingSpec(x_u=x)
        with pytest.raises(ValueError, match="rbf"):
            sample_prior(spec, kernel=kernel, n_samples=1)

    @pytest.mark.parametrize("kernel", U_UNSUPPORTED_KERNELS)
    def test_posterior_raises_for_unsupported_kernel_pred(self, kernel):
        obs = Observations(x_f=np.array([1.0]), y_f=np.array([0.0]))
        spec = SamplingSpec(x_u=np.linspace(0, 2, 5))
        with pytest.raises(ValueError, match="rbf"):
            sample_posterior(obs, spec, kernel=kernel, n_samples=1)

    @pytest.mark.parametrize("kernel", U_UNSUPPORTED_KERNELS)
    def test_posterior_raises_for_unsupported_kernel_obs(self, kernel):
        obs = Observations(x_u=np.array([1.0]), y_u=np.array([0.0]))
        spec = SamplingSpec(x_f=np.linspace(0, 2, 5))
        with pytest.raises(ValueError, match="rbf"):
            sample_posterior(obs, spec, kernel=kernel, n_samples=1)


# =============================================================================
# 6. Numerical stability
# =============================================================================


class TestSecondDerivativeNumericalStability:
    """u sampling should remain finite across a wide range of parameters."""

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    @pytest.mark.parametrize("ell", [0.1, 0.5, 1.0, 2.0, 5.0])
    def test_prior_finite_at_various_length_scales(self, kernel, ell):
        x = np.linspace(0, 5, 30)
        spec = SamplingSpec(x_f=x, x_u=x)
        result = sample_prior(spec, sigma2=1.0, ell=ell, kernel=kernel, n_samples=5, seed=0)

        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.u))

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    @pytest.mark.parametrize("sigma2", [0.1, 1.0, 10.0])
    def test_prior_finite_at_various_variances(self, kernel, sigma2):
        x = np.linspace(0, 5, 30)
        spec = SamplingSpec(x_f=x, x_u=x)
        result = sample_prior(spec, sigma2=sigma2, ell=1.0, kernel=kernel, n_samples=5, seed=0)

        assert np.all(np.isfinite(result.f))
        assert np.all(np.isfinite(result.u))

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_u_variance_scales_with_sigma2_over_ell4(self, kernel):
        """Var[u(x)] matches analytical K_uu diagonal — verify scaling w.r.t. σ² and ℓ."""
        x = np.array([2.5])
        n = 3000
        kuu_diag = KUU_DIAGONAL[kernel]
        extra = KERNEL_KWARGS[kernel]  # e.g. {"period": 2.5} for periodic kernels

        for sigma2, ell in [(1.0, 1.0), (2.0, 1.0), (1.0, 0.5)]:
            spec = SamplingSpec(x_u=x)
            result = sample_prior(spec, sigma2=sigma2, ell=ell, kernel=kernel, n_samples=n,
                                  seed=0, **extra)
            empirical = result.u.var()
            # locally_periodic needs ell_p; use ell_p=ell (default) and period from extra
            period = extra.get("period", 1.0)
            if kernel == "locally_periodic":
                expected = kuu_diag(sigma2, ell, ell_p=ell, period=period)
            elif kernel == "periodic":
                expected = kuu_diag(sigma2, ell, period=period)
            else:
                expected = kuu_diag(sigma2, ell)
            assert abs(empirical / expected - 1.0) < 0.15, (
                f"σ²={sigma2}, ℓ={ell}: Var[u]={empirical:.4f}, expected ≈{expected:.4f}"
            )

    @pytest.mark.parametrize("kernel", U_SUPPORTED_KERNELS)
    def test_posterior_finite_at_various_length_scales(self, kernel):
        obs = Observations(x_u=np.array([1.0, 2.0, 3.0]), y_u=np.zeros(3))
        x_test = np.linspace(0, 4, 20)
        spec = SamplingSpec(x_f=x_test, x_u=x_test)

        for ell in [0.3, 1.0, 3.0]:
            result = sample_posterior(obs, spec, sigma2=1.0, ell=ell, kernel=kernel, n_samples=3, seed=0)
            assert np.all(np.isfinite(result.f)), f"f non-finite at ell={ell}"
            assert np.all(np.isfinite(result.u)), f"u non-finite at ell={ell}"
            assert np.all(result.u_std >= 0), f"u_std negative at ell={ell}"
