from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import logging
import math
import torch
from .kernel import Kernel

logger = logging.getLogger()


class SpectralMixtureKernel(Kernel):
    r"""
    Computes a covariance matrix based on the Spectral Mixture Kernel
    between inputs :math:`mathbf{x_1}` and :math:`mathbf{x_2}`:
    It was proposed in `Gaussian Process Kernels for Pattern Discovery and Extrapolation`_.

    .. note::

        Unlike other kernels,
        * :attr:`ard_num_dums` **must equal** the number of dimensions of the data
        * :attr:`batch_size` **must equal** the batch size of the data (1 if the data is not batched)
        * This kernel should not be combined with a :class:`gpytorch.kernels.ScaleKernel`.

    Args:
        :attr:`num_mixtures` (int, optional):
            The number of components in the mixture.
        :attr:`ard_num_dims` (int, optional):
            Set this to match the dimensionality of the input.
            It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `1`
        :attr:`batch_size` (int, optional):
            Set this if the data is
            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `1`
        :attr:`active_dims` (tuple of ints, optional):
            Set this if you want to
            compute the covariance of only a few input dimensions. The ints
            corresponds to the indices of the dimensions. Default: `None`.
        :attr:`eps` (float):
            The minimum value that the lengthscale can take
            (prevents divide by zero errors). Default: `1e-6`.

    Attributes:
        :attr:`mixture_lengthscale` (Tensor):
            The lengthscale parameter. Given `k` mixture components,
            and `b x n x d` data, this will be of size `b x k x 1 x d`.
        :attr:`mixture_means` (Tensor):
            The mixture mean parameters (`b x k x 1 x d`).
        :attr:`mixture_weights` (Tensor):
            The mixture weight parameters (`b x k`).

    Example:
        >>> # Non-batch
        >>> x = torch.randn(10, 5)
        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, ard_dum_dims=5)
        >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
        >>>
        >>> # Batch
        >>> batch_x = torch.randn(2, 10, 5)
        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, batch_size=2, ard_dum_dims=5)
        >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)


    .. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
        https://arxiv.org/pdf/1302.4245.pdf
    """
    # TODO: add equation to docs

    def __init__(
        self,
        num_mixtures=None,
        ard_num_dims=1,
        batch_size=1,
        active_dims=None,
        eps=1e-6,
        log_mixture_scales_prior=None,
        log_mixture_means_prior=None,
        log_mixture_weights_prior=None,
    ):
        if num_mixtures is None:
            raise RuntimeError("num_mixtures is a required argument")
        if (
            log_mixture_means_prior is not None
            or log_mixture_scales_prior is not None
            or log_mixture_weights_prior is not None
        ):
            logger.warning("Priors not implemented for SpectralMixtureKernel")

        # This kernel does not use the default lengthscale
        super(SpectralMixtureKernel, self).__init__(active_dims=active_dims)
        self.num_mixtures = num_mixtures
        self.batch_size = batch_size
        self.ard_num_dims = ard_num_dims
        self.eps = eps

        self.register_parameter(
            name="log_mixture_weights", parameter=torch.nn.Parameter(torch.zeros(self.batch_size, self.num_mixtures))
        )
        self.register_parameter(
            name="log_mixture_means",
            parameter=torch.nn.Parameter(torch.zeros(self.batch_size, self.num_mixtures, 1, self.ard_num_dims)),
        )
        self.register_parameter(
            name="log_mixture_scales",
            parameter=torch.nn.Parameter(torch.zeros(self.batch_size, self.num_mixtures, 1, self.ard_num_dims)),
        )

    @property
    def mixture_scales(self):
        return self.log_mixture_scales.exp().clamp(self.eps, 1e5)

    @property
    def mixture_means(self):
        return self.log_mixture_means.exp().clamp(self.eps, 1e5)

    @property
    def mixture_weights(self):
        return self.log_mixture_weights.exp().clamp(self.eps, 1e5)

    def initialize_from_data(self, train_x, train_y, **kwargs):
        if not torch.is_tensor(train_x) or not torch.is_tensor(train_y):
            raise RuntimeError("train_x and train_y should be tensors")
        if train_x.ndimension() == 1:
            train_x = train_x.unsqueeze(-1)
        if train_x.ndimension() == 2:
            train_x = train_x.unsqueeze(0)

        train_x_sort = train_x.sort(1)[0]
        max_dist = train_x_sort[:, -1, :] - train_x_sort[:, 0, :]
        min_dist_sort = (train_x_sort[:, 1:, :] - train_x_sort[:, :-1, :]).squeeze(0)
        min_dist = torch.zeros(1, self.ard_num_dims)
        for ind in range(self.ard_num_dims):
            min_dist[:, ind] = min_dist_sort[(torch.nonzero(min_dist_sort[:, ind]))[0], ind]

        # Inverse of lengthscales should be drawn from truncated Gaussian | N(0, max_dist^2) |
        self.log_mixture_scales.data.normal_().mul_(max_dist).abs_().pow_(-1).log_()
        # Draw means from Unif(0, 0.5 / minimum distance between two points)
        self.log_mixture_means.data.uniform_().mul_(0.5).div_(min_dist).log_()
        # Mixture weights should be roughly the stdv of the y values divided by the number of mixtures
        self.log_mixture_weights.data.fill_(train_y.std() / self.num_mixtures).log_()

    def forward(self, x1, x2, **params):
        batch_size, n, num_dims = x1.size()
        _, m, _ = x2.size()

        if not num_dims == self.ard_num_dims:
            raise RuntimeError(
                "The SpectralMixtureKernel expected the input to have {} dimensionality "
                "(based on the ard_num_dims argument). Got {}.".format(self.ard_num_dims, num_dims)
            )
        if not batch_size == self.batch_size:
            raise RuntimeError(
                "The SpectralMixtureKernel expected the input to have a batch_size of {} "
                "(based on the batch_size argument). Got {}.".format(self.batch_size, batch_size)
            )

        # Expand x1 and x2 to account for the number of mixtures
        # Should make x1/x2 (b x k x n x d) for k mixtures
        x1_ = x1.unsqueeze(1)
        x2_ = x2.unsqueeze(1)

        # Compute distances - scaled by appropriate parameters
        x1_exp = x1_ * self.mixture_scales
        x2_exp = x2_ * self.mixture_scales
        x1_cos = x1_ * self.mixture_means
        x2_cos = x2_ * self.mixture_means

        # Create grids
        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, **params)
        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, **params)

        # Compute the exponential and cosine terms
        exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
        cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
        res = exp_term.exp_() * cos_term.cos_()

        # Product omer dimensions
        res = res.prod(-1)

        # Sum over mixtures
        mixture_weights = self.mixture_weights
        while mixture_weights.dim() < res.dim():
            mixture_weights.unsqueeze_(-1)
        res = (res * mixture_weights).sum(1)
        return res
