"""Tests for hyperparameter optimization module."""

import numpy as np
import pytest

import gp4c
from gp4c.optimize import compute_log_marginal_likelihood


class TestLogMarginalLikelihood:
    """Test the log marginal likelihood function (C implementation)."""

    @pytest.mark.parametrize("kernel", ["rbf", "matern52", "matern32", "periodic", "locally_periodic"])
    def test_all_kernels(self, kernel):
        """Log marginal likelihood should work for all kernels."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        val = gp4c.log_marginal_likelihood(
            obs, sigma2=1.0, ell=1.0, period=1.0, kernel=kernel
        )

        assert isinstance(val, float)
        assert not np.isnan(val)
        assert not np.isinf(val)

    @pytest.mark.parametrize(
        "params",
        [
            {"sigma2": 1.0, "ell": 1.0, "period": 1.0},
            {"sigma2": 0.5, "ell": 2.0, "period": 0.5},
            {"sigma2": 2.0, "ell": 0.5, "period": 2.0},
        ],
    )
    def test_various_params(self, params):
        """Should work across various parameter values."""
        np.random.seed(42)
        x = np.linspace(0, 5, 15)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        val = gp4c.log_marginal_likelihood(obs, kernel="rbf", **params)
        assert isinstance(val, float)
        assert not np.isnan(val)

    def test_lml_alias(self):
        """LML alias should work identically to log_marginal_likelihood."""
        np.random.seed(42)
        x = np.linspace(0, 5, 15)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        val1 = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="rbf")
        val2 = gp4c.LML(obs, sigma2=1.0, ell=1.0, kernel="rbf")

        assert val1 == val2

    def test_scalar_vs_diagonal_noise(self):
        """Scalar noise 0.1 should match diagonal noise [0.1, 0.1, 0.1]."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.0, 1.0, 0.5])

        obs_scalar = gp4c.Observations(x_f=x, y_f=y, noise_f=0.1)
        obs_diag = gp4c.Observations(x_f=x, y_f=y, noise_f=np.array([0.1, 0.1, 0.1]))

        val_scalar = gp4c.log_marginal_likelihood(obs_scalar, sigma2=1.0, ell=1.0, kernel="rbf")
        val_diag = gp4c.log_marginal_likelihood(obs_diag, sigma2=1.0, ell=1.0, kernel="rbf")

        assert np.isclose(val_scalar, val_diag, rtol=1e-5)

    def test_invalid_sigma2_raises(self):
        """Should raise error for non-positive sigma2."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        with pytest.raises(ValueError, match="sigma2 must be positive"):
            gp4c.log_marginal_likelihood(obs, sigma2=-1.0, ell=1.0, kernel="rbf")

    def test_invalid_ell_raises(self):
        """Should raise error for non-positive ell."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        with pytest.raises(ValueError, match="ell must be positive"):
            gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=-1.0, kernel="rbf")

    def test_invalid_period_raises(self):
        """Should raise error for non-positive period."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        with pytest.raises(ValueError, match="period must be positive"):
            gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, period=-1.0, kernel="rbf")

    def test_matern32_with_derivative_works(self):
        """Matérn 3/2 supports derivative observations (once mean-square differentiable)."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(
            x_f=x, y_f=y, noise_f=0.01, x_h=x, y_h=y, noise_h=0.01
        )

        lml = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="matern32")
        assert np.isfinite(lml)

    def test_periodic_with_integral_raises(self):
        """Should raise error for periodic kernel with integral observations."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(
            x_f=x, y_f=y, noise_f=0.01, x_g=x, y_g=y, noise_g=0.01
        )

        with pytest.raises(ValueError, match="does not support integral"):
            gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="periodic")

    def test_derivative_observations(self):
        """Should handle derivative observations."""
        np.random.seed(42)
        x = np.linspace(0, 5, 15)
        y = np.sin(x) + np.random.normal(0, 0.05, len(x))
        dy = np.cos(x) + np.random.normal(0, 0.05, len(x))
        obs = gp4c.Observations(
            x_f=x, y_f=y, noise_f=0.01, x_h=x, y_h=dy, noise_h=0.01
        )

        val = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="matern52")
        assert isinstance(val, float)
        assert not np.isnan(val)

    def test_h_only_observations(self):
        """Should handle derivative-only observations."""
        np.random.seed(42)
        x = np.linspace(0, 5, 15)
        dy = np.cos(x) + np.random.normal(0, 0.05, len(x))
        obs = gp4c.Observations(x_h=x, y_h=dy, noise_h=0.01)

        val = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="rbf")
        assert isinstance(val, float)
        assert not np.isnan(val)

    def test_not_observations_raises(self):
        """Should raise error for non-Observations input."""
        with pytest.raises(TypeError, match="must be an Observations instance"):
            gp4c.log_marginal_likelihood(
                {"x": [1, 2, 3], "y": [1, 2, 3]}, sigma2=1.0, ell=1.0, kernel="rbf"
            )

    def test_return_type(self):
        """Should return a float."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="rbf")
        assert isinstance(result, float)
        assert not np.isnan(result)
        assert not np.isinf(result)

    def test_integral_observations_rbf(self):
        """C implementation handles integral observations for RBF."""
        x = np.array([0.0, 1.0, 2.0])
        y_f = np.array([0.1, 0.2, 0.3])
        y_g = np.array([0.05, 0.15, 0.35])  # Approximate integrals
        obs = gp4c.Observations(x_f=x, y_f=y_f, noise_f=0.01, x_g=x, y_g=y_g, noise_g=0.01)

        result = gp4c.log_marginal_likelihood(obs, sigma2=1.0, ell=1.0, kernel="rbf")
        assert isinstance(result, float)
        assert not np.isnan(result)
        assert not np.isinf(result)


