# Written by Jacob I. Monroe, NIST employee

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

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import gpflow
import numpy as np
import pytest
import sympy as sp
from scipy import linalg

from thermoextrap import idealgas
from thermoextrap.gpr_active import active_utils, gp_models, ig_active

if TYPE_CHECKING:
    from cmomy.resample.typing import SamplerType
    from numpy.random import Generator

# For now, not testing DataWrapper or SimWrapper objects
# These will likely change in future iterations and be highly dependent on individual problems


@pytest.fixture
def rng() -> Generator:
    return np.random.default_rng(42)


@pytest.fixture
def sampler(rng: Generator) -> SamplerType:
    return {"nrep": 100, "rng": rng}


# First test some simple utility functions, starting with sympy expressions for RBF
def test_rbf_expr() -> None:
    l_sym = sp.symbols("l_0", real=True)
    var_sym = sp.symbols("var", real=True)
    x1_sym = sp.symbols("x1_0", real=True)
    x2_sym = sp.symbols("x2_0", real=True)

    check_expr, check_params = active_utils.make_rbf_expr()
    assert l_sym in check_expr.free_symbols
    assert var_sym in check_expr.free_symbols
    assert x1_sym in check_expr.free_symbols
    assert x2_sym in check_expr.free_symbols
    assert "l_0" in check_params
    assert "var" in check_params

    def rbf(var, volume, x1, x2):
        return var * np.exp(-0.5 * (x1 / volume - x2 / volume) ** 2)

    check_var = np.array([0.0, 1.0, 2.0, -1.0])
    check_l = np.array([1.0, 2.0])
    check_x = np.array([0.0, 0.5, 1.0, 2.0, -0.5, -1.0, -2.0])
    for v in check_var:
        for volume in check_l:
            for x1 in check_x:
                for x2 in check_x:
                    out_expr = check_expr.subs(
                        [(var_sym, v), (l_sym, volume), (x1_sym, x1), (x2_sym, x2)]
                    )
                    np.testing.assert_allclose(float(out_expr), rbf(v, volume, x1, x2))


# Next test function for building GP inputs from thermoextrap ExtrapModel objects
# (representing thermodynamic states)
def test_make_gp_input(rng: Generator, sampler: SamplerType) -> None:
    beta = 5.6
    state = ig_active.extrap_IG(beta, rng=rng)

    # Test without logarithm on x
    check_x, check_y, check_cov = active_utils.input_GP_from_state(
        state, sampler=sampler
    )
    # Check shapes
    assert check_x.shape == (state.order + 1, 2)
    assert check_y.shape == (state.order + 1, 1)
    assert check_cov.shape == (1, state.order + 1, state.order + 1)
    # Check derivatives calculated
    for k in range(3):
        np.testing.assert_allclose(
            check_y[k], idealgas.dbeta_xave(k)(beta, 1.0), rtol=(10**k) * 1e-02
        )

    # Test with logarithm on x
    # Note that using base 10 logarithm
    # Using sympy to compare because have analytical expression for ideal gas
    def dxave_dlogbeta(k):
        logbeta_sym = sp.symbols("lnbeta_sym")
        log_expr = idealgas.xave_sym.subs([(idealgas.beta_sym, 10 ** (logbeta_sym))])
        deriv = sp.diff(log_expr, logbeta_sym, k)
        return sp.lambdify([logbeta_sym, idealgas.vol_sym], deriv, modules="numpy")

    check_x_log, check_y_log, check_cov_log = active_utils.input_GP_from_state(
        state,
        log_scale=True,
        sampler=sampler,
    )
    np.testing.assert_equal(np.log10(check_x[:, 0]), check_x_log[:, 0])
    assert check_x_log.shape == (state.order + 1, 2)
    assert check_y_log.shape == (state.order + 1, 1)
    assert check_cov_log.shape == (1, state.order + 1, state.order + 1)
    for k in range(3):
        np.testing.assert_allclose(
            check_y_log[k],
            dxave_dlogbeta(k)(np.log10(beta), 1.0),
            rtol=(10**k) * 1e-02,
        )

    # Test multiple output dimensions
    def dbeta_xsqave(k):
        beta_sym = sp.symbols("beta_sym")
        xsqave_sym = (
            (2 / (beta_sym**2))
            - (2 / (beta_sym * (sp.exp(beta_sym) - 1)))
            - (1 / (sp.exp(beta_sym) - 1))
        )
        deriv = sp.diff(xsqave_sym, beta_sym, k)
        return sp.lambdify(
            [
                beta_sym,
            ],
            deriv,
            modules="numpy",
        )

    state_mult = ig_active.multiOutput_extrap_IG(beta, rng=rng)
    check_x_mult, check_y_mult, check_cov_mult = active_utils.input_GP_from_state(
        state_mult,
        sampler=sampler,
    )
    assert check_x_mult.shape == (state.order + 1, 2)
    assert check_y_mult.shape == (state.order + 1, 2)
    assert check_cov_mult.shape == (2, state.order + 1, state.order + 1)
    for k in range(3):
        np.testing.assert_allclose(
            check_y_mult[k, 0],
            idealgas.dbeta_xave(k)(beta, 1.0),
            rtol=(10**k) * 1e-02,
        )
        np.testing.assert_allclose(
            check_y_mult[k, 1], dbeta_xsqave(k)(beta), rtol=(10**k) * 1e-02
        )


