'''All credit for this function goes to <NAME> at
stackoverflow.com/questions/21100716/fast-arbitrary-distribution-random-sampling'''

import numpy as np
import scipy.misc
import scipy.special as special

class Distribution(object):
    """
    draws samples from a one dimensional probability distribution,
    by means of inversion of a discrete inverstion of a cumulative density function

    the pdf can be sorted first to prevent numerical error in the cumulative sum
    this is set as default; for big density functions with high contrast,
    it is absolutely necessary, and for small density functions,
    the overhead is minimal

    a call to this distibution object returns indices into density array
    """
    def __init__(self, pdf, sort = True, interpolation = True, transform = lambda x: x):
        self.shape          = pdf.shape
        self.pdf            = pdf.ravel()
        self.sort           = sort
        self.interpolation  = interpolation
        self.transform      = transform

        # a pdf can not be negative
        assert(np.all(pdf>=0))

        # sort the pdf by magnitude
        if self.sort:
            self.sortindex = np.argsort(self.pdf, axis=None)
            self.pdf = self.pdf[self.sortindex]
        # construct the cumulative distribution function
        self.cdf = np.cumsum(self.pdf)

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def sum(self):
        """cached sum of all pdf values; the pdf need not sum to one, and is imlpicitly normalized"""
        return self.cdf[-1]

    def __call__(self, N):
        """draw """
        # pick numbers which are uniformly random over the cumulative distribution function
        # print N, self.ndim, self.sum
        choice = np.random.uniform(high = self.sum, size = N)
        # find the indices corresponding to this point on the CDF
        index = np.searchsorted(self.cdf, choice)
        # if necessary, map the indices back to their original ordering
        if self.sort:
            index = self.sortindex[index]
        # map back to multi-dimensional indexing
        index = np.unravel_index(index, self.shape)
        index = np.vstack(index)
        # is this a discrete or piecewise continuous distribution?
        if self.interpolation:
            index = np.float_(index)
            # index[0] += np.random.uniform(size=index.shape[1])
            index[:int(self.interpolation)] += np.random.uniform(size=(int(self.interpolation), index.shape[1]))

        return self.transform(index)


def planck(T, l):
    from scipy.constants import codata
    import numpy as np
    import matplotlib.pyplot as plt

    D = codata.physical_constants

    h = D['Planck constant'][0]
    k = D['Boltzmann constant'][0]
    c = D['speed of light in vacuum'][0]
    wienConstant = 2.897e-3
    # calculate the Planck Law for a specific temperature and an array of wavelengths
    p = c*h/(k*l*T)
    result = np.zeros(np.shape(l))+1e-99
    # prevent underflow - compute only when p is "not too big"
    calcMe = np.where(p<700)
    result[calcMe] = (h*c*c)/(np.power(l[calcMe], 5.0) * (np.exp(p[calcMe])-1))
    return result


def poisson(lamda, k):
    pdf = (lamda ** k * np.exp(-lamda)) / scipy.misc.factorial(k)
    return pdf


def bessel(k):
    '''modified zero order besset=l'''
    pdf = scipy.special.i0(k)
    return pdf


def MR(I,Ic,Is):
    '''modified rician distribution'''
    # print np.exp(-1.0*(I+Ic)/Is)[::10], special.iv(0,2.0*np.sqrt(I*Ic)/Is)[::10] 
    # print Ic, Is
    pdf = 1.0/Is * np.exp(-1.0*(I+Ic)/Is)* special.iv(0,2.0*np.sqrt(I*Ic)/Is)
    #pdf = (1./Is) * np.exp(-(I+Ic)/Is) * bessel(2*np.sqrt(I*Ic)/Is)
    pdf = pdf/sum(pdf)
    return pdf


def gaussian(mu, sig, x):
    pdf = np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
    return pdf


def gaussian2(x, sig,mu):
    pdf = np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
    pdf = pdf/sum(pdf)
    return pdf


def lognorm(x,mu, sigma):
    pdf = (np.exp(-(np.log10(x - mu))**2 / (2 * sigma**2))/ (x * sigma))
    pdf = pdf/sum(pdf)
    return pdf


if __name__=='__main__':
    res_elements = 1000
    I = np.linspace(0.1, 200, res_elements)

    ratios = [0,0.5,1,2,7]
    for ratio in ratios:
        Is = 10
        pdf = MR(Is*ratio,Is,I)
        pdf = pdf/ np.sum(pdf)
        # plt.plot(pdf)


        dist = Distribution(pdf)#, transform=lambda i: i - res_elements/2)
        TOAs = dist(100)
        print(TOAs, np.shape(TOAs))
        hist, bins = np.histogram(TOAs, bins ='auto')
        plt.plot(bins[:-1], hist)

    # plt.plot(dist(1000),'o')
    plt.show()

    # shape = 3,3
    # pdf = np.ones(shape)
    # pdf[1]=0
    # dist = Distribution(pdf, transform=lambda i:i-1.5)
    # print dist(10)
    # import matplotlib.pyplot as pp
    # pp.scatter(*dist(1000))
    # pp.show()