import importlib

import pandas as pd
import pytest
import torch
from botorch.models import MixedSingleTaskGP, SingleTaskGP
from botorch.models.kernels.categorical import CategoricalKernel
from botorch.models.robust_relevance_pursuit_model import (
    RobustRelevancePursuitSingleTaskGP,
)
from botorch.models.transforms.input import (
    ChainedInputTransform,
    InputStandardize,
    Normalize,
    NumericToCategoricalEncoding,
)
from botorch.models.transforms.outcome import Standardize
from gpytorch.constraints import GreaterThan
from pandas.testing import assert_frame_equal
from torch.nn import Module

import bofire.surrogates.api as surrogates
from bofire.benchmarks.api import Hartmann, Himmelblau
from bofire.data_models.domain.api import EngineeredFeatures, Inputs, Outputs
from bofire.data_models.enum import CategoricalEncodingEnum, RegressionMetricsEnum
from bofire.data_models.features.api import (
    CategoricalInput,
    CategoricalMolecularInput,
    ContinuousInput,
    ContinuousOutput,
    SumFeature,
)
from bofire.data_models.kernels.api import (
    AdditiveKernel,
    HammingDistanceKernel,
    MaternKernel,
    RBFKernel,
    ScaleKernel,
    TanimotoKernel,
)
from bofire.data_models.molfeatures.api import Fingerprints, MordredDescriptors
from bofire.data_models.priors.api import (
    HVARFNER_LENGTHSCALE_PRIOR,
    HVARFNER_NOISE_PRIOR,
    MBO_LENGTHSCALE_PRIOR,
    MBO_NOISE_PRIOR,
    MBO_OUTPUTSCALE_PRIOR,
    ROBUSTGP_LENGTHSCALE_CONSTRAINT,
    THREESIX_LENGTHSCALE_PRIOR,
    THREESIX_NOISE_PRIOR,
    THREESIX_SCALE_PRIOR,
)
from bofire.data_models.surrogates.api import (
    MixedSingleTaskGPSurrogate,
    RobustSingleTaskGPSurrogate,
    ScalerEnum,
    SingleTaskGPHyperconfig,
    SingleTaskGPSurrogate,
)
from bofire.data_models.surrogates.trainable import metrics2objectives
from bofire.utils.torch_tools import tkwargs


RDKIT_AVAILABLE = importlib.util.find_spec("rdkit") is not None


@pytest.mark.parametrize(
    "kernel, scaler, output_scaler",
    [
        (
            ScaleKernel(base_kernel=RBFKernel(ard=True)),
            ScalerEnum.NORMALIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.STANDARDIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.IDENTITY,
            ScalerEnum.IDENTITY,
        ),
    ],
)
def test_SingleTaskGPModel(kernel, scaler, output_scaler):
    inputs = Inputs(
        features=[
            ContinuousInput(
                key=f"x_{i + 1}",
                bounds=(-4, 4),
            )
            for i in range(2)
        ],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = inputs.sample(n=10)
    experiments.eval("y=((x_1**2 + x_2 - 11)**2+(x_1 + x_2**2 -7)**2)", inplace=True)
    experiments["valid_y"] = 1
    model = SingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
    )
    model = surrogates.map(model)
    samples = inputs.sample(5)
    # test error on non fitted model
    with pytest.raises(ValueError):
        model.predict(samples)
    model.fit(experiments)
    # dump the model
    dump = model.dumps()
    # make predictions
    samples2 = samples.copy()
    samples2 = samples2.astype({"x_1": "object"})
    preds = model.predict(samples2)
    assert preds.shape == (5, 2)
    # check that model is composed correctly
    assert isinstance(model.model, SingleTaskGP)
    if output_scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.outcome_transform, Standardize)
    elif output_scaler == ScalerEnum.IDENTITY:
        assert not hasattr(model.model, "outcome_transform")
    if scaler == ScalerEnum.NORMALIZE:
        assert isinstance(model.model.input_transform, Normalize)
    elif scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.input_transform, InputStandardize)
    else:
        with pytest.raises(AttributeError):
            assert model.model.input_transform is None
    assert model.is_compatibilized is False
    # reload the model from dump and check for equality in predictions
    model2 = SingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
    )
    model2 = surrogates.map(model2)
    model2.loads(dump)
    preds2 = model2.predict(samples)
    assert_frame_equal(preds, preds2)


