# -*- coding: utf-8 -*-
from numpy import * 
from scipy import special as sp
from scipy import stats as st
from matplotlib.pyplot import *


def Gamma(alpha):

    # Gamma inverse 
    ret = exp(sp.gammaln(alpha))
    return ret


def RandomdBar(alpha, beta):
    # ppf here is the inverse of gamma CDF
    ret = st.gamma.ppf(random.uniform(0,1), alpha + 1, 0,beta)
    # st.gamma.ppf(probability,alpha,0,beta)
    # st.gamma.ppf(0.3,a=alpha,loc=0,scale=beta)
    ret = random.uniform(0,1) * ret
    return ret


def DiGamma(x, DiGamma_alpha, DiGamma_beta, tau):
    '''
This distribution is a double-side symmetric distribution with a mixture of two gamma distribution
Weight is the mixture weight of the first gamma distribution
'''
    # Gamma CDF fucntion 
    # st.gamma.cdf(x,alpha,0,beta)
    
    if x < 0:
        ret = 0.5 - 0.5 * st.gamma.cdf(( x )  ** tau, DiGamma_alpha,0, DiGamma_beta ** tau)
    elif x > 0:
        ret = 0.5 + 0.5 * st.gamma.cdf(x ** tau, DiGamma_alpha, 0,DiGamma_beta ** tau)
    else:
        ret = 0.5
        
    return ret


# This might be an inverse of DiGamma function defined above
def DiGammaInv(x, DiGamma_alpha, DiGamma_beta, tau):
    
    Value = 0
    if x < 0.5:
        Value = st.gamma.ppf(2 *  ( 0.5 - x ), DiGamma_alpha,0, DiGamma_beta ** tau)
        _ret = - Value **  ( 1 / tau )
    elif x > 0.5:
        Value = st.gamma.ppf(2 *  ( x - 0.5 ), DiGamma_alpha, 0,DiGamma_beta ** tau)
        _ret = Value **  ( 1 / tau )
    else:
        _ret = 0
    return _ret


'''
I skipped DiGammaInv2 function as I might not need it other than plotting
'''



def DiGammaMoment(n, DiGamma_alpha, DiGamma_beta, tau):
    n=float(n)
    _ret = DiGamma_beta ** n * Gamma(DiGamma_alpha + n / tau) / Gamma(DiGamma_alpha)
    return _ret



# this Laplace function is still not working  :'(
# starts from here


