#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""RNGs do not necessary only apply on the interval (0, 1).

For this reason, this file also allows numerous different
distributions from which we can obtain random streams.
Based on the simple proof of:

.. math:: Pr[X <= x] = Pr[F^{-1}(U) <= x] = Pr[U <= F(x)] = F(x)

we know that all we need is an inverse cumulative distribution
function for all distributions we like to offer.

In this module, there are multiple inverse cumulative distribution functions implemented
in such a way that you can use any RNG in the range (0, 1) for generating the data. It
is therefore completely standalone from some builtin implementation for a RNG.

Besides Wikipedia, I refer to the following link when it comes to the actual distributions:
http://www2.isikun.edu.tr/mustafahekimoglu/simulation/Lecture5-ProbabilityReview.pdf

Also, all functions return numpy-classes!
"""

from pypdevsbbl.extra.mathutils import binprob, harmonic, harmonicinv, binomial as bn
import scipy.special as sp
import numpy as np

def beta(y, a, b):
    """Inverse of the beta cdf on the interval [0, 1] with shapes a1 and a2.

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The first shape parameter.
        b (numeric):    The second shape parameter.

    Returns:
        A numpy.float64
    """
    assert a > 0 and b > 0 and 0 < y < 1
    return sp.betaincinv(a, b, y)


def binomial(y, n, p):
    """Inverse of the discrete binomial cdf with n experiments that have p chance of success.

    Since a Bernoulli distribution is a special case of the binomial distribution,
    you can also use this function for those usecases.

    .. math:: Bernouilli(y, p) = binomial(y, 1, p)

    Args:
        y (numeric):    The y-value of the function.
        n (numeric):    The number of experiments.
        p (numeric):    The chance for success in an experiment.

    Returns:
        A numpy.float64
    """
    assert 0 <= p <= 1 and n > 0 and 0 < y < 1
    x = 0
    s = 0
    while y > s:
        x += 1
        s += binprob(x, n, p)
    return x


def cauchy(y, x0=0, scale=1):
    """The inverse of the cauchy cdf.

    This distribution is also known as the Lorentz distribution or the
    Breit-Wigner distribution. It is often used as the canonical example
    of a pathological distribution.

    Args:
        y (numeric):        The y-value of the function.
        x0 (numeric):       The location of the distribution. Defaults to 0.
        scale (numeric):    The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert 0 < scale and 0 < y < 1
    return scale * np.tan(np.pi * (y - 0.5)) + x0


def chiSquare(y, k):
    """Inverse of the chi square cdf with k degrees of freedom.

    This distribution is a special case of the gamma distribution and occurs
    the most often in inferential statistics.

    Technically, this is a special case of the gamma distribution where:
        shape = k / 2
        rate  = 1

    Args:
        y (numeric):    The y-value of the function.
        k (numeric):    The degrees of freedom to use.

    Returns:
        A numpy.float64
    """
    assert k > 0 and 0 < y < 1
    return gamma(y, k / 2, 1)


def exponential(y, l):
    """The inverse of the exponential cdf.

    This distribution is often used to predict the wait time until the
    first event.
    The mean of this distribution is 1/l.

    The exponential function is "memoryless". In other words:
    Pr[X > s+t | X > s] = Pr[X > t]

    Args:
        y (numeric):    The y-value of the function.
        l (numeric):    The rate parameter (lambda),
                        i.e. mean number of occurrences per time unit.

    Returns:
        A numpy.float64
    """
    assert l > 0 and 0 < y < 1
    # np.log(x) does the natural logarithm of x
    return -np.log(y) / l


def F(y, d1, d2):
    """Inverse of the Fisher-Snedecor cdf with d1 and d2 its degrees of freedom.

    Args:
        y (numeric):    The y-value of the function.
        d1 (numeric):   The degrees of freedom to use.
        d2 (numeric):   The degrees of freedom to use.

    Returns:
        A numpy.float64
    """
    assert d1 > 0 and d2 > 0 and 0 < y < 1
    I = sp.betaincinv(d1 / 2, d2 / 2, y)
    return d2 * I / (d1 * (1 - I))


