# Written by Jacob I. Monroe, NIST employee

"""Tests for GP models with derivatives and active learning based on those models."""

import copy

import gpflow
import numpy as np
import pytest

from thermoextrap.gpr_active import sine_active
from thermoextrap.gpr_active.active_utils import make_rbf_expr, train_GPR
from thermoextrap.gpr_active.gp_models import (
    ConstantMeanWithDerivs,
    DerivativeKernel,
    HeteroscedasticGPR,
    HetGaussianDeriv,
    LinearWithDerivs,
    SympyMeanFunc,
    multioutput_multivariate_normal,
)


# For derivative kernel, test if reproduces derivatives for known functional form and values
# Work with RBF and test up to second order derivatives
# Will build covariance matrix with derivative orders indexing, i.e., cov[0, 1] = dk/dx2
class RBFCovs:
    def __init__(self, var, volume) -> None:
        self.var = var
        self.volume = volume

    def d(self, x1, x2):
        return x1 / self.volume - x2 / self.volume

    # k(x1, x2) = var * exp(-0.5*(x1/volume - x2/volume)**2)
    def cov00(self, x1, x2):
        return self.var * np.exp(-0.5 * self.d(x1, x2) ** 2)

    # dk/dx2 = var * exp(-0.5*(x1/volume - x2/volume)**2) * ((1/volume)*(x1/volume - x2/volume)) = k(x1, x2) * ((1/volume)*(x1/volume - x2/volume))
    def cov01(self, x1, x2):
        return self.cov00(x1, x2) * ((1 / self.volume) * self.d(x1, x2))

    # dk/dx1 = -dk/dx2
    def cov10(self, x1, x2):
        return -self.cov01(x1, x2)

    # d^2k/dx1*dx2 = -dk/dx2 * ((1/volume)*(x1/volume - x2/volume)) + k(x1, x2) * (1/volume^2)
    def cov11(self, x1, x2):
        return -self.cov01(x1, x2) * ((1 / self.volume) * self.d(x1, x2)) + self.cov00(
            x1, x2
        ) * (1 / self.volume**2)

    # d^2k/dx1^2 = -dk/dx1 * ((1/volume)*(x1/volume - x2/volume)) - k(x1, x2) * (1/volume^2) = -d^2k/dx1*dx2
    def cov02(self, x1, x2):
        return -self.cov11(x1, x2)

    # d^2k/dx2^2 = dk/dx2 * ((1/volume)*(x1/volume - x2/volume)) - k(x1, x2) * (1/volume^2) = -d^2k/dx1*dx2
    def cov20(self, x1, x2):
        return -self.cov11(x1, x2)

    # d^3k/dx2^2*dx1 = d^2k/dx2*dx1 * ((1/volume)*(x1/volume - x2/volume)) - 2*dk/dx1 * (1/volume^2)
    def cov12(self, x1, x2):
        return self.cov11(x1, x2) * (
            (1 / self.volume) * self.d(x1, x2)
        ) + 2 * self.cov01(x1, x2) * (1 / self.volume**2)

    # d^3k/dx1^2*dx2 = -d^2k/dx1*dx2 * ((1/volume)*(x1/volume - x2/volume)) + 2*dk/dx1 * (1/volume^2) = -d^3k/dx2^2*dx1
    def cov21(self, x1, x2):
        return -self.cov12(x1, x2)

    # d^4k/dx1^2*dx2^2 = d^3k/dx1^2*dx2 * ((1/volume)*(x1/volume - x2/volume)) + d^2k/dx1*dx2 * (1/volume^2) - 2*d^2k/dx1^2 * (1/volume^2)
    def cov22(self, x1, x2):
        return -self.cov12(x1, x2) * (
            (1 / self.volume) * self.d(x1, x2)
        ) + 3 * self.cov11(x1, x2) * (1 / self.volume**2)

    def __call__(self, x1, x2):
        cov_mat = [
            [self.cov00, self.cov01, self.cov02],
            [self.cov10, self.cov11, self.cov12],
            [self.cov20, self.cov21, self.cov22],
        ]

        out = np.zeros((3, 3))
        for i in range(3):
            for j in range(3):
                out[i, j] = cov_mat[i][j](x1, x2)

        return out