def Laplace(random_variable, F, n, alpha, beta, t=None):

    # Calculate the value of laplace transform
    # N is the coeff in the exp including the sign
    # t is the optional parameter here
    n=float(n)
    alpha=float(alpha)
    beta=float(beta)
    F=float(F)

    
    
    ePsilon = 0.0001
    TheFactor = 1
    x = - 1 * n * math.log(1 - F)
    
    if (random_variable == 'Ld'):
        #In (3.3) page 13 paper
        _ret = ( 1 + beta * x )  **  ( -1 * alpha )
 
        
    elif (random_variable == 'Ld-'):
        if x == 0:
            #Following the newest correction in the note
            _ret = 1
        elif  ( x < 0 )  and  ( x > -1 / beta ) :
            #e.5 in the correction, page 11
            _ret = ( 1 -  ( 1 + beta * x )  **  ( -1 * alpha ) )  /  ( alpha * beta * x )
        elif  ( x <= -1 / beta ) :
            #Added case ( I don't really know what this means ) 
            _ret = PInf
        else:
            #For x > 0, section 3.3 is ok as mentioned in the note
            _ret = ( 1 -  ( 1 + beta * x )  **  ( -1 * alpha ) )  /  ( alpha * beta * x )
    elif (random_variable == 'Ld-t'):
        if t== None:
            _ret = 'Missing t'
        else:
            Prob_dt = 1 - st.gamma.cdf(t, alpha + 1,0,beta) - t *  ( 1 - st.gamma.cdf(t, alpha,0, beta) )  /  ( alpha * beta )
            if x == 0:
                #In the note (newest)
                _ret = 1
            elif  ( x < 0 )  and  ( x > -1 / beta ) :
                #e.5 in the correction, page 11
                _ret = ( 1 - exp(- x * t) *  ( 1 - st.gamma.cdf(t, alpha,0, beta) )  -  ( 1 + beta * x )  **  (-alpha )  * st.gamma.cdf(t, alpha,0,beta /  ( 1 + beta * x )) )  /  ( alpha * beta * x )  + exp(- x * t) * Prob_dt
            elif x == -1 / beta:
                _ret = exp(t / beta) * Prob_dt -  ( 1 - exp(t / beta) *  ( 1 - st.gamma.cdf(t, alpha, 0,beta) )  - t ** alpha /  ( beta ** alpha * Gamma(alpha + 1) ) )  / alpha
            elif x < -1 / beta:
                #[Updated 7/23/2012]
                #Added two conditions to determine convergency. See (e.7.2) and (e.7.3) in "more corrections July 2012"
                #[Updated 11/06/2012]
                #When t = 0, TheSum = 0
                TheSum = 0
                if t != 0:
                    k = 0
                    Cond1 = k * math.log(- ( x + 1 / beta )  * t) - sp.gammaln(k + 1) - math.log(alpha + k) < math.log(- alpha * beta * x * Gamma(alpha) * beta ** alpha * ePsilon /  ( 2 * t ** alpha ))
                    Cond2 = k > 2 *  ( ( x + 1 / beta )  * t )  - 1
                    while not (Cond1 and Cond2):
                        TheSum = TheSum + exp((alpha * math.log(t) - (sp.gammaln(alpha) + alpha * math.log(beta)) + k * math.log(-t * (x + 1 / beta)) - sp.gammaln(k + 1) - math.log(alpha + k)))
                        k = k + 1
                        Cond1 = k * math.log(-(x + 1 / beta) * t) - sp.gammaln(k + 1) - math.log(alpha + k) < math.log(-alpha * beta * x * Gamma(alpha) * beta ** alpha * ePsilon / (2 * t ** alpha))
                        Cond2 = k > 2 * (-(x + 1 / beta) * t) - 1
    
                _ret = (1 - exp(-x * t) * (1 - st.gamma.cdf(t, alpha, 0,beta)) - TheSum) / (alpha * beta * x) + exp(-x * t) * Prob_dt

            else:
                # In (3.3) in the paper, page 14
                _ret = (1 - exp(-x * t) * (1 - st.gamma.cdf(t, alpha,0, beta)) - (1 + beta * x) ** (-alpha) * st.gamma.cdf(t, alpha, 0,beta / (1 + beta * x))) / (alpha * beta * x) + exp(-x * t) * Prob_dt
    else:
        _ret = 'Wrong input'
    return _ret


def LaplaceChain(random_variable, F, n, alpha, beta, t=None):
    _ret = None
    #Calulate the value of laplace transform Chain
    #Consult the note
    F=float(F)
 
    alpha=float(alpha)
    beta=float(beta)
    
    k = 1
    Result = 1
    
    for i in range(1, n+1):
        k = k *  ( n - i + 1 )  / i

      #  print 'the result before is ' + str(Result)
        Result = Result + k * (-1)** i * Laplace(random_variable, F, i, alpha, beta, t)
        
       # print 'the variable is '+str( k * (-1)** i * Laplace(random_variable, F, i, alpha, beta, t)) 
      #  print 'the result after is ' + str(Result)

    _ret = Result
    return _ret

def iseven(a):
        return a%2==0