def frechet(y, alpha, m=0, s=1):
    """The inverse of the Fréchet cdf.

    This distribution is also known as the Generalized Extreme Value Type B
    distribution or the log-Weibull distribution. It is commonly used in
    hydrology.

    Args:
        y (numeric):        The y-value of the function.
        alpha (numeric):    The shape of the distribution.
        m (numeric):        The location of the distribution. Defaults to 0.
        s (numeric):        The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert alpha > 0 and s > 0 and 0 < y < 1
    return s * np.power(-np.log(y), -1.0 / alpha) + m


_gamma_cache = {} # cache of some intermediary values of the gamma function
def gamma(y, a, b):
    """Inverse of the gamma cdf with shape a and rate b.

    This distribution is used to predict the wait time until the a-th event.

    This formula is based upon the shape-rate system, so if you want to use
    the shape-scale system, use the property tao = 1/b

    Seeing as the erlang distribution is a special case of the gamma
    distribution (i.e. b is an integer), you can also use this function
    for those purposes.

    This function makes use of the _gamma_cache dictionary for efficiency.

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The shape of the curve.
        b (numeric):    The rate parameter.

    Returns:
        A numpy.float64
    """
    assert a > 0 and b > 0 and 0 < y < 1
    if 0 < a < 1:
        if a not in _gamma_cache:
            _gamma_cache[a] = (np.e + a) / np.e
        _b = _gamma_cache[a]
        if y <= 1/_b:
            return b * np.power(_b*y, 1/a)
        else:
            return -b * np.log((_b * (1 - y)) / a)
    elif a == 1:
        return b * exponential(y, 1)
    else:
        l = np.sqrt(2*a-1)
        mu = np.power(a, l)
        return b * np.power((mu * y) / (1 - y), 1 / l)


def geometric(y, p):
    """Inverse of the discrete geometric cdf with p chance for success.

    :math:`x` would represent the amount of faillures.

    Args:
        y (numeric):    The y-value of the function.
        p (numeric):    The chance for success in an experiment.

    Returns:
        A numpy.float64
    """
    assert 0 < p < 1 and 0 < y < 1
    return np.log(1 - y) / np.log(1 - p) - 1


def gumbel(y, mu=0, beta=1):
    """The inverse of the Gumbel cdf.

    This distribution is also known as the Generalized Extreme Value Type A
    distribution or the log-Weibull distribution. It is commonly used in
    hydrology.

    Args:
        y (numeric):        The y-value of the function.
        mu (numeric):       The location of the distribution. Defaults to 0.
        beta (numeric):     The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert beta > 0 and 0 < y < 1
    return -beta * np.log(-np.log(y)) + mu


def JSB(y, a1, a2, a, b):
    """Inverse of the Johnson Bounded cdf.

    Args:
        y (numeric):    The y-value of the function.
        a1 (numeric):   The shape of the distribution.
        a2 (numeric):   The shape of the distribution.
        a (numeric):    The location of the distribution.
        b (numeric):    The scale of the distribution.

    Returns:
        A numpy.float64
    """
    assert a1 > 0 and a2 > 0 and a > 0 and b > 0 and 0 < y < 1
    z = normal(y)
    y = np.exp((z - a1) / a2)
    return (a + b * y) / (y + 1)

johnsonB = JSB
"""Alias for :meth:`JSB` for readability."""


def JSU(y, a1, a2, g, b):
    """Inverse of the Johnson Unbounded cdf.

    Args:
        y (numeric):    The y-value of the function.
        a1 (numeric):   The shape of the distribution.
        a2 (numeric):   The shape of the distribution.
        g (numeric):    The location of the distribution.
        b (numeric):    The scale of the distribution.

    Returns:
        A numpy.float64
    """
    assert a1 > 0 and a2 > 0 and g > 0 and b > 0 and 0 < y < 1
    z = normal(y)
    y = np.exp((z - a1) / a2)
    return g + (b / 2) * (y - 1 / y)

johnsonU = JSU
"""Alias for :meth:`JSU` for readability."""


def laplace(y, mu=0, b=1):
    """Inverse of the Laplace cdf.

    The distribution is used in speech recognition and hydrology.

    Args:
        y (numeric):    The y-value of the function.
        mu (numeric):   The location of the distribution. Defaults to 0.
        b (numeric):    The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert b > 0 and 0 < y < 1
    u = y - 0.5
    return mu - b * np.sign(u) * np.log(1 - 2 * np.abs(u))


def logistic(y, mu=0, s=1):
    """Inverse of the Logistic cdf.

    Also known as the Sech-squared distribution. It appears in logistic
    regression and feedforward neural networks.

    Args:
        y (numeric):    The y-value of the function.
        mu (numeric):   The location of the distribution. Defaults to 0.
        s (numeric):    The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert s > 0 and 0 < y < 1
    return -s * np.log(np.power(y, -1.0) - 1) + mu


