"""
pymc.distributions

A collection of common probability distributions. The objects associated
with a distribution called 'dist' are:

  dist_like : function
    The log-likelihood function corresponding to dist. PyMC's convention
    is to sum the log-likelihoods of multiple input values, so all
    log-likelihood functions return a single float.
  rdist : function
    The random variate generator corresponding to dist. These take a
    'size' argument indicating how many variates should be generated.
  dist_expval : function
    Computes the expected value of a dist-distributed variable.
  Dist : Stochastic subclass
    Instances have dist_like as their log-probability function
    and rdist as their random function.
"""

#-------------------------------------------------------------------
# Decorate fortran functions from pymc.flib to ease argument passing
#-------------------------------------------------------------------
# TODO: Add exponweib_expval
# TODO: categorical, mvhypergeometric
# TODO: __all__

__docformat__ = 'reStructuredText'

from . import flib, utils
import numpy as np
from scipy.stats.kde import gaussian_kde
from .Node import ZeroProbability
from .PyMCObjects import Stochastic, Deterministic
from .CommonDeterministics import Lambda
from numpy import pi, inf
import itertools
import pdb
from . import utils
import warnings

from . import six
from .six import print_
xrange = six.moves.xrange

def poiscdf(a, x):
    x = np.atleast_1d(x)
    a = np.resize(a, x.shape)
    values = np.array([flib.gammq(b, y) for b, y in zip(a.ravel(), x.ravel())])
    return values.reshape(x.shape)

# Import utility functions
import inspect, types
from copy import copy
random_number = np.random.random
inverse = np.linalg.pinv

sc_continuous_distributions = ['beta', 'cauchy', 'chi2',
                               'degenerate', 'exponential', 'exponweib',
                               'gamma', 'half_cauchy', 'half_normal',
                               'inverse_gamma', 'laplace', 'logistic',
                               'lognormal', 'noncentral_t', 'normal',
                               'pareto', 't', 'truncated_pareto', 'uniform',
                               'weibull', 'skew_normal', 'truncated_normal',
                               'von_mises']
sc_bool_distributions = ['bernoulli']
sc_discrete_distributions = ['binomial', 'betabin', 'geometric', 'poisson',
                             'negative_binomial', 'categorical', 'hypergeometric',
                             'discrete_uniform', 'truncated_poisson']

sc_nonnegative_distributions = ['bernoulli', 'beta', 'betabin', 'binomial', 'chi2', 'exponential',
                                'exponweib', 'gamma', 'half_cauchy', 'half_normal',
                                'hypergeometric', 'inverse_gamma', 'lognormal',
                                'weibull']

mv_continuous_distributions = ['dirichlet', 'mv_normal',
                               'mv_normal_cov', 'mv_normal_chol', 'wishart',
                               'wishart_cov']

mv_discrete_distributions = ['multivariate_hypergeometric', 'multinomial']

mv_nonnegative_distributions = ['dirichlet', 'wishart',
                                'wishart_cov', 'multivariate_hypergeometric',
                                'multinomial']

availabledistributions = (sc_continuous_distributions +
                          sc_bool_distributions +
                          sc_discrete_distributions +
                          mv_continuous_distributions +
                          mv_discrete_distributions)

# Changes lower case, underscore-separated names into "Class style"
# capitalized names For example, 'negative_binomial' becomes
# 'NegativeBinomial'
capitalize = lambda name: ''.join([s.capitalize() for s in name.split('_')])


# ==============================================================================
# User-accessible function to convert a logp and random function to a
# Stochastic subclass.
# ==============================================================================

# TODO Document this function
def bind_size(randfun, shape):
    def newfun(*args, **kwargs):
        try:
            return np.reshape(randfun(size=shape, *args, **kwargs), shape)
        except ValueError:
            # Account for non-array return values
            return randfun(size=shape, *args, **kwargs)
    newfun.scalar_version = randfun
    return newfun

def new_dist_class(*new_class_args):
    """
    Returns a new class from a distribution.

    :Parameters:
      dtype : numpy dtype
        The dtype values of instances of this class.
      name : string
        Name of the new class.
      parent_names : list of strings
        The labels of the parents of this class.
      parents_default : list
        The default values of parents.
      docstr : string
        The docstring of this class.
      logp : function
        The log-probability function for this class.
      random : function
        The random function for this class.
      mv : boolean
        A flag indicating whether this class represents array-valued
        variables.

      .. note::
        stochastic_from_dist provides a higher-level version.
        stochastic_from_data is suited for non-parametric distributions.

      :SeeAlso:
        stochastic_from_dist, stochastic_from_data
    """

    (dtype, name, parent_names, parents_default, docstr, logp, random, mv, logp_partial_gradients) = new_class_args

    class new_class(Stochastic):
        __doc__ = docstr

        def __init__(self, *args, **kwds):
            (dtype, name, parent_names, parents_default, docstr, logp, random, mv, logp_partial_gradients) = new_class_args
            parents = parents_default

            # Figure out what argument names are needed.
            arg_keys = ['name', 'parents', 'value', 'observed', 'size', 'trace', 'rseed', 'doc', 'debug', 'plot', 'verbose']
            arg_vals = [None, parents, None, False, None, True, True, None, False, None, -1]
            if 'isdata' in kwds:
                warnings.warn('"isdata" is deprecated, please use "observed" instead.')
                kwds['observed'] = kwds['isdata']
                pass


            # No size argument allowed for multivariate distributions.
            if mv:
                arg_keys.pop(4)
                arg_vals.pop(4)

            arg_dict_out = dict(zip(arg_keys, arg_vals))
            args_needed = ['name'] + parent_names + arg_keys[2:]

            # Sort positional arguments
            for i in xrange(len(args)):
                try:
                    k = args_needed.pop(0)
                    if k in parent_names:
                        parents[k] = args[i]
                    else:
                        arg_dict_out[k] = args[i]
                except:
                    raise ValueError('Too many positional arguments provided. Arguments for class ' + self.__class__.__name__ + ' are: ' + str(args_needed))


            # Sort keyword arguments
            for k in args_needed:
                if k in parent_names:
                    try:
                        parents[k] = kwds.pop(k)
                    except:
                        if k in parents_default:
                            parents[k] = parents_default[k]
                        else:
                            raise ValueError('No value given for parent ' + k)
                elif k in arg_dict_out.keys():
                    try:
                        arg_dict_out[k] = kwds.pop(k)
                    except:
                        pass

            # Remaining unrecognized arguments raise an error.
            if len(kwds) > 0:
                raise TypeError('Keywords ' + str(kwds.keys()) + ' not recognized. Arguments recognized are ' + str(args_needed))

        # Determine size desired for scalar variables.
        # Notes
        # -----
        # Case | init_val     | parents       | size | value.shape | bind size
        # ------------------------------------------------------------------
        # 1.1  | None         | scalars       | None | 1           | 1
        # 1.2  | None         | scalars       | n    | n           | n
        # 1.3  | None         | n             | None | n           | 1
        # 1.4  | None         | n             | n(m) | n (Error)   | 1 (-)
        # 2.1  | scalar       | scalars       | None | 1           | 1
        # 2.2  | scalar       | scalars       | n    | n           | n
        # 2.3  | scalar       | n             | None | n           | 1
        # 2.4  | scalar       | n             | n(m) | n (Error)   | 1 (-)
        # 3.1  | n            | scalars       | None | n           | n
        # 3.2  | n            | scalars       | n(m) | n (Error)   | n (-)
        # 3.3  | n            | n             | None | n           | 1
        # 3.4  | n            | n             | n(m) | n (Error)   | 1 (-)

            if not mv:

                shape = arg_dict_out.pop('size')
                shape = None if shape is None else tuple(np.atleast_1d(shape))

                init_val = arg_dict_out['value']
                init_val_shape = None if init_val is None else np.shape(init_val)

                if len(parents) > 0:
                    pv = [np.shape(utils.value(v)) for v in parents.values()]
                    biggest_parent = np.argmax([(np.prod(v) if v else 0) for v in pv])
                    parents_shape = pv[biggest_parent]

                    # Scalar parents can support any shape.
                    if np.prod(parents_shape) < 1:
                        parents_shape = None

                else:
                    parents_shape = None

                def shape_error():
                    raise ValueError('Shapes are incompatible: value %s, largest parent %s, shape argument %s'%(shape, init_val_shape, parents_shape))

                if init_val_shape is not None and shape is not None and init_val_shape != shape:
                    shape_error()

                given_shape = init_val_shape or shape
                bindshape = given_shape or parents_shape

                # Check consistency of bindshape and parents_shape
                if parents_shape is not None:
                    # Uncomment to leave broadcasting completely up to NumPy's random functions
                    # if bindshape[-np.alen(parents_shape):]!=parents_shape:
                    # Uncomment to limit broadcasting flexibility to what the Fortran likelihoods can handle.
                    if bindshape < parents_shape:
                        shape_error()

                if random is not None:
                    random = bind_size(random, bindshape)


            elif 'size' in kwds.keys():
                raise ValueError('No size argument allowed for multivariate stochastic variables.')


            # Call base class initialization method
            if arg_dict_out.pop('debug'):
                logp = debug_wrapper(logp)
                random = debug_wrapper(random)
            else:
                Stochastic.__init__(self, logp=logp, random=random, logp_partial_gradients = logp_partial_gradients, dtype=dtype, **arg_dict_out)

    new_class.__name__ = name
    new_class.parent_names = parent_names
    new_class.parents_default = parents_default
    new_class.dtype = dtype
    new_class.mv = mv
    new_class.raw_fns = {'logp': logp, 'random': random}

    return new_class


def stochastic_from_dist(name, logp, random=None, logp_partial_gradients={}, dtype=np.float, mv=False):
    """
    Return a Stochastic subclass made from a particular distribution.

    :Parameters:
      name : string
        The name of the new class.
      logp : function
        The log-probability function.
      random : function
        The random function
      dtype : numpy dtype
        The dtype of values of instances.
      mv : boolean
        A flag indicating whether this class represents
        array-valued variables.

    :Example:
      >>> Exponential = stochastic_from_dist('exponential',
                                              logp=exponential_like,
                                              random=rexponential,
                                              dtype=np.float,
                                              mv=False)
      >>> A = Exponential(self_name, value, beta)

    .. note::
      new_dist_class is a more flexible class factory. Also consider
      subclassing Stochastic directly.
      stochastic_from_data is suited for non-parametric distributions.

    :SeeAlso:
      new_dist_class, stochastic_from_data
    """

    (args, varargs, varkw, defaults) = inspect.getargspec(logp)
    parent_names = args[1:]
    try:
        parents_default = dict(zip(args[-len(defaults):], defaults))
    except TypeError: # No parents at all.
        parents_default = {}

    name = capitalize(name)

    # Build docstring from distribution
    parents_str = ''
    if parent_names:
        parents_str = ', '.join(parent_names) + ', '
    docstr = name[0]+' = '+name + '(name, '+parents_str+'value=None, observed=False,'
    if not mv:
        docstr += ' size=1,'
    docstr += ' trace=True, rseed=True, doc=None, verbose=-1, debug=False)\n\n'
    docstr += 'Stochastic variable with '+name+' distribution.\nParents are: '+', '.join(parent_names) + '.\n\n'
    docstr += 'Docstring of log-probability function:\n'
    try:
        docstr += logp.__doc__
    except TypeError:
        pass  # This will happen when logp doesn't have a docstring

    logp = valuewrapper(logp)
    distribution_arguments = logp.__dict__

    wrapped_logp_partial_gradients = {}

    for parameter, func in six.iteritems(logp_partial_gradients):
        wrapped_logp_partial_gradients[parameter] = valuewrapper(logp_partial_gradients[parameter], arguments = distribution_arguments)

    return new_dist_class(dtype, name, parent_names, parents_default, docstr,
                                                 logp, random, mv, wrapped_logp_partial_gradients)


