#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools

import numpy as np

from ... import opcodes as OperandDef
from ...serialization.serializables import NDArrayField, StringField, Float64Field
from ...config import options
from ..utils import decide_chunk_sizes, gen_random_seeds
from ..array_utils import array_module, device
from .core import TensorRandomOperandMixin, TensorDistribution, TENSOR_CHUNK_TYPE


class TensorMultivariateNormal(TensorDistribution, TensorRandomOperandMixin):
    _op_type_ = OperandDef.RAND_MULTIVARIATE_NORMAL

    _fields_ = "mean", "cov", "size", "check_valid", "tol"
    mean = NDArrayField("mean")
    cov = NDArrayField("cov")
    check_valid = StringField("check_valid")
    tol = Float64Field("tol")
    _func_name = "multivariate_normal"

    def __call__(self, chunk_size=None):
        N = self.mean.size
        if self.size is None:
            shape = (N,)
        else:
            try:
                shape = tuple(self.size) + (N,)
            except TypeError:
                shape = (self.size, N)

        return self.new_tensor(None, shape, raw_chunk_size=chunk_size)

    @classmethod
    def tile(cls, op):
        tensor = op.outputs[0]
        chunk_size = tensor.extra_params.raw_chunk_size or options.chunk_size
        nsplits = decide_chunk_sizes(
            tensor.shape[:-1], chunk_size, tensor.dtype.itemsize
        ) + ((tensor.shape[-1],),)

        mean_chunk = op.mean.chunks[0] if hasattr(op.mean, "chunks") else op.mean
        cov_chunk = op.cov.chunks[0] if hasattr(op.cov, "chunks") else op.cov

        idxes = list(itertools.product(*[range(len(s)) for s in nsplits]))
        seeds = gen_random_seeds(len(idxes), np.random.RandomState(op.seed))

        out_chunks = []
        for seed, out_idx, shape in zip(seeds, idxes, itertools.product(*nsplits)):
            chunk_op = op.copy().reset_key()
            chunk_op._state = None
            chunk_op.seed = seed
            chunk_op.size = shape[:-1]
            out_chunk = chunk_op.new_chunk(
                [mean_chunk, cov_chunk], shape=shape, index=out_idx
            )
            out_chunks.append(out_chunk)

        new_op = op.copy()
        return new_op.new_tensors(
            op.inputs, tensor.shape, chunks=out_chunks, nsplits=nsplits
        )

    @classmethod
    def execute(cls, ctx, op):
        xp = array_module(op.gpu)
        if xp is np:
            device_id = -1
        else:
            device_id = op.device or 0

        with device(device_id):
            rs = xp.random.RandomState(op.seed)

            args = []
            for k in op.args:
                val = getattr(op, k, None)
                if isinstance(val, TENSOR_CHUNK_TYPE):
                    args.append(ctx[val.key])
                else:
                    args.append(val)
            mean, cov = args[:2]
            kw = {}
            if args[2] is not None:
                kw["size"] = args[2]
            if args[3] is not None:
                kw["check_valid"] = args[3]
            if args[4] is not None:
                kw["tol"] = args[4]

            try:
                res = rs.multivariate_normal(mean, cov, **kw)
                if xp is not np:
                    ctx[op.outputs[0].key] = xp.asarray(res)
                else:
                    ctx[op.outputs[0].key] = res
            except AttributeError:
                if xp is not np:
                    # cupy cannot generate data, fallback to numpy first
                    rs = np.random.RandomState(op.seed)
                    res = rs.multivariate_normal(mean, cov, **kw)
                    ctx[op.outputs[0].key] = xp.asarray(res)
                else:
                    raise