# Next want to test creation of a GP model
def test_base_gp_creation(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    raw_x_data = []
    raw_y_data = []
    raw_cov_data = []
    for beta in (1.0, 5.6, 9.0):
        s = ig_active.extrap_IG(beta, rng=rng)
        this_x, this_y, this_cov = active_utils.input_GP_from_state(s, sampler=sampler)
        raw_x_data.append(this_x)
        raw_y_data.append(this_y)
        raw_cov_data.append(this_cov)

    # Mainly testing options, so testing decision structure
    # Test using 2 data points
    x_data = np.vstack([raw_x_data[0], raw_x_data[-1]])
    y_data = np.vstack([raw_y_data[0], raw_y_data[-1]])
    cov_data = linalg.block_diag(*[raw_cov_data[0][0, ...], raw_cov_data[-1][0, ...]])
    data_input_2p = (x_data, y_data, cov_data)

    order = int(np.max(x_data[:, 1]))

    check_gp = active_utils.create_base_GP_model(
        data_input_2p, d_order_ref=0, shared_kernel=True
    )
    # Should have constant mean function
    assert isinstance(check_gp.mean_function, gp_models.ConstantMeanWithDerivs)
    # And should scale by the std of the y inputs - with 2 data points, that's just the distance to the mean
    # Also taking zeroth order data because should be 0 order for d_order_ref
    np.testing.assert_allclose(
        check_gp.scale_fac, np.abs(y_data[0, :] - np.mean(y_data[:: (order + 1), :]))
    )

    # Test using 3 data points
    data_input_3p = (
        np.vstack(raw_x_data),
        np.vstack(raw_y_data),
        linalg.block_diag(*[dat[0, ...] for dat in raw_cov_data]),
    )
    check_gp_3p = active_utils.create_base_GP_model(
        data_input_3p, d_order_ref=0, shared_kernel=True
    )
    # Should be same as for 2 points but use linear mean function
    assert isinstance(check_gp_3p.mean_function, gp_models.LinearWithDerivs)
    # And check that scaling factor is std of zeroth-order data
    np.testing.assert_allclose(
        check_gp_3p.scale_fac,
        np.std(
            np.vstack(raw_y_data)[:: (order + 1), :]
            - check_gp_3p.mean_function(np.vstack(raw_x_data)[:: (order + 1), :]),
            axis=0,
        ),
    )

    # Or with 1 data point
    data_input_1p = (
        np.vstack([raw_x_data[0]]),
        np.vstack([raw_y_data[0]]),
        linalg.block_diag(*[raw_cov_data[0][0, ...]]),
    )
    check_gp_1p = active_utils.create_base_GP_model(
        data_input_1p, d_order_ref=0, shared_kernel=True
    )
    assert isinstance(check_gp_1p.mean_function, gp_models.ConstantMeanWithDerivs)
    # And scaling factor should be 1
    np.testing.assert_allclose(check_gp_1p.scale_fac, 1.0)

    # Test if d_order_ref is 1, not 0
    check_gp_d1 = active_utils.create_base_GP_model(
        data_input_2p, d_order_ref=1, shared_kernel=True
    )
    assert isinstance(check_gp_d1.mean_function, gp_models.ConstantMeanWithDerivs)
    # And should use 1st derivative to determine scaling instead of zeroth order
    np.testing.assert_allclose(
        check_gp_d1.scale_fac,
        np.abs(y_data[1, :] - np.mean(y_data[1 :: (order + 1), :])),
    )

    # Check properties of kernel
    # Make sure using shared kernel if shared_kernel=True
    assert isinstance(check_gp.kernel, gpflow.kernels.SharedIndependent)
    # Should be an RBF kernel (for default inputs, but just check here)
    # (note default kernel is SharedIndependent, so referencing kernel.kernel)
    assert (
        sp.simplify(
            check_gp.kernel.kernel.kernel_expr - active_utils.make_rbf_expr()[0]  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
        )
        == 0
    )

    # If shared_kernel is False, check to make sure uses SeparateIndependent
    check_gp_sep = active_utils.create_base_GP_model(
        data_input_2p, d_order_ref=0, shared_kernel=False
    )
    assert isinstance(check_gp_sep.kernel, gpflow.kernels.SeparateIndependent)
    assert (
        sp.simplify(
            check_gp_sep.kernel.kernels[0].kernel_expr - active_utils.make_rbf_expr()[0]  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
        )
        == 0
    )
    for k in check_gp_sep.kernel.kernels[1:]:
        for j, t_par in enumerate(k.trainable_parameters):
            assert check_gp_sep.kernel.kernels[0].trainable_parameters[j] != t_par

    # Test if can pass in full kernel
    # Note instantiation of kernel so is object, not class
    # And note that checking to make sure gets passed through to HeteroscedasticGPR
    # and wrapped as SharedIndependent, even though shared_kernel is False
    # Warning should be printed, so should come up with some way to check for that...
    k_rbf = active_utils.RBFDerivKernel()
    # Changing parameter values to check if passed faithfully
    # pylint: disable=no-member
    k_rbf.l_0.assign(0.5)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
    k_rbf.var.assign(5.0)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
    check_kernel = active_utils.create_base_GP_model(
        data_input_1p, d_order_ref=0, shared_kernel=False, kernel=k_rbf
    )
    assert isinstance(check_kernel.kernel, gpflow.kernels.SharedIndependent)
    assert check_kernel.kernel.kernel == k_rbf
    np.testing.assert_allclose(check_kernel.kernel.kernel.l_0.numpy(), 0.5)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
    np.testing.assert_allclose(check_kernel.kernel.kernel.var.numpy(), 5.0)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]


