# Copyright 2024 The GPJax Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Tests that core GPJax operations produce correct dtypes under float32 and
float64.  The global ``jax_enable_x64`` flag is toggled per-test via the
*enable_x64* fixture so these tests exercise both precision modes.
"""

from gpjax.dataset import Dataset
from gpjax.gps import Prior
from gpjax.kernels.stationary import RBF, Matern32
from gpjax.likelihoods import Gaussian
from gpjax.mean_functions import Constant
from gpjax.objectives import conjugate_mll
from jax import config
import jax.numpy as jnp
import pytest

config.update("jax_enable_x64", True)

KERNELS = [RBF, Matern32]


@pytest.fixture(params=[jnp.float32, jnp.float64], ids=["float32", "float64"])
def enable_x64(request):
    """Toggle x64 mode and yield the target dtype."""
    dtype = request.param
    use_x64 = dtype == jnp.float64
    prev = config.jax_enable_x64
    config.update("jax_enable_x64", use_x64)
    try:
        yield dtype
    finally:
        config.update("jax_enable_x64", prev)


# --------------------------------------------------------------------------- #
#  Kernel dtype tests                                                          #
# --------------------------------------------------------------------------- #


@pytest.mark.parametrize("KernelClass", KERNELS)
def test_kernel_gram_dtype(KernelClass, enable_x64):
    dtype = enable_x64
    x = jnp.linspace(0, 1, 10, dtype=dtype)[:, None]
    kernel = KernelClass()
    Kxx = kernel.gram(x).to_dense()
    assert Kxx.dtype == dtype, f"Expected {dtype}, got {Kxx.dtype}"


@pytest.mark.parametrize("KernelClass", KERNELS)
def test_kernel_cross_covariance_dtype(KernelClass, enable_x64):
    dtype = enable_x64
    x = jnp.linspace(0, 1, 10, dtype=dtype)[:, None]
    y = jnp.linspace(0.5, 1.5, 5, dtype=dtype)[:, None]
    kernel = KernelClass()
    Kxy = kernel.cross_covariance(x, y)
    assert Kxy.dtype == dtype, f"Expected {dtype}, got {Kxy.dtype}"


# --------------------------------------------------------------------------- #
#  GP prior dtype tests                                                        #
# --------------------------------------------------------------------------- #


def test_prior_predict_dtype(enable_x64):
    dtype = enable_x64
    x = jnp.linspace(0, 1, 10, dtype=dtype)[:, None]
    prior = Prior(mean_function=Constant(), kernel=RBF())
    dist = prior.predict(x)
    assert dist.mean.dtype == dtype, (
        f"Mean dtype: expected {dtype}, got {dist.mean.dtype}"
    )


# --------------------------------------------------------------------------- #
#  GP posterior dtype tests                                                     #
# --------------------------------------------------------------------------- #


@pytest.mark.filterwarnings("ignore:X is not of type float64:UserWarning")
@pytest.mark.filterwarnings("ignore:y is not of type float64:UserWarning")
def test_posterior_predict_dtype(enable_x64):
    dtype = enable_x64
    x_train = jnp.linspace(0, 1, 20, dtype=dtype)[:, None]
    y_train = jnp.sin(x_train)
    D = Dataset(X=x_train, y=y_train)

    x_test = jnp.linspace(0, 1, 5, dtype=dtype)[:, None]
    prior = Prior(mean_function=Constant(), kernel=RBF())
    likelihood = Gaussian(num_datapoints=D.n)
    posterior = prior * likelihood

    dist = posterior.predict(x_test, D)
    assert dist.mean.dtype == dtype, (
        f"Mean dtype: expected {dtype}, got {dist.mean.dtype}"
    )


# --------------------------------------------------------------------------- #
#  Objective dtype tests                                                       #
# --------------------------------------------------------------------------- #


@pytest.mark.filterwarnings("ignore:X is not of type float64:UserWarning")
@pytest.mark.filterwarnings("ignore:y is not of type float64:UserWarning")
def test_conjugate_mll_dtype(enable_x64):
    dtype = enable_x64
    x = jnp.linspace(0, 1, 20, dtype=dtype)[:, None]
    y = jnp.sin(x)
    D = Dataset(X=x, y=y)

    prior = Prior(mean_function=Constant(), kernel=RBF())
    likelihood = Gaussian(num_datapoints=D.n)
    posterior = prior * likelihood

    mll = conjugate_mll(posterior, D)
    assert mll.dtype == dtype, f"MLL dtype: expected {dtype}, got {mll.dtype}"
