# Need to accomodate for SciPy's recent changes in organization
# of statistical distributuions
try:
    from scipy.stats import _continuous_distns as cdists
    from scipy.stats import _discrete_distns as ddists
    rv_discrete = ddists.rv_discrete
    rv_continuous = cdists.rv_continuous
    sp_old = False
except ImportError:
    import scipy.stats.distributions as sc_dst
    rv_discrete = sc_dst.rv_discrete
    rv_continuous = sc_dst.rv_continuous
    sp_old = True
import inspect
import numpy as np
from . import Stochastic
from copy import copy
from .distributions import *
from .utils import get_signature

"""
Wraps a SciPy rv object in a PyMC random variable object.
Needs to wrap the following methods:

generic.rvs(<shape(s)>,loc=0,scale=1)
    - random variates

generic.pdf(x,<shape(s)>,loc=0,scale=1)
    - probability density function

generic.cdf(x,<shape(s)>,loc=0,scale=1)
    - cumulative density function

generic.sf(x,<shape(s)>,loc=0,scale=1)
    - survival function (1-cdf --- sometimes more accurate)

generic.ppf(q,<shape(s)>,loc=0,scale=1)
    - percent point function (inverse of cdf --- percentiles)

generic.isf(q,<shape(s)>,loc=0,scale=1)
    - inverse survival function (inverse of sf)

generic.stats(<shape(s)>,loc=0,scale=1,moments='mv')
    - mean('m',axis=0), variance('v'), skew('s'), and/or kurtosis('k')

generic.entropy(<shape(s)>,loc=0,scale=1)
    - (differential) entropy of the RV.
"""
__all__ = ['stochastic_from_scipy_dist']


def separate_shape_args(kwds, shape_args):

    new_kwds = {}
    for name in kwds:
        new_kwds[name] = kwds[name]

    args = [new_kwds.pop(name) for name in shape_args]

    return args, new_kwds


def remove_size_arg(fun):
    def wrapper(size=1, *args, **kwds):
        return fun(*args, **kwds)