# Simple test for checking training of GP model
# Test is a bit slow, though
@pytest.mark.slow
def test_train_gp(sampler: SamplerType) -> None:
    # Will compare training results to a reference
    # Need data to work with
    raw_x_data = []
    raw_y_data = []
    raw_cov_data = []
    for beta in (1.0, 5.6, 9.0):
        s = ig_active.extrap_IG(
            beta, rng=np.random.default_rng(42)
        )  # NOTE: need this to produce results below...
        this_x, this_y, this_cov = active_utils.input_GP_from_state(s, sampler=sampler)
        raw_x_data.append(this_x)
        raw_y_data.append(this_y)
        raw_cov_data.append(this_cov)

    # Mainly testing options, so testing decision structure
    # Test using 2 data points
    x_data = np.vstack(raw_x_data)
    y_data = np.vstack(raw_y_data)
    cov_data = linalg.block_diag(*[dat[0, ...] for dat in raw_cov_data])

    gp = active_utils.create_base_GP_model((x_data, y_data, cov_data))
    output = active_utils.train_GPR(gp, record_loss=False)
    assert output is None

    gp = active_utils.create_base_GP_model((x_data, y_data, cov_data))
    output = active_utils.train_GPR(gp, record_loss=True)
    assert output is not None

    # ref_params = [4.9320208698871015, 17.07650174925236, 0.0, 4.3460319697777543e-16]  # with L-BFGS-B minimize
    ref_params = [
        5.0146512540878225,
        18.083488371624117,
        0.0,
        1.0677842958152827e-09,
    ]  # with SLSQP minimize

    # print(gp.parameters)
    # raise ValueError(f"hello {gp.parameters}")
    for p, rp in zip(gp.parameters, ref_params, strict=True):
        # np.testing.assert_allclose(p.numpy(), rp, rtol=1e-1, atol=1e-12)  # for L-BFGS-B
        np.testing.assert_allclose(p.numpy(), rp, rtol=1e-7, atol=0)

    # Also test training from a different starting point
    # Not checking for specific behavior here, just seeing if runs
    output = active_utils.train_GPR(
        gp, record_loss=False, start_params=[ref_params[ind] for ind in (0, 1, 3)]
    )


