from __future__ import division
from stats_toolkit import _
from numpy import random, zeros, isnan, arange
from scipy import stats
from stats_toolkit.errors import InvalidParamsError
from base import UncertaintyBase, one_row_params_array


class BetaUncertainty(UncertaintyBase):
    """
The Beta distribution has the probability distribution function:

.. math:: f(x; a,b) = \frac{1}{B(\alpha, \beta)} x^{\alpha - 1}(1 - x)^{\beta - 1},

where the normalisation, B, is the beta function,

.. math:: B(\alpha, \beta) = \int_0^1 t^{\alpha - 1}(1 - t)^{\beta - 1} dt.

See http://en.wikipedia.org/wiki/Beta_distribution.                                 
    """
    id = 10
    description = _("Beta uncertainty")

    @classmethod
    def validate(cls, params):
        if (params['amount'] > 0).sum() != params.shape[0]:
            raise InvalidParamsError("Real, positive alpha values are" + \
                " required for Beta uncertainties.")
        if (params['sigma'] > 0).sum() != params.shape[0]:
            raise InvalidParamsError("Real, positive beta values are" + \
                " required for Beta uncertainties.")
        if (params['maximum'] <= 0).sum():
            raise InvalidParamsError("Scale value must be positive or NaN")

    @classmethod
    def random_variables(cls, params, size, seeded_random=None, 
            transform=False):
        if not seeded_random:
            seeded_random = random
        scale = params['maximum']
        scale[isnan(scale)] = 1
        return scale.reshape((-1, 1)) * seeded_random.beta(
            params['amount'],
            params['sigma'],
            size=(size, params.shape[0])).T

    @classmethod
    def cdf(cls, params, vector):
        vector = cls.check_2d_inputs(params, vector)
        results = zeros(vector.shape)
        scale = params['maximum']
        scale[isnan(scale)] = 1
        for row in range(params.shape[0]):
            results[row, :] = stats.beta.cdf(vector[row,:], 
                params['amount'][row], params['sigma'][row], 
                scale=scale[row])
        return results

    @classmethod
    def ppf(cls, params, percentages):
        percentages = cls.check_2d_inputs(params, percentages)
        results = zeros(percentages.shape)
        scale = params['maximum']
        scale[isnan(scale)] = 1
        for row in range(percentages.shape[0]):
            results[row, :] = stats.beta.ppf(percentages[row, :], 
                params['amount'][row], params['sigma'][row], 
                scale=scale[row])
        return results

    @classmethod
    @one_row_params_array
    def statistics(cls, params):
        alpha = float(params['amount'])
        beta = float(params['sigma'])
        scale = 1 if isnan(params['maximum'])[0] else float(params['maximum'])
        if alpha <= 1 or beta <= 1:
            mode = "Undefined"
        else:
            mode = (alpha - 1) / (alpha + beta - 2)
        return {
            'mean': alpha / (alpha + beta), 
            'mode': mode, 
            'median': "Not Implemented", 
            'lower': "Not Implemented", 
            'upper': "Not Implemented"
            }

    @classmethod
    @one_row_params_array
    def pdf(cls, params, xs=None):
        scale = 1 if isnan(params['maximum'])[0] else float(params['maximum'])
        if xs == None:
            xs = arange(0, scale, scale/cls.default_number_points_in_pdf)
        ys = stats.beta.pdf(xs, params['amount'], params['sigma'], 
            scale=scale)
        return xs, ys.reshape(ys.shape[1])