def ET(F, n,alpha,beta, Arguments=None):
    
 # After looking original VBA code, the whole bussiness is funny
 # Arguments is orginally not an array when they first put as parameter
 # Then, he made it an array of one element in the code for ET function 
    
     
    if Arguments== None:
        #The only case when ET is called without any parameter array is from Drift function when t equals to infinity.
        #More corrections Page 7
        print 'the first step'
        answer = Laplace('Ld-', F, n, alpha, beta) * LaplaceChain('Ld', F, n, alpha, beta) /  ( 1 - Laplace('Ld', F, n, alpha, beta) )  + LaplaceChain('Ld-', F, n, alpha, beta)
    else:
##        # This block follows the original piece of code
##        # But, I am going to change it because I checked with Professor Bridgeman. We should be checking if the whole Arguments is an array.
##        
##        if type(Arguments) is int or type(Arguments) is float:  
##            
##            temp=[]
##            temp.append(Arguments)
##            Arguments=temp
##
##        # the following block makes no sense, But,  I put it there just to be the same as the original VBA code   
##        if type(Arguments[0]) is list:
##            Args = Arguments[0]
##        else:
##            Args = Arguments




####         I checked with Professor Bridgeman
####         We should be checking if the whole Arguments is an array.
        if type(Arguments) is list:
            Args = []
            Args.append(Arguments[0])
        else:
            Args = []
            Args.append(Arguments)

                            
        # Look at paper (3.4) (3.5)
        
        FiniteT = 0
        if iseven(len(Args)-1):
            FiniteT = 1
            FirstFactor = Laplace('Ld-t', F, n, alpha, beta, Args[0])
            Truncation = LaplaceChain('Ld-t', F, n, alpha, beta, Args[0])
        #    print 'FirstFactor is '+str(FirstFactor)
         #   print 'Truncation is '+str(Truncation)
            
            
            x = ( 1 - F )  ** Args[0]
            Y = 1 - x
            VX = x ** n
            VXChain = Y ** n
            #This is consistent with subtracting 1 from the highest value of n_k
            #See (3.5) and (3.5a) in more corrections 2011
        else:
            FirstFactor = Laplace('Ld-', F, n, alpha, beta)
            Truncation = LaplaceChain('Ld-', F, n, alpha, beta)
            
        ComFactor = LaplaceChain('Ld', F, n, alpha, beta)
        GeoSeries = Laplace('Ld', F, n, alpha, beta)
      #  print 'ComFactor is '+str(ComFactor)
      #  print 'GeoSeries is  ' + str(GeoSeries)
        
        # I am going to remove these error terms



        # The original function they used over here is UBound(Args) which basically gives the highest subscript of the array, Args
        # for example UBound([2])=0

        for i in range(FiniteT, len(Args), 2):
              #  if IsError(Args(i)) == True:
           #     Debug.Print("ERROR In ParamArray At Element: " + CStr(i + 1))
          #  else:
               
            Arg = int(Args[i])
            ComFactor = ComFactor * LaplaceChain('Ld', F, Arg, alpha, beta) ** Args[i + 1]
            GeoSeries = GeoSeries * Laplace('Ld', F, Arg, alpha, beta) ** Args[i + 1]
            if FiniteT == 0:
                FirstFactor = FirstFactor * Laplace('Ld-', F, Arg, alpha, beta) ** Args[i + 1]
                Truncation = Truncation * LaplaceChain('Ld-', F, Arg, alpha, beta) ** Args[i + 1]
            else:
                FirstFactor = FirstFactor * Laplace('Ld-t', F, Arg, alpha, beta, Args[0]) ** Args[i + 1]
                Truncation = Truncation * LaplaceChain('Ld-t', F, Arg, alpha, beta, Args[0]) ** Args[i + 1]
                VX = VX * x **  ( Arg * Args[i + 1] )
                VXChain = VXChain * Y **  ( Arg * Args[i + 1] )
                
        if FiniteT == 0:
         #   print 'the second step'
            answer = ( FirstFactor * ComFactor )  /  ( 1 - GeoSeries )  + Truncation
        else:
            # Look at page 14 paper
            Prob_dt = 1 - st.gamma.cdf(Args[0], alpha + 1,0, beta) - Args[0] *  ( 1 - st.gamma.cdf(Args[0], alpha, 0,beta) )  /  ( alpha * beta )
            # Look at paper page 13 for inverse Laplace transform, and 1st formula in (3.2), g here is G in the paper
            g = 1 - exp(( 1 - GeoSeries **  ( -1 / alpha ) )  / beta)
            # K in page 12 paper
            k = 1 -  ( 1 - g )  ** Args[0] *  ( Laplace('Ld-t', g, - 1, alpha, beta, Args[0]) - Prob_dt *  ( 1 - g )  **  ( -Args[0] ) )  /  ( Laplace('Ld-t', g, 1, alpha, beta, Args[0]) - Prob_dt *  ( 1 - g )  ** Args[0] )
            # (3.5) in paper page 12
     
            answer = k *  ( FirstFactor - Prob_dt * VX )  * ComFactor /  ( 1 - GeoSeries )  + Truncation - Prob_dt * VXChain