def stochastic_from_data(name, data, lower=-np.inf, upper=np.inf,
                         value=None, observed=False, trace=True, verbose=-1, debug=False):
    """
    Return a Stochastic subclass made from arbitrary data.

    The histogram for the data is fitted with Kernel Density Estimation.

    :Parameters:
      - `data`  : An array with samples (e.g. trace[:])
      - `lower` : Lower bound on possible outcomes
      - `upper` : Upper bound on possible outcomes

    :Example:
       >>> from pymc import stochastic_from_data
       >>> pos = stochastic_from_data('posterior', posterior_samples)
       >>> prior = pos # update the prior with arbitrary distributions

    :Alias:
      Histogram
    """
    pdf = gaussian_kde(data) # automatic bandwidth selection

    # account for tail contribution
    lower_tail = upper_tail = 0.
    if lower > -np.inf:
        lower_tail = pdf.integrate_box(-np.inf, lower)
    if upper <  np.inf:
        upper_tail = pdf.integrate_box(upper, np.inf)
    factor = 1./(1. - (lower_tail + upper_tail))

    def logp(value):
        prob = factor*pdf(value)
        if value < lower or value > upper:
            return -np.inf
        elif prob <= 0.:
            return -np.inf
        else:
            return np.log(prob)

    def random():
        res = pdf.resample(1)[0][0]
        while res < lower or res > upper:
            res = pdf.resample(1)[0][0]
        return res

    if value == None:
        value = random()

    return Stochastic(logp     = logp,
                      doc      = 'Non-parametric density with Gaussian Kernels.',
                      name     = name,
                      parents  = {},
                      random   = random,
                      trace    = trace,
                      value    = value,
                      dtype    = float,
                      observed = observed,
                      verbose  = verbose)


# Alias following Stochastics naming convention
Histogram = stochastic_from_data

#-------------------------------------------------------------
# Light decorators
#-------------------------------------------------------------

def randomwrap(func):
    """
    Decorator for random value generators

    Allows passing of sequence of parameters, as well as a size argument.


    Convention:

      - If size=1 and the parameters are all scalars, return a scalar.
      - If size=1, the random variates are 1D.
      - If the parameters are scalars and size > 1, the random variates are 1D.
      - If size > 1 and the parameters are sequences, the random variates are
        aligned as (size, max(length)), where length is the parameters size.


    :Example:
      >>> rbernoulli(.1)
      0
      >>> rbernoulli([.1,.9])
      np.asarray([0, 1])
      >>> rbernoulli(.9, size=2)
      np.asarray([1, 1])
      >>> rbernoulli([.1,.9], 2)
      np.asarray([[0, 1],
             [0, 1]])
    """

    # Find the order of the arguments.
    refargs, varargs, varkw, defaults = inspect.getargspec(func)
    # vfunc = np.vectorize(self.func)
    npos = len(refargs)-len(defaults) # Number of pos. arg.
    nkwds = len(defaults) # Number of kwds args.
    mv = func.__name__[1:] in mv_continuous_distributions + mv_discrete_distributions

    # Use the NumPy random function directly if this is not a multivariate distribution
    if not mv:
        return func

    def wrapper(*args, **kwds):
        # First transform keyword arguments into positional arguments.
        n = len(args)
        if nkwds > 0:
            args = list(args)
            for i, k in enumerate(refargs[n:]):
                if k in kwds.keys():
                    args.append(kwds[k])
                else:
                    args.append(defaults[n-npos+i])

        r = []
        s = []
        largs = []
        nr = args[-1]
        length = [np.atleast_1d(a).shape[0] for a in args]
        dimension = [np.atleast_1d(a).ndim for a in args]
        N = max(length)
        if len(set(dimension)) > 2:
            raise('Dimensions do not agree.')
        # Make sure all elements are iterable and have consistent lengths, ie
        # 1 or n, but not m and n.

        for arg, s in zip(args, length):
            t = type(arg)
            arr = np.empty(N, type)
            if s == 1:
                arr.fill(arg)
            elif s == N:
                arr = np.asarray(arg)
            else:
                raise RuntimeError('Arguments size not allowed: %s.' % s)
            largs.append(arr)

        if mv and N > 1 and max(dimension) > 1 and nr > 1:
            raise ValueError('Multivariate distributions cannot take s>1 and multiple values.')

        if mv:
            for i, arg in enumerate(largs[:-1]):
                largs[0] = np.atleast_2d(arg)

        for arg in zip(*largs):
            r.append(func(*arg))

        size = arg[-1]
        vec_stochastics = len(r) > 1
        if mv:
            if nr == 1:
                return r[0]
            else:
                return np.vstack(r)
        else:
            if size > 1 and vec_stochastics:
                return np.atleast_2d(r).T
            elif vec_stochastics or size > 1:
                return np.concatenate(r)
            else: # Scalar case
                return r[0][0]

    wrapper.__doc__ = func.__doc__
    wrapper.__name__ = func.__name__
    return wrapper

def debug_wrapper(func, name):
    # Wrapper to debug distributions

    import pdb

    def wrapper(*args, **kwargs):

        print_('Debugging inside %s:' % name)
        print_('\tPress \'s\' to step into function for debugging')
        print_('\tCall \'args\' to list function arguments')

        # Set debugging trace
        pdb.set_trace()

        # Call function
        return func(*args, **kwargs)

    return wrapper


#-------------------------------------------------------------
# Utility functions
#-------------------------------------------------------------

def constrain(value, lower=-np.Inf, upper=np.Inf, allow_equal=False):
    """
    Apply interval constraint on stochastic value.
    """

    ok = flib.constrain(value, lower, upper, allow_equal)
    if ok == 0:
        raise ZeroProbability

def standardize(x, loc=0, scale=1):
    """
    Standardize x

    Return (x-loc)/scale
    """

    return flib.standardize(x, loc, scale)

# ==================================
# = vectorize causes memory leaks. =
# ==================================
# @Vectorize
def gammaln(x):
    """
    Logarithm of the Gamma function
    """

    return flib.gamfun(x)

def expand_triangular(X, k):
    """
    Expand flattened triangular matrix.
    """

    X = X.tolist()
    # Unflatten matrix
    Y = np.asarray([[0] * i + X[i * k - (i * (i - 1)) / 2 : i * k + (k - i)] for i in range(k)])
    # Loop over rows
    for i in range(k):
        # Loop over columns
        for j in range(k):
            Y[j, i] = Y[i, j]
    return Y


# Loss functions

absolute_loss = lambda o, e: absolute(o - e)

squared_loss = lambda o, e: (o - e)**2

chi_square_loss = lambda o, e: (1.*(o - e)**2)/e

loss_functions = {'absolute': absolute_loss, 'squared': squared_loss, 'chi_square': chi_square_loss}

def GOFpoints(x, y, expval, loss):
    # Return pairs of points for GOF calculation
    return np.sum(np.transpose([loss(x, expval), loss(y, expval)]), 0)

def gofwrapper(f, loss_function='squared'):
    """
    Goodness-of-fit decorator function for likelihoods
    ==================================================
    Generates goodness-of-fit points for data likelihoods.

    Wrap function f(*args, **kwds) where f is a likelihood.

    Assume args = (x, parameter1, parameter2, ...)
    Before passing the arguments to the function, the wrapper makes sure that
    the parameters have the same shape as x.
    """

    name = f.__name__[:-5]
    # Take a snapshot of the main namespace.

    # Find the functions needed to compute the gof points.
    expval_func = eval(name+'_expval')
    random_func = eval('r'+name)

    def wrapper(*args, **kwds):
        """
        This wraps a likelihood.
        """

        """Return gof points."""

        # Calculate loss
        loss = kwds.pop('gof', loss_functions[loss_function])

        # Expected value, given parameters
        expval = expval_func(*args[1:], **kwds)
        y = random_func(size=len(args[0]), *args[1:], **kwds)
        f.gof_points = GOFpoints(args[0], y, expval, loss)

        """Return likelihood."""

        return f(*args, **kwds)


    # Assign function attributes to wrapper.
    wrapper.__doc__ = f.__doc__
    wrapper.__name__ = f.__name__
    wrapper.name = name

    return wrapper

#--------------------------------------------------------
# Statistical distributions
# random generator, expval, log-likelihood
#--------------------------------------------------------

# Autoregressive lognormal
def rarlognormal(a, sigma, rho, size=1):
    R"""
    Autoregressive normal random variates.

    If a is a scalar, generates one series of length size.
    If a is a sequence, generates size series of the same length
    as a.
    """
    f = utils.ar1
    if np.isscalar(a):
        r = f(rho, 0, sigma, size)
    else:
        n = len(a)
        r = [f(rho, 0, sigma, n) for i in range(size)]
        if size == 1:
            r = r[0]
    return a*np.exp(r)


def arlognormal_like(x, a, sigma, rho):
    R"""
    Autoregressive lognormal log-likelihood.

    .. math::
        x_i & = a_i \exp(e_i) \\
        e_i & = \rho e_{i-1} + \epsilon_i

    where :math:`\epsilon_i \sim N(0,\sigma)`.
    """
    return flib.arlognormal(x, np.log(a), sigma, rho, beta=1)


# Bernoulli----------------------------------------------
@randomwrap
def rbernoulli(p, size=None):
    """
    Random Bernoulli variates.
    """

    return np.random.random(size) < p

def bernoulli_expval(p):
    """
    Expected value of bernoulli distribution.
    """

    return p


def bernoulli_like(x, p):
    R"""Bernoulli log-likelihood

    The Bernoulli distribution describes the probability of successes (x=1) and
    failures (x=0).

    .. math::  f(x \mid p) = p^{x} (1-p)^{1-x}

    :Parameters:
      - `x` : Series of successes (1) and failures (0). :math:`x=0,1`
      - `p` : Probability of success. :math:`0 < p < 1`.

    :Example:
       >>> from pymc import bernoulli_like
       >>> bernoulli_like([0,1,0,1], .4)
       -2.854232711280291

    .. note::
      - :math:`E(x)= p`
      - :math:`Var(x)= p(1-p)`

    """

    return flib.bernoulli(x, p)

bernoulli_grad_like = {'p' : flib.bern_grad_p}

# Beta----------------------------------------------
@randomwrap
def rbeta(alpha, beta, size=None):
    """
    Random beta variates.
    """
    from scipy.stats.distributions import beta as sbeta
    return sbeta.ppf(np.random.random(size), alpha, beta)
    # return np.random.beta(alpha, beta, size)

def beta_expval(alpha, beta):
    """
    Expected value of beta distribution.
    """

    return 1.0 * alpha / (alpha + beta)


def beta_like(x, alpha, beta):
    R"""
    Beta log-likelihood. The conjugate prior for the parameter
    :math:`p` of the binomial distribution.

    .. math::
        f(x \mid \alpha, \beta) = \frac{\Gamma(\alpha + \beta)}{\Gamma(\alpha) \Gamma(\beta)} x^{\alpha - 1} (1 - x)^{\beta - 1}

    :Parameters:
      - `x` : 0 < x < 1
      - `alpha` : alpha > 0
      - `beta` : beta > 0

    :Example:
      >>> from pymc import beta_like
      >>> beta_like(.4,1,2)
      0.182321556793954

    .. note::
      - :math:`E(X)=\frac{\alpha}{\alpha+\beta}`
      - :math:`Var(X)=\frac{\alpha \beta}{(\alpha+\beta)^2(\alpha+\beta+1)}`

    """
    # try:
    #     constrain(alpha, lower=0, allow_equal=True)
    #     constrain(beta, lower=0, allow_equal=True)
    #     constrain(x, 0, 1, allow_equal=True)
    # except ZeroProbability:
    #     return -np.Inf
    return flib.beta_like(x, alpha, beta)

beta_grad_like = {'value' : flib.beta_grad_x,
                  'alpha' : flib.beta_grad_a,
                  'beta' : flib.beta_grad_b}


# Binomial----------------------------------------------
@randomwrap
def rbinomial(n, p, size=None):
    """
    Random binomial variates.
    """
    # return np.random.binomial(n,p,size)
    return np.random.binomial(np.ravel(n), np.ravel(p), size)

def binomial_expval(n, p):
    """
    Expected value of binomial distribution.
    """

    return p*n