def test_SingleTaskGPModel_with_engineered_features():
    bench = Himmelblau()
    experiments = bench.f(bench.domain.inputs.sample(20), return_complete=True)

    surrogate_data = SingleTaskGPSurrogate(
        inputs=bench.domain.inputs,
        outputs=bench.domain.outputs,
        engineered_features=EngineeredFeatures(
            features=[
                SumFeature(key="sum", features=["x_1", "x_2"], keep_features=True)
            ]
        ),
    )

    surrogate = surrogates.map(surrogate_data)

    with pytest.raises(KeyError, match="Feature with key 'sum123' not found."):
        surrogate.get_feature_indices(["sum123", "x_1", "x_2"])

    surrogate.engineered_features[0].keep_features = False
    with pytest.raises(
        NotImplementedError,
        match="Cannot get feature indices if original features are filtered.",
    ):
        surrogate.get_feature_indices(["sum", "x_1", "x_2"])

    surrogate.engineered_features[0].keep_features = True

    assert surrogate.get_feature_indices(["sum", "x_1", "x_2"]) == [0, 1, 2]
    assert surrogate.get_feature_indices(["sum", "x_2"]) == [1, 2]

    surrogate.fit(experiments)
    assert surrogate.model.covar_module.active_dims.tolist() == [0, 1, 2]

    surrogate_data = SingleTaskGPSurrogate(
        inputs=bench.domain.inputs,
        outputs=bench.domain.outputs,
        engineered_features=EngineeredFeatures(
            features=[
                SumFeature(key="sum", features=["x_1", "x_2"], keep_features=True)
            ]
        ),
        kernel=RBFKernel(ard=True, features=["x_1", "sum"]),
    )
    surrogate = surrogates.map(surrogate_data)
    surrogate.fit(experiments)
    assert surrogate.model.covar_module.active_dims.tolist() == [0, 2]


@pytest.mark.skipif(not RDKIT_AVAILABLE, reason="requires rdkit")
@pytest.mark.parametrize(
    "kernel, scaler, output_scaler",
    [
        (
            ScaleKernel(base_kernel=RBFKernel(ard=True)),
            ScalerEnum.NORMALIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.STANDARDIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.IDENTITY,
            ScalerEnum.IDENTITY,
        ),
    ],
)
def test_SingleTaskGPModel_mordred(kernel, scaler, output_scaler):
    inputs = Inputs(
        features=[
            CategoricalMolecularInput(
                key="x_mol",
                categories=[
                    "CC(=O)Oc1ccccc1C(=O)O",
                    "c1ccccc1",
                    "[CH3][CH2][OH]",
                    "N[C@](C)(F)C(=O)O",
                ],
            )
        ]
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = [
        ["CC(=O)Oc1ccccc1C(=O)O", 88.0],
        ["c1ccccc1", 35.0],
        ["[CH3][CH2][OH]", 69.0],
        ["N[C@](C)(F)C(=O)O", 20.0],
    ]
    experiments = pd.DataFrame(experiments, columns=["x_mol", "y"])
    experiments["valid_y"] = 1
    model = SingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
        categorical_encodings={
            "x_mol": MordredDescriptors(descriptors=["NssCH2", "ATSC2d"]),
        },
        # input_preprocessing_specs={
        #     "x_mol": MordredDescriptors(descriptors=["NssCH2", "ATSC2d"]),
        # },
    )
    model = surrogates.map(model)
    model.fit(experiments)
    # dump the model
    dump = model.dumps()
    # make predictions
    preds = model.predict(experiments.iloc[:-1])
    assert preds.shape == (3, 2)
    # check that model is composed correctly
    assert isinstance(model.model, SingleTaskGP)
    if output_scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.outcome_transform, Standardize)
    elif output_scaler == ScalerEnum.IDENTITY:
        assert not hasattr(model.model, "outcome_transform")
    if scaler == ScalerEnum.NORMALIZE:
        assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform.scaler, Normalize)
    elif scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform.scaler, InputStandardize)
    else:
        assert isinstance(model.model.input_transform, NumericToCategoricalEncoding)
    assert model.is_compatibilized is False
    # reload the model from dump and check for equality in predictions
    model2 = SingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
        categorical_encodings={
            "x_mol": MordredDescriptors(descriptors=["NssCH2", "ATSC2d"]),
        },
    )
    model2 = surrogates.map(model2)
    model2.loads(dump)
    preds2 = model2.predict(experiments.iloc[:-1])
    assert_frame_equal(preds, preds2)