# For multiple input dimensions with separate lengthscales, implement as product of kernels
# Derivatives will also be products since dimensions separable as products of exponentials
class MultiDRBFCovs:
    def __init__(self, rbf_cov_list) -> None:
        self.rbf_cov_list = rbf_cov_list
        self.n_dims = len(rbf_cov_list)

    def __call__(self, x1, x2):
        # Loop over dimensions and collect covariance matrices for each
        cov_mats = [self.rbf_cov_list[i](x1[i], x2[i]) for i in range(self.n_dims)]

        # Put together all combinations of covariances
        out = copy.deepcopy(cov_mats[0])
        for i in range(1, self.n_dims):
            new_out = [
                [c * cov_mats[i] for c in out[j, :]] for j in range(out.shape[0])
            ]
            out = np.block(new_out)

        return out


def test_deriv_kernel_manual_1d() -> None:
    # Create manually implemented RBF class to use as reference
    rbf_check = RBFCovs(1.0, 2.0)

    # Create a derivative kernel to check
    kern_expr, kern_params = make_rbf_expr()
    # Overwrite transforms on kernel parameters
    with pytest.raises(ValueError):
        deriv_kern = DerivativeKernel(
            kern_expr, 1, {"var": [1.0, {}], "volume": [2.0, {}]}
        )
    kern_params = {"var": [1.0, {}], "l_0": [2.0, {}]}
    deriv_kern = DerivativeKernel(kern_expr, 1, kernel_params=kern_params)

    # Want to work with a few points to check...
    # (0, 0), (1, 1), (1, 0), (0, 1), and (-1, 0)
    for x_pair in [(0.0, 0.0), (1.0, 1.0), (1.0, 0.0), (0.0, 1.0), (-1.0, 0.0)]:
        ref = rbf_check(x_pair[0], x_pair[1])
        this_x1 = np.vstack([x_pair[0] * np.ones(3), np.arange(3)]).T
        this_x2 = np.vstack([x_pair[1] * np.ones(3), np.arange(3)]).T
        to_check = deriv_kern.K(this_x1, this_x2).numpy()
        np.testing.assert_allclose(ref, to_check)

    # Next test with multiple points provided at once
    x_check = np.block(
        [
            [0.0 * np.ones((3, 1)), np.arange(3)[:, None]],
            [1.0 * np.ones((3, 1)), np.arange(3)[:, None]],
        ]
    )
    to_check = deriv_kern.K(x_check, x_check).numpy()
    ref = np.block(
        [
            [rbf_check(0.0, 0.0), rbf_check(0.0, 1.0)],
            [rbf_check(1.0, 0.0), rbf_check(1.0, 1.0)],
        ]
    )
    np.testing.assert_allclose(ref, to_check)


def test_deriv_kernel_manual_multi_d() -> None:
    # Create manually implemented RBF class to use as reference
    rbf_check = MultiDRBFCovs(
        [
            RBFCovs(1.0, 2.0),
            RBFCovs(1.0, 1.0),
        ]
    )

    # Create a derivative kernel to check
    kern_expr, kern_params = make_rbf_expr(2)
    # Overwrite transforms on kernel parameters
    kern_params = {
        "var": [1.0, {}],
        "l_0": [2.0, {}],
        "l_1": [1.0, {}],
    }
    deriv_kern = DerivativeKernel(kern_expr, 2, kernel_params=kern_params)

    # Want to work with a few points to check...
    # (0.0, 0.0) and (0.0, 0.0)
    # (1.0, 1.0) and (1.0, 1.0)
    # (0.0, 0.0) and (0.0, 1.0)
    # (1.0, 0.0) and (0.0, 0.0)
    # (0.0, 0.0) and (-1.0, -1.0)
    # Helpful to put together all combos of derivatives for each set of points
    d_combos = np.array([(i, j) for i in range(3) for j in range(3)])

    for x_pair in [
        [np.array([0.0, 0.0]), np.array([0.0, 0.0])],
        [np.array([1.0, 1.0]), np.array([1.0, 1.0])],
        [np.array([0.0, 0.0]), np.array([0.0, 1.0])],
        [np.array([1.0, 0.0]), np.array([0.0, 0.0])],
        [np.array([0.0, 0.0]), np.array([-1.0, -1.0])],
    ]:
        ref = rbf_check(x_pair[0], x_pair[1])
        this_x1 = np.hstack([np.tile(x_pair[0], (9, 1)), d_combos])
        this_x2 = np.hstack([np.tile(x_pair[1], (9, 1)), d_combos])
        to_check = deriv_kern.K(this_x1, this_x2).numpy()
        np.testing.assert_allclose(ref, to_check)

    # Next test with multiple points provided at once
    pt1 = np.array([0.0, 0.0])
    pt2 = np.array([1.0, 1.0])
    x_check = np.block(
        [
            [np.tile(pt1, (9, 1)), d_combos],
            [np.tile(pt2, (9, 1)), d_combos],
        ]
    )
    to_check = deriv_kern.K(x_check, x_check).numpy()
    ref = np.block(
        [
            [rbf_check(pt1, pt1), rbf_check(pt1, pt2)],
            [rbf_check(pt2, pt1), rbf_check(pt2, pt2)],
        ]
    )
    np.testing.assert_allclose(ref, to_check)