def binomial_like(x, n, p):
    R"""
    Binomial log-likelihood.  The discrete probability distribution of the
    number of successes in a sequence of n independent yes/no experiments,
    each of which yields success with probability p.

    .. math::
        f(x \mid n, p) = \frac{n!}{x!(n-x)!} p^x (1-p)^{n-x}

    :Parameters:
      - `x` : [int] Number of successes, > 0.
      - `n` : [int] Number of Bernoulli trials, > x.
      - `p` : Probability of success in each trial, :math:`p \in [0,1]`.

    .. note::
       - :math:`E(X)=np`
       - :math:`Var(X)=np(1-p)`

    """

    return flib.binomial(x, n, p)


binomial_grad_like = {'p' : flib.binomial_gp}

# Beta----------------------------------------------
@randomwrap
def rbetabin(alpha, beta, n, size=None):
    """
    Random beta-binomial variates.
    """

    phi = np.random.beta(alpha, beta, size)
    return np.random.binomial(n, phi)

def betabin_expval(alpha, beta, n):
    """
    Expected value of beta-binomial distribution.
    """

    return n * alpha / (alpha + beta)


def betabin_like(x, alpha, beta, n):
    R"""
    Beta-binomial log-likelihood. Equivalent to binomial random
    variables with probabilities drawn from a
    :math:`\texttt{Beta}(\alpha,\beta)` distribution.

    .. math::
        f(x \mid \alpha, \beta, n) = \frac{\Gamma(\alpha + \beta)}{\Gamma(\alpha)} \frac{\Gamma(n+1)}{\Gamma(x+1)\Gamma(n-x+1)} \frac{\Gamma(\alpha + x)\Gamma(n+\beta-x)}{\Gamma(\alpha+\beta+n)}

    :Parameters:
      - `x` : x=0,1,\ldots,n
      - `alpha` : alpha > 0
      - `beta` : beta > 0
      - `n` : n=x,x+1,\ldots

    :Example:
      >>> betabin_like(3,1,1,10)
      -2.3978952727989

    .. note::
      - :math:`E(X)=n\frac{\alpha}{\alpha+\beta}`
      - :math:`Var(X)=n\frac{\alpha \beta}{(\alpha+\beta)^2(\alpha+\beta+1)}`

    """
    return flib.betabin_like(x, alpha, beta, n)

betabin_grad_like = {'alpha' : flib.betabin_ga,
                 'beta' : flib.betabin_gb}

# Categorical----------------------------------------------
# Note that because categorical elements are not ordinal, there
# is no expected value.

#@randomwrap
def rcategorical(p, size=None):
    """
    Categorical random variates.
    """
    out = flib.rcat(p, np.random.random(size=size))
    if sum(out.shape) == 1:
        return out.squeeze()
    else:
        return out

def categorical_like(x, p):
    R"""
    Categorical log-likelihood. The most general discrete distribution.

    .. math::  f(x=i \mid p) = p_i

    for :math:`i \in 0 \ldots k-1`.

    :Parameters:
      - `x` : [int] :math:`x \in 0\ldots k-1`
      - `p` : [float] :math:`p > 0`, :math:`\sum p = 1`

    """

    p = np.atleast_2d(p)
    if np.any(abs(np.sum(p, 1)-1) > 0.0001):
        print_("Probabilities in categorical_like sum to", np.sum(p, 1))
    if np.array(x).dtype != int:
        # print_("Non-integer values in categorical_like")
        return -inf
    return flib.categorical(x, p)


# Cauchy----------------------------------------------
@randomwrap
def rcauchy(alpha, beta, size=None):
    """
    Returns Cauchy random variates.
    """

    return alpha + beta*np.tan(pi*random_number(size) - pi/2.0)

def cauchy_expval(alpha, beta):
    """
    Expected value of cauchy distribution.
    """

    return alpha

# In wikipedia, the arguments name are k, x0.
def cauchy_like(x, alpha, beta):
    R"""
    Cauchy log-likelihood. The Cauchy distribution is also known as the
    Lorentz or the Breit-Wigner distribution.

    .. math::
        f(x \mid \alpha, \beta) = \frac{1}{\pi \beta [1 + (\frac{x-\alpha}{\beta})^2]}

    :Parameters:
      - `alpha` : Location parameter.
      - `beta` : Scale parameter > 0.

    .. note::
       - Mode and median are at alpha.

    """

    return flib.cauchy(x, alpha, beta)

cauchy_grad_like = {'value' : flib.cauchy_grad_x,
                 'alpha' : flib.cauchy_grad_a,
                 'beta' : flib.cauchy_grad_b}

# Chi square----------------------------------------------
@randomwrap
def rchi2(nu, size=None):
    """
    Random :math:`\chi^2` variates.
    """

    return np.random.chisquare(nu, size)

def chi2_expval(nu):
    """
    Expected value of Chi-squared distribution.
    """

    return nu

def chi2_like(x, nu):
    R"""
    Chi-squared :math:`\chi^2` log-likelihood.

    .. math::
        f(x \mid \nu) = \frac{x^{(\nu-2)/2}e^{-x/2}}{2^{\nu/2}\Gamma(\nu/2)}

    :Parameters:
      - `x` : > 0
      - `nu` : [int] Degrees of freedom ( nu > 0 )

    .. note::
      - :math:`E(X)=\nu`
      - :math:`Var(X)=2\nu`

    """

    return flib.gamma(x, 0.5*nu, 1./2)

chi2_grad_like = {'value'  : lambda x, nu : flib.gamma_grad_x    (x, 0.5 * nu, 1./2),
                  'nu' : lambda x, nu : flib.gamma_grad_alpha(x, 0.5 * nu, 1./2) * .5}

# chi2_grad_like = {'x'  : lambda x, nu : (nu / 2 - 1) / x -.5,
#                  'nu' : flib.chi2_grad_nu }

# Degenerate---------------------------------------------
@randomwrap
def rdegenerate(k, size=1):
    """
    Random degenerate variates.
    """
    return np.ones(size)*k

def degenerate_expval(k):
    """
    Expected value of degenerate distribution.
    """
    return k

def degenerate_like(x, k):
    R"""
    Degenerate log-likelihood.

    .. math::
        f(x \mid k) = \left\{ \begin{matrix} 1 \text{ if } x = k \\ 0 \text{ if } x \ne k\end{matrix} \right.

    :Parameters:
      - `x` : Input value.
      - `k` : Degenerate value.

    """
    x = np.atleast_1d(x)
    return sum(np.log([i == k for i in x]))

# def degenerate_grad_like(x, k):
#    R"""
#    degenerate_grad_like(x, k)
#
#    Degenerate gradient log-likelihood.
#
#    .. math::
#        f(x \mid k) = \left\{ \begin{matrix} 1 \text{ if } x = k \\ 0 \text{ if } x \ne k\end{matrix} \right.
#
#    :Parameters:
#      - `x` : Input value.
#      - `k` : Degenerate value.
#    """
#    return np.zeros(np.size(x))*k

# Dirichlet----------------------------------------------
@randomwrap
def rdirichlet(theta, size=1):
    """
    Dirichlet random variates.
    """
    gammas = np.vstack([rgamma(theta, 1) for i in xrange(size)])
    if size > 1 and np.size(theta) > 1:
        return (gammas.T/gammas.sum(1))[:-1].T
    elif np.size(theta) > 1:
        return (gammas[0]/gammas[0].sum())[:-1]
    else:
        return 1.

def dirichlet_expval(theta):
    """
    Expected value of Dirichlet distribution.
    """
    return theta/np.sum(theta).astype(float)

def dirichlet_like(x, theta):
    R"""
    Dirichlet log-likelihood.

    This is a multivariate continuous distribution.

    .. math::
        f(\mathbf{x}) = \frac{\Gamma(\sum_{i=1}^k \theta_i)}{\prod \Gamma(\theta_i)}\prod_{i=1}^{k-1} x_i^{\theta_i - 1}
        \cdot\left(1-\sum_{i=1}^{k-1}x_i\right)^\theta_k

    :Parameters:
      x : (n, k-1) array
        Array of shape (n, k-1) where `n` is the number of samples
        and `k` the dimension.
        :math:`0 < x_i < 1`,  :math:`\sum_{i=1}^{k-1} x_i < 1`
      theta : array
        An (n,k) or (1,k) array > 0.

    .. note::
        Only the first `k-1` elements of `x` are expected. Can be used
        as a parent of Multinomial and Categorical nevertheless.

    """
    x = np.atleast_2d(x)
    theta = np.atleast_2d(theta)
    if (np.shape(x)[-1]+1) != np.shape(theta)[-1]:
        raise ValueError('The dimension of x in dirichlet_like must be k-1.')
    return flib.dirichlet(x, theta)

# Exponential----------------------------------------------
@randomwrap
def rexponential(beta, size=None):
    """
    Exponential random variates.
    """

    return np.random.exponential(1./beta, size)

def exponential_expval(beta):
    """
    Expected value of exponential distribution.
    """
    return 1./beta


def exponential_like(x, beta):
    R"""
    Exponential log-likelihood.

    The exponential distribution is a special case of the gamma distribution
    with alpha=1. It often describes the time until an event.

    .. math:: f(x \mid \beta) = \beta e^{-\beta x}

    :Parameters:
      - `x` : x > 0
      - `beta` : Rate parameter (beta > 0).

    .. note::
      - :math:`E(X) = 1/\beta`
      - :math:`Var(X) = 1/\beta^2`
      - PyMC's beta is named 'lambda' by Wikipedia, SciPy, Wolfram MathWorld and other sources.

    """

    return flib.gamma(x, 1, beta)

exponential_grad_like = {'value' : lambda x, beta : flib.gamma_grad_x(x, 1.0, beta),
                         'beta' : lambda x, beta : flib.gamma_grad_beta(x, 1.0, beta)}

# Exponentiated Weibull-----------------------------------
@randomwrap
def rexponweib(alpha, k, loc=0, scale=1, size=None):
    """
    Random exponentiated Weibull variates.
    """

    q = np.random.uniform(size=size)
    r = flib.exponweib_ppf(q, alpha, k)
    return loc + r*scale

def exponweib_expval(alpha, k, loc, scale):
    # Not sure how we can do this, since the first moment is only
    # tractable at particular values of k
    raise NotImplementedError('exponweib_expval has not been implemented yet.')

def exponweib_like(x, alpha, k, loc=0, scale=1):
    R"""
    Exponentiated Weibull log-likelihood.

    The exponentiated Weibull distribution is a generalization of the Weibull
    family. Its value lies in being able to model monotone and non-monotone
    failure rates.

    .. math::
        f(x \mid \alpha,k,loc,scale)  & = \frac{\alpha k}{scale} (1-e^{-z^k})^{\alpha-1} e^{-z^k} z^{k-1} \\
        z & = \frac{x-loc}{scale}

    :Parameters:
      - `x` : x > 0
      - `alpha` : Shape parameter
      - `k` : k > 0
      - `loc` : Location parameter
      - `scale` : Scale parameter (scale > 0).

    """
    return flib.exponweib(x, alpha, k, loc, scale)

"""
commented out because tests fail
exponweib_grad_like = {'value' : flib.exponweib_gx,
                   'alpha' : flib.exponweib_ga,
                   'k' : flib.exponweib_gk,
                   'loc' : flib.exponweib_gl,
                   'scale' : flib.exponweib_gs}
"""
# Gamma----------------------------------------------
@randomwrap
def rgamma(alpha, beta, size=None):
    """
    Random gamma variates.
    """

    return np.random.gamma(shape=alpha, scale=1./beta, size=size)

def gamma_expval(alpha, beta):
    """
    Expected value of gamma distribution.
    """
    return 1. * np.asarray(alpha) / beta