def loglogistic(y, a, b):
    """Inverse of the Log-Logistic / Fisk cdf.

    In economics, this distribution is known as the Fisk distribution.

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The shape of the distribution.
        b (numeric):    The scale of the distribution.

    Returns:
        A numpy.float64
    """
    assert a > 0 and b > 0 and 0 < y < 1
    return b * np.power(y / (1 - y), 1 / a)


def lognormal(y, mu, s):
    """Inverse of the log-normal cdf with mean mu and standard derivation s.

    Generates a variable whose natural logarithm is normally distributed with
    mean mu and standard derivation s.

    Args:
        y (numeric):    The y-value of the function.
        mu (numeric):   The mean of the distribution.
        s (numeric):    The standard derivation of the distribution.

    Returns:
        A numpy.float64
    """
    assert s > 0 and 0 < y < 1
    Y = normal(y, mu, s)
    return np.exp(Y)


def negbinomial(y, r, p):
    """Inverse of the discrete negative binomial cdf with p chance of success foreach of s dist.

    Args:
        y (numeric):    The y-value of the function.
        r (numeric):    The number of faillures.
        p (numeric):    The chance for success in an experiment.

    Returns:
        A numpy.float64
    """
    assert 0 <= p <= 1 and 0 < r
    return _sample(y, lambda x: bn(x + r - 1, x) * (1 - p) ** r * p ** x, 1)


