#   Copyright 2020 The PyMC Developers
#
#   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 warnings

import aesara.tensor as at
import numpy as np

from aesara.compile import SharedVariable
from aesara.tensor.slinalg import (  # noqa: W0611; pylint: disable=unused-import
    cholesky,
    solve,
)
from aesara.tensor.slinalg import (  # noqa: W0611; pylint: disable=unused-import
    solve_lower_triangular as solve_lower,
)
from aesara.tensor.slinalg import (  # noqa: W0611; pylint: disable=unused-import
    solve_upper_triangular as solve_upper,
)
from aesara.tensor.var import TensorConstant
from scipy.cluster.vq import kmeans

from pymc.aesaraf import compile_pymc, walk_model

# Avoid circular dependency when importing modelcontext
from pymc.distributions.distribution import NoDistribution

assert NoDistribution  # keep both pylint and black happy
from pymc.model import modelcontext

JITTER_DEFAULT = 1e-6


def replace_with_values(vars_needed, replacements=None, model=None):
    R"""
    Replace random variable nodes in the graph with values given by the replacements dict.
    Uses untransformed versions of the inputs, performs some basic input validation.

    Parameters
    ----------
    vars_needed: list of TensorVariables
        A list of variable outputs
    replacements: dict with string keys, numeric values
        The variable name and values to be replaced in the model graph.
    model: Model
        A PyMC model object
    """
    model = modelcontext(model)

    inputs, input_names = [], []
    for rv in walk_model(vars_needed, walk_past_rvs=True):
        if rv in model.named_vars.values() and not isinstance(rv, SharedVariable):
            inputs.append(rv)
            input_names.append(rv.name)

    # Then it's deterministic, no inputs are required, can eval and return
    if len(inputs) == 0:
        return tuple(v.eval() for v in vars_needed)

    fn = compile_pymc(
        inputs,
        vars_needed,
        allow_input_downcast=True,
        accept_inplace=True,
        on_unused_input="ignore",
    )

    # Remove unneeded inputs
    replacements = {name: val for name, val in replacements.items() if name in input_names}
    missing = set(input_names) - set(replacements.keys())

    # Error if more inputs are needed
    if len(missing) > 0:
        missing_str = ", ".join(missing)
        raise ValueError(f"Values for {missing_str} must be included in `replacements`.")

    return fn(**replacements)


def infer_size(X, n_points=None):
    R"""
    Maybe attempt to infer the size, or N, of a Gaussian process input matrix.

    If a specific shape cannot be inferred, for instance if X is symbolic, then an
    error is raised.

    Parameters
    ----------
    X: array-like
        Gaussian process input matrix.
    n_points: None or int
        The number of rows of `X`.  If `None`, the number of rows of `X` is
        calculated from `X` if possible.
    """
    if n_points is None:
        try:
            n_points = int(X.shape[0])
        except TypeError:
            raise TypeError("Cannot infer 'shape', provide as an argument")
    return n_points


def stabilize(K, jitter=JITTER_DEFAULT):
    R"""
    Adds small diagonal to a covariance matrix.

    Often the matrices calculated from covariance functions, `K = cov_func(X)`
    do not appear numerically to be positive semi-definite.  Adding a small
    correction, `jitter`, to the diagonal is usually enough to fix this.

    Parameters
    ----------
    K: array-like
        A square covariance or kernel matrix.
    jitter: float
        A small constant.
    """
    return K + jitter * at.identity_like(K)