def gamma_like(x, alpha, beta):
    R"""
    Gamma log-likelihood.

    Represents the sum of alpha exponentially distributed random variables, each
    of which has mean beta.

    .. math::
        f(x \mid \alpha, \beta) = \frac{\beta^{\alpha}x^{\alpha-1}e^{-\beta x}}{\Gamma(\alpha)}

    :Parameters:
      - `x` : math:`x \ge 0`
      - `alpha` : Shape parameter (alpha > 0).
      - `beta` : Rate parameter (beta > 0).

    .. note::
       - :math:`E(X) = \frac{\alpha}{\beta}`
       - :math:`Var(X) = \frac{\alpha}{\beta^2}`

    """

    return flib.gamma(x, alpha, beta)


gamma_grad_like = {'value'     : flib.gamma_grad_x,
                   'alpha' : flib.gamma_grad_alpha,
                   'beta'  : flib.gamma_grad_beta}

# GEV Generalized Extreme Value ------------------------
# Modify parameterization -> Hosking (kappa, xi, alpha)
@randomwrap
def rgev(xi, mu=0, sigma=1, size=None):
    """
    Random generalized extreme value (GEV) variates.
    """

    q = np.random.uniform(size=size)
    z = flib.gev_ppf(q, xi)
    return z*sigma + mu

def gev_expval(xi, mu=0, sigma=1):
    """
    Expected value of generalized extreme value distribution.
    """
    return mu - (sigma / xi) + (sigma / xi) * flib.gamfun(1 - xi)

def gev_like(x, xi, mu=0, sigma=1):
    R"""
    Generalized Extreme Value log-likelihood

    .. math::
        pdf(x \mid \xi,\mu,\sigma) = \frac{1}{\sigma}(1 + \xi \left[\frac{x-\mu}{\sigma}\right])^{-1/\xi-1}\exp{-(1+\xi \left[\frac{x-\mu}{\sigma}\right])^{-1/\xi}}

    .. math::
        \sigma & > 0,\\
        x & > \mu-\sigma/\xi \text{ if } \xi > 0,\\
        x & < \mu-\sigma/\xi \text{ if } \xi < 0\\
        x & \in [-\infty,\infty] \text{ if } \xi = 0

    """

    return flib.gev(x, xi, mu, sigma)

# Geometric----------------------------------------------
# Changed the return value
@randomwrap
def rgeometric(p, size=None):
    """
    Random geometric variates.
    """

    return np.random.geometric(p, size)

def geometric_expval(p):
    """
    Expected value of geometric distribution.
    """
    return 1. / p

def geometric_like(x, p):
    R"""
    Geometric log-likelihood. The probability that the first success in a
    sequence of Bernoulli trials occurs on the x'th trial.

    .. math::
        f(x \mid p) = p(1-p)^{x-1}

    :Parameters:
      - `x` : [int] Number of trials before first success (x > 0).
      - `p` : Probability of success on an individual trial, :math:`p \in [0,1]`

    .. note::
      - :math:`E(X)=1/p`
      - :math:`Var(X)=\frac{1-p}{p^2}`

    """

    return flib.geometric(x, p)

geometric_grad_like = {'p' : flib.geometric_gp}

# Half Cauchy----------------------------------------------
@randomwrap
def rhalf_cauchy(alpha, beta, size=None):
    """
    Returns half-Cauchy random variates.
    """

    return abs(alpha + beta*np.tan(pi*random_number(size) - pi/2.0))

def half_cauchy_expval(alpha, beta):
    """
    Expected value of cauchy distribution is undefined.
    """

    return inf

# In wikipedia, the arguments name are k, x0.
def half_cauchy_like(x, alpha, beta):
    R"""
    Half-Cauchy log-likelihood. Simply the absolute value of Cauchy.

    .. math::
        f(x \mid \alpha, \beta) = \frac{2}{\pi \beta [1 + (\frac{x-\alpha}{\beta})^2]}

    :Parameters:
      - `alpha` : Location parameter.
      - `beta` : Scale parameter (beta > 0).

    .. note::
      - x must be non-negative.
    """

    x = np.atleast_1d(x)
    if sum(x < 0):
        return -inf
    return flib.cauchy(x, alpha, beta) + len(x)*np.log(2)

# Half-normal----------------------------------------------
@randomwrap
def rhalf_normal(tau, size=None):
    """
    Random half-normal variates.
    """

    return abs(np.random.normal(0, np.sqrt(1/tau), size))

def half_normal_expval(tau):
    """
    Expected value of half normal distribution.
    """

    return np.sqrt(2. * pi / np.asarray(tau))

def half_normal_like(x, tau):
    R"""
    Half-normal log-likelihood, a normal distribution with mean 0 limited
    to the domain :math:`x \in [0, \infty)`.

    .. math::
        f(x \mid \tau) = \sqrt{\frac{2\tau}{\pi}}\exp\left\{ {\frac{-x^2 \tau}{2}}\right\}

    :Parameters:
      - `x` : :math:`x \ge 0`
      - `tau` : tau > 0

    """

    return flib.hnormal(x, tau)

half_normal_grad_like = {'value'   : flib.hnormal_gradx,
                 'tau' : flib.hnormal_gradtau}

# Hypergeometric----------------------------------------------
def rhypergeometric(n, m, N, size=None):
    """
    Returns hypergeometric random variates.
    """
    if n == 0:
        return np.zeros(size, dtype=int)
    elif n == N:
        out = np.empty(size, dtype=int)
        out.fill(m)
        return out
    return np.random.hypergeometric(n, N-n, m, size)

def hypergeometric_expval(n, m, N):
    """
    Expected value of hypergeometric distribution.
    """
    return 1. * n * m / N

def hypergeometric_like(x, n, m, N):
    R"""
    Hypergeometric log-likelihood.

    Discrete probability distribution that describes the number of successes in
    a sequence of draws from a finite population without replacement.

    .. math::

        f(x \mid n, m, N) = \frac{\left({ \begin{array}{c} {m} \\ {x} \\
        \end{array} }\right)\left({ \begin{array}{c} {N-m} \\ {n-x} \\
        \end{array}}\right)}{\left({ \begin{array}{c} {N} \\ {n} \\
        \end{array}}\right)}


    :Parameters:
      - `x` : [int] Number of successes in a sample drawn from a population.
      - `n` : [int] Size of sample drawn from the population.
      - `m` : [int] Number of successes in the population.
      - `N` : [int] Total number of units in the population.

    .. note::

        :math:`E(X) = \frac{n n}{N}`
    """

    return flib.hyperg(x, n, m, N)

# Inverse gamma----------------------------------------------
@randomwrap
def rinverse_gamma(alpha, beta, size=None):
    """
    Random inverse gamma variates.
    """

    return 1. / np.random.gamma(shape=alpha, scale=1./beta, size=size)

def inverse_gamma_expval(alpha, beta):
    """
    Expected value of inverse gamma distribution.
    """
    return 1. * np.asarray(beta) / (alpha-1.)

def inverse_gamma_like(x, alpha, beta):
    R"""
    Inverse gamma log-likelihood, the reciprocal of the gamma distribution.

    .. math::
        f(x \mid \alpha, \beta) = \frac{\beta^{\alpha}}{\Gamma(\alpha)} x^{-\alpha - 1} \exp\left(\frac{-\beta}{x}\right)

    :Parameters:
      - `x` : x > 0
      - `alpha` : Shape parameter (alpha > 0).
      - `beta` : Scale parameter (beta > 0).

    .. note::

       :math:`E(X)=\frac{\beta}{\alpha-1}`  for :math:`\alpha > 1`
       :math:`Var(X)=\frac{\beta^2}{(\alpha-1)^2(\alpha)}`  for :math:`\alpha > 2`

    """

    return flib.igamma(x, alpha, beta)

inverse_gamma_grad_like = {'value' : flib.igamma_grad_x,
             'alpha' : flib.igamma_grad_alpha,
             'beta' : flib.igamma_grad_beta}

# Inverse Wishart---------------------------------------------------

# def rinverse_wishart(n, C):
#     """
#     Return an inverse Wishart random matrix.
#
#     :Parameters:
#       - `n` : [int] Degrees of freedom (n > 0).
#       - `C` : Symmetric and positive definite scale matrix
#     """
#     wi = rwishart(n, np.asmatrix(C).I).I
#     flib.symmetrize(wi)
#     return wi
#
# def inverse_wishart_expval(n, C):
#     """
#     Expected value of inverse Wishart distribution.
#
#     :Parameters:
#       - `n` : [int] Degrees of freedom (n > 0).
#       - `C` : Symmetric and positive definite scale matrix
#
#     """
#     return np.asarray(C)/(n-len(C)-1)
#
# def inverse_wishart_like(X, n, C):
#     R"""
#     Inverse Wishart log-likelihood. The inverse Wishart distribution
#     is the conjugate prior for the covariance matrix of a multivariate
#     normal distribution.
#
#     .. math::
#         f(X \mid n, T) = \frac{{\mid T \mid}^{n/2}{\mid X
#         \mid}^{-(n+k+1)/2} \exp\left\{ -\frac{1}{2} Tr(TX^{-1})
#         \right\}}{2^{nk/2} \Gamma_p(n/2)}
#
#     where :math:`k` is the rank of X.
#
#     :Parameters:
#       - `X` : Symmetric, positive definite matrix.
#       - `n` : [int] Degrees of freedom (n > 0).
#       - `C` : Symmetric and positive definite scale matrix
#
#     .. note::
#        Step method MatrixMetropolis will preserve the symmetry of
#        Wishart variables.
#
#     """
#     return flib.blas_wishart(X.I, n, C.I)

# Double exponential (Laplace)--------------------------------------------
@randomwrap
def rlaplace(mu, tau, size=None):
    """
    Laplace (double exponential) random variates.
    """

    u = np.random.uniform(-0.5, 0.5, size)
    return mu - np.sign(u)*np.log(1 - 2*np.abs(u))/tau

rdexponential = rlaplace

def laplace_expval(mu, tau):
    """
    Expected value of Laplace (double exponential) distribution.
    """
    return mu

dexponential_expval = laplace_expval

def laplace_like(x, mu, tau):
    R"""
    Laplace (double exponential) log-likelihood.

    The Laplace (or double exponential) distribution describes the
    difference between two independent, identically distributed exponential
    events. It is often used as a heavier-tailed alternative to the normal.

    .. math::
        f(x \mid \mu, \tau) = \frac{\tau}{2}e^{-\tau |x-\mu|}

    :Parameters:
      - `x` : :math:`-\infty < x < \infty`
      - `mu` : Location parameter :math: `-\infty < mu < \infty`
      - `tau` : Scale parameter :math:`\tau > 0`

    .. note::
      - :math:`E(X) = \mu`
      - :math:`Var(X) = \frac{2}{\tau^2}`
    """

    return flib.gamma(np.abs(x-mu), 1, tau) - np.log(2)

laplace_grad_like = {'value'   : lambda x, mu, tau: flib.gamma_grad_x(np.abs(x - mu), 1, tau) * np.sign(x - mu),
                     'mu'  : lambda x, mu, tau: -flib.gamma_grad_x(np.abs(x - mu), 1, tau) * np.sign(x - mu),
                     'tau' : lambda x, mu, tau: flib.gamma_grad_beta(np.abs(x - mu), 1, tau)}


dexponential_like = laplace_like
dexponential_grad_like = laplace_grad_like

# Logistic-----------------------------------
@randomwrap
def rlogistic(mu, tau, size=None):
    """
    Logistic random variates.
    """

    u = np.random.random(size)
    return mu + np.log(u/(1-u))/tau


def logistic_expval(mu, tau):
    """
    Expected value of logistic distribution.
    """
    return mu


def logistic_like(x, mu, tau):
    R"""
    Logistic log-likelihood.

    The logistic distribution is often used as a growth model; for example,
    populations, markets. Resembles a heavy-tailed normal distribution.

    .. math::
        f(x \mid \mu, tau) = \frac{\tau \exp(-\tau[x-\mu])}{[1 + \exp(-\tau[x-\mu])]^2}

    :Parameters:
      - `x` : :math:`-\infty < x < \infty`
      - `mu` : Location parameter :math:`-\infty < mu < \infty`
      - `tau` : Scale parameter (tau > 0)

    .. note::
      - :math:`E(X) = \mu`
      - :math:`Var(X) = \frac{\pi^2}{3\tau^2}`
    """

    return flib.logistic(x, mu, tau)