class TestFit:
    """Test gp4c.fit() hyperparameter optimization."""

    def test_simple_rbf_optimization(self):
        """Test optimization with RBF kernel."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=42, maxiter=50
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0
        assert result.period is None  # RBF doesn't have period
        assert isinstance(result.log_likelihood, float)

    def test_matern52_optimization(self):
        """Test optimization with Matérn 5/2 kernel."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.fit(
            obs, kernel="matern52", method="differential_evolution", seed=42, maxiter=50
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0
        assert result.period is None

    def test_periodic_optimization(self):
        """Test optimization with periodic kernel."""
        np.random.seed(42)
        x = np.linspace(0, 10, 30)
        y = np.sin(2 * np.pi * x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.fit(
            obs,
            kernel="periodic",
            method="differential_evolution",
            seed=42,
            maxiter=100,
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0
        assert result.period is not None
        assert result.period > 0
        # Period should be close to 1.0 (true period)
        assert 0.5 < result.period < 2.0

    def test_derivative_observations_optimization(self):
        """Test optimization with derivative observations."""
        np.random.seed(42)
        x = np.linspace(0, 5, 15)
        y = np.sin(x) + np.random.normal(0, 0.05, len(x))
        dy = np.cos(x) + np.random.normal(0, 0.05, len(x))
        obs = gp4c.Observations(
            x_f=x, y_f=y, noise_f=0.01, x_h=x, y_h=dy, noise_h=0.01
        )

        result = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=42, maxiter=50
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0

    def test_custom_bounds(self):
        """Test optimization with custom bounds."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        custom_bounds = {"sigma2": (0.5, 2.0), "ell": (0.5, 2.0)}
        result = gp4c.fit(
            obs,
            kernel="rbf",
            bounds=custom_bounds,
            method="differential_evolution",
            seed=42,
            maxiter=50,
        )

        assert result.success
        # Check that results are within custom bounds
        assert 0.5 <= result.sigma2 <= 2.0
        assert 0.5 <= result.ell <= 2.0

    def test_local_optimizer_lbfgsb(self):
        """Test local optimization with L-BFGS-B."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        x0 = {"sigma2": 1.0, "ell": 1.0}
        result = gp4c.fit(
            obs, kernel="rbf", method="L-BFGS-B", x0=x0
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0

    def test_local_optimizer_nelder_mead(self):
        """Test local optimization with Nelder-Mead."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        x0 = {"sigma2": 1.0, "ell": 1.0}
        result = gp4c.fit(
            obs, kernel="rbf", method="Nelder-Mead", x0=x0
        )

        # Nelder-Mead might not always succeed with default settings
        assert result.sigma2 > 0
        assert result.ell > 0

    def test_local_optimizer_without_x0_raises(self):
        """Test that local optimizer without x0 raises error."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        with pytest.raises(ValueError, match="requires x0"):
            gp4c.fit(obs, kernel="rbf", method="L-BFGS-B")

    def test_invalid_kernel_raises(self):
        """Test that invalid kernel raises error."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        with pytest.raises(ValueError, match="Unknown kernel"):
            gp4c.fit(obs, kernel="invalid_kernel")

    def test_matern32_with_derivative_works(self):
        """Test that Matérn 3/2 with derivative observations works (once mean-square differentiable)."""
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.1, 0.2, 0.3])
        dy = np.array([0.1, 0.1, 0.1])
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01, x_h=x, y_h=dy, noise_h=0.01)

        result = gp4c.fit(obs, kernel="matern32")
        assert result is not None

    def test_not_observations_raises(self):
        """Test that non-Observations input raises error."""
        with pytest.raises(TypeError, match="must be an Observations instance"):
            gp4c.fit(
                {"x": [1, 2, 3], "y": [1, 2, 3]}, kernel="rbf"
            )

    def test_optimization_result_attributes(self):
        """Test that OptimizationResult has all required attributes."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=42, maxiter=50
        )

        assert hasattr(result, "sigma2")
        assert hasattr(result, "ell")
        assert hasattr(result, "period")
        assert hasattr(result, "log_likelihood")
        assert hasattr(result, "success")
        assert hasattr(result, "scipy_result")

    def test_reproducibility_with_seed(self):
        """Test that results are reproducible with seed."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result1 = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=123, maxiter=50
        )
        result2 = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=123, maxiter=50
        )

        assert np.allclose(result1.sigma2, result2.sigma2)
        assert np.allclose(result1.ell, result2.ell)
        assert np.allclose(result1.log_likelihood, result2.log_likelihood)

    def test_optimization_improves_likelihood(self):
        """Test that optimization improves likelihood over initial guess."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        # Initial guess
        sigma2_init = 0.5
        ell_init = 0.5
        log_lik_init = compute_log_marginal_likelihood(
            obs, sigma2_init, ell_init, 1.0, "rbf"
        )

        # Optimized
        result = gp4c.fit(
            obs, kernel="rbf", method="differential_evolution", seed=42, maxiter=100
        )

        # Optimized likelihood should be better (more positive)
        assert result.log_likelihood >= log_lik_init

    def test_dual_annealing_method(self):
        """Test optimization with dual_annealing method."""
        np.random.seed(42)
        x = np.linspace(0, 5, 20)
        y = np.sin(x) + np.random.normal(0, 0.1, len(x))
        obs = gp4c.Observations(x_f=x, y_f=y, noise_f=0.01)

        result = gp4c.fit(
            obs, kernel="rbf", method="dual_annealing", seed=42, maxiter=100
        )

        assert result.success
        assert result.sigma2 > 0
        assert result.ell > 0