def test_deriv_kernel_self_1d() -> None:
    # Now test for self-consistency within the derivative kernel
    kern_expr, kern_params = make_rbf_expr()
    kern_params = {"var": [1.0, {}], "l_0": [2.0, {}]}
    deriv_kern = DerivativeKernel(kern_expr, 1, kernel_params=kern_params)

    # Create test input
    x_check = np.block(
        [
            [0.0 * np.ones((3, 1)), np.arange(3)[:, None]],
            [1.0 * np.ones((3, 1)), np.arange(3)[:, None]],
            [-1.0 * np.ones((3, 1)), np.arange(3)[:, None]],
        ]
    )

    # Test that splits x data correctly (told it that observation dims is 1)
    check_locs, check_deriv_order = deriv_kern._split_x_into_locs_and_deriv_info(
        x_check
    )
    np.testing.assert_allclose(check_locs, x_check[:, :1])
    np.testing.assert_allclose(check_deriv_order, x_check[:, 1:])

    # Generate full output to check against various equivalent scenarios
    output = deriv_kern.K(x_check, x_check).numpy()

    # Check that K_diag function same as diagonal of K
    output_diag = deriv_kern.K_diag(x_check).numpy()
    np.testing.assert_allclose(output_diag, np.diag(output))

    # Check that passing single x is same as passing twice
    output_single = deriv_kern.K(x_check)
    np.testing.assert_allclose(output_single, output)

    # Make sure that if re-order x inputs get same thing with different order
    x_check_reorder = np.roll(x_check, 3, axis=0)
    output_reorder = deriv_kern.K(x_check_reorder, x_check_reorder).numpy()
    output_reorder = np.roll(np.roll(output_reorder, -3, axis=0), -3, axis=1)
    np.testing.assert_allclose(output_reorder, output)

    # And if switch around derivatives and points
    x_check_deriv_order = x_check_reorder = np.roll(x_check, 1, axis=0)
    output_deriv_order = deriv_kern.K(x_check_deriv_order, x_check_deriv_order).numpy()
    output_deriv_order = np.roll(np.roll(output_deriv_order, -1, axis=0), -1, axis=1)
    np.testing.assert_allclose(output_deriv_order, output)


