import numpy as np
import scipy
import scipy.special
import scipy.stats




# Function to find log-density functions matching given values for the 2.5, 50 and 97.5th percentiles
# 'prctiles' is a vector of the [2.5, 50, 97.5]th percentiles
# 'distribution' specifies whether log normal or beta 
# 'visualising' should be set to 'true' if wanting to see the estimated pdf

def logncdf(x,y,z):
    return 0

def betacdf(x,y,z):
    
    return 0

    
def get_distribution_fns(prctiles, distribution, visualising):
    aux = {}
    val=1
    dat  = np.sort(prctiles)
    opts = []#optimset('TolX',1e-12,'TolFun',1e-12,'MaxIter',1e6,'MaxFunEvals',1e6)

    # Estimators to help setting up initial guesses
    mn = dat[1] 
    var = (dat[2]-dat[0])**(2/4)

    if distribution == 'lognorm':
        
        ff  = lambda x : [scipy.stats.lognorm.cdf(dat[0],x[0],x[1]), scipy.stats.lognorm.cdf(dat[1],x[0],x[1]), scipy.stats.lognorm.cdf(dat[2],x[0],x[1])]

        # Set up initial guess
        sig = np.sqrt(np.log(var/(mn**2)+1))
        mu  = np.log(mn)-(sig**2)/2 
        init = [mu, sig]

        # Do the optimisation
        perc = np.array([2.5, 50, 97.5])
        obj = lambda x : sum((ff(x)/(perc/100) - 1)**2)
        out = scipy.optimize.fmin(obj, init)# fminsearch(obj, init, opts)
        #if val > 1e-2:
        #    raise  Exception('Calibration setup not converged')
        
        # Get the log-pdf
        mu = out[0]
        sig = out[1]
        logfn = lambda x : -(np.log(x)-mu)**2/(2*sig**2) - np.log(x*sig*np.sqrt(2*np.pi))
        
    elif distribution == 'beta':
        ff  = lambda x : [scipy.stats.beta(dat[0],x[0],x[1]), scipy.stats.beta(dat[1],x[0],x[1]), scipy.stats.beta(dat[2],x[0],x[1])]

        # Set up initial guess
        tmp = mn*(1-mn)/var-1
        a = tmp*mn 
        b = tmp*(1-mn) 
        init = [a, b]

        # Do the optimisation
        obj = lambda x: sum((ff(x)/([2.5, 50, 97.5]/100) - 1)**2)
        [out, val] = scipy.optimize.fmin(obj, init)#fminsearch(obj,init,opts);
        if (val > 1e-2):
            raise Exception('Calibration setup not converged')
        
        # Get the log-pdf
        a = out[0] 
        b = out[1]
        logfn = lambda x : (a-1)*np.log(x) + (b-1)*np.log(1-x) - scipy.special.betaln(a,b)
    

    aux['sim'] = ff(out)                                                         # Simulated values of CDF at given percentile points
    aux['val'] = val                                                             # The final objective function

    '''
    if nargin == 3 && visualising
        figure;
        x = linspace(dat(1)*.8, dat(3)*1.2);
        if strcmp(distribution,'lognorm')
            y = lognpdf(x,out(1),out(2));
        elseif strcmp(distribution,'beta')
            y = betapdf(x,out(1),out(2));
      
        plot(x,y); hold on;
        # Show the percentiles
        for id = 1:3
            line(dat(id)*[1 1],ylim,'linestyle','--');
    '''
      

    return logfn, out, aux