def multivariate_normal(
    random_state,
    mean,
    cov,
    size=None,
    check_valid=None,
    tol=None,
    chunk_size=None,
    gpu=None,
    dtype=None,
):
    """
    Draw random samples from a multivariate normal distribution.

    The multivariate normal, multinormal or Gaussian distribution is a
    generalization of the one-dimensional normal distribution to higher
    dimensions.  Such a distribution is specified by its mean and
    covariance matrix.  These parameters are analogous to the mean
    (average or "center") and variance (standard deviation, or "width,"
    squared) of the one-dimensional normal distribution.

    Parameters
    ----------
    mean : 1-D array_like, of length N
        Mean of the N-dimensional distribution.
    cov : 2-D array_like, of shape (N, N)
        Covariance matrix of the distribution. It must be symmetric and
        positive-semidefinite for proper sampling.
    size : int or tuple of ints, optional
        Given a shape of, for example, ``(m,n,k)``, ``m*n*k`` samples are
        generated, and packed in an `m`-by-`n`-by-`k` arrangement.  Because
        each sample is `N`-dimensional, the output shape is ``(m,n,k,N)``.
        If no shape is specified, a single (`N`-D) sample is returned.
    check_valid : { 'warn', 'raise', 'ignore' }, optional
        Behavior when the covariance matrix is not positive semidefinite.
    tol : float, optional
        Tolerance when checking the singular values in covariance matrix.
    chunk_size : int or tuple of int or tuple of ints, optional
        Desired chunk size on each dimension
    gpu : bool, optional
        Allocate the tensor on GPU if True, False as default
    dtype : data-type, optional
      Data-type of the returned tensor.

    Returns
    -------
    out : Tensor
        The drawn samples, of shape *size*, if that was provided.  If not,
        the shape is ``(N,)``.

        In other words, each entry ``out[i,j,...,:]`` is an N-dimensional
        value drawn from the distribution.

    Notes
    -----
    The mean is a coordinate in N-dimensional space, which represents the
    location where samples are most likely to be generated.  This is
    analogous to the peak of the bell curve for the one-dimensional or
    univariate normal distribution.

    Covariance indicates the level to which two variables vary together.
    From the multivariate normal distribution, we draw N-dimensional
    samples, :math:`X = [x_1, x_2, ... x_N]`.  The covariance matrix
    element :math:`C_{ij}` is the covariance of :math:`x_i` and :math:`x_j`.
    The element :math:`C_{ii}` is the variance of :math:`x_i` (i.e. its
    "spread").

    Instead of specifying the full covariance matrix, popular
    approximations include:

      - Spherical covariance (`cov` is a multiple of the identity matrix)
      - Diagonal covariance (`cov` has non-negative elements, and only on
        the diagonal)

    This geometrical property can be seen in two dimensions by plotting
    generated data-points:

    >>> mean = [0, 0]
    >>> cov = [[1, 0], [0, 100]]  # diagonal covariance

    Diagonal covariance means that points are oriented along x or y-axis:

    >>> import matplotlib.pyplot as plt
    >>> import mars.tensor as mt
    >>> x, y = mt.random.multivariate_normal(mean, cov, 5000).T
    >>> plt.plot(x.execute(), y.execute(), 'x')
    >>> plt.axis('equal')
    >>> plt.show()

    Note that the covariance matrix must be positive semidefinite (a.k.a.
    nonnegative-definite). Otherwise, the behavior of this method is
    undefined and backwards compatibility is not guaranteed.

    References
    ----------
    .. [1] Papoulis, A., "Probability, Random Variables, and Stochastic
           Processes," 3rd ed., New York: McGraw-Hill, 1991.
    .. [2] Duda, R. O., Hart, P. E., and Stork, D. G., "Pattern
           Classification," 2nd ed., New York: Wiley, 2001.

    Examples
    --------
    >>> mean = (1, 2)
    >>> cov = [[1, 0], [0, 1]]
    >>> x = mt.random.multivariate_normal(mean, cov, (3, 3))
    >>> x.shape
    (3, 3, 2)

    The following is probably true, given that 0.6 is roughly twice the
    standard deviation:

    >>> list(((x[0,0,:] - mean) < 0.6).execute())
    [True, True]
    """
    mean = np.asarray(mean)
    cov = np.asarray(cov)

    if mean.ndim != 1:
        raise ValueError("mean must be 1 dimensional")
    if cov.ndim != 2:
        raise ValueError("cov must be 1 dimensional")
    if len(set(mean.shape + cov.shape)) != 1:
        raise ValueError("mean and cov must have same length")

    if dtype is None:
        small_kw = {}
        if check_valid:
            small_kw["check_valid"] = check_valid
        if tol:
            small_kw["tol"] = tol
        dtype = np.random.multivariate_normal(mean, cov, size=(0,), **small_kw).dtype

    size = random_state._handle_size(size)
    seed = gen_random_seeds(1, random_state.to_numpy())[0]
    op = TensorMultivariateNormal(
        mean=mean,
        cov=cov,
        size=size,
        check_valid=check_valid,
        tol=tol,
        seed=seed,
        gpu=gpu,
        dtype=dtype,
    )
    return op(chunk_size=chunk_size)