def test_deriv_kernel_self_multi_d() -> None:
    # Now test for self-consistency within the derivative kernel
    kern_expr, kern_params = make_rbf_expr(2)
    kern_params = {
        "var": [1.0, {}],
        "l_0": [2.0, {}],
        "l_1": [1.0, {}],
    }
    deriv_kern = DerivativeKernel(kern_expr, 2, kernel_params=kern_params)

    # Helpful to put together all combos of derivatives for each set of points
    d_combos = np.array([(i, j) for i in range(3) for j in range(3)])

    # Create test input
    pt1 = np.array([0.0, 0.0])
    pt2 = np.array([1.0, 1.0])
    pt3 = np.array([-1.0, 0.0])
    x_check = np.block(
        [
            [np.tile(pt1, (9, 1)), d_combos],
            [np.tile(pt2, (9, 1)), d_combos],
            [np.tile(pt3, (9, 1)), d_combos],
        ]
    )

    # Test that splits x data correctly (told it that observation dims is 2)
    check_locs, check_deriv_order = deriv_kern._split_x_into_locs_and_deriv_info(
        x_check
    )
    np.testing.assert_allclose(check_locs, x_check[:, :2])
    np.testing.assert_allclose(check_deriv_order, x_check[:, 2:])

    # Generate full output to check against various equivalent scenarios
    output = deriv_kern.K(x_check, x_check).numpy()

    # Check that K_diag function same as diagonal of K
    output_diag = deriv_kern.K_diag(x_check).numpy()
    np.testing.assert_allclose(output_diag, np.diag(output))

    # Check that passing single x is same as passing twice
    output_single = deriv_kern.K(x_check)
    np.testing.assert_allclose(output_single, output)

    # Make sure that if re-order x inputs get same thing with different order
    x_check_reorder = np.roll(x_check, 9, axis=0)
    output_reorder = deriv_kern.K(x_check_reorder, x_check_reorder).numpy()
    output_reorder = np.roll(np.roll(output_reorder, -9, axis=0), -9, axis=1)
    np.testing.assert_allclose(output_reorder, output)

    # And if switch around derivatives and points
    x_check_deriv_order = np.roll(x_check, 1, axis=0)
    output_deriv_order = deriv_kern.K(x_check_deriv_order, x_check_deriv_order).numpy()
    output_deriv_order = np.roll(np.roll(output_deriv_order, -1, axis=0), -1, axis=1)
    np.testing.assert_allclose(output_deriv_order, output)


# Testing mean functions is straightforward
# First define some data to use for testing
def mean_funcs_check_data():
    # Define points to check at
    x_check = np.block(
        [
            [0.0 * np.ones((3, 1)), np.arange(3)[:, None]],
            [0.5 * np.ones((3, 1)), np.arange(3)[:, None]],
            [1.0 * np.ones((3, 1)), np.arange(3)[:, None]],
        ]
    )

    # Will use y=x^2, so just have that at zero order
    # Then for first order derivative, it's just 2*x
    # Then just 2
    y_check_sq = np.zeros((x_check.shape[0], 1))
    for i in range(3):
        this_inds = np.where(x_check[:, 1] == i)[0]
        this_vals = x_check[this_inds, 0] ** (2 - i)
        if i > 0:
            this_vals *= 2.0
        y_check_sq[this_inds, 0] = this_vals

    # But to check multidimensional data, also provide just linear function of y
    # (as another column)
    y_check_lin = np.zeros((x_check.shape[0], 1))
    y_check_lin[::3, 0] = x_check[::3, 0]
    y_check_lin[1::3, 0] = 1.0

    y_check = np.hstack([y_check_lin, y_check_sq])

    return x_check, y_check


# For constant, just make sure that zeroth order derivatives have constant and all others are zero
def test_constant_mean_func() -> None:
    x_check, y_check = mean_funcs_check_data()

    # Generate constant mean function to test
    m_func = ConstantMeanWithDerivs(y_check[::3])

    # Make sure mean is computed correctly
    ref_avg = np.average(y_check[::3, :], axis=0)
    np.testing.assert_allclose(m_func.c, ref_avg)

    # And when call make sure get constant at order zero and zero otherwise
    check_out = m_func(x_check).numpy()
    zero_bool = x_check[:, 1] == 0
    nonzero_bool = x_check[:, 1] != 0
    np.testing.assert_array_equal(
        check_out[zero_bool], ref_avg * np.ones_like(check_out[zero_bool])
    )
    np.testing.assert_array_equal(
        check_out[nonzero_bool], np.zeros_like(check_out[nonzero_bool])
    )