# Lognormal----------------------------------------------
@randomwrap
def rlognormal(mu, tau, size=None):
    """
    Return random lognormal variates.
    """

    return np.random.lognormal(mu, np.sqrt(1./tau), size)

def lognormal_expval(mu, tau):
    """
    Expected value of log-normal distribution.
    """
    return np.exp(mu + 1./2/tau)

def lognormal_like(x, mu, tau):
    R"""
    Log-normal log-likelihood.

    Distribution of any random variable whose logarithm is normally
    distributed. A variable might be modeled as log-normal if it can be thought
    of as the multiplicative product of many small independent factors.

    .. math::
        f(x \mid \mu, \tau) = \sqrt{\frac{\tau}{2\pi}}\frac{
        \exp\left\{ -\frac{\tau}{2} (\ln(x)-\mu)^2 \right\}}{x}

    :Parameters:
      - `x` : x > 0
      - `mu` : Location parameter.
      - `tau` : Scale parameter (tau > 0).

    .. note::

       :math:`E(X)=e^{\mu+\frac{1}{2\tau}}`
       :math:`Var(X)=(e^{1/\tau}-1)e^{2\mu+\frac{1}{\tau}}`

    """
    return flib.lognormal(x, mu, tau)

lognormal_grad_like = {'value'   : flib.lognormal_gradx,
                       'mu'  : flib.lognormal_gradmu,
                       'tau' : flib.lognormal_gradtau}


# Multinomial----------------------------------------------
#@randomwrap
def rmultinomial(n, p, size=None):
    """
    Random multinomial variates.
    """
    # Leaving size=None as the default means return value is 1d array
    # if not specified-- nicer.

    # Single value for p:
    if len(np.shape(p)) == 1:
        return np.random.multinomial(n, p, size)

    # Multiple values for p:
    if np.isscalar(n):
        n = n * np.ones(np.shape(p)[0], dtype=np.int)
    out = np.empty(np.shape(p))
    for i in xrange(np.shape(p)[0]):
        out[i,:] = np.random.multinomial(n[i], p[i,:], size)
    return out

def multinomial_expval(n, p):
    """
    Expected value of multinomial distribution.
    """
    return np.asarray([pr * n for pr in p])

def multinomial_like(x, n, p):
    R"""
    Multinomial log-likelihood.

    Generalization of the binomial
    distribution, but instead of each trial resulting in "success" or
    "failure", each one results in exactly one of some fixed finite number k
    of possible outcomes over n independent trials. 'x[i]' indicates the number
    of times outcome number i was observed over the n trials.

    .. math::
        f(x \mid n, p) = \frac{n!}{\prod_{i=1}^k x_i!} \prod_{i=1}^k p_i^{x_i}

    :Parameters:
      x : (ns, k) int
          Random variable indicating the number of time outcome i is
          observed. :math:`\sum_{i=1}^k x_i=n`, :math:`x_i \ge 0`.
      n : int
          Number of trials.
      p : (k,)
          Probability of each one of the different outcomes.
          :math:`\sum_{i=1}^k p_i = 1)`, :math:`p_i \ge 0`.

    .. note::
       - :math:`E(X_i)=n p_i`
       - :math:`Var(X_i)=n p_i(1-p_i)`
       - :math:`Cov(X_i,X_j) = -n p_i p_j`
       - If :math: `\sum_i p_i < 0.999999` a log-likelihood value of -inf
       will be returned.

    """
    # flib expects 2d arguments. Do we still want to support multiple p
    # values along realizations ?
    x = np.atleast_2d(x)
    p = np.atleast_2d(p)

    return flib.multinomial(x, n, p)

# Multivariate hypergeometric------------------------------
def rmultivariate_hypergeometric(n, m, size=None):
    """
    Random multivariate hypergeometric variates.

    Parameters:
      - `n` : Number of draws.
      - `m` : Number of items in each categoy.
    """

    N = len(m)
    urn = np.repeat(np.arange(N), m)

    if size:
        draw = np.array([[urn[i] for i in np.random.permutation(len(urn))[:n]]
                         for j in range(size)])

        r = [[np.sum(draw[j] == i) for i in range(len(m))]
             for j in range(size)]
    else:
        draw = np.array([urn[i] for i in np.random.permutation(len(urn))[:n]])

        r = [np.sum(draw == i) for i in range(len(m))]
    return np.asarray(r)

def multivariate_hypergeometric_expval(n, m):
    """
    Expected value of multivariate hypergeometric distribution.

    Parameters:
      - `n` : Number of draws.
      - `m` : Number of items in each categoy.
    """
    m = np.asarray(m, float)
    return n * (m / m.sum())


def multivariate_hypergeometric_like(x, m):
    R"""
    Multivariate hypergeometric log-likelihood

    Describes the probability of drawing x[i] elements of the ith category,
    when the number of items in each category is given by m.

    .. math::
        \frac{\prod_i \left({ \begin{array}{c} {m_i} \\ {x_i} \\
        \end{array}}\right)}{\left({ \begin{array}{c} {N} \\ {n} \\
        \end{array}}\right)}


    where :math:`N = \sum_i m_i` and :math:`n = \sum_i x_i`.

    :Parameters:
      - `x` : [int sequence] Number of draws from each category, (x < m).
      - `m` : [int sequence] Number of items in each categoy.

    """
    return flib.mvhyperg(x, m)


# Multivariate normal--------------------------------------
def rmv_normal(mu, tau, size=1):
    """
    Random multivariate normal variates.
    """

    sig = np.linalg.cholesky(tau)
    mu_size = np.shape(mu)

    if size == 1:
        out = np.random.normal(size=mu_size)
        try:
            flib.dtrsm_wrap(sig, out, 'L', 'T', 'L', 1.)
        except:
            out = np.linalg.solve(sig, out)
        out += mu
        return out
    else:
        if not hasattr(size, '__iter__'):
            size = (size,)
        tot_size = np.prod(size)
        out = np.random.normal(size = (tot_size,) + mu_size)
        for i in xrange(tot_size):
            try:
                flib.dtrsm_wrap(sig, out[i,:], 'L', 'T', 'L', 1.)
            except:
                out[i,:] = np.linalg.solve(sig, out[i,:])
            out[i,:] += mu
        return out.reshape(size+mu_size)

def mv_normal_expval(mu, tau):
    """
    Expected value of multivariate normal distribution.
    """
    return mu

def mv_normal_like(x, mu, tau):
    R"""
    Multivariate normal log-likelihood

    .. math::
        f(x \mid \pi, T) = \frac{|T|^{1/2}}{(2\pi)^{1/2}} \exp\left\{ -\frac{1}{2} (x-\mu)^{\prime}T(x-\mu) \right\}

    :Parameters:
      - `x` : (n,k)
      - `mu` : (k) Location parameter sequence.
      - `Tau` : (k,k) Positive definite precision matrix.

    .. seealso:: :func:`mv_normal_chol_like`, :func:`mv_normal_cov_like`

    """
    # TODO: Vectorize in Fortran
    if len(np.shape(x)) > 1:
        return np.sum([flib.prec_mvnorm(r, mu, tau) for r in x])
    else:
        return flib.prec_mvnorm(x, mu, tau)

# Multivariate normal, parametrized with covariance---------------------------
def rmv_normal_cov(mu, C, size=1):
    """
    Random multivariate normal variates.
    """
    mu_size = np.shape(mu)
    if size == 1:
        return np.random.multivariate_normal(mu, C, size).reshape(mu_size)
    else:
        return np.random.multivariate_normal(mu, C, size).reshape((size,)+mu_size)

def mv_normal_cov_expval(mu, C):
    """
    Expected value of multivariate normal distribution.
    """
    return mu

def mv_normal_cov_like(x, mu, C):
    R"""
    Multivariate normal log-likelihood parameterized by a covariance
    matrix.

    .. math::
        f(x \mid \pi, C) = \frac{1}{(2\pi|C|)^{1/2}} \exp\left\{ -\frac{1}{2} (x-\mu)^{\prime}C^{-1}(x-\mu) \right\}

    :Parameters:
      - `x` : (n,k)
      - `mu` : (k) Location parameter.
      - `C` : (k,k) Positive definite covariance matrix.

    .. seealso:: :func:`mv_normal_like`, :func:`mv_normal_chol_like`

    """
    # TODO: Vectorize in Fortran
    if len(np.shape(x)) > 1:
        return np.sum([flib.cov_mvnorm(r, mu, C) for r in x])
    else:
        return flib.cov_mvnorm(x, mu, C)


# Multivariate normal, parametrized with Cholesky factorization.----------
def rmv_normal_chol(mu, sig, size=1):
    """
    Random multivariate normal variates.
    """
    mu_size = np.shape(mu)

    if size == 1:
        out = np.random.normal(size=mu_size)
        try:
            flib.dtrmm_wrap(sig, out, 'L', 'N', 'L', 1.)
        except:
            out = np.dot(sig, out)
        out += mu
        return out
    else:
        if not hasattr(size, '__iter__'):
            size = (size,)
        tot_size = np.prod(size)
        out = np.random.normal(size = (tot_size,) + mu_size)
        for i in xrange(tot_size):
            try:
                flib.dtrmm_wrap(sig, out[i,:], 'L', 'N', 'L', 1.)
            except:
                out[i,:] = np.dot(sig, out[i,:])
            out[i,:] += mu
        return out.reshape(size+mu_size)


def mv_normal_chol_expval(mu, sig):
    """
    Expected value of multivariate normal distribution.
    """
    return mu

def mv_normal_chol_like(x, mu, sig):
    R"""
    Multivariate normal log-likelihood.

    .. math::
        f(x \mid \pi, \sigma) = \frac{1}{(2\pi)^{1/2}|\sigma|)} \exp\left\{ -\frac{1}{2} (x-\mu)^{\prime}(\sigma \sigma^{\prime})^{-1}(x-\mu) \right\}

    :Parameters:
      - `x` : (n,k)
      - `mu` : (k) Location parameter.
      - `sigma` : (k,k) Lower triangular matrix.

    .. seealso:: :func:`mv_normal_like`, :func:`mv_normal_cov_like`

    """
    # TODO: Vectorize in Fortran
    if len(np.shape(x)) > 1:
        return np.sum([flib.chol_mvnorm(r, mu, sig) for r in x])
    else:
        return flib.chol_mvnorm(x, mu, sig)



# Negative binomial----------------------------------------
@randomwrap
def rnegative_binomial(mu, alpha, size=None):
    """
    Random negative binomial variates.
    """
    # Using gamma-poisson mixture rather than numpy directly
    # because numpy apparently rounds
    mu = np.asarray(mu, dtype=float)
    pois_mu = np.random.gamma(alpha, mu/alpha, size)
    return np.random.poisson(pois_mu, size)
    # return np.random.negative_binomial(alpha, alpha / (mu + alpha), size)

def negative_binomial_expval(mu, alpha):
    """
    Expected value of negative binomial distribution.
    """
    return mu


def negative_binomial_like(x, mu, alpha):
    R"""
    Negative binomial log-likelihood.

    The negative binomial
    distribution describes a Poisson random variable whose rate
    parameter is gamma distributed. PyMC's chosen parameterization is
    based on this mixture interpretation.

    .. math::
        f(x \mid \mu, \alpha) = \frac{\Gamma(x+\alpha)}{x! \Gamma(\alpha)} (\alpha/(\mu+\alpha))^\alpha (\mu/(\mu+\alpha))^x

    :Parameters:
      - `x` : Input data (x > 0).
      - `mu` : mu > 0
      - `alpha` : alpha > 0

    .. note::
      - :math:`E[x]=\mu`
      - In Wikipedia's parameterization,
        :math:`r=\alpha`
        :math:`p=\alpha/(\mu+\alpha)`
        :math:`\mu=r(1-p)/p`

    """
    alpha = np.array(alpha)
    if (alpha > 1e10).any():
        if (alpha > 1e10).all():
            # Return Poisson when alpha gets very large
            return flib.poisson(x, mu)

        # Split big and small dispersion values
        big = alpha > 1e10
        return flib.poisson(x[big], mu[big]) + flib.negbin2(x[big-True],
                mu[big-True], alpha[big-True])
    return flib.negbin2(x, mu, alpha)