def kmeans_inducing_points(n_inducing, X, **kmeans_kwargs):
    R"""
    Use the K-means algorithm to initialize the locations `X` for the inducing
    points `fu`.

    Parameters
    ----------
    n_inducing: int
        The number of inducing points (or k, the number of clusters)
    X: array-like
        Gaussian process input matrix.
    **kmeans_kwargs:
        Extra keyword arguments that are passed to `scipy.cluster.vq.kmeans`
    """
    # first whiten X
    if isinstance(X, TensorConstant):
        X = X.value
    elif isinstance(X, (np.ndarray, tuple, list)):
        X = np.asarray(X)
    else:
        raise TypeError(
            "To use K-means initialization, "
            "please provide X as a type that "
            "can be cast to np.ndarray, instead "
            "of {}".format(type(X))
        )
    scaling = np.std(X, 0)
    # if std of a column is very small (zero), don't normalize that column
    scaling[scaling <= 1e-6] = 1.0
    Xw = X / scaling

    if "k_or_guess" in kmeans_kwargs:
        warn.UserWarning("Use `n_inducing` to set the `k_or_guess` parameter instead.")

    Xu, distortion = kmeans(Xw, k_or_guess=n_inducing, **kmeans_kwargs)
    return Xu * scaling


def conditioned_vars(varnames):
    """Decorator for validating attrs that are conditioned on."""

    def gp_wrapper(cls):
        def make_getter(name):
            def getter(self):
                value = getattr(self, name, None)
                if value is None:
                    raise AttributeError(
                        "'{}' not set.  Provide as argument "
                        "to condition, or call 'prior' "
                        "first".format(name.lstrip("_"))
                    )
                else:
                    return value
                return getattr(self, name)

            return getter

        def make_setter(name):
            def setter(self, val):
                setattr(self, name, val)

            return setter

        for name in varnames:
            getter = make_getter("_" + name)
            setter = make_setter("_" + name)
            setattr(cls, name, property(getter, setter))
        return cls

    return gp_wrapper


def plot_gp_dist(
    ax,
    samples: np.ndarray,
    x: np.ndarray,
    plot_samples=True,
    palette="Reds",
    fill_alpha=0.8,
    samples_alpha=0.1,
    fill_kwargs=None,
    samples_kwargs=None,
):
    """A helper function for plotting 1D GP posteriors from trace

        Parameters
    ----------
    ax: axes
        Matplotlib axes.
    samples: numpy.ndarray
        Array of S posterior predictive sample from a GP.
        Expected shape: (S, X)
    x: numpy.ndarray
        Grid of X values corresponding to the samples.
        Expected shape: (X,) or (X, 1), or (1, X)
    plot_samples: bool
        Plot the GP samples along with posterior (defaults True).
    palette: str
        Palette for coloring output (defaults to "Reds").
    fill_alpha: float
        Alpha value for the posterior interval fill (defaults to 0.8).
    samples_alpha: float
        Alpha value for the sample lines (defaults to 0.1).
    fill_kwargs: dict
        Additional arguments for posterior interval fill (fill_between).
    samples_kwargs: dict
        Additional keyword arguments for samples plot.

    Returns
    -------

    ax: Matplotlib axes
    """
    import matplotlib.pyplot as plt

    if fill_kwargs is None:
        fill_kwargs = {}
    if samples_kwargs is None:
        samples_kwargs = {}
    if np.any(np.isnan(samples)):
        warnings.warn(
            "There are `nan` entries in the [samples] arguments. "
            "The plot will not contain a band!",
            UserWarning,
        )

    cmap = plt.get_cmap(palette)
    percs = np.linspace(51, 99, 40)
    colors = (percs - np.min(percs)) / (np.max(percs) - np.min(percs))
    samples = samples.T
    x = x.flatten()
    for i, p in enumerate(percs[::-1]):
        upper = np.percentile(samples, p, axis=1)
        lower = np.percentile(samples, 100 - p, axis=1)
        color_val = colors[i]
        ax.fill_between(x, upper, lower, color=cmap(color_val), alpha=fill_alpha, **fill_kwargs)
    if plot_samples:
        # plot a few samples
        idx = np.random.randint(0, samples.shape[1], 30)
        ax.plot(x, samples[:, idx], color=cmap(0.9), lw=1, alpha=samples_alpha, **samples_kwargs)

    return ax