# For linear, just make sure it returns correct linear fit for provided test case
def test_linear_mean_func() -> None:
    x_check, y_check = mean_funcs_check_data()

    # Generate linear mean function to test
    m_func = LinearWithDerivs(x_check[::3, :1], y_check[::3])

    # Will check over whole range at each order
    x_full_range = np.linspace(np.min(x_check), np.max(x_check), 100)

    # For checking against quadratic data (second column of y_check), compare to scipy
    from scipy.stats import linregress

    sp_regress = linregress(x_check[::3, 0], y_check[::3, 1])
    sp_regress_check = sp_regress.slope * x_full_range + sp_regress.intercept

    # Loop over orders
    for i in range(3):
        this_x = np.vstack([x_full_range, i * np.ones_like(x_full_range)]).T
        this_output = m_func(this_x)
        if i == 0:
            # For linear function input, fit should be exact (so for first column of y_check)
            np.testing.assert_allclose(this_output[:, 0], x_full_range)
            np.testing.assert_allclose(this_output[:, 1], sp_regress_check)
        elif i == 1:
            np.testing.assert_array_equal(
                this_output, m_func.slope * np.ones_like(this_output)
            )
        else:
            np.testing.assert_array_equal(this_output, np.zeros_like(this_output))


# Testing sympy mean function is more difficult - mainly testing sympy stuff
def test_sympy_mean_func() -> None:
    import sympy as sp

    # Will work with logistic because relevant derivatives can be computed in terms of 0 order
    x = sp.symbols("x")
    m = sp.symbols("m")
    b = sp.symbols("b")
    sig_expr = 1.0 / (1.0 + sp.exp(-m * (x + b)))
    params = {
        "m": 1.0,
        "b": 0.0,
    }

    # Create some x data
    x_vals = np.linspace(-10.0, 10.0, 10)
    x_check = [np.vstack([x_vals, d_o * np.ones_like(x_vals)]).T for d_o in range(3)]
    x_check = np.vstack(x_check)

    # And some y data
    m_check = 0.26
    b_check = 3.4
    y_vals = 1.0 / (1.0 + np.exp(-m_check * (x_vals + b_check)))
    # Will not use higher order info, so just set to zeros
    y_check = np.zeros((x_check.shape[0], 1))
    y_check[: len(x_vals), 0] = y_vals

    # Create our mean function to test
    check_sym = SympyMeanFunc(
        sig_expr, x_check[: len(x_vals), :1], y_check[: len(x_vals)], params=params
    )

    # Check that expression matches input
    assert sp.simplify(check_sym.expr - sig_expr) == 0
    # Check that found optimal parameters
    np.testing.assert_allclose(check_sym.m, m_check, rtol=1e-06)
    np.testing.assert_allclose(check_sym.b, b_check, rtol=1e-06)
    # And check values and derivatives
    output = check_sym(x_check)
    np.testing.assert_allclose(
        output[:10].numpy(),
        sp.lambdify(x, sig_expr.subs(((m, m_check), (b, b_check))))(x_check[:10, :1]),
        rtol=1e-04,
    )
    np.testing.assert_allclose(
        output[10:20].numpy(), check_sym.m * (output[:10] - output[:10] ** 2).numpy()
    )
    np.testing.assert_allclose(
        output[20:].numpy(),
        (check_sym.m**2)
        * (output[:10] - 3.0 * output[:10] ** 2 + 2.0 * output[:10] ** 3),
    )


# Need to test all components of the GP model
# Start with testing the multioutput_multivariate_normal function
def test_multiout_multivar_normal() -> None:
    # Idea of this function is to parallelize assessment of multivariate Gaussians in tensorflow
    # First need valid covariance matrices
    cov1 = np.array([[1.0, 0.0], [0.0, 1.0]])
    cov2 = np.array([[0.5, 0.0], [0.0, 2.0]])
    cov3 = np.array([[0.5, 0.0], [-2.5, 3.5]])
    cov_mats = np.array([cov1, cov2, cov3.T @ cov3])

    # Generate Cholesky decomposition, which is required input for
    cov_cholesky = np.array([np.linalg.cholesky(c) for c in cov_mats])

    # And need means and points to compute log probability at
    means = np.array([[0.0, 0.5, -1.0], [1.0, 1.0, 1.0]])
    points = np.array([[0.5, 1.0, 1.0], [1.0, 0.5, -1.0]])

    # Generate output to check
    check_out = multioutput_multivariate_normal(points, means, cov_cholesky).numpy()

    # Will compare to a loop with scipy's implementation of log probabilities for Gaussians
    from scipy.stats import multivariate_normal

    for i in range(cov_mats.shape[0]):
        this_dist = multivariate_normal(mean=means[:, i], cov=cov_mats[i, ...])
        ref_logp = this_dist.logpdf(points[:, i])
        np.testing.assert_allclose(ref_logp, check_out[i])