@pytest.mark.parametrize("target_metric", list(RegressionMetricsEnum))
def test_hyperconfig_domain(target_metric: RegressionMetricsEnum):
    # we test here also the abstract methods from the corresponding base class
    # should be move somewhere else when tidying up all tests
    hy = SingleTaskGPHyperconfig(target_metric=target_metric)
    assert hy.domain.inputs == hy.inputs
    assert hy.domain.outputs.get_keys() == [target_metric.name]
    assert hy.domain.outputs[0].objective == metrics2objectives[target_metric]()


def test_hyperconfig_invalid():
    with pytest.raises(
        ValueError,
        match="It is not allowed to specify the number of its for FractionalFactorialStrategy",
    ):
        SingleTaskGPHyperconfig(n_iterations=5)
    with pytest.raises(
        ValueError,
        match="At least number of hyperparams plus 2 iterations has to be specified",
    ):
        SingleTaskGPHyperconfig(n_iterations=3, hyperstrategy="RandomStrategy")
    hy = SingleTaskGPHyperconfig(n_iterations=None, hyperstrategy="RandomStrategy")
    assert hy.n_iterations == 14


def test_SingleTaskGPHyperconfig():
    # we test here also the basic trainable
    benchmark = Himmelblau()
    surrogate_data_no_hy = SingleTaskGPSurrogate(
        inputs=benchmark.domain.inputs,
        outputs=benchmark.domain.outputs,
        hyperconfig=None,
    )
    with pytest.raises(ValueError, match="No hyperconfig available."):
        surrogate_data_no_hy.update_hyperparameters(
            benchmark.domain.inputs.sample(1).loc[0],
        )
    # test that correct stuff is written
    surrogate_data = SingleTaskGPSurrogate(
        inputs=benchmark.domain.inputs,
        outputs=benchmark.domain.outputs,
    )
    candidate = surrogate_data.hyperconfig.inputs.sample(1).loc[0]
    surrogate_data.update_hyperparameters(candidate)
    # if hasattr(surrogate_data.kernel, "base_kernel"):
    base_kernel = (
        surrogate_data.kernel.base_kernel
        if hasattr(surrogate_data.kernel, "base_kernel")
        else surrogate_data.kernel
    )
    if candidate.scalekernel == "True":
        assert hasattr(surrogate_data.kernel, "base_kernel")
    else:
        assert not hasattr(surrogate_data.kernel, "base_kernel")
    if candidate.kernel == "matern_1.5":
        assert isinstance(base_kernel, MaternKernel)
        assert base_kernel.nu == 1.5
    elif candidate.kernel == "matern_2.5":
        assert isinstance(base_kernel, MaternKernel)
        assert base_kernel.nu == 2.5
    else:
        assert isinstance(base_kernel, RBFKernel)
    if candidate.prior == "mbo":
        assert surrogate_data.noise_prior == MBO_NOISE_PRIOR()
        if candidate.scalekernel == "True":
            assert surrogate_data.kernel.outputscale_prior == MBO_OUTPUTSCALE_PRIOR()
        assert base_kernel.lengthscale_prior == MBO_LENGTHSCALE_PRIOR()
    elif candidate.prior == "threesix":
        assert surrogate_data.noise_prior == THREESIX_NOISE_PRIOR()
        if candidate.scalekernel == "True":
            assert surrogate_data.kernel.outputscale_prior == THREESIX_SCALE_PRIOR()
        assert base_kernel.lengthscale_prior == THREESIX_LENGTHSCALE_PRIOR()
    else:
        assert surrogate_data.noise_prior == HVARFNER_NOISE_PRIOR()
        if candidate.scalekernel == "True":
            assert surrogate_data.kernel.outputscale_prior == THREESIX_SCALE_PRIOR()
        assert base_kernel.lengthscale_prior == HVARFNER_LENGTHSCALE_PRIOR()