# Simple test for creating a GP model from list of states
# Also a bit slow
@pytest.mark.slow
def test_create_gp_from_states(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    states = [ig_active.extrap_IG(beta, rng=rng) for beta in (1.0, 5.6, 9.0)]

    n_gp_points = np.sum([s.order + 1 for s in states])

    gp = active_utils.create_GPR(states, sampler=sampler)
    assert gp.data[0].shape == (n_gp_points, 2)
    assert gp.data[1].shape == (n_gp_points, 1)
    assert gp.likelihood.cov.shape == (gp.data[1].shape[1], n_gp_points, n_gp_points)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]

    # Also test with multidimensional data
    states_mult = [
        ig_active.multiOutput_extrap_IG(beta, rng=rng) for beta in (1.0, 5.6, 9.0)
    ]
    gp_mult = active_utils.create_GPR(states_mult, sampler=sampler)
    assert gp_mult.data[0].shape == (n_gp_points, 2)
    assert gp_mult.data[1].shape == (n_gp_points, 2)
    assert gp_mult.likelihood.cov.shape == (  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
        gp_mult.data[1].shape[1],
        n_gp_points,
        n_gp_points,
    )


# Testing update and stopping function classes
@pytest.mark.slow
def test_update_stop_abc(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    beta_list = [1.0, 5.6, 9.0]
    states = [ig_active.extrap_IG(beta, rng=rng) for beta in beta_list]
    # And trained GP
    gp = active_utils.create_GPR(states, sampler=sampler)

    # Start by checking with default inputs
    check_default = active_utils.UpdateStopABC()

    # Check creation of alpha grid
    out_default_grid, out_default_sel = check_default.create_alpha_grid(beta_list)
    np.testing.assert_allclose(
        out_default_grid, np.linspace(beta_list[0], beta_list[-1], 1000)
    )
    np.testing.assert_equal(out_default_grid, out_default_sel)

    # With log scale and avoiding repeats
    check_logrepeats = active_utils.UpdateStopABC(log_scale=True, avoid_repeats=True)
    out_logrepeats_grid, out_logrepeats_sel = check_logrepeats.create_alpha_grid(
        beta_list
    )
    np.testing.assert_equal(
        out_logrepeats_grid,
        np.linspace(np.log10(beta_list[0]), np.log10(beta_list[-1]), 1000),
    )
    assert not np.array_equal(out_logrepeats_grid, out_logrepeats_sel)
    # Make sure randomization stays within grid
    assert np.all(
        abs(out_logrepeats_sel - out_logrepeats_grid[1:-1])
        <= (out_logrepeats_grid[1] - out_logrepeats_grid[0])
    )

    # Check transformation for default - should just return same thing
    x = np.arange(10.0)
    y = x**3
    out_notrans_y, out_notrans_std, out_notrans_conf = check_default.transform_func(
        x, y, 1.0
    )
    np.testing.assert_equal(out_notrans_y, y)
    np.testing.assert_allclose(out_notrans_std, 1.0)
    np.testing.assert_allclose(out_notrans_conf[0], y - 2.0 * 1.0)
    np.testing.assert_allclose(out_notrans_conf[1], y + 2.0 * 1.0)

    # And if provide fancier transformation
    def fancy_transform(x, y, y_var):
        # Transformation itself being checked, resulting uncertainty and confidence
        # interval are nonsense
        y_std = np.sqrt(y_var)
        out = x * (y ** (1 / 3))
        conf_int = [x * (y - 2.0 * y_std), x * (y + 2.0 * y_std)]
        return out, y_std, conf_int

    check_trans = active_utils.UpdateStopABC(transform_func=fancy_transform)
    out_trans_y = check_trans.transform_func(x, y, 1.0)[0]
    np.testing.assert_allclose(out_trans_y, x**2)

    # Check to make sure GP used properly for default case (just runs the function)
    (
        _out_default_mu,
        _out_default_std,
        _out_default_conf,
    ) = check_default.get_transformed_GP_output(gp, out_default_grid)


# For update classes, all have different update criteria
# Rather than check all of these (hard for random...)
# just check to make sure satisfy correct input/output structure
@pytest.mark.slow
def test_update_classes(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    beta_list = [1.0, 5.6, 9.0]
    states = [ig_active.extrap_IG(beta, rng=rng) for beta in beta_list]
    # And trained GP
    gp = active_utils.create_GPR(states, sampler=sampler)

    # Start with base, which should throw an error
    check_base = active_utils.UpdateFuncBase()
    # Should be able to call, which just wraps do_update()
    np.testing.assert_raises(NotImplementedError, check_base, gp, beta_list)
    np.testing.assert_raises(NotImplementedError, check_base.do_update, gp, beta_list)

    # Random update function
    check_rand = active_utils.UpdateRandom()
    out_rand = check_rand(gp, beta_list)
    assert len(out_rand) == 3
    assert isinstance(out_rand[0], float)

    # Space filling update function (furthest point)
    check_space = active_utils.UpdateSpaceFill()
    out_space = check_space(gp, beta_list)
    assert len(out_space) == 3
    assert isinstance(out_space[0], float)

    # ALM (max variance) update function
    check_alm = active_utils.UpdateALMbrute()
    out_alm = check_alm(gp, beta_list)
    assert len(out_alm) == 3
    assert isinstance(out_alm[0], float)

    # Adaptive integration updates
    # (furthest point satisfying error tolerance)
    check_adapt = active_utils.UpdateAdaptiveIntegrate(tol=0.005)
    out_adapt = check_adapt(gp, beta_list)
    assert len(out_adapt) == 3
    assert isinstance(out_adapt[0], float)


@pytest.mark.slow
def test_update_classes_multioutput(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    beta_list = [1.0, 5.6, 9.0]
    states = [ig_active.multiOutput_extrap_IG(beta, rng=rng) for beta in beta_list]
    # And trained GP
    gp = active_utils.create_GPR(states, sampler=sampler)

    # Random update function
    check_rand = active_utils.UpdateRandom()
    out_rand = check_rand(gp, beta_list)
    assert len(out_rand) == 3
    assert isinstance(out_rand[0], float)
    assert out_rand[1].shape[0] == 2

    # Space filling update function (furthest point)
    check_space = active_utils.UpdateSpaceFill()
    out_space = check_space(gp, beta_list)
    assert len(out_space) == 3
    assert isinstance(out_space[0], float)
    assert out_space[1].shape[0] == 2

    # ALM (max variance) update function
    check_alm = active_utils.UpdateALMbrute()
    out_alm = check_alm(gp, beta_list)
    assert len(out_alm) == 3
    assert isinstance(out_alm[0], float)
    assert out_alm[1].shape[0] == 2

    # Adaptive integration updates
    # (furthest point satisfying error tolerance)
    check_adapt = active_utils.UpdateAdaptiveIntegrate()
    out_adapt = check_adapt(gp, beta_list)
    assert len(out_adapt) == 3
    assert isinstance(out_adapt[0], float)
    assert out_adapt[1].shape[0] == 2


# Same for metric classes
# just check mechanics for taking inputs and generic features of outputs
@pytest.mark.slow
def test_metrics(rng: Generator, sampler: SamplerType) -> None:
    # Need to create inputs to work with
    # Expects "history" which is a list of array-likes
    # where index zero is GP means over time and index one is vars
    x = np.arange(-10.0, 10.0, 1.0)
    hist = [np.vstack([x, x**3]), np.vstack([0.01 * x**2, 0.01 * x**2])]

    # Base class
    tol = 1e-03
    check_base = active_utils.MetricBase("Base", tol)
    assert check_base.tol == tol
    # Use to test _check_history function for all inheriting classes
    # pylint: disable=protected-access
    np.testing.assert_raises(ValueError, check_base._check_history, None)
    np.testing.assert_raises(ValueError, check_base._check_history, x)  # type: ignore[call-overload]
    check_base._check_history(hist)
    # And generic call and calc_metric both work
    np.testing.assert_raises(NotImplementedError, check_base, hist, x, None)
    np.testing.assert_raises(NotImplementedError, check_base.calc_metric, hist, x, None)

    # Other classes
    check_maxvar = active_utils.MaxVar(tol)
    assert isinstance(check_maxvar(hist, x, None), float)

    check_avgvar = active_utils.AvgVar(tol)
    assert isinstance(check_avgvar(hist, x, None), float)

    check_maxrelvar = active_utils.MaxRelVar(tol)
    assert isinstance(check_maxrelvar(hist, x, None), float)

    check_avgrelvar = active_utils.AvgRelVar(tol)
    assert isinstance(check_avgrelvar(hist, x, None), float)

    check_msd = active_utils.MSD(tol)
    assert isinstance(check_msd(hist, x, None), float)

    check_maxabsreldev = active_utils.MaxAbsRelDeviation(tol)
    assert isinstance(check_maxabsreldev(hist, x, None), float)

    check_avgabsreldev = active_utils.AvgAbsRelDeviation(tol)
    assert isinstance(check_avgabsreldev(hist, x, None), float)

    # For metric ensuring maximum iterations, do check criteria since important and easy (never reaches tol)
    check_maxiter = active_utils.MaxIter()
    assert isinstance(check_maxiter(hist, x, None), float)
    assert check_maxiter(hist, x, None) > check_maxiter.tol

    # For special ErrorStability class, actually need a GP model
    beta_list = [1.0, 2.3, 5.6, 9.0]
    states = [ig_active.extrap_IG(beta, rng=rng) for beta in beta_list]
    # And trained GP
    gp_2 = active_utils.create_GPR([states[0], states[-1]], sampler=sampler)
    gp_3 = active_utils.create_GPR([states[0], states[2], states[-1]], sampler=sampler)
    gp_4 = active_utils.create_GPR(states, sampler=sampler)

    check_errorstability = active_utils.ErrorStability(0.05)
    # Check that returns 1 for just 2 points
    assert check_errorstability(hist, x, gp_2) == 1
    # Check that sets normalization once provide more points
    assert check_errorstability.r1 is None

    out_errorstability = check_errorstability(hist, x, gp_3)
    assert isinstance(check_errorstability.r1, float)
    assert isinstance(out_errorstability, float)  # type: ignore[unreachable]
    np.testing.assert_allclose(out_errorstability, 1.0)
    # Check with 4 points
    out_errorstability_4 = check_errorstability(hist, x, gp_4)
    assert isinstance(out_errorstability, float)
    assert out_errorstability_4 < out_errorstability

    # Also need GP for relative quantities compared to variance of GP data
    check_maxrelglobalvar = active_utils.MaxRelGlobalVar(tol)
    assert isinstance(check_maxrelglobalvar(hist, x, gp_4), float)

    check_maxabsrelglobaldev = active_utils.MaxAbsRelGlobalDeviation(tol)
    assert isinstance(check_maxabsrelglobaldev(hist, x, gp_4), float)


# Test class for implementing stopping criteria
@pytest.mark.slow
def test_stop_criteria(rng: Generator, sampler: SamplerType) -> None:
    # Need data to work with
    beta_list = [1.0, 5.6, 9.0]
    states = [ig_active.extrap_IG(beta, rng=rng) for beta in beta_list]
    # And trained GP
    gp_prev = active_utils.create_GPR([states[0], states[-1]], sampler=sampler)
    gp = active_utils.create_GPR(states, sampler=sampler)

    tol = 1e-03
    m_var = active_utils.MaxVar(tol)
    m_msd = active_utils.MSD(tol)
    m_es = active_utils.ErrorStability(tol)

    # With one metric
    check_stop_1m = active_utils.StopCriteria([m_var])
    # Test call for first round without history
    assert check_stop_1m.history is None
    out_1m_bool, out_1m_info = check_stop_1m(gp_prev, beta_list)
    assert check_stop_1m.history is not None
    assert len(check_stop_1m.history) == 2  # type: ignore[unreachable]
    assert check_stop_1m.history[0].shape == (1, 1000, 1)
    assert check_stop_1m.history[1].shape == (1, 1000, 1)
    assert isinstance(out_1m_bool, bool)
    assert isinstance(out_1m_info, dict)
    assert len(out_1m_info.keys()) == 2  # Also have key for the tolerance
    # Call again and check if updated correctly
    out_1m_bool, out_1m_info = check_stop_1m(gp, beta_list)
    assert len(check_stop_1m.history) == 2
    assert check_stop_1m.history[0].shape == (2, 1000, 1)
    assert check_stop_1m.history[1].shape == (2, 1000, 1)
    assert isinstance(out_1m_bool, bool)
    assert isinstance(out_1m_info, dict)
    assert len(out_1m_info.keys()) == 2

    # With multiple metrics
    check_stop_2m = active_utils.StopCriteria([m_var, m_msd])
    assert check_stop_2m.history is None
    out_2m_bool, out_2m_info = check_stop_2m(gp_prev, beta_list)
    assert check_stop_2m.history is not None
    assert len(check_stop_2m.history) == 2
    assert check_stop_2m.history[0].shape == (1, 1000, 1)
    assert check_stop_2m.history[1].shape == (1, 1000, 1)
    assert isinstance(out_2m_bool, bool)
    assert isinstance(out_2m_info, dict)
    assert len(out_2m_info.keys()) == 4
    out_2m_bool, out_2m_info = check_stop_2m(gp, beta_list)
    assert len(check_stop_2m.history) == 2
    assert check_stop_2m.history[0].shape == (2, 1000, 1)
    assert check_stop_2m.history[1].shape == (2, 1000, 1)
    assert isinstance(out_2m_bool, bool)
    assert isinstance(out_2m_info, dict)
    assert len(out_2m_info.keys()) == 4

    # Check that correctly manipulates ErrorStability metric attributes
    def check_transform(x, y, y_var):
        y_std = x * np.sqrt(y_var)
        out = x * y
        conf = [x * (y - 2.0 * y_std), x * (y + 2.0 * y_std)]
        return out, y_std, conf

    check_stop_with_es = active_utils.StopCriteria(
        [m_var, m_msd, m_es],
        d_order_pred=1,
        transform_func=check_transform,
        log_scale=True,
    )
    assert m_es.d_order_pred == 1
    assert m_es.transform_func == check_stop_with_es.transform_func
    assert m_es.log_scale is True


# Test full active learning routine on ideal gas
# Need way to make this fully reproducible...
# Specifically, need to be able to set random seed for all parts of process
# Includes bootstrapping of data in thermoextrap/cmomy
# Just need to carefully check where calling np.random throughout all code
# import io
# from contextlib import redirect_stdout

# def test_thing_1(caplog) -> None:
#     with caplog.at_level(logging.INFO):
#         active_utils.thing("max")
#     assert "Reached maximum iterations" in caplog.text


# def test_thing_2(caplog) -> None:
#     with caplog.at_level(logging.INFO):
#         active_utils.thing("stop")
#     assert "Stopping criteria satisfied" in caplog.text


@pytest.mark.slow
def test_active_learning(caplog, sampler: SamplerType) -> None:
    # Starting beta values
    init_states = [1.0, 9.6]
    sims = ig_active.SimulateIG()
    updates = active_utils.UpdateALMbrute()
    metrics = [
        active_utils.MaxVar(1e-03),
        active_utils.MaxRelVar(1e-02),
        active_utils.MSD(1.0),
        active_utils.MaxAbsRelDeviation(1e-02),
        active_utils.MaxIter(),
    ]

    with caplog.at_level(logging.INFO):
        # Once where should reach maximum iterations
        stops = active_utils.StopCriteria(metrics)
        active_utils.active_learning(
            init_states,
            sims,
            updates,
            stop_criteria=stops,
            max_iter=4,
            sampler=sampler,
        )

    assert "Reached maximum iterations" in caplog.text


@pytest.mark.slow
def test_active_learning_2(caplog, sampler: SamplerType) -> None:
    # Starting beta values
    init_states = [1.0, 9.6]
    sims = ig_active.SimulateIG()
    updates = active_utils.UpdateALMbrute()
    metrics = [
        active_utils.MaxVar(1e-03),
        active_utils.MaxRelVar(1e-02),
        active_utils.MSD(1.0),
        active_utils.MaxAbsRelDeviation(1e-02),
        active_utils.MaxIter(),
    ]

    with caplog.at_level(logging.INFO):
        # And once where expect to stop early
        stops_early = active_utils.StopCriteria(metrics[:-1])
        active_utils.active_learning(
            init_states,
            sims,
            updates,
            sampler=sampler,
            stop_criteria=stops_early,
            max_iter=4,
        )
    assert "Stopping criteria satisfied" in caplog.text