##            print 'Prob_dt is '+str(Prob_dt)
##            print 'g is  ' +str(g)
##            print 'k is ' +str(k)
##            print 'answer is '+str(answer)


            
        return answer
    
        





def Drift(F, alpha, beta, dt, sigmaT, adjVar, t=None):
    _ret = None
    # Look at page 10 correction
    Factor1 = -0.5 * adjVar * (1 - F) ** (dt) / (1 + (1 - F) ** (dt))

    
    Factor2 = 0.5 * sigmaT ** 2 / (dt) / (1 - F) ** (dt)
    if t == None:
        # Look at page 10 correction
        _ret = Factor1 + Factor2 * (1 - (1 - F) ** (dt)) * (1 - ET(F, 2, alpha, beta))
    else:
        #   Drift = Factor1 * (1 + (1 - F) ^ (2 * t - dt)) + Factor2 * (1 - (1 - F) ^ (dt) - (1 - F) ^ t * (Laplace("Ld-t", F, -1, alpha, beta, t) - Laplace("Ld-", F, -1, alpha, beta, t - dt)) - (ET(F, 2, alpha, beta, t) - (1 - F) ^ (dt) * ET(F, 2, alpha, beta, t - dt)))
        #   Need a solution for case ET with t = 0
        if t != dt:
            # Look at page 10 correction
            _ret = Factor1 * (1 + (1 - F) ** (2 * t - dt)) + Factor2 * (1 - (1 - F) ** (dt) - (1 - F) ** t * (Laplace("Ld-t", F, -1, alpha, beta, t) - Laplace("Ld-t", F, -1, alpha, beta, t - dt)) - (ET(F, 2, alpha, beta, t) - (1 - F) ** (dt) * ET(F, 2, alpha, beta, t - dt)))
        else:
            # Look at page 10 correction, ET is 0 when t is 0
            _ret = Factor1 * (1 + (1 - F) ** (2 * t - dt)) + Factor2 * (1 - (1 - F) ** (dt) - (1 - F) ** t * (Laplace("Ld-t", F, -1, alpha, beta, t) - Laplace("Ld-t", F, -1, alpha, beta, t - dt)) - (ET(F, 2, alpha, beta, t)))
    return _ret

# Functions to clear screen
def clc():
    print('\n'*50)


'''
# Trying to calculate with using numpy array in the beginning
drift_terms=zeros(731)
for i in range(1,732):
    drift_terms[i]=Drift(0.386806339, 3.519950456, 2.317331572, 1.0/12, 0.459514051, 0.045470** 2, i/12.0)

'''