def test_SingleTaskGPModel_feature_subsets():
    """make an additive kernel using feature subsets for each kernel in the sum"""
    benchmark = Hartmann()
    bench_x = benchmark.domain.inputs.sample(12)
    bench_expts = pd.concat([bench_x, benchmark.f(bench_x)], axis=1)

    input_names = benchmark.domain.inputs.get_keys()
    inputs_kernel_1 = input_names[:2]
    inputs_kernel_2 = input_names[2:]

    gp_data = SingleTaskGPSurrogate(
        inputs=benchmark.domain.inputs,
        outputs=benchmark.domain.outputs,
        kernel=AdditiveKernel(
            kernels=[
                RBFKernel(
                    ard=True,
                    lengthscale_prior=HVARFNER_LENGTHSCALE_PRIOR(),
                    features=inputs_kernel_1,
                ),
                RBFKernel(
                    ard=True,
                    lengthscale_prior=HVARFNER_LENGTHSCALE_PRIOR(),
                    features=inputs_kernel_2,
                ),
            ]
        ),
    )

    gp_mapped = surrogates.map(gp_data)
    assert hasattr(gp_mapped, "fit")
    assert len(gp_mapped.kernel.kernels) == 2
    assert gp_mapped.kernel.kernels[0].features == ["x_0", "x_1"]
    assert gp_mapped.kernel.kernels[1].features == ["x_2", "x_3", "x_4", "x_5"]
    gp_mapped.fit(bench_expts)
    pred = gp_mapped.predict(bench_expts)
    assert pred.shape == (12, 2)
    assert gp_mapped.model.covar_module.kernels[0].active_dims.tolist() == [0, 1]
    assert gp_mapped.model.covar_module.kernels[1].active_dims.tolist() == [2, 3, 4, 5]


@pytest.mark.skipif(not RDKIT_AVAILABLE, reason="requires rdkit")
def test_SingleTaskGPModel_mixed_features():
    """test that we can use a single task gp with mixed features"""
    inputs = Inputs(
        features=[
            ContinuousInput(key="x_1", bounds=(-4, 4)),
            ContinuousInput(key="x_2", bounds=(-4, 4)),
            CategoricalInput(key="x_cat_1", categories=["mama", "papa"]),
            CategoricalInput(key="x_cat_2", categories=["cat", "dog"]),
            CategoricalMolecularInput(
                key="x_mol",
                categories=[
                    "CC(=O)Oc1ccccc1C(=O)O",
                    "c1ccccc1",
                    "[CH3][CH2][OH]",
                    "N[C@](C)(F)C(=O)O",
                ],
            ),
        ],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])

    experiment_values = [
        [2.56, -1.42, "papa", "dog", -3.98, 1, "CC(=O)Oc1ccccc1C(=O)O"],
        [3.84, -2.73, "mama", "cat", -197.46, 1, "c1ccccc1"],
        [3.57, 3.23, "papa", "cat", -74.55, 1, "[CH3][CH2][OH]"],
        [-0.07, -1.55, "mama", "dog", -179.14, 1, "N[C@](C)(F)C(=O)O"],
    ]
    experiments = pd.DataFrame(
        experiment_values,
        columns=["x_1", "x_2", "x_cat_1", "x_cat_2", "y", "valid_y", "x_mol"],
    )

    gp_data = SingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=AdditiveKernel(
            kernels=[
                HammingDistanceKernel(
                    ard=True,
                    features=["x_cat_1", "x_cat_2"],
                ),
                RBFKernel(
                    ard=True,
                    lengthscale_prior=HVARFNER_LENGTHSCALE_PRIOR(),
                    features=["x_1", "x_2"],
                ),
                TanimotoKernel(features=["x_mol"]),
            ]
        ),
        categorical_encodings={
            "x_cat_1": CategoricalEncodingEnum.ORDINAL,
            "x_cat_2": CategoricalEncodingEnum.ORDINAL,
            "x_mol": Fingerprints(n_bits=2048),
        },
    )

    gp_mapped = surrogates.map(gp_data)
    gp_mapped.fit(experiments)
    pred = gp_mapped.predict(experiments)
    assert pred.shape == (4, 2)
    assert gp_mapped.model.covar_module.kernels[0].active_dims.tolist() == [
        2050,
        2051,
    ]
    assert gp_mapped.model.covar_module.kernels[1].active_dims.tolist() == [0, 1]
    assert gp_mapped.model.covar_module.kernels[2].active_dims.tolist() == list(
        range(2, 2050)
    )
    # assert (pred['y_pred'] - experiments['y']).abs().mean() < 0.4


def test_MixedSingleTaskGP_only_categorical():
    inputs = Inputs(
        features=[
            CategoricalInput(key="x_cat_1", categories=["mama", "papa"]),
            CategoricalInput(key="x_cat_2", categories=["cat", "dog"]),
        ]
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    surrogate_data = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
    )
    experiments = [
        ["mama", "cat", 1.0],
        ["papa", "dog", 2.0],
        ["mama", "dog", 3.0],
        ["papa", "cat", 4.0],
    ]
    experiments = pd.DataFrame(experiments, columns=["x_cat_1", "x_cat_2", "y"])
    experiments["valid_y"] = 1

    assert surrogate_data.categorical_kernel.features == ["x_cat_1", "x_cat_2"]
    assert surrogate_data.continuous_kernel.features == []
    surrogate = surrogates.map(surrogate_data)
    assert isinstance(surrogate, surrogates.SingleTaskGPSurrogate)
    assert isinstance(surrogate.kernel, ScaleKernel)
    assert isinstance(surrogate.kernel.base_kernel, HammingDistanceKernel)

    surrogate.fit(experiments)
    assert surrogate.model.covar_module.base_kernel.active_dims.tolist() == [0, 1]
    assert surrogate.model.covar_module.base_kernel.ard_num_dims == 2