def test_gp_likelihood() -> None:
    # Check that handles input correctly
    cov_id = np.eye(3)
    d_orders = np.arange(3.0)[:, None]  # Add dimension since now handles n-d inputs
    x = np.hstack([d_orders, d_orders])
    check = HetGaussianDeriv(cov_id, 1)
    check_flat = HetGaussianDeriv(np.diag(cov_id), 1)
    np.testing.assert_allclose(check.cov, check_flat.cov)

    # Including if input is 1D
    check_1d = HetGaussianDeriv(cov_id[:1, :1], 1)
    check_1d_flat = HetGaussianDeriv(cov_id[0, :1], 1)
    np.testing.assert_allclose(check_1d.cov, check_1d_flat.cov)

    # Check that applies scaling appropriately for different scenarios above
    np.testing.assert_allclose(
        check.build_scaled_cov_mat(x).numpy(),
        check_flat.build_scaled_cov_mat(x).numpy(),
    )
    np.testing.assert_allclose(
        check_1d.build_scaled_cov_mat(x[:1]).numpy(),
        check_1d_flat.build_scaled_cov_mat(x[:1]).numpy(),
    )

    # Now check specifics of building covariance matrix
    # Check using alternative calculation
    def log_cov_model(cov0, p, s, d_o):
        d_o = d_o + 1  # noqa: PLR6104
        return np.log(cov0) + p * np.add(*np.meshgrid(d_o, d_o)) + s

    # Default has p=10.0, s=0.0
    cov_ = np.array([[0.5, 0.0, 0.0], [-1.5, 2.5, 0.0], [3.5, 4.5, 5.5]])
    cov = cov_.T @ cov_
    check = HetGaussianDeriv(cov, 1)
    np.testing.assert_allclose(
        log_cov_model(cov, 10.0, 0.0, d_orders),
        np.log(check.build_scaled_cov_mat(x).numpy()),
    )
    # Also do for new model with p and s changed
    check_newps = HetGaussianDeriv(cov, 1, p=2.0, s=-2.0)
    np.testing.assert_allclose(
        log_cov_model(cov, 2.0, -2.0, d_orders),
        np.log(check_newps.build_scaled_cov_mat(x).numpy()),
    )