def ScenarioStats(numScenarios, ScenarioStat):
    # this is to print out the final values
    third = 0
    fourth = 0
    sixth = 0

    
    finalMatrix=zeros((5,9))
    for j in range(0, 9):

        Ave = ScenarioStat[j]
        Appaverage = average(Ave)
        AppStDev_P = std(Ave)
        

        third =0
        fourth =0
        sixth =0
        
        for i in range(0, numScenarios):
            third = third +  ( Ave[i] - Appaverage )  ** 3
            fourth = fourth +  ( Ave[i] - Appaverage )  ** 4
            sixth = sixth +  ( Ave[i] - Appaverage )  ** 6
            
        #Output Stats of Scenario Stats
        finalMatrix[0,j] = Appaverage
        finalMatrix[1,j] = AppStDev_P
        finalMatrix[2,j] = third / numScenarios / AppStDev_P ** 3
        finalMatrix[3,j] = fourth / numScenarios / AppStDev_P ** 4
        finalMatrix[4,j] = sixth / numScenarios / AppStDev_P ** 6


    return finalMatrix



def generating_spread(long_rates_matrix,a,b,c,target_spread_at_t, initial_spread,long_rate_targets_matrix,theta,volatility_matrix,DiGamma_alpha,DiGamma_beta,tau,rho):
    [row,column]=shape(long_rates_matrix)

    sigma = sqrt(  DiGamma_beta*DiGamma_beta * (math.gamma(DiGamma_alpha+(2/tau))/math.gamma(DiGamma_alpha) ) )
    # this sigma is used to transform the random volatility of our digamma distribution so that it looks like Normal Distribution
    # a in this program is beta2 in AAA
    # b in this function is phi in AAA
    # c in this function is sigma2
    # After transformation you get 1Zt from AAA
    # To get 2Zt from 1Zt
    # you have to use 1Zt = ρ* 1Zt + ( (  (1- ( ρ^2 ))^0.5)*2Nt)
    # rho (ρ) is the correlation coefficient  between 1Zt and 2Zt, which is used in calculating spread
    spread=zeros((row,column))
    OneZt=volatility_matrix/sigma
    
   # TwoZt is basically  rho*OneZt[i,j]+  ( ( 1- (rho*rho) )**0.5  *   random.normal(0,1))
    
    for i in range(0,row):
        spread[i,0] = initial_spread 
        for j in range(1,column):
            spread[i,j]= spread[i,j-1]+ a* (target_spread_at_t-spread[i,j-1]) + b*(log(long_rates_matrix[i,j])- log(long_rate_targets_matrix[i,j]) )+ (c* (long_rates_matrix[i,j-1]**theta) *  ( rho*OneZt[i,j]+  ( ( 1- (rho*rho) )**0.5  *   random.normal(0,1)))) 
    


    return spread


def short_rate_generator(long_rates_matrix,spread_matrix,kappa, lower_bound):
    [row,column]=shape(long_rates_matrix)
    short=zeros((row,column))
    for i in range(0,row):
        current_short_rate=0
        for j in range(0,column):
            current_short_rate= long_rates_matrix[i,j]- spread_matrix[i,j]
            if current_short_rate < lower_bound:
                current_short_rate= kappa*long_rates_matrix[i,j]
            short[i,j]=current_short_rate


    return short
    

def Plotting_rates(scenario_number,Months,Long_rates_matrix,short_rates_matrix,spread_matrix):
    m=range(1,Months+1)
    plot(m,Long_rates_matrix[scenario_number,:],label='Long Rate')
    plot(m,short_rates_matrix[scenario_number,:],label='Short Rate')
    plot(m,spread_matrix[scenario_number,:],label='Spread')
    xlabel('Months')
    ylabel('Interest Rates')
    title(' Scenario '+ str(scenario_number))
    legend()
    show()
    return None



def Saving_Plots(numScenarios,Months,Long_rates_matrix,short_rates_matrix,spread_matrix):
    m=range(1,Months+1)
    plot(m,Long_rates_matrix[scenario_number,:],label='Long Rate')
    plot(m,short_rates_matrix[scenario_number,:],label='Short Rate')
    plot(m,spread_matrix[scenario_number,:],label='Spread')
    xlabel('Months')
    ylabel('Interest Rates')
    title(' Scenario '+ str(scenario_number))
    legend()
    show()
    return None