negative_binomial_grad_like = {'mu'    : flib.negbin2_gmu,
                               'alpha' : flib.negbin2_ga}

# Normal---------------------------------------------------
@randomwrap
def rnormal(mu, tau,size=None):
    """
    Random normal variates.
    """
    return np.random.normal(mu, 1./np.sqrt(tau), size)

def normal_expval(mu, tau):
    """
    Expected value of normal distribution.
    """
    return mu

def normal_like(x, mu, tau):
    R"""
    Normal log-likelihood.

    .. math::
        f(x \mid \mu, \tau) = \sqrt{\frac{\tau}{2\pi}} \exp\left\{ -\frac{\tau}{2} (x-\mu)^2 \right\}

    :Parameters:
      - `x` : Input data.
      - `mu` : Mean of the distribution.
      - `tau` : Precision of the distribution, which corresponds to
        :math:`1/\sigma^2` (tau > 0).

    .. note::
       - :math:`E(X) = \mu`
       - :math:`Var(X) = 1/\tau`

    """
    # try:
    #     constrain(tau, lower=0)
    # except ZeroProbability:
    #     return -np.Inf

    return flib.normal(x, mu, tau)

def t_normal_grad_x(x, mu, tau):
    return flib.normal_grad_x(x, mu, tau)

def t_normal_grad_mu(x, mu, tau):
    return flib.normal_grad_mu(x, mu, tau)
def t_normal_grad_tau(x, mu, tau):
    return flib.normal_grad_tau(x, mu, tau)

normal_grad_like = {'value' : t_normal_grad_x,
             'mu' : t_normal_grad_mu,
             'tau' : t_normal_grad_tau}

# normal_grad_like = {'x' : flib.normal_grad_x,
#             'mu' : flib.normal_grad_mu,
#             'tau' : flib.normal_grad_tau}

# von Mises--------------------------------------------------
@randomwrap
def rvon_mises(mu, kappa, size=None):
    """
    Random von Mises variates.
    """
    # TODO: Just return straight from numpy after release 1.3
    return (np.random.mtrand.vonmises(mu, kappa, size) + np.pi)%(2.*np.pi)-np.pi

def von_mises_expval(mu, kappa):
    """
    Expected value of von Mises distribution.
    """
    return mu

def von_mises_like(x, mu, kappa):
    R"""
    von Mises log-likelihood.

    .. math::
        f(x \mid \mu, k) = \frac{e^{k \cos(x - \mu)}}{2 \pi I_0(k)}

    where `I_0` is the modified Bessel function of order 0.

    :Parameters:
      - `x` : Input data.
      - `mu` : Mean of the distribution.
      - `kappa` : Dispersion of the distribution

    .. note::
       - :math:`E(X) = \mu`

    """
    return flib.vonmises(x, mu, kappa)

# Pareto---------------------------------------------------
@randomwrap
def rpareto(alpha, m, size=None):
    """
    Random Pareto variates.
    """
    return m / (random_number(size)**(1./alpha))

def pareto_expval(alpha, m):
    """
    Expected value of Pareto distribution.
    """

    if alpha <= 1:
        return inf
    return alpha*m/(alpha-1)

def pareto_like(x, alpha, m):
    R"""
    Pareto log-likelihood. The Pareto is a continuous, positive
    probability distribution with two parameters. It is often used
    to characterize wealth distribution, or other examples of the
    80/20 rule.

    .. math::
        f(x \mid \alpha, m) = \frac{\alpha m^{\alpha}}{x^{\alpha+1}}

    :Parameters:
      - `x` : Input data (x > m)
      - `alpha` : Shape parameter (alpha>0)
      - `m` : Scale parameter (m>0)

    .. note::
       - :math:`E(x)=\frac{\alpha m}{\alpha-1} if \alpha > 1`
       - :math:`Var(x)=\frac{m^2 \alpha}{(\alpha-1)^2(\alpha-2)} if \alpha > 2`

    """
    return flib.pareto(x, alpha, m)

# Truncated Pareto---------------------------------------------------
@randomwrap
def rtruncated_pareto(alpha, m, b, size=None):
    """
    Random bounded Pareto variates.
    """
    u = random_number(size)
    return (-(u*b**alpha - u*m**alpha - b**alpha)/(b**alpha * m**alpha))**(-1./alpha)

def truncated_pareto_expval(alpha, m, b):
    """
    Expected value of truncated Pareto distribution.
    """

    if alpha <= 1:
        return inf
    part1 = (m**alpha)/(1. - (m/b)**alpha)
    part2 = 1.*alpha/(alpha-1)
    part3 = (1./(m**(alpha-1)) - 1./(b**(alpha-1.)))
    return part1*part2*part3

def truncated_pareto_like(x, alpha, m, b):
    R"""
    Truncated Pareto log-likelihood. The Pareto is a continuous, positive
    probability distribution with two parameters. It is often used
    to characterize wealth distribution, or other examples of the
    80/20 rule.

    .. math::
        f(x \mid \alpha, m, b) = \frac{\alpha m^{\alpha} x^{-\alpha}}{1-(m/b)**{\alpha}}

    :Parameters:
      - `x` : Input data (x > m)
      - `alpha` : Shape parameter (alpha>0)
      - `m` : Scale parameter (m>0)
      - `b` : Upper bound (b>m)

    """
    return flib.truncated_pareto(x, alpha, m, b)

# Poisson--------------------------------------------------
@randomwrap
def rpoisson(mu, size=None):
    """
    Random poisson variates.
    """

    return np.random.poisson(mu, size)


def poisson_expval(mu):
    """
    Expected value of Poisson distribution.
    """

    return mu


def poisson_like(x, mu):
    R"""
    Poisson log-likelihood.

    The Poisson is a discrete probability
    distribution.  It is often used to model the number of events
    occurring in a fixed period of time when the times at which events
    occur are independent. The Poisson distribution can be derived as
    a limiting case of the binomial distribution.

    .. math::
        f(x \mid \mu) = \frac{e^{-\mu}\mu^x}{x!}

    :Parameters:
      - `x` : [int] :math:`x \in {0,1,2,...}`
      - `mu` : Expected number of occurrences during the given interval, :math:`\mu \geq 0`.

    .. note::
       - :math:`E(x)=\mu`
       - :math:`Var(x)=\mu`

    """
    return flib.poisson(x, mu)

poisson_grad_like = {'mu' : flib.poisson_gmu}

# Truncated Poisson--------------------------------------------------
@randomwrap
def rtruncated_poisson(mu, k, size=None):
    """
    Random truncated Poisson variates with minimum value k, generated
    using rejection sampling.
    """

    # Calculate m
    try:
        m = max(0, np.floor(k-mu))
    except (TypeError, ValueError):
        # More than one mu
        return np.array([rtruncated_poisson(x, i)
                         for x, i in zip(mu, np.resize(k, np.size(mu)))]).squeeze()
    k -= 1

    # Calculate constant for acceptance probability
    C = np.exp(flib.factln(k+1)-flib.factln(k+1-m))

    # Empty array to hold random variates
    rvs = np.empty(0, int)

    total_size = np.prod(size or 1)

    while(len(rvs) < total_size):

        # Propose values by sampling from untruncated Poisson with mean mu + m
        proposals = np.random.poisson(mu+m, (total_size*4, np.size(m))).squeeze()

        # Acceptance probability
        a = C * np.array([np.exp(flib.factln(y-m)-flib.factln(y))
                          for y in proposals])
        a *= proposals > k

        # Uniform random variates
        u = np.random.random(total_size*4)

        rvs = np.append(rvs, proposals[u < a])

    return np.reshape(rvs[:total_size], size)


def truncated_poisson_expval(mu, k):
    """
    Expected value of Poisson distribution truncated to be no smaller than k.
    """

    return mu/(1.-poiscdf(k, mu))


def truncated_poisson_like(x, mu, k):
    R"""
    Truncated Poisson log-likelihood.

    The Truncated Poisson is a
    discrete probability distribution that is arbitrarily truncated to
    be greater than some minimum value k. For example, zero-truncated
    Poisson distributions can be used to model counts that are
    constrained to be non-negative.

    .. math::
        f(x \mid \mu, k) = \frac{e^{-\mu}\mu^x}{x!(1-F(k|\mu))}

    :Parameters:
      - `x` : [int] :math:`x \in {0,1,2,...}`
      - `mu` : Expected number of occurrences during the given interval,
               :math:`\mu \geq 0`.
      - `k` : Truncation point representing the minimum allowable value.

    .. note::
       - :math:`E(x)=\frac{\mu}{1-F(k|\mu)}`
       - :math:`Var(x)=\frac{\mu}{1-F(k|\mu)}`

    """
    return flib.trpoisson(x, mu, k)

truncated_poisson_grad_like = {'mu' : flib.trpoisson_gmu}

# Truncated normal distribution--------------------------
@randomwrap
def rtruncated_normal(mu, tau, a=-np.inf, b=np.inf, size=None):
    """
    Random truncated normal variates.
    """

    sigma = 1./np.sqrt(tau)
    na = utils.normcdf((a-mu)/sigma)
    nb = utils.normcdf((b-mu)/sigma)

    # Use the inverse CDF generation method.
    U = np.random.mtrand.uniform(size=size)
    q = U * nb + (1-U)*na
    R = utils.invcdf(q)

    # Unnormalize
    return R*sigma + mu

rtruncnorm = rtruncated_normal

def truncated_normal_expval(mu, tau, a, b):
    """Expected value of the truncated normal distribution.

    .. math::
       E(X) =\mu + \frac{\sigma(\varphi_1-\varphi_2)}{T}


    where

    .. math::
       T & =\Phi\left(\frac{B-\mu}{\sigma}\right)-\Phi
       \left(\frac{A-\mu}{\sigma}\right)\text \\
       \varphi_1 &=
       \varphi\left(\frac{A-\mu}{\sigma}\right) \\
       \varphi_2 &=
       \varphi\left(\frac{B-\mu}{\sigma}\right) \\

    and :math:`\varphi = N(0,1)` and :math:`tau & 1/sigma**2`.

    :Parameters:
      - `mu` : Mean of the distribution.
      - `tau` : Precision of the distribution, which corresponds to 1/sigma**2 (tau > 0).
      - `a` : Left bound of the distribution.
      - `b` : Right bound of the distribution.

    """
    phia = np.exp(normal_like(a, mu, tau))
    phib = np.exp(normal_like(b, mu, tau))
    sigma = 1./np.sqrt(tau)
    Phia = utils.normcdf((a-mu)/sigma)
    if b == np.inf:
        Phib = 1.0
    else:
        Phib = utils.normcdf((b-mu)/sigma)
    return (mu + (phia-phib)/(Phib - Phia))[0]

truncnorm_expval = truncated_normal_expval

def truncated_normal_like(x, mu, tau, a=None, b=None):
    R"""
    Truncated normal log-likelihood.

    .. math::
        f(x \mid \mu, \tau, a, b) = \frac{\phi(\frac{x-\mu}{\sigma})} {\Phi(\frac{b-\mu}{\sigma}) - \Phi(\frac{a-\mu}{\sigma})},

    where :math:`\sigma^2=1/\tau`, `\phi` is the standard normal PDF and `\Phi` is the standard normal CDF.

    :Parameters:
      - `x` : Input data.
      - `mu` : Mean of the distribution.
      - `tau` : Precision of the distribution, which corresponds to 1/sigma**2 (tau > 0).
      - `a` : Left bound of the distribution.
      - `b` : Right bound of the distribution.
    """
    x = np.atleast_1d(x)
    if a is None:
        a = -np.inf
    a = np.atleast_1d(a)
    if b is None:
        b = np.inf
    b = np.atleast_1d(b)
    mu = np.atleast_1d(mu)
    sigma = (1./np.atleast_1d(np.sqrt(tau)))
    if (x < a).any() or (x > b).any():
        return -np.inf
    else:
        n = len(x)
        phi = normal_like(x, mu, tau)
        lPhia = utils.normcdf((a-mu)/sigma, log=True)
        lPhib = utils.normcdf((b-mu)/sigma, log=True)
        try:
            d = utils.log_difference(lPhib, lPhia)
        except ValueError:
            return -np.inf
        # d = np.log(Phib-Phia)
        if len(d) == n:
            Phi = d.sum()
        else:
            Phi = n*d
        if np.isnan(Phi) or np.isinf(Phi):
            return -np.inf
        return phi - Phi