def stochastic_from_scipy_dist(scipy_dist):
    """
    Return a Stochastic subclass made from a particular SciPy distribution.
    """

    name = scipy_dist.__class__.__name__.replace('_gen', '').capitalize()

    (args, defaults) = get_signature(scipy_dist._cdf)
    shape_args = args[2:]
    if isinstance(scipy_dist, rv_continuous):
        dtype = float

        def logp(value, **kwds):
            args, kwds = separate_shape_args(kwds, shape_args)
            return np.sum(scipy_dist.logpdf(value, *args, **kwds))

        parent_names = shape_args + ['loc', 'scale']
        defaults = [None] * (len(parent_names) - 2) + [0., 1.]

    elif isinstance(scipy_dist, rv_discrete):
        dtype = int

        def logp(value, **kwds):
            args, kwds = separate_shape_args(kwds, shape_args)
            return np.sum(scipy_dist.logpmf(value, *args, **kwds))

        parent_names = shape_args + ['loc']
        defaults = [None] * (len(parent_names) - 1) + [0]
    else:
        return None

    parents_default = dict(zip(parent_names, defaults))

    def random(shape=None, **kwds):
        args, kwds = separate_shape_args(kwds, shape_args)

        if shape is None:
            return scipy_dist.rvs(*args, **kwds)
        else:
            return np.reshape(scipy_dist.rvs(*args, **kwds), shape)

    # Build docstring from distribution
    docstr = name[0] + ' = ' + name + '(name, ' + ', '.join(
        parent_names) + ', value=None, shape=None, trace=True, rseed=True, doc=None)\n\n'
    docstr += 'Stochastic variable with ' + name + \
        ' distribution.\nParents are: ' + ', '.join(parent_names) + '.\n\n'
    docstr += """
Methods:

    random()
        - draws random value
          sets value to return value

    ppf(q)
        - percent point function (inverse of cdf --- percentiles)
          sets value to return value

    isf(q)
        - inverse survival function (inverse of sf)
          sets value to return value

    stats(moments='mv')
        - mean('m',axis=0), variance('v'), skew('s'), and/or kurtosis('k')


Attributes:

    logp
        - sum(log(pdf())) or sum(log(pmf()))

    cdf
        - cumulative distribution function

    sf
        - survival function (1-cdf --- sometimes more accurate)

    entropy
        - (differential) entropy of the RV.


NOTE: If you encounter difficulties with this object, please try the analogous
computation using the rv objects in scipy.stats.distributions directly before
reporting the bug.
    """

    new_class = new_dist_class(
        dtype,
        name,
        parent_names,
        parents_default,
        docstr,
        logp,
        random,
        True,
        None)

    class newer_class(new_class):
        __doc__ = docstr
        rv = scipy_dist

        def __init__(self, *args, **kwds):
            new_class.__init__(self, *args, **kwds)
            self.args, self.kwds = separate_shape_args(
                self.parents, shape_args)
            self.frozen_rv = self.rv(self.args, self.kwds)

        def _pymc_dists_to_value(self, args):
            """Replace arguments that are a pymc.Node with their value."""
            # This is needed because the scipy rv function transforms
            # every input argument which causes new pymc lambda
            # functions to be generated. Thus, when calling this many
            # many times, excessive amounts of RAM are used.
            new_args = []
            for arg in args:
                if isinstance(arg, pm.Node):
                    new_args.append(arg.value)
                else:
                    new_args.append(arg)
            return new_args

        def _cdf(self):
            """
            The cumulative distribution function of self conditional on parents
            evaluated at self's current value
            """
            return self.rv.cdf(
                self.value, *self._pymc_dists_to_value(self.args), **self.kwds)
        cdf = property(_cdf, doc=_cdf.__doc__)

        def _sf(self):
            """
            The survival function of self conditional on parents
            evaluated at self's current value
            """
            return self.rv.sf(
                self.value, *self._pymc_dists_to_value(self.args), **self.kwds)
        sf = property(_sf, doc=_sf.__doc__)

        def ppf(self, q):
            """
            The percentile point function (inverse cdf) of self conditional on parents.
            Self's value will be set to the return value.
            """
            self.value = self.rv.ppf(
                q,
                *
                self._pymc_dists_to_value(
                    self.args),
                **self.kwds)
            return self.value

        def isf(self, q):
            """
            The inverse survival function of self conditional on parents.
            Self's value will be set to the return value.
            """
            self.value = self.rv.isf(
                q,
                *self._pymc_dists_to_value(
                    self.args),
                **self.kwds)
            return self.value

        def stats(self, moments='mv'):
            """The first few moments of self's distribution conditional on parents"""
            return self.rv.stats(
                moments=moments, *self._pymc_dists_to_value(self.args), **self.kwds)

        def _entropy(self):
            """The entropy of self's distribution conditional on its parents"""
            return self.rv.entropy(
                *self._pymc_dists_to_value(self.args), **self.kwds)
        entropy = property(_entropy, doc=_entropy.__doc__)

    newer_class.__name__ = new_class.__name__
    return newer_class

if sp_old:

    for scipy_dist_name in sc_dst.__all__:
        scipy_dist = sc_dst.__dict__[scipy_dist_name]
        if isinstance(scipy_dist, rv_continuous) or isinstance(
                scipy_dist, rv_discrete):
            new_dist = stochastic_from_scipy_dist(scipy_dist)
            locals()[new_dist.__name__] = new_dist
            __all__.append(new_dist.__name__)

else:

    for scipy_dist_name in ddists.__all__:
        scipy_dist = ddists.__dict__[scipy_dist_name]
        if isinstance(scipy_dist, rv_discrete):
            new_dist = stochastic_from_scipy_dist(scipy_dist)
            locals()[new_dist.__name__] = new_dist
            __all__.append(new_dist.__name__)

    for scipy_dist_name in cdists.__all__:
        scipy_dist = cdists.__dict__[scipy_dist_name]
        if isinstance(scipy_dist, rv_continuous):
            new_dist = stochastic_from_scipy_dist(scipy_dist)
            locals()[new_dist.__name__] = new_dist
            __all__.append(new_dist.__name__)