def normal(y, mu=0, s=1):
    """Inverse of the normal cdf with mean mu and standard derivation s.

    The implementation of this function is based on the quantile function
    and the properties of the errorfunction erf.

    Args:
        y (numeric):    The y-value of the function.
        mu (numeric):   The mean of the distribution.
                        Defaults to 0.
        s (numeric):    The standard derivation of the distribution.
                        Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert s ** 2 > 0 and 0 < y < 1
    return mu + s * np.sqrt(2) * sp.erfinv(2 * y - 1)


def pareto(y, alpha, xm):
    """Inverse of the Pareto cdf.

    The distribution is used to describe social, scientifical, geophysical,
    actuatial... phenomena. It originally was used to describe the distribution
    of wealth.

    Colloquially, the distribution is refered to as the `80-20 rule`.

    Args:
        y (numeric):        The y-value of the function.
        alpha (numeric):    The shape of the distribution.
        xm (numeric):       The scale of the distribution. Defaults to 1.

    Returns:
        A numpy.float64
    """
    assert xm > 0 and alpha > 0 and 0 < y < 1
    return xm * np.power(1 - y, -1.0 / alpha)


def poisson(y, l):
    """Inverse of the discrete Poisson cdf with lambda l.

    This distribution is a discrete distribution that expresses the probability
    that any given number of intervals occurs in a fixed interval of time.

    The implemented algorithm here is based on the inverse transform sampling algorithm
    on Wikipedia: https://en.wikipedia.org/wiki/Poisson_distribution

    This has been chosen, because it requires only a single random value for each call.
    Cumulative probabilities are examined in turn until one exceeds the random value.

    Args:
        y (numeric):    The y-value of the function.
        l (numeric):    The average number of events that occur (i.e. lambda).

    Returns:
        A numpy.int
    """
    assert 0 < y < 1
    x = 0
    p = np.exp(-l)
    s = p
    while y > s:
        x += 1
        p *= l / x
        s += p
    return np.int(x)


def PT5(y, a, b):
    """Inverse of the Pearson Type V cdf.

    This method makes use of the gamma distribution.

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The shape of the distribution.
        b (numeric):    The scale of the distribution.

    Returns:
        A numpy.float64
    """
    assert a > 0 and b > 0 and 0 < y < 1
    return 1 / gamma(y, a, 1 / b)

pearson5 = PT5
"""Alias of :meth:`PT5` for readability."""


def PT6(y, a1, a2, b):
    """Inverse of the Pearson Type VI cdf.

    Args:
        y (numeric):    The y-value of the function.
        a1 (numeric):   The shape of the first distribution.
        a2 (numeric):   The shape of the second distribution.
        b (numeric):    The scale of the distribution.

    Returns:
        A numpy.float64
    """
    assert b > 0
    X = beta(y, a1, a2)
    return b * X / (1 - X)

pearson6 = PT6
"""Alias of :meth:`PT6` for readability."""


def student(y, v, step=0.001):
    """Approximation of the inverse student t-cdf with v degrees of freedom.

    This distribution is used for identifying unknowns in a distribution and occurs
    most often when estimating unknowns. When combined with RNG, these often arise
    in multi-dimensional applications of copula-dependency (dependence between variables).

    This implementation uses a simple sampling method, but requires discretization.
    For this to be accurate, minimize the step size to use.

    Args:
        y (numeric):    The y-value of the function.
        v (numeric):    The degrees of freedom to use.
        step (float):   The stepsize to use when sampling.

    Returns:
        A numpy.float64
    """
    assert v > 0 and 0 < y < 1
    f = lambda x: sp.gamma((v + 1) / 2) / (np.sqrt(v * np.pi) * sp.gamma(v / 2)) *\
                  np.power(1 + np.power(x, 2) / v, -(v + 1) / 2)
    return _sample(y, f, step)


def triangular(y, a, b, c):
    """Inverse of the triangular distribution with minimum a, mode b and maximum c.

    This distribution is typically used as a subjective description of
    a population for which there is limited sample data, especially if
    the relationship between variables is known, but the data is scarse.

    Please be aware of the differences between b (mode) and c (maximum). Some sources
    use b and c the other way around (i.e. they use c as mode and b as maximum).

    There are also resources that encode these values in a single number q, which can
    be obtained via (a, b, c w.r.t. the parameters of *this* function):

    .. math:: q = (b - a) / (c - a)

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The minimum.
        b (numeric):    The mode.
        c (numeric):    The maximum.

    Returns:
        A numpy.float64
    """
    assert a < b <= c and 0 < y < 1
    f = (c - a) / (b - a)
    if y < f:
        return a + np.sqrt(y * (b - a) * (c - a))
    return c - np.sqrt((1 - y) * (c - a) * (c - b))


def uniform(y, a, b):
    """Inverse of the uniform cdf over the range (a, b).

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The lower bound of the distribution.
        b (numeric):    The upper bound of the distribution.

    Returns:
        A numpy.float64 in the range (a, b).
    """
    assert 0 < y < 1
    return np.float64(a + ((b - a) * y))


def uniformInt(y, a, b):
    """Inverse of the rounded uniform cdf over the range [a, b].

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The lower bound of the distribution.
        b (numeric):    The upper bound of the distribution.

    Returns:
        A numpy.float64 in the range [a, b].
    """
    assert 0 < y < 1
    return np.float64(a + np.floor((b - a + 1.0) * y))


def wald(y, mu, l, step=0.001):
    """Approximation of the inverse Wald cdf.

    Also known as the Inverse Gaussian Distribution, the Wald distribution
    is commonly used to compute the properties of the Brownian Motion.

    Args:
        y (numeric):    The y-value of the function.
        mu (numeric):   The mean of the distribution.
        l (numeric):    The shape of the distribution.
        step (float):   The stepsize to use when sampling.

    Returns:
        A numpy.float64
    """
    assert mu > 0 and l > 0 and 0 < y < 1
    f = lambda x: np.sqrt(l / (2 * np.pi * np.power(x, 3))) * \
                  np.exp(-(l * np.power(x - mu, 2)) / (2 * np.power(mu, 2) * x))
    return _sample(y, f, step)


def weibull(y, a, b, v=0):
    """Inverse of the weibell cdf with scale a, shape b and location v.

    This distribution is most commonly used to predict the lifetime
    of technical equipment. The location is not often used, in which
    cases it is to be set at 0 (also the default).

    Args:
        y (numeric):    The y-value of the function.
        a (numeric):    The scale of the curve.
        b (numeric):    The shape parameter.
        v (numeric):    The location parameter.
                        Defaults to 0 (not often used).

    Returns:
        A numpy.float64
    """
    assert a > 0 and b > 0 and 0 < y < 1
    l = -np.log(y)
    x = v + np.power(l, 1/b) * a
    if x >= v:
        return x
    else:
        return 0.0


def zipf(y, s, N:int):
    """Inverse of a discrete cdf w.r.t. Zipf's law.

    Args:
        y (numeric):   The y-value of the function.
        s (numeric):   The degrees of freedom to use.
        N (int):       The degrees of freedom to use.

    Returns:
        A numpy.float64
    """
    assert s >= 0 and N > 0 and 0 < y < 1
    return harmonicinv(y * harmonic(N, s), s)


def _sample(y, f, step):
    """Sampling method."""
    x = 0
    p = f(x)
    while y > p:
        x += step
        p += f(x)
    return x


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from pypdevsbbl.extra.rngstreams import RngStream, U01
    gen = RngStream.RngStream_CreateStream("test")
    y = sorted([U01(gen) for i in range(10000)])
    x = [zipf(i, 4, 10) for i in y]
    fig, ax = plt.subplots()
    ax.plot(x, y)
    plt.show()