truncnorm_like = truncated_normal_like

# Azzalini's skew-normal-----------------------------------
@randomwrap
def rskew_normal(mu,tau,alpha,size=()):
    """
    Skew-normal random variates.
    """
    size_ = size or (1,)
    len_ = np.prod(size_)
    return flib.rskewnorm(len_, mu, tau, alpha, np.random.normal(size=2*len_)).reshape(size)

def skew_normal_expval(mu, tau, alpha):
    """
    Expectation of skew-normal random variables.
    """
    delta = alpha / np.sqrt(1.+alpha**2)
    return mu + np.sqrt(2/pi/tau) * delta

def skew_normal_like(x, mu, tau, alpha):
    R"""
    Azzalini's skew-normal log-likelihood

    .. math::
        f(x \mid \mu, \tau, \alpha) = 2 \Phi((x-\mu)\sqrt{\tau}\alpha) \phi(x,\mu,\tau)

    where :math: \Phi is the normal CDF and :math: \phi is the normal PDF.

    :Parameters:
      - `x` : Input data.
      - `mu` : Mean of the distribution.
      - `tau` : Precision of the distribution (> 0).
      - `alpha` : Shape parameter of the distribution.

    .. note::
      See http://azzalini.stat.unipd.it/SN/
    """
    return flib.sn_like(x, mu, tau, alpha)


# Student's t-----------------------------------
@randomwrap
def rt(nu, size=None):
    """
    Student's t random variates.
    """
    return rnormal(0, 1, size) / np.sqrt(rchi2(nu, size)/nu)

def t_expval(nu):
    """
    Expectation of Student's t random variables.
    """
    return 0

def t_like(x, nu):
    R"""
    Student's T log-likelihood.

    Describes a zero-mean normal variable
    whose precision is gamma distributed. Alternatively, describes the
    mean of several zero-mean normal random variables divided by their
    sample standard deviation.

    .. math::
        f(x \mid \nu) = \frac{\Gamma(\frac{\nu+1}{2})}{\Gamma(\frac{\nu}{2}) \sqrt{\nu\pi}} \left( 1 + \frac{x^2}{\nu} \right)^{-\frac{\nu+1}{2}}

    :Parameters:
      - `x` : Input data.
      - `nu` : Degrees of freedom.

    """
    nu = np.asarray(nu)
    return flib.t(x, nu)

# Non-central Student's t-----------------------------------
@randomwrap
def rnoncentral_t(mu, lam, nu, size=None):
    """
    Non-central Student's t random variates.
    """
    tau = rgamma(nu/2., nu/(2.*lam), size)
    return rnormal(mu, tau)

def noncentral_t_expval(mu, lam, nu):
    """noncentral_t_expval(mu, lam, nu)

    Expectation of non-central Student's t random variables. Only defined
    for nu>1.
    """
    if nu > 1:
        return mu
    return inf

def noncentral_t_like(x, mu, lam, nu):
    R"""
    Non-central Student's T log-likelihood.

    Describes a normal variable whose precision is gamma distributed.

    .. math::
        f(x|\mu,\lambda,\nu) = \frac{\Gamma(\frac{\nu +
        1}{2})}{\Gamma(\frac{\nu}{2})}
        \left(\frac{\lambda}{\pi\nu}\right)^{\frac{1}{2}}
        \left[1+\frac{\lambda(x-\mu)^2}{\nu}\right]^{-\frac{\nu+1}{2}}

    :Parameters:
      - `x` : Input data.
      - `mu` : Location parameter.
      - `lambda` : Scale parameter.
      - `nu` : Degrees of freedom.

    """
    mu = np.asarray(mu)
    lam = np.asarray(lam)
    nu = np.asarray(nu)
    return flib.nct(x, mu, lam, nu)

def t_grad_setup(x, nu, f):
    nu = np.asarray(nu)

    return f(x, nu)

t_grad_like = {'value'  : lambda x, nu : t_grad_setup(x, nu, flib.t_grad_x),
               'nu' : lambda x, nu : t_grad_setup(x, nu, flib.t_grad_nu)}


# DiscreteUniform--------------------------------------------------
@randomwrap
def rdiscrete_uniform(lower, upper, size=None):
    """
    Random discrete_uniform variates.
    """
    return np.random.randint(lower, upper+1, size)

def discrete_uniform_expval(lower, upper):
    """
    Expected value of discrete_uniform distribution.
    """
    return (upper - lower) / 2.

def discrete_uniform_like(x, lower, upper):
    R"""
    Discrete uniform log-likelihood.

    .. math::
        f(x \mid lower, upper) = \frac{1}{upper-lower}

    :Parameters:
      - `x` : [int] :math:`lower \leq x \leq upper`
      - `lower` : Lower limit.
      - `upper` : Upper limit (upper > lower).

    """
    return flib.duniform_like(x, lower, upper)


# Uniform--------------------------------------------------
@randomwrap
def runiform(lower, upper, size=None):
    """
    Random uniform variates.
    """
    return np.random.uniform(lower, upper, size)

def uniform_expval(lower, upper):
    """
    Expected value of uniform distribution.
    """
    return (upper - lower) / 2.

def uniform_like(x, lower, upper):
    R"""
    Uniform log-likelihood.

    .. math::
        f(x \mid lower, upper) = \frac{1}{upper-lower}

    :Parameters:
      - `x` : :math:`lower \leq x \leq upper`
      - `lower` : Lower limit.
      - `upper` : Upper limit (upper > lower).

    """

    return flib.uniform_like(x, lower, upper)

uniform_grad_like = {'value' : flib.uniform_grad_x,
             'lower' : flib.uniform_grad_l,
             'upper' : flib.uniform_grad_u}

# Weibull--------------------------------------------------
@randomwrap
def rweibull(alpha, beta,size=None):
    """
    Weibull random variates.
    """
    tmp = -np.log(runiform(0, 1, size))
    return beta * (tmp ** (1. / alpha))

def weibull_expval(alpha, beta):
    """
    Expected value of weibull distribution.
    """
    return beta * np.exp(gammaln((alpha + 1.) / alpha))

def weibull_like(x, alpha, beta):
    R"""
    Weibull log-likelihood

    .. math::
        f(x \mid \alpha, \beta) = \frac{\alpha x^{\alpha - 1}
        \exp(-(\frac{x}{\beta})^{\alpha})}{\beta^\alpha}

    :Parameters:
      - `x` : :math:`x \ge 0`
      - `alpha` : alpha > 0
      - `beta` : beta > 0

    .. note::
      - :math:`E(x)=\beta \Gamma(1+\frac{1}{\alpha})`
      - :math:`Var(x)=\beta^2 \Gamma(1+\frac{2}{\alpha} - \mu^2)`

    """
    return flib.weibull(x, alpha, beta)

weibull_grad_like = {'value' : flib.weibull_gx,
                 'alpha' : flib.weibull_ga,
                 'beta' : flib.weibull_gb}

# Wishart---------------------------------------------------

def rwishart(n, Tau):
    """
    Return a Wishart random matrix.

    Tau is the inverse of the 'covariance' matrix :math:`C`.
    """
    p = np.shape(Tau)[0]
    sig = np.linalg.cholesky(Tau)
    if n < p:
        raise ValueError('Wishart parameter n must be greater '
                         'than size of matrix.')
    norms = np.random.normal(size=p*(p-1)/2)
    chi_sqs = np.sqrt(np.random.chisquare(df=np.arange(n, n-p, -1)))
    A = flib.expand_triangular(chi_sqs, norms)

    flib.dtrsm_wrap(sig, A, side='L', uplo='L', transa='T', alpha=1.)
    w = np.asmatrix(np.dot(A, A.T))
    flib.symmetrize(w)
    return w

def wishart_expval(n, Tau):
    """
    Expected value of wishart distribution.
    """
    return n * np.asarray(Tau.I)

def wishart_like(X, n, Tau):
    R"""
    Wishart log-likelihood.

    The Wishart distribution is the probability
    distribution of the maximum-likelihood estimator (MLE) of the precision
    matrix of a multivariate normal distribution. If Tau=1, the distribution
    is identical to the chi-square distribution with n degrees of freedom.

    For an alternative parameterization based on :math:`C=T{-1}`, see
    `wishart_cov_like`.

    .. math::
        f(X \mid n, T) = \frac{{\mid T \mid}^{n/2}{\mid X \mid}^{(n-k-1)/2}}{2^{nk/2}
        \Gamma_p(n/2)} \exp\left\{ -\frac{1}{2} Tr(TX) \right\}

    where :math:`k` is the rank of X.

    :Parameters:
      X : matrix
        Symmetric, positive definite.
      n : int
        Degrees of freedom, > 0.
      Tau : matrix
        Symmetric and positive definite

    .. note::
      Step method MatrixMetropolis will preserve the symmetry of Wishart variables.

    """
    return flib.blas_wishart(X, n, Tau)

# Wishart, parametrized by covariance ------------------------------------
def rwishart_cov(n, C):
    """
    Return a Wishart random matrix.

    :Parameters:
      n : int
        Degrees of freedom, > 0.
      C : matrix
        Symmetric and positive definite
    """
    # return rwishart(n, np.linalg.inv(C))

    p = np.shape(C)[0]
    # Need cholesky decomposition of precision matrix C^-1?
    sig = np.linalg.cholesky(C)

    if n < p:
        raise ValueError('Wishart parameter n must be greater '
                         'than size of matrix.')

    norms = np.random.normal(size=p*(p-1)/2)
    chi_sqs = np.sqrt(np.random.chisquare(df=np.arange(n, n-p, -1)))
    A = flib.expand_triangular(chi_sqs, norms)

    flib.dtrmm_wrap(sig, A, side='L', uplo='L', transa='N', alpha=1.)
    w = np.asmatrix(np.dot(A, A.T))
    flib.symmetrize(w)
    return w

def wishart_cov_expval(n, C):
    """
    Expected value of wishart distribution.

    :Parameters:
      n : int
        Degrees of freedom, > 0.
      C : matrix
        Symmetric and positive definite
    """
    return n * np.asarray(C)

def wishart_cov_like(X, n, C):
    R"""
    wishart_like(X, n, C)

    Wishart log-likelihood. The Wishart distribution is the probability
    distribution of the maximum-likelihood estimator (MLE) of the covariance
    matrix of a multivariate normal distribution. If C=1, the distribution
    is identical to the chi-square distribution with n degrees of freedom.

    For an alternative parameterization based on :math:`T=C^{-1}`, see
    `wishart_like`.

    .. math::
        f(X \mid n, C) = {\mid C^{-1} \mid}^{n/2}{\mid X \mid}^{(n-k-1)/2} \exp\left\{ -\frac{1}{2} Tr(C^{-1}X) \right\}

    where :math:`k` is the rank of X.

    :Parameters:
      X : matrix
        Symmetric, positive definite.
      n : int
        Degrees of freedom, > 0.
      C : matrix
        Symmetric and positive definite

    """
    return flib.blas_wishart_cov(X, n, C)


# -----------------------------------------------------------
# DECORATORS
# -----------------------------------------------------------

def name_to_funcs(name, module):
    if isinstance(module, types.ModuleType):
        module = copy(module.__dict__)
    elif isinstance(module, dict):
        module = copy(module)
    else:
        raise AttributeError

    try:
        logp = module[name+"_like"]
    except:
        raise KeyError("No likelihood found with this name ", name+"_like")

    try:
        random = module['r'+name]
    except:
        random = None

    try:
        grad_logp = module[name + "_grad_like"]
    except:
        grad_logp = {}

    return logp, random, grad_logp

