import warnings

import numpy as np
import theano.tensor as tt

from scipy import stats
from scipy.interpolate import InterpolatedUnivariateSpline
from scipy.special import expit

from pymc3.distributions import transforms
from pymc3.distributions.dist_math import (
    SplineWrapper,
    betaln,
    bound,
    clipped_beta_rvs,
    gammaln,
    i0e,
    incomplete_beta,
    log_normal,
    logpow,
    normal_lccdf,
    normal_lcdf,
    zvalue,
)
from pymc3.distributions.distribution import Continuous, draw_values, generate_samples
from pymc3.distributions.special import log_i0
from pymc3.math import invlogit, log1mexp, log1pexp, logdiffexp, logit
from pymc3.theanof import floatX

import theano.tensor as tt
from theano.compile.ops import as_op

def assert_negative_support(var, label, distname, value=-1e-6):
    # Checks for evidence of positive support for a variable
    if var is None:
        return
    try:
        # Transformed distribution
        support = np.isfinite(var.transformed.distribution.dist.logp(value).tag.test_value)
    except AttributeError:
        try:
            # Untransformed distribution
            support = np.isfinite(var.distribution.logp(value).tag.test_value)
        except AttributeError:
            # Otherwise no direct evidence of non-positive support
            support = False

    if np.any(support):
        msg = f"The variable specified for {label} has negative support for {distname}, "
        msg += "likely making it unsuitable for this parameter."
        warnings.warn(msg)

def within_pi_op(k):
    while ((k<-np.pi).any() or (k>np.pi).any()):
        k[k>np.pi] -= np.pi*2
        k[k<-np.pi] += np.pi*2
        return k

class CustomVonMises(Continuous):
    r"""
    Univariate VonMises log-likelihood.
    The pdf of this distribution is
    .. math::
        f(x \mid \mu, \kappa) =
            \frac{e^{\kappa\cos(x-\mu)}}{2\pi I_0(\kappa)}
    where :math:`I_0` is the modified Bessel function of order 0.
    .. plot::
        import matplotlib.pyplot as plt
        import numpy as np
        import scipy.stats as st
        plt.style.use('seaborn-darkgrid')
        x = np.linspace(-np.pi, np.pi, 200)
        mus = [0., 0., 0.,  -2.5]
        kappas = [.01, 0.5,  4., 2.]
        for mu, kappa in zip(mus, kappas):
            pdf = st.vonmises.pdf(x, kappa, loc=mu)
            plt.plot(x, pdf, label=r'$\mu$ = {}, $\kappa$ = {}'.format(mu, kappa))
        plt.xlabel('x', fontsize=12)
        plt.ylabel('f(x)', fontsize=12)
        plt.legend(loc=1)
        plt.show()
    ========  ==========================================
    Support   :math:`x \in [-\pi, \pi]`
    Mean      :math:`\mu`
    Variance  :math:`1-\frac{I_1(\kappa)}{I_0(\kappa)}`
    ========  ==========================================
    Parameters
    ----------
    mu: float
        Mean.
    kappa: float
        Concentration (\frac{1}{kappa} is analogous to \sigma^2).
    """

    def __init__(self, mu=0.0, kappa=None, margin=None, transform="circular", *args, **kwargs):
        if transform == "circular":
            transform = transforms.Circular()
        super().__init__(transform=transform, *args, **kwargs)

        self.mean = self.median = self.mode = self.mu = mu = tt.as_tensor_variable(floatX(mu))
        self.kappa = kappa = tt.as_tensor_variable(floatX(kappa))

        self._margin = margin

        assert_negative_support(kappa, "kappa", "VonMises")

    def random(self, point=None, size=None):
        """
        Draw random values from VonMises distribution.
        Parameters
        ----------
        point: dict, optional
            Dict of variable values on which random values are to be
            conditioned (uses default point if not specified).
        size: int, optional
            Desired size of random sample (returns one sample if not
            specified).
        Returns
        -------
        array
        """

        mu, kappa = draw_values([self.mu, self.kappa], point=point, size=size)
        return within_pi_op(generate_samples(
            stats.vonmises.rvs, loc=mu, kappa=kappa, dist_shape=self.shape, size=size
        ))

    def logp(self, value):
        """
        Calculate log-probability of VonMises distribution at specified value.
        Parameters
        ----------
        value: numeric
            Value(s) for which log-probability is calculated. If the log probabilities for multiple
            values are desired the values must be provided in a numpy array or theano tensor
        Returns
        -------
        TensorVariable
        """

        mu = self.mu
        kappa = self.kappa

        return bound(
            kappa * tt.cos(mu - value) - (tt.log(2 * np.pi) + log_i0(kappa)),
            kappa > 0,
            value >= -np.pi,
            value <= np.pi,
        )

    def _distr_parameters_for_repr(self):
        return ["mu", "kappa"]