def test_MixedSingleTaskGPHyperconfig():
    inputs = Inputs(
        features=[
            ContinuousInput(
                key=f"x_{i + 1}",
                bounds=(-4, 4),
            )
            for i in range(2)
        ]
        + [CategoricalInput(key="x_cat", categories=["mama", "papa"])],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    surrogate_data = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
    )
    candidate = surrogate_data.hyperconfig.inputs.sample(1).loc[0]
    surrogate_data.update_hyperparameters(candidate)
    assert surrogate_data.continuous_kernel.ard == (candidate["ard"] == "True")
    if candidate.continuous_kernel == "matern_1.5":
        assert isinstance(surrogate_data.continuous_kernel, MaternKernel)
        assert surrogate_data.continuous_kernel.nu == 1.5
    elif candidate.continuous_kernel == "matern_2.5":
        assert isinstance(surrogate_data.continuous_kernel, MaternKernel)
        assert surrogate_data.continuous_kernel.nu == 2.5
    else:
        assert isinstance(surrogate_data.continuous_kernel, RBFKernel)
    if candidate.prior == "mbo":
        assert surrogate_data.noise_prior == MBO_NOISE_PRIOR()
        assert (
            surrogate_data.continuous_kernel.lengthscale_prior
            == MBO_LENGTHSCALE_PRIOR()
        )
    if candidate.prior == "threesix":
        assert surrogate_data.noise_prior == THREESIX_NOISE_PRIOR()
        assert (
            surrogate_data.continuous_kernel.lengthscale_prior
            == THREESIX_LENGTHSCALE_PRIOR()
        )
    if candidate.prior == "hvarfner":
        assert surrogate_data.noise_prior == HVARFNER_NOISE_PRIOR()
        assert (
            surrogate_data.continuous_kernel.lengthscale_prior
            == HVARFNER_LENGTHSCALE_PRIOR()
        )


def test_MixedSingletaskGPModel_with_botorch():
    inputs = Inputs(
        features=[
            ContinuousInput(
                key=f"x_{i + 1}",
                bounds=(-4, 4),
            )
            for i in range(2)
        ]
        + [CategoricalInput(key="x_cat", categories=["mama", "papa"])],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = inputs.sample(n=10)
    experiments.eval("y=((x_1**2 + x_2 - 11)**2+(x_1 + x_2**2 -7)**2)", inplace=True)
    experiments.loc[experiments.x_cat == "mama", "y"] *= 5.0
    experiments.loc[experiments.x_cat == "papa", "y"] /= 2.0
    experiments["valid_y"] = 1
    surrogate_data = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
    )
    model = surrogates.map(surrogate_data)
    assert isinstance(model, surrogates.SingleTaskGPSurrogate)
    model.fit(experiments)

    transformed_X = surrogate_data.inputs.transform(
        experiments[surrogate_data.inputs.get_keys()],
        surrogate_data.input_preprocessing_specs,
    )
    Y = experiments[outputs.get_keys()]
    tX, tY = (
        torch.from_numpy(transformed_X.values).to(**tkwargs),
        torch.from_numpy(Y.values).to(**tkwargs),
    )
    pure_botorch_model = MixedSingleTaskGP(
        train_X=tX,
        train_Y=tY,
        cat_dims=[2],
        input_transform=Normalize(
            d=3, indices=[0, 1], bounds=torch.tensor([(-4, -4), (4, 4)])
        ).to(**tkwargs),
    )

    def compare_modules(module1: Module, module2: Module) -> bool:
        if type(module1) is not type(module2):
            return False
        if isinstance(module1, RBFKernel) or isinstance(module1, CategoricalKernel):
            if not torch.equal(module1.active_dims, module2.active_dims):
                return False
        if not isinstance(module1, GreaterThan):
            if hasattr(module1, "children") and hasattr(module2, "children"):
                children1 = list(module1.children())
                children2 = list(module2.children())
                if len(children1) != len(children2):
                    return False
                return all(
                    compare_modules(c1, c2) for c1, c2 in zip(children1, children2)
                )
        return True

    assert (
        compare_modules(pure_botorch_model.covar_module, model.model.covar_module)
        is True
    )
    assert torch.allclose(model.model.input_transform.indices, torch.tensor([0, 1]))


