#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import itertools
import warnings

import torch
from botorch.exceptions.errors import BotorchTensorDimensionError
from botorch.exceptions.warnings import OptimizationWarning
from botorch.fit import fit_gpytorch_model
from botorch.models import ModelListGP
from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
from botorch.models.transforms import Standardize
from botorch.models.transforms.input import Normalize
from botorch.posteriors import GPyTorchPosterior
from botorch.utils.testing import BotorchTestCase, _get_random_data
from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
from gpytorch.kernels import MaternKernel, ScaleKernel
from gpytorch.likelihoods import LikelihoodList
from gpytorch.means import ConstantMean
from gpytorch.mlls import SumMarginalLogLikelihood
from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
from gpytorch.priors import GammaPrior


def _get_model(fixed_noise=False, use_octf=False, use_intf=False, **tkwargs):
    train_x1, train_y1 = _get_random_data(
        batch_shape=torch.Size(), m=1, n=10, **tkwargs
    )
    train_x2, train_y2 = _get_random_data(
        batch_shape=torch.Size(), m=1, n=11, **tkwargs
    )
    octfs = [Standardize(m=1), Standardize(m=1)] if use_octf else [None, None]
    intfs = [Normalize(d=1), Normalize(d=1)] if use_intf else [None, None]
    if fixed_noise:
        train_y1_var = 0.1 + 0.1 * torch.rand_like(train_y1, **tkwargs)
        train_y2_var = 0.1 + 0.1 * torch.rand_like(train_y2, **tkwargs)
        model1 = FixedNoiseGP(
            train_X=train_x1,
            train_Y=train_y1,
            train_Yvar=train_y1_var,
            outcome_transform=octfs[0],
            input_transform=intfs[0],
        )
        model2 = FixedNoiseGP(
            train_X=train_x2,
            train_Y=train_y2,
            train_Yvar=train_y2_var,
            outcome_transform=octfs[1],
            input_transform=intfs[1],
        )
    else:
        model1 = SingleTaskGP(
            train_X=train_x1,
            train_Y=train_y1,
            outcome_transform=octfs[0],
            input_transform=intfs[0],
        )
        model2 = SingleTaskGP(
            train_X=train_x2,
            train_Y=train_y2,
            outcome_transform=octfs[1],
            input_transform=intfs[1],
        )
    model = ModelListGP(model1, model2)
    return model.to(**tkwargs)


