import scipy.stats as stats
from numpy import rint, asscalar

# For now we just wrap all the scipy.stats distributions. Since the scipy ones
# have annoying abbreviated names, we provide more descriptive local aliases.
# Since we need to ability to both generate random variates drawn from a
# distribution, *and* the ability to calculate the probability for a given
# value, we use scipy rather than the numpy.random utilities.
_scipy_dists = {'Alpha': 'alpha', 'Anglit': 'anglit', 'Arcsine': 'arcsine',
                'Beta': 'beta', 'BetaPrime': 'betaprime',
                'Bradford': 'bradford', 'Burr3': 'burr',
                'Burr12': 'burr12', 'Cauchy': 'cauchy', 'Chi': 'chi',
                'ChiSquared': 'chi2', 'Cosine': 'cosine',
                'DoubleGamma': 'dgamma', 'DoubleWeibull': 'dweibull',
                'Erlang': 'erlang', 'Exponential': 'expon',
                'ExponentialNormal': 'exponnorm',
                'ExponentialWeibull': 'exponweib',
                'ExponentialPower': 'exponpow', 'F': 'f',
                'FatigueLife': 'fatiguelife', 'Fisk': 'fisk',
                'FoldedCauchy': 'foldcauchy', 'FoldedNormal': 'foldnorm',
                'GeneralLogistic': 'genlogistic', 'GeneralNormal': 'gennorm',
                'GeneralPareto': 'genpareto', 'GeneralExponential': 'genexpon',
                'GeneralExtreme': 'genextreme',
                'GaussHypergeometric': 'gausshyper', 'Gamma': 'gamma',
                'GeneralGamma': 'gengamma',
                'GeneralHalfLogistic': 'genhalflogistic', 'Gilbrat': 'gilbrat',
                'Gompertz': 'gompertz', 'GumbelRight': 'gumbel_r',
                'GumbelLeft': 'gumbel_l', 'HalfCauchy': 'halfcauchy',
                'HalfLogistic': 'halflogistic', 'HalfNormal': 'halfnorm',
                'HalfGeneralNormal': 'halfgennorm',
                'HyperbolicSecant': 'hypsecant', 'InverseGamma': 'invgamma',
                'InverseGaussian': 'invgauss', 'InverseWeibull': 'invweibull',
                'JohnsonSB': 'johnsonsb', 'JohnsonSU': 'johnsonsu',
                'Kappa4': 'kappa4', 'Kappa3': 'kappa3', 'KSOneSided': 'ksone',
                'KSTwoSided': 'kstwobign', 'Laplace': 'laplace',
                'Levy': 'levy', 'LevyLeft': 'levy_l',
                'LevyStable': 'levy_stable', 'Logistic': 'logistic',
                'LogGamma': 'loggamma', 'LogLaplace': 'loglaplace',
                'LogNormal': 'lognorm', 'Lomax': 'lomax', 'Maxwell': 'maxwell',
                'Mielke': 'mielke', 'Nakagami': 'nakagami',
                'NonCentralChiSquared': 'ncx2', 'NonCentralF': 'ncf',
                'NonCentralT': 'nct', 'Normal': 'norm', 'Pareto': 'pareto',
                'PearsonType3': 'pearson3', 'PowerLaw': 'powerlaw',
                'PowerLogNormal': 'powerlognorm', 'PowerNormal': 'powernorm',
                'RDistributed': 'rdist', 'Reciprocal': 'reciprocal',
                'Rayleigh': 'rayleigh', 'Rice': 'rice',
                'ReciprocalInverseGaussian': 'recipinvgauss',
                'Semicircular': 'semicircular', 'SkewNormal': 'skewnorm',
                'T': 't', 'Trapezoidal': 'trapz', 'Triangular': 'triang',
                'TruncatedExponential': 'truncexpon',
                'TruncatedNormal': 'truncnorm',
                'TukeyLambda': 'tukeylambda', 'Uniform': 'uniform',
                'VonMises': 'vonmises', 'VonMisesLine': 'vonmises_line',
                'Wald': 'wald', 'WeibullMinimum': 'weibull_min',
                'WeibullMaximum': 'weibull_max', 'WrappedCauchy': 'wrapcauchy',
                # Now discrete distributions
                'Bernoulli': 'bernoulli', 'Binomial': 'binom',
                'Boltzmann': 'boltzmann', 'DiscreteLaplace': 'dlaplace',
                'Geometric': 'geom', 'Hypergeometric': 'hypergeom',
                'LogSeries': 'logser', 'NegativeBinomial': 'nbinom',
                'Planck': 'planck', 'Poisson': 'poisson',
                'DiscreteUniform': 'randint', 'Skellam': 'skellam',
                'Zipf': 'zipf'}


#
class Distribution(object):
    """
    Template version of Distribution class. Can directly subclass or use
    one of the scipy distributions generated by the _class_from_scipy_dist
    factory
    """
    def __init__(self):
        super(Distribution, self).__init__()
        # Initialize with randomly-drawn value
        self._value = self.random()
        self.name = ''
        self.fitsname = ''

    def random(self):
        return 0

    def logp(self, x):
        return 0

    def median(self):
        return 0

    def get_value(self):
        return self._value

    def set_value(self, val):
        self._value = val

    value = property(fget=get_value, fset=set_value)


def _class_from_scipy_dist(class_name, scipy_rv):
    """
    Create a Distribution subclass from a given scipy dist_gen class
    :param class_name: Name of returned class
    :param scipy_rv: scipy.stats.<dist>_gen class to wrap
    :return: subclass of Distribution
    """
    class NewDist(Distribution):
        """
        {} probability distribution. Wraps scipy.stats.{}, whose docstring is
        reproduced below:

        {}
        """
        rv_class = getattr(stats, scipy_rv)
        __doc__ = str(__doc__).format(class_name, scipy_rv, rv_class.__doc__)

        def __init__(self, *args, **kwargs):
            self.rv_frozen = self.rv_class(*args, **kwargs)

            # Set logp and random functions
            if isinstance(self.rv_frozen.dist, stats.rv_continuous):
                self.logp = self.rv_frozen.logpdf
            elif isinstance(self.rv_frozen.dist, stats.rv_discrete):
                self.logp = self.rv_frozen.logpmf
            else:
                raise TypeError('Only rv_continuous and rv_discete '
                                'distributions are supported')
            self.random = self.rv_frozen.rvs
            self.median = self.rv_frozen.median
            super(NewDist, self).__init__()

        def set_value(self, val):
            if isinstance(self.rv_frozen.dist, stats.rv_discrete):
                new_val = rint(val).astype(int)
            else:
                new_val = val
            try:
                self._value = asscalar(new_val)
            except ValueError:
                self._value = new_val

        value = property(fget=Distribution.get_value, fset=set_value)

    NewDist.__name__ = class_name
    return NewDist

# Create Distribution classes for every scipy distribution
__all__ = list(_scipy_dists.keys()) + ['Distribution']
for _class_name, _dist in _scipy_dists.items():
    dist_class = _class_from_scipy_dist(_class_name, _dist)
    locals()[_class_name] = dist_class