@pytest.mark.parametrize(
    "kernel, scaler, output_scaler",
    [
        (RBFKernel(ard=True), ScalerEnum.NORMALIZE, ScalerEnum.STANDARDIZE),
        (RBFKernel(ard=False), ScalerEnum.STANDARDIZE, ScalerEnum.STANDARDIZE),
        (RBFKernel(ard=False), ScalerEnum.IDENTITY, ScalerEnum.IDENTITY),
    ],
)
def test_MixedSingleTaskGPModel(kernel, scaler, output_scaler):
    inputs = Inputs(
        features=[
            ContinuousInput(
                key=f"x_{i + 1}",
                bounds=(-4, 4),
            )
            for i in range(2)
        ]
        + [CategoricalInput(key="x_cat", categories=["mama", "papa"])],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = inputs.sample(n=10)
    experiments.eval("y=((x_1**2 + x_2 - 11)**2+(x_1 + x_2**2 -7)**2)", inplace=True)
    experiments.loc[experiments.x_cat == "mama", "y"] *= 5.0
    experiments.loc[experiments.x_cat == "papa", "y"] /= 2.0
    experiments["valid_y"] = 1

    surrogate_data = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        # input_preprocessing_specs={"x_cat": CategoricalEncodingEnum.ONE_HOT},
        scaler=scaler,
        output_scaler=output_scaler,
        continuous_kernel=kernel,
        categorical_kernel=HammingDistanceKernel(),
    )
    model = surrogates.map(surrogate_data)
    assert isinstance(model, surrogates.SingleTaskGPSurrogate)
    with pytest.raises(ValueError):
        model.dumps()
    model.fit(experiments)

    # dump the model
    dump = model.dumps()
    # make predictions
    samples = inputs.sample(5)
    preds = model.predict(samples)
    assert preds.shape == (5, 2)
    # check that model is composed correctly
    assert isinstance(model.model, SingleTaskGP)

    if output_scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.outcome_transform, Standardize)
    elif output_scaler == ScalerEnum.IDENTITY:
        assert not hasattr(model.model, "outcome_transform")
    if scaler == ScalerEnum.NORMALIZE:
        # assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform, Normalize)
        assert torch.eq(
            model.model.input_transform.indices,
            torch.tensor([0, 1], dtype=torch.int64),
        ).all()
        # assert isinstance(model.model.input_transform.tf2, OneHotToNumeric)
    elif scaler == ScalerEnum.STANDARDIZE:
        # assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform, InputStandardize)
        assert torch.eq(
            model.model.input_transform.indices,
            torch.tensor([0, 1], dtype=torch.int64),
        ).all()
    else:
        assert not hasattr(model.model, "input_transform")
        # assert isinstance(model.model.input_transform, OneHotToNumeric)
    assert model.is_compatibilized is False
    # reload the model from dump and check for equality in predictions
    model2 = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        continuous_kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
    )
    model2 = surrogates.map(model2)
    model2.loads(dump)
    preds2 = model2.predict(samples)
    assert_frame_equal(preds, preds2)