# Testing full GP model with derivative information
# This test method is quite large and slow
# Can consider turning into a test class with methods inside
# Would allow for most comparisons to a "base" model created in __init__
# But would make testing more modular and specific
# Manually parsing and running all methods in class is a pain for testing, though
@pytest.mark.slow
def test_gp() -> None:  # noqa: PLR0915
    # First create data we can use
    rng = np.random.default_rng(42)
    x_data, y_data, y_var_data = sine_active.make_data(
        np.linspace(-np.pi, np.pi, 5), noise=0.01, max_order=2, rng=rng
    )
    cov_data = np.diag(np.squeeze(y_var_data))[
        None, ...
    ]  # Give first dimension of 1 since 1D output

    # Next specify parameters for GP model, with kernel produced by a function
    # Want to make sure no parameters shared across models we test
    def make_kern():
        kern_expr, kern_params = make_rbf_expr()
        return DerivativeKernel(kern_expr, 1, kernel_params=kern_params)

    like_kwargs = {"p": 1.0, "s": -2.0}

    # Create GP model with just the 1D data
    check_1d = HeteroscedasticGPR(
        (x_data, y_data, cov_data), kernel=make_kern(), likelihood_kwargs=like_kwargs
    )

    # Check to make sure likelihood handled correctly
    ref_like = HetGaussianDeriv(
        cov_data,
        1,
        **like_kwargs,
    )
    np.testing.assert_allclose(ref_like.cov, check_1d.likelihood.cov)
    np.testing.assert_allclose(
        ref_like.build_scaled_cov_mat(x_data),
        check_1d.likelihood.build_scaled_cov_mat(x_data),
    )

    # Since working with 1D data, scaling should not change the trained model
    # Also should not change the kernel lengthscale
    # And variance should change according to scale factor squared
    # Or for x scaling, changes lengthscale proportionally, but not variance
    check_scale = HeteroscedasticGPR(
        (x_data, y_data, cov_data),
        kernel=make_kern(),
        scale_fac=2.0,
    )
    # Removing tests with x scaling factor
    # Essentially unused previously and not helpful, just confusing
    # check_xscale = HeteroscedasticGPR(
    #     (x_data, y_data, cov_data),
    #     kernel=make_kern(),
    #     x_scale_fac=2.0,
    # )
    check_base = HeteroscedasticGPR(
        (x_data, y_data, cov_data),
        kernel=make_kern(),
    )
    train_GPR(check_base)
    train_GPR(check_scale)
    # train_GPR(check_xscale)

    x_test = np.linspace(-np.pi, np.pi, 100)[:, None]
    x_test = np.block([[x_test, np.zeros_like(x_test)], [x_test, np.ones_like(x_test)]])

    pred_base = check_base.predict_f(x_test)
    pred_scale = check_base.predict_f(x_test)
    # pred_xscale = check_base.predict_f(x_test)

    # Comparing predictions
    np.testing.assert_allclose(pred_base[0], pred_scale[0])
    np.testing.assert_allclose(pred_base[1], pred_scale[1])
    # np.testing.assert_allclose(pred_base[0], pred_xscale[0])
    # np.testing.assert_allclose(pred_base[1], pred_xscale[1])

    # Comparing parameters
    np.testing.assert_allclose(
        check_base.kernel.kernel.l_0.numpy(),
        check_scale.kernel.kernel.l_0.numpy(),
        rtol=1e-03,
    )
    np.testing.assert_allclose(
        check_base.kernel.kernel.var.numpy(),
        check_scale.kernel.kernel.var.numpy() * (check_scale.scale_fac**2),
        rtol=1e-03,
    )
    # np.testing.assert_allclose(
    #     check_base.kernel.kernel.volume.numpy(),
    #     check_xscale.kernel.kernel.volume.numpy() / (check_xscale.x_scale_fac),
    #     rtol=1e-03,
    # )
    # np.testing.assert_allclose(
    #     check_base.kernel.kernel.var.numpy(),
    #     check_xscale.kernel.kernel.var.numpy(),
    #     rtol=1e-03,
    # )

    # Define points for testing model prediction and compare to reference
    # This will test for correct behavior across operating systems and software versions
    x_val = np.array([-3.0, 0.5, 3.0])[:, None]
    x_val = np.block([[x_val, np.zeros_like(x_val)], [x_val, np.ones_like(x_val)]])

    pred_val_base = check_base.predict_f(x_val)
    ref_pred_means = np.array(
        [
            [-0.12116922],
            [0.54235242],
            [0.02270773],
            [-1.12612318],
            [0.87582276],
            [-1.10413224],
        ]
    )
    ref_pred_vars = np.array(
        [
            [0.00460669],
            [0.00413819],
            [0.01032148],
            [0.00841806],
            [0.00238915],
            [0.01527506],
        ]
    )
    np.testing.assert_allclose(pred_val_base[0].numpy(), ref_pred_means, rtol=1e-03)
    np.testing.assert_allclose(pred_val_base[1].numpy(), ref_pred_vars, rtol=1e-03)

    # Test handling of multiple outputs (extra y dimensions)
    _x_shift, y_shift, y_var_shift = sine_active.make_data(
        np.linspace(-np.pi, np.pi, 5),
        phase_shift=np.pi / 3.0,
        noise=0.01,
        max_order=2,
        rng=rng,
    )
    cov_shift = np.diag(np.squeeze(y_var_shift))[None, ...]
    check_multi_d = HeteroscedasticGPR(
        (
            x_data,
            np.hstack([y_data, y_shift]),
            np.concatenate([cov_data, cov_shift], axis=0),
        ),
        kernel=make_kern(),
    )

    # Make sure wrapping works into SharedIndependent kernel works correctly
    assert isinstance(check_multi_d.kernel, gpflow.kernels.SharedIndependent)

    # Train
    train_GPR(check_multi_d)

    # Produce output and compare to below - will depend on optimization procedure
    # But better than nothing?
    pred_val_multi_d = check_multi_d.predict_f(x_val)
    ref_multi_d_means = np.array(
        [
            [-0.12177132, -0.8197274],
            [0.54288101, 0.97107731],
            [0.02439859, -0.74359773],
            [-1.11681105, -0.42663647],
            [0.87410287, -0.01394911],
            [-1.0930425, -0.48817702],
        ]
    )
    ref_multi_d_vars = np.array(
        [
            [0.00461916, 0.00477108],
            [0.00430657, 0.00396219],
            [0.01033009, 0.00807202],
            [0.00865764, 0.00895764],
            [0.00256747, 0.00244668],
            [0.01557182, 0.01329149],
        ]
    )
    np.testing.assert_allclose(
        pred_val_multi_d[0].numpy(), ref_multi_d_means, rtol=1e-03
    )
    np.testing.assert_allclose(
        pred_val_multi_d[1].numpy(), ref_multi_d_vars, rtol=1e-03
    )

    # Make sure multiple outputs also work with separate independent kernels
    check_sep_ind = HeteroscedasticGPR(
        (
            x_data,
            np.hstack([y_data, y_shift]),
            np.concatenate([cov_data, cov_shift], axis=0),
        ),
        kernel=gpflow.kernels.SeparateIndependent([make_kern(), make_kern()]),
    )

    # With independent kernels, can compare first dimension to model trained on 1D data only
    train_GPR(check_sep_ind)
    pred_sep_ind = check_sep_ind.predict_f(x_test)
    np.testing.assert_allclose(pred_base[0], pred_sep_ind[0][:, :1], rtol=1e-03)
    np.testing.assert_allclose(pred_base[1], pred_sep_ind[1][:, :1], rtol=1e-03)
    np.testing.assert_allclose(
        check_base.kernel.kernel.l_0.numpy(),
        check_sep_ind.kernel.kernels[0].l_0.numpy(),
        rtol=1e-03,
    )
    np.testing.assert_allclose(
        check_base.kernel.kernel.var.numpy(),
        check_sep_ind.kernel.kernels[0].var.numpy(),
        rtol=1e-03,
    )

    # Test output with full_cov (specifically variance, second output)
    # Best done with multidimensional output data
    pred_ff = check_sep_ind.predict_f(x_data, full_cov=False)[1]  # Default
    pred_tf = check_sep_ind.predict_f(x_data, full_cov=True)[1]

    # Assert correct shapes
    assert pred_ff.shape == (x_data.shape[0], 2)
    assert pred_tf.shape == (2, x_data.shape[0], x_data.shape[0])

    # And assert that diagonal of full_cov=True matches full_cov=False
    np.testing.assert_allclose(
        pred_ff.numpy(), np.vstack([np.diag(v) for v in pred_tf]).T
    )

    # Check proper handling of mean functions
    # Will assume mean function itself correct, so just checking how model uses it
    # Easy enough to pass in constant mean function and make sure model output very similar
    # Does change predictions, just shouldn't change by much
    # To work, need to make sure no scaling happens to covariance matrix
    mean_func = ConstantMeanWithDerivs(np.ones((3, 1)))
    check_meanf = HeteroscedasticGPR(
        (x_data, y_data, cov_data),
        kernel=make_kern(),
        mean_function=mean_func,
        likelihood_kwargs={"p": 0.0, "transform_p": None},
    )

    check_meanf.kernel.kernel.l_0.assign(1e-06)
    np.testing.assert_allclose(y_data, check_meanf.predict_f(x_data)[0], atol=1e-01)

    train_GPR(check_meanf)
    pred_meanf = check_meanf.predict_f(x_test)

    # Can compare predictions and lengthscale parameter, but not kernel variance parameter
    np.testing.assert_allclose(pred_base[0], pred_meanf[0], atol=1e-01)
    np.testing.assert_allclose(pred_base[1], pred_meanf[1], atol=1e-01)
    np.testing.assert_allclose(
        check_base.kernel.kernel.l_0.numpy(),
        check_meanf.kernel.kernel.l_0.numpy(),
        atol=2e-01,
    )