class TestModelListGP(BotorchTestCase):
    def test_ModelListGP(self):
        for dtype, use_octf in itertools.product(
            (torch.float, torch.double), (False, True)
        ):
            tkwargs = {"device": self.device, "dtype": dtype}
            model = _get_model(use_octf=use_octf, **tkwargs)
            self.assertIsInstance(model, ModelListGP)
            self.assertIsInstance(model.likelihood, LikelihoodList)
            for m in model.models:
                self.assertIsInstance(m.mean_module, ConstantMean)
                self.assertIsInstance(m.covar_module, ScaleKernel)
                matern_kernel = m.covar_module.base_kernel
                self.assertIsInstance(matern_kernel, MaternKernel)
                self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)
                if use_octf:
                    self.assertIsInstance(m.outcome_transform, Standardize)

            # test constructing likelihood wrapper
            mll = SumMarginalLogLikelihood(model.likelihood, model)
            for mll_ in mll.mlls:
                self.assertIsInstance(mll_, ExactMarginalLogLikelihood)

            # test model fitting (sequential)
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=OptimizationWarning)
                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=OptimizationWarning)
                # test model fitting (joint)
                mll = fit_gpytorch_model(
                    mll, options={"maxiter": 1}, max_retries=1, sequential=False
                )

            # test subset outputs
            subset_model = model.subset_output([1])
            self.assertIsInstance(subset_model, ModelListGP)
            self.assertEqual(len(subset_model.models), 1)
            sd_subset = subset_model.models[0].state_dict()
            sd = model.models[1].state_dict()
            self.assertTrue(set(sd_subset.keys()) == set(sd.keys()))
            self.assertTrue(all(torch.equal(v, sd[k]) for k, v in sd_subset.items()))

            # test posterior
            test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
            posterior = model.posterior(test_x)
            self.assertIsInstance(posterior, GPyTorchPosterior)
            self.assertIsInstance(posterior.mvn, MultitaskMultivariateNormal)
            if use_octf:
                # ensure un-transformation is applied
                submodel = model.models[0]
                p0 = submodel.posterior(test_x)
                tmp_tf = submodel.outcome_transform
                del submodel.outcome_transform
                p0_tf = submodel.posterior(test_x)
                submodel.outcome_transform = tmp_tf
                expected_var = tmp_tf.untransform_posterior(p0_tf).variance
                self.assertTrue(torch.allclose(p0.variance, expected_var))

            # test observation_noise
            posterior = model.posterior(test_x, observation_noise=True)
            self.assertIsInstance(posterior, GPyTorchPosterior)
            self.assertIsInstance(posterior.mvn, MultitaskMultivariateNormal)

            # test output_indices
            posterior = model.posterior(
                test_x, output_indices=[0], observation_noise=True
            )
            self.assertIsInstance(posterior, GPyTorchPosterior)
            self.assertIsInstance(posterior.mvn, MultivariateNormal)

            # test condition_on_observations
            f_x = torch.rand(2, 1, **tkwargs)
            f_y = torch.rand(2, 2, **tkwargs)
            cm = model.condition_on_observations(f_x, f_y)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations batched
            f_x = torch.rand(3, 2, 1, **tkwargs)
            f_y = torch.rand(3, 2, 2, **tkwargs)
            cm = model.condition_on_observations(f_x, f_y)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations batched (fast fantasies)
            f_x = torch.rand(2, 1, **tkwargs)
            f_y = torch.rand(3, 2, 2, **tkwargs)
            cm = model.condition_on_observations(f_x, f_y)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations (incorrect input shape error)
            with self.assertRaises(BotorchTensorDimensionError):
                model.condition_on_observations(f_x, torch.rand(3, 2, 3, **tkwargs))

    def test_ModelListGP_fixed_noise(self):
        for dtype, use_octf in itertools.product(
            (torch.float, torch.double), (False, True)
        ):
            tkwargs = {"device": self.device, "dtype": dtype}
            model = _get_model(fixed_noise=True, use_octf=use_octf, **tkwargs)
            self.assertIsInstance(model, ModelListGP)
            self.assertIsInstance(model.likelihood, LikelihoodList)
            for m in model.models:
                self.assertIsInstance(m.mean_module, ConstantMean)
                self.assertIsInstance(m.covar_module, ScaleKernel)
                matern_kernel = m.covar_module.base_kernel
                self.assertIsInstance(matern_kernel, MaternKernel)
                self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)

            # test model fitting
            mll = SumMarginalLogLikelihood(model.likelihood, model)
            for mll_ in mll.mlls:
                self.assertIsInstance(mll_, ExactMarginalLogLikelihood)
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=OptimizationWarning)
                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)

            # test posterior
            test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
            posterior = model.posterior(test_x)
            self.assertIsInstance(posterior, GPyTorchPosterior)
            self.assertIsInstance(posterior.mvn, MultitaskMultivariateNormal)
            if use_octf:
                # ensure un-transformation is applied
                submodel = model.models[0]
                p0 = submodel.posterior(test_x)
                tmp_tf = submodel.outcome_transform
                del submodel.outcome_transform
                p0_tf = submodel.posterior(test_x)
                submodel.outcome_transform = tmp_tf
                expected_var = tmp_tf.untransform_posterior(p0_tf).variance
                self.assertTrue(torch.allclose(p0.variance, expected_var))

            # test output_indices
            posterior = model.posterior(
                test_x, output_indices=[0], observation_noise=True
            )
            self.assertIsInstance(posterior, GPyTorchPosterior)
            self.assertIsInstance(posterior.mvn, MultivariateNormal)

            # test condition_on_observations
            f_x = torch.rand(2, 1, **tkwargs)
            f_y = torch.rand(2, 2, **tkwargs)
            noise = 0.1 + 0.1 * torch.rand_like(f_y)
            cm = model.condition_on_observations(f_x, f_y, noise=noise)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations batched
            f_x = torch.rand(3, 2, 1, **tkwargs)
            f_y = torch.rand(3, 2, 2, **tkwargs)
            noise = 0.1 + 0.1 * torch.rand_like(f_y)
            cm = model.condition_on_observations(f_x, f_y, noise=noise)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations batched (fast fantasies)
            f_x = torch.rand(2, 1, **tkwargs)
            f_y = torch.rand(3, 2, 2, **tkwargs)
            noise = 0.1 + 0.1 * torch.rand(2, 2, **tkwargs)
            cm = model.condition_on_observations(f_x, f_y, noise=noise)
            self.assertIsInstance(cm, ModelListGP)

            # test condition_on_observations (incorrect input shape error)
            with self.assertRaises(BotorchTensorDimensionError):
                model.condition_on_observations(
                    f_x, torch.rand(3, 2, 3, **tkwargs), noise=noise
                )
            # test condition_on_observations (incorrect noise shape error)
            f_y = torch.rand(2, 2, **tkwargs)
            with self.assertRaises(BotorchTensorDimensionError):
                model.condition_on_observations(
                    f_x, f_y, noise=torch.rand(2, 3, **tkwargs)
                )

    def test_ModelListGP_single(self):
        tkwargs = {"device": self.device, "dtype": torch.float}
        train_x1, train_y1 = _get_random_data(
            batch_shape=torch.Size(), m=1, n=10, **tkwargs
        )
        model1 = SingleTaskGP(train_X=train_x1, train_Y=train_y1)
        model = ModelListGP(model1)
        model.to(**tkwargs)
        test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
        posterior = model.posterior(test_x)
        self.assertIsInstance(posterior, GPyTorchPosterior)
        self.assertIsInstance(posterior.mvn, MultivariateNormal)

    def test_transform_revert_train_inputs(self):
        tkwargs = {"device": self.device, "dtype": torch.float}
        model_list = _get_model(use_intf=True, **tkwargs)
        org_inputs = [m.train_inputs[0] for m in model_list.models]
        model_list.eval()
        for i, m in enumerate(model_list.models):
            self.assertTrue(
                torch.allclose(
                    m.train_inputs[0],
                    m.input_transform.preprocess_transform(org_inputs[i]),
                )
            )
            self.assertTrue(m._has_transformed_inputs)
            self.assertTrue(torch.equal(m._original_train_inputs, org_inputs[i]))
        model_list.train(mode=True)
        for i, m in enumerate(model_list.models):
            self.assertTrue(torch.equal(m.train_inputs[0], org_inputs[i]))
            self.assertFalse(m._has_transformed_inputs)
        model_list.train(mode=False)
        for i, m in enumerate(model_list.models):
            self.assertTrue(
                torch.allclose(
                    m.train_inputs[0],
                    m.input_transform.preprocess_transform(org_inputs[i]),
                )
            )
            self.assertTrue(m._has_transformed_inputs)
            self.assertTrue(torch.equal(m._original_train_inputs, org_inputs[i]))