@pytest.mark.skipif(not RDKIT_AVAILABLE, reason="requires rdkit")
@pytest.mark.parametrize(
    "kernel, scaler, output_scaler",
    [
        (RBFKernel(ard=True), ScalerEnum.NORMALIZE, ScalerEnum.STANDARDIZE),
        (RBFKernel(ard=False), ScalerEnum.STANDARDIZE, ScalerEnum.STANDARDIZE),
        (RBFKernel(ard=False), ScalerEnum.IDENTITY, ScalerEnum.IDENTITY),
    ],
)
def test_MixedSingleTaskGPModel_mordred(kernel, scaler, output_scaler):
    inputs = Inputs(
        features=[
            CategoricalMolecularInput(
                key="x_mol",
                categories=[
                    "CC(=O)Oc1ccccc1C(=O)O",
                    "c1ccccc1",
                    "[CH3][CH2][OH]",
                    "N[C@](C)(F)C(=O)O",
                ],
            ),
            CategoricalInput(key="x_cat", categories=["a", "b"]),
        ],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = [
        ["CC(=O)Oc1ccccc1C(=O)O", "a", 88.0],
        ["c1ccccc1", "a", 35.0],
        ["[CH3][CH2][OH]", "b", 69.0],
        ["N[C@](C)(F)C(=O)O", "b", 20.0],
    ]
    experiments = pd.DataFrame(experiments, columns=["x_mol", "x_cat", "y"])
    experiments["valid_y"] = 1
    model = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        scaler=scaler,
        output_scaler=output_scaler,
        continuous_kernel=kernel,
        categorical_encodings={
            "x_mol": MordredDescriptors(descriptors=["NssCH2", "ATSC2d"]),
        },
    )
    model = surrogates.map(model)
    model.fit(experiments)
    # dump the model
    dump = model.dumps()
    # make predictions
    preds = model.predict(experiments.iloc[:-1])
    assert preds.shape == (3, 2)
    # check that model is composed correctly
    assert isinstance(model.model, SingleTaskGP)
    if output_scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.outcome_transform, Standardize)
    elif output_scaler == ScalerEnum.IDENTITY:
        assert not hasattr(model.model, "outcome_transform")

    if scaler == ScalerEnum.NORMALIZE:
        assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform.scaler, Normalize)
        assert torch.eq(
            model.model.input_transform.scaler.indices,
            torch.tensor([0, 1], dtype=torch.int64),
        ).all()
        assert isinstance(model.model.input_transform.cat, NumericToCategoricalEncoding)
    elif scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.input_transform, ChainedInputTransform)
        assert isinstance(model.model.input_transform.scaler, InputStandardize)
        assert torch.eq(
            model.model.input_transform.scaler.indices,
            torch.tensor([0, 1], dtype=torch.int64),
        ).all()
        assert isinstance(model.model.input_transform.cat, NumericToCategoricalEncoding)
    else:
        assert isinstance(model.model.input_transform, NumericToCategoricalEncoding)

    # check correct indices
    assert torch.allclose(
        model.model.covar_module.kernels[0].base_kernel.kernels[0].active_dims,
        torch.tensor([0, 1]),
    )

    assert torch.allclose(
        model.model.covar_module.kernels[0]
        .base_kernel.kernels[1]
        .base_kernel.active_dims,
        torch.tensor([2]),
    )

    assert model.is_compatibilized is False
    # reload the model from dump and check for equality in predictions
    model2 = MixedSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        continuous_kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
        categorical_encodings={
            "x_mol": MordredDescriptors(descriptors=["NssCH2", "ATSC2d"]),
        },
    )
    model2 = surrogates.map(model2)
    model2.loads(dump)
    preds2 = model2.predict(experiments.iloc[:-1])
    assert_frame_equal(preds, preds2)


@pytest.mark.parametrize(
    "kernel, scaler, output_scaler",
    [
        (
            ScaleKernel(base_kernel=RBFKernel(ard=True)),
            ScalerEnum.NORMALIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.STANDARDIZE,
            ScalerEnum.STANDARDIZE,
        ),
        (
            ScaleKernel(base_kernel=RBFKernel(ard=False)),
            ScalerEnum.IDENTITY,
            ScalerEnum.IDENTITY,
        ),
    ],
)
def test_RobustSingleTaskGPModel(kernel, scaler, output_scaler):
    inputs = Inputs(
        features=[
            ContinuousInput(
                key=f"x_{i + 1}",
                bounds=(-4, 4),
            )
            for i in range(2)
        ],
    )
    outputs = Outputs(features=[ContinuousOutput(key="y")])
    experiments = inputs.sample(n=10)
    experiments.eval("y=((x_1**2 + x_2 - 11)**2+(x_1 + x_2**2 -7)**2)", inplace=True)
    experiments["valid_y"] = 1
    model = RobustSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
        output_scaler=output_scaler,
    )
    model = surrogates.map(model)
    samples = inputs.sample(5)
    # test error on non fitted model
    with pytest.raises(ValueError):
        model.predict(samples)
    model.fit(experiments)
    # dump the model
    dump = model.dumps()
    # make predictions
    samples2 = samples.copy()
    samples2 = samples2.astype({"x_1": "object"})
    preds = model.predict(samples2)
    assert preds.shape == (5, 2)
    # check that model is composed correctly
    assert isinstance(model.model, RobustRelevancePursuitSingleTaskGP)
    if output_scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.outcome_transform, Standardize)
    elif output_scaler == ScalerEnum.IDENTITY:
        assert not hasattr(model.model, "outcome_transform")
    if scaler == ScalerEnum.NORMALIZE:
        assert isinstance(model.model.input_transform, Normalize)
    elif scaler == ScalerEnum.STANDARDIZE:
        assert isinstance(model.model.input_transform, InputStandardize)
    else:
        with pytest.raises(AttributeError):
            assert model.model.input_transform is None
    assert model.is_compatibilized is False
    # reload the model from dump and check for equality in predictions
    model2 = RobustSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
    )
    model2 = surrogates.map(model2)
    model2.loads(dump)
    preds2 = model2.predict(samples)

    assert_frame_equal(preds, preds2)

    model3 = RobustSingleTaskGPSurrogate(
        inputs=inputs,
        outputs=outputs,
        kernel=kernel,
        scaler=scaler,
    )

    model3 = surrogates.map(model3)
    # test predict outliers
    preds_outliers = model3.predict_outliers(experiments)

    # assert that preds_outliers dataframe had the same length as experiments
    assert len(preds_outliers) == len(experiments)

    # check for the correct columns
    assert set(preds_outliers.columns) == {"y_pred", "y_sd", "y_rho"}