def valuewrapper(f, arguments = None):
    """Return a likelihood accepting value instead of x as a keyword argument.
    This is specifically intended for the instantiator above.
    """
    def wrapper(**kwds):
        value = kwds.pop('value')
        return f(value, **kwds)

    if arguments is None:
        wrapper.__dict__.update(f.__dict__)
    else :
        wrapper.__dict__.update(arguments)

    return wrapper

"""
Decorate the likelihoods
"""

snapshot = locals().copy()
likelihoods = {}
for name, obj in six.iteritems(snapshot):
    if name[-5:] == '_like' and name[:-5] in availabledistributions:
        likelihoods[name[:-5]] = snapshot[name]

def local_decorated_likelihoods(obj):
    """
    New interface likelihoods
    """

    for name, like in six.iteritems(likelihoods):
        obj[name+'_like'] = gofwrapper(like, snapshot)


# local_decorated_likelihoods(locals())
# Decorating the likelihoods breaks the creation of distribution
# instantiators -DH.

# Create Stochastic instantiators

def _inject_dist(distname, kwargs={}, ns=locals()):
    """
    Reusable function to inject Stochastic subclasses into module
    namespace
    """

    dist_logp, dist_random, grad_logp = name_to_funcs(distname, ns)
    classname = capitalize(distname)
    ns[classname] = stochastic_from_dist(distname, dist_logp,
                                        dist_random,
                                                                                grad_logp, **kwargs)


for dist in sc_continuous_distributions:
    _inject_dist(dist)

for dist in mv_continuous_distributions:
    _inject_dist(dist, kwargs={'mv' : True})

for dist in sc_bool_distributions:
    _inject_dist(dist, kwargs={'dtype' : np.bool})

for dist in sc_discrete_distributions:
    _inject_dist(dist, kwargs={'dtype' : np.int})

for dist in mv_discrete_distributions:
    _inject_dist(dist, kwargs={'dtype' : np.int, 'mv' : True})




def uninformative_like(x):
    """
    Uninformative log-likelihood. Returns 0 regardless of the value of x.
    """
    return 0.


def one_over_x_like(x):
    """
    returns -np.Inf if x<0, -np.log(x) otherwise.
    """
    if np.any(x < 0):
        return -np.Inf
    else:
        return -np.sum(np.log(x))


Uninformative = stochastic_from_dist('uninformative', logp=uninformative_like)
DiscreteUninformative = stochastic_from_dist('uninformative', logp=uninformative_like, dtype=np.int)
DiscreteUninformative.__name__ = 'DiscreteUninformative'
OneOverX = stochastic_from_dist('one_over_x', logp = one_over_x_like)

# Conjugates of Dirichlet get special treatment, can be parametrized
# by first k-1 'p' values

def extend_dirichlet(p):
    """
    extend_dirichlet(p)

    Concatenates 1-sum(p) to the end of p and returns.
    """
    if len(np.shape(p)) > 1:
        return np.hstack((p, np.atleast_2d(1.-np.sum(p))))
    else:
        return np.hstack((p, 1.-np.sum(p)))


def mod_categorical_like(x, p):
    """
    Categorical log-likelihood with parent p of length k-1.

    An implicit k'th category  is assumed to exist with associated
    probability 1-sum(p).

    ..math::
        f(x=i \mid p, m, s) = p_i,
    ..math::
        i \in 0\ldots k-1

    :Parameters:
      x : integer
        :math: `x \in 0\ldots k-1`
      p : (k-1) float
        :math: `p > 0`
        :math: `\sum p < 1`
      minval : integer
      step : integer
        :math: `s \ge 1`
    """
    return categorical_like(x, extend_dirichlet(p))


def mod_categorical_expval(p):
    """
    Expected value of categorical distribution with parent p of length k-1.

    An implicit k'th category  is assumed to exist with associated
    probability 1-sum(p).
    """
    p = extend_dirichlet(p)
    return np.sum([p*i for i, p in enumerate(p)])


def rmod_categor(p,size=None):
    """
    Categorical random variates with parent p of length k-1.

    An implicit k'th category  is assumed to exist with associated
    probability 1-sum(p).
    """
    return rcategorical(extend_dirichlet(p), size)

class Categorical(Stochastic):
    __doc__ = """
C = Categorical(name, p, value=None, dtype=np.int, observed=False,
size=1, trace=True, rseed=False, cache_depth=2, plot=None)

Stochastic variable with Categorical distribution.
Parent is: p

If parent p is Dirichlet and has length k-1, an implicit k'th
category is assumed to exist with associated probability 1-sum(p.value).

Otherwise parent p's value should sum to 1.

Docstring of categorical_like (case where P is not a Dirichlet):
    """\
    + categorical_like.__doc__ +\
    """
Docstring of categorical_like (case where P is a Dirichlet):
    """\
    + categorical_like.__doc__


    parent_names = ['p', 'minval', 'step']

    def __init__(self, name, p, value=None, dtype=np.int, observed=False,
                 size=None, trace=True, rseed=False, cache_depth=2, plot=None,
                 verbose=-1,**kwds):

        if value is not None:
            if np.isscalar(value):
                self.size = None
            else:
                self.size = len(value)
        else:
            self.size = size

        if isinstance(p, Dirichlet):
            Stochastic.__init__(self, logp=valuewrapper(mod_categorical_like),
                                doc='A Categorical random variable', name=name,
                                parents={'p':p}, random=bind_size(rmod_categor, self.size),
                                trace=trace, value=value, dtype=dtype,
                                rseed=rseed, observed=observed,
                                cache_depth=cache_depth, plot=plot,
                                verbose=verbose, **kwds)
        else:
            Stochastic.__init__(self, logp=valuewrapper(categorical_like),
                                doc='A Categorical random variable', name=name,
                                parents={'p':p},
                                random=bind_size(rcategorical, self.size),
                                trace=trace, value=value, dtype=dtype,
                                rseed=rseed, observed=observed,
                                cache_depth=cache_depth, plot=plot,
                                verbose=verbose, **kwds)

# class ModCategorical(Stochastic):
#     __doc__ = """
# C = ModCategorical(name, p, minval, step[, trace=True, value=None,
#    rseed=False, observed=False, cache_depth=2, plot=None, verbose=0])
#
# Stochastic variable with ModCategorical distribution.
# Parents are: p, minval, step.
#
# If parent p is Dirichlet and has length k-1, an implicit k'th
# category is assumed to exist with associated probability 1-sum(p.value).
#
# Otherwise parent p's value should sum to 1.
#
# Docstring of mod_categorical_like (case where P is not a Dirichlet):
#     """\
#     + mod_categorical_like.__doc__ +\
#     """
# Docstring of mod_categorical_like (case where P is a Dirichlet):
#     """\
#     + mod_categorical_like.__doc__
#
#
#     parent_names = ['p', 'minval', 'step']
#
#     def __init__(self, name, p, minval=0, step=1, value=None, dtype=np.float, observed=False, size=1, trace=True, rseed=False, cache_depth=2, plot=None, verbose=0, **kwds):
#
#         if value is not None:
#             if np.isscalar(value):
#                 self.size = 1
#             else:
#                 self.size = len(value)
#         else:
#             self.size = size
#
#         if isinstance(p, Dirichlet):
#             Stochastic.__init__(self, logp=valuewrapper(mod_categorical_like), doc='A ModCategorical random variable', name=name,
#                 parents={'p':p,'minval':minval,'step':step}, random=bind_size(rmod_categor, self.size), trace=trace, value=value, dtype=dtype,
#                 rseed=rseed, observed=observed, cache_depth=cache_depth, plot=plot, verbose=verbose, **kwds)
#         else:
#             Stochastic.__init__(self, logp=valuewrapper(mod_categorical_like), doc='A ModCategorical random variable', name=name,
#                 parents={'p':p,'minval':minval,'step':step}, random=bind_size(rmod_categorical, self.size), trace=trace, value=value, dtype=dtype,
#                 rseed=rseed, observed=observed, cache_depth=cache_depth, plot=plot, verbose=verbose, **kwds)

def mod_rmultinom(n, p):
    return rmultinomial(n, extend_dirichlet(p))


def mod_multinom_like(x, n, p):
    return multinomial_like(x, n, extend_dirichlet(p))

class Multinomial(Stochastic):
    """
M = Multinomial(name, n, p, trace=True, value=None,
   rseed=False, observed=False, cache_depth=2, plot=None])

A multinomial random variable. Parents are p, minval, step.

If parent p is Dirichlet and has length k-1, an implicit k'th
category is assumed to exist with associated probability 1-sum(p.value).

Otherwise parent p's value should sum to 1.
    """

    parent_names = ['n', 'p']

    def __init__(self, name, n, p, trace=True, value=None, rseed=False,
                 observed=False, cache_depth=2, plot=None, verbose=-1,
                 **kwds):

        if isinstance(p, Dirichlet):
            Stochastic.__init__(self, logp=valuewrapper(mod_multinom_like),
                                doc='A Multinomial random variable', name=name,
                                parents={'n':n,'p':p}, random=mod_rmultinom,
                                trace=trace, value=value, dtype=np.int,
                                rseed=rseed, observed=observed,
                                cache_depth=cache_depth, plot=plot,
                                verbose=verbose, **kwds)
        else:
            Stochastic.__init__(self, logp=valuewrapper(multinomial_like),
                                doc='A Multinomial random variable', name=name,
                                parents={'n':n,'p':p}, random=rmultinomial,
                                trace=trace, value=value, dtype=np.int,
                                rseed=rseed, observed=observed,
                                cache_depth=cache_depth, plot=plot,
                                verbose=verbose, **kwds)

def Impute(name, dist_class, imputable, **parents):
    """
    This function accomodates missing elements for the data of simple
    Stochastic distribution subclasses. The masked_values argument is an
    object of type numpy.ma.MaskedArray, which contains the raw data and
    a boolean mask indicating missing values. The resulting list contains
    a list of stochastics of type dist_class, with the extant values as data
    stochastics and the missing values as variable stochastics.

    :Arguments:
      - name : string
        Name of the data stochastic
      - dist_class : Stochastic
        Stochastic subclass such as Poisson, Normal, etc.
      - imputable : numpy.ma.core.MaskedArray or iterable
        A masked array with missing elements (where mask=True, value
        is assumed missing), or any iterable that contains None
        elements that will be imputed.
      - parents (optional): dict
        Arbitrary keyword arguments.
    """

    dims = np.shape(imputable)
    masked_values = np.ravel(imputable)

    if not isinstance(masked_values, np.ma.core.MaskedArray):
        # Generate mask

        mask = [v is None or np.isnan(v) for v in masked_values]
        # Generate masked array
        masked_values = np.ma.masked_array(masked_values, mask)

    # Initialise list
    vars = []
    for i in xrange(len(masked_values)):

        # Name of element
        this_name = name + '[%i]'%i
        # Dictionary to hold parents
        these_parents = {}
        # Parse parents
        for key, parent in six.iteritems(parents):

            try:
                # If parent is a PyMCObject
                shape = np.shape(parent.value)
            except AttributeError:
                shape = np.shape(parent)

            if shape == dims:
                these_parents[key] = Lambda(key + '[%i]'%i,
                                            lambda p=np.ravel(parent),
                                            i=i: p[i])
            elif shape == np.shape(masked_values):
                these_parents[key] = Lambda(key + '[%i]'%i, lambda p=parent,
                                            i=i: p[i])
            else:
                these_parents[key] = parent

        if masked_values.mask[i]:
            # Missing values
            vars.append(dist_class(this_name, **these_parents))
        else:
            # Observed values
            vars.append(dist_class(this_name, value=masked_values[i],
                                   observed=True, **these_parents))
    return np.reshape(vars, dims)

ImputeMissing = Impute

if __name__ == "__main__":
    import doctest
    doctest.testmod()