def test_RobustSingleTaskGPHyperconfig():
    # we test here also the basic trainable
    benchmark = Himmelblau()
    surrogate_data_no_hy = RobustSingleTaskGPSurrogate(
        inputs=benchmark.domain.inputs,
        outputs=benchmark.domain.outputs,
        hyperconfig=None,
    )
    with pytest.raises(ValueError, match="No hyperconfig available."):
        surrogate_data_no_hy.update_hyperparameters(
            benchmark.domain.inputs.sample(1).loc[0],
        )
    # test that correct stuff is written
    surrogate_data = RobustSingleTaskGPSurrogate(
        inputs=benchmark.domain.inputs,
        outputs=benchmark.domain.outputs,
    )

    assert (
        surrogate_data.kernel.lengthscale_constraint
        == ROBUSTGP_LENGTHSCALE_CONSTRAINT()
    )

    candidate = surrogate_data.hyperconfig.inputs.sample(1).loc[0]
    # surrogate_data.update_hyperparameters(candidate, lengthscale_constraint=ROBUSTGP_LENGTHSCALE_CONSTRAINT(), outputscale_constraint=ROBUSTGP_OUTPUTSCALE_CONSTRAINT())
    surrogate_data.update_hyperparameters(candidate)
    if hasattr(surrogate_data.kernel, "base_kernel"):
        # if surrogate_data.kernel == ScaleKernel():
        #     assert surrogate_data.kernel.outputscale_constraint == ROBUSTGP_OUTPUTSCALE_CONSTRAINT()
        assert surrogate_data.kernel.base_kernel.ard == (candidate["ard"] == "True")
        # assert surrogate_data.kernel.base_kernel.lengthscale_constraint == ROBUSTGP_LENGTHSCALE_CONSTRAINT()
        if candidate.kernel == "matern_1.5":
            assert isinstance(surrogate_data.kernel.base_kernel, MaternKernel)
            assert surrogate_data.kernel.base_kernel.nu == 1.5
        elif candidate.kernel == "matern_2.5":
            assert isinstance(surrogate_data.kernel.base_kernel, MaternKernel)
            assert surrogate_data.kernel.base_kernel.nu == 2.5
        else:
            assert isinstance(surrogate_data.kernel.base_kernel, RBFKernel)
        if candidate.prior == "mbo":
            assert surrogate_data.noise_prior == MBO_NOISE_PRIOR()
            assert surrogate_data.kernel.outputscale_prior == MBO_OUTPUTSCALE_PRIOR()
            assert (
                surrogate_data.kernel.base_kernel.lengthscale_prior
                == MBO_LENGTHSCALE_PRIOR()
            )
        elif candidate.prior == "threesix":
            assert surrogate_data.noise_prior == THREESIX_NOISE_PRIOR()
            assert surrogate_data.kernel.outputscale_prior == THREESIX_SCALE_PRIOR()
            assert (
                surrogate_data.kernel.base_kernel.lengthscale_prior
                == THREESIX_LENGTHSCALE_PRIOR()
            )
        else:
            assert surrogate_data.noise_prior == HVARFNER_NOISE_PRIOR()
            assert surrogate_data.kernel.outputscale_prior == THREESIX_SCALE_PRIOR()
            assert (
                surrogate_data.kernel.base_kernel.lengthscale_prior
                == HVARFNER_LENGTHSCALE_PRIOR()
            )
