import scipy.interpolate as si  
  
  
from OptionPricer import *  
  
class PDECalculator(OptionPricer):  
    _NumberSteps = 0  
    _NumberAssetSteps = 0  
      
    _StepAdapt = False  
    _Regul = False  
  
    @staticmethod  
    def preparedivs(divi, T, dt):  
        dividends = [[], []]   
        #2 columns vector, one for amount one for time.  
        #we make sure we don't take into account dividend happening after expiration  
     
        if (np.size(divi)>0 and divi[0][0]<T) :  
            lastdiv = np.nonzero(np.array(divi[0][:])<= T)[0][-1]  
            dividends[0] = divi[0][:lastdiv+1]          
            dividends[1] = divi[1][:lastdiv+1]   
      
        #Transform the dividend date into a step  
        if np.size(dividends)>0:  
            dividendsStep = np.floor(np.multiply(dividends[0], 1/dt))  
        else:  
            dividendsStep = []   
        return dividends, dividendsStep  
  
    @staticmethod  
    def preparerates(r2, T, dt, NTS):  
        r = [[], []]  
          
        #similar structure for rates  
        # rates with date d1 is use until date d1 so we need the first date post maturity.     
        if np.size(r2)>0 :  
            lastrate = np.nonzero(np.array(r2[0][:])>= T)[0]+1  
            #i take the one just after expiration  
            r[0] = r2[0][:lastrate]          
            r[1] = r2[1][:lastrate]         
        if np.size(r)>0:  
            rateChangeStep = np.floor(np.multiply(r[0], 1/dt))  
        else:  
            rateChangeStep = []        
          
        #prepare the array for rates used for the PDE  
        rArray = np.zeros((NTS+2, 1), float)     
        DFArray = np.zeros((NTS+2, 1), float)  
      
        #get the steps where the rate changes  
        prevstep = 0  
          
        for step in range(len(rateChangeStep)):  
              
            DFArray[prevstep:min(rateChangeStep[step], NTS+2)] = np.exp(-r[1][step]*dt)  
            rArray[prevstep:min(rateChangeStep[step], NTS+2)] = r[1][step]  
            prevstep = rateChangeStep[step]  
  
        return r, rArray, rateChangeStep, DFArray  
  
    @staticmethod  
    def PDEEUDuo(pricingParameters, optionContract, NTS, NAS, regul, stepAdapt):  
              
        #parameters  
        S0 = pricingParameters._S0  
        r2 = pricingParameters._r2  
        T = pricingParameters._T         
        sigma = optionContract._Sigma  
        divi = pricingParameters._div  
        K = optionContract._K  
          
        #calculate delta T      
        deltaT = float(T) / NTS  
   
        Zmin = -7.5  
        Zmax = 7.5  
        
        dZ = float((Zmax-Zmin)/NAS)     
          
        #time step is constraint by stability to be a function  of dZ  
        #but if the user wants a smaller time step, that is fine too  
        dt = min(dZ**2/(3*sigma*sigma), float(T)/NTS)  
          
        dividends, dividendsStep = PDEPricer.preparedivs(divi, T, deltaT)  
        r, rArray, rateChangeStep, DFArray = PDEPricer.preparerates(r2, T, dt, NTS)  
       
              
        #we get the rate  
        rr = rArray[-1][0]  
        NTS = int(T/dt)+1  
          
        a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
        b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
        c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
             
         
        VC = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        VP = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        Z = np.array([(Zmin+i*dZ) for i in xrange(int(NAS)+1)])  
          
          
        if regul:  
            VC = optionContract.expiCR(np.exp(Z), rr, sigma, dt)  
            VP = optionContract.expiPR(np.exp(Z), rr, sigma, dt)  
            NTS = NTS-1  
        else:  
            VC = optionContract.expiC(np.exp(Z))  
            VP = optionContract.expiP(np.exp(Z))  
      
        a1 = (Z[NAS]-Z[NAS-1])  
        b1 = (Z[NAS-1]-Z[NAS-2])  
              
      
        for k in range(NTS, -1, -1):  
              
            #we need to change the a, b, c only when rate changes  
            if (k in rateChangeStep[:]):  
                rr = (r[1][np.nonzero(rateChangeStep == (k))[0]])  
                a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
            if (k in dividendsStep):  
                if stepAdapt:  
                      
                    #we perform a fraction of step for the pde, place the dividend and finish the step  
                    #so no matter the step size, the div is at the right place.  
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                    fraction = -(k*dt-dividends[0][np.nonzero(dividendsStep == (k))[0]])  
                    dtf = dt-fraction  
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                    VP[1:-1] = af*VP[:-2]+bf*VP[1:-1]+cf*VP[2:]  
                 
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    tckC = si.splrep(Z, VC, k = 1)  
                    tckP = si.splrep(Z, VP, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckC, der = 0)              
                    VP = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckP, der = 0)             
                      
                                   
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    #we finish the step  
                    dtf = fraction  
                     
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                    VP[1:-1] = af*VP[:-2]+bf*VP[1:-1]+cf*VP[2:]  
                 
                else:  
                    #otherwise perform the step and then place the dividend  
                     #call  
                    VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                    VC[0] = VC[0]*(1-rr*dt)  
                      
                    #put  
                    VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
                      
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                      
                    tckC = si.splrep(np.exp(Z), VC, k = 1)  
                    tckP = si.splrep(np.exp(Z), VP, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckC, der = 0)              
                    VP = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckP, der = 0)    
                     
                      
            else:  
                #there is no dividends  
                #call  
                VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
               #put  
                VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
                   
            VC[0] = 0  
            VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
            VP[0] = K  
            VP[NAS] = 0  
            #check early exercise  
            #VC = earlyCall(VC, np.exp(Z), Kk)  
            #VP = earlyPut(VP, np.exp(Z), Kk)  
              
        #get the prices and the grid from the vector      
        tck = si.splrep(Z, VC)  
        optionContract.call.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        optionContract.call.delta = (p1-p2)/2  
        optionContract.call.gamma = p1+p2-2*optionContract.call.price  
        VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
        tck = si.splrep(Z, VC)  
        optionContract.call.theta = (optionContract.call.price-si.splev(np.log(S0), tck, der = 0) )/dt  
          
          
        tck = si.splrep(Z, VP)  
        optionContract.put.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        optionContract.put.delta = (p1-p2)/2  
        optionContract.put.gamma = p1+p2-2*optionContract.put.price  
        VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
        tck = si.splrep(Z, VP)  
        optionContract.put.theta = (optionContract.put.price-si.splev(np.log(S0), tck, der = 0) )/dt  
          
          
        return VC        
          
    @staticmethod  
    def PDEEUSolo(pricingParameters, optionContract, NTS, NAS, regul, stepAdapt, option, payoff):  
              
        #parameters  
        S0 = pricingParameters._S0  
        r2 = pricingParameters._r2  
        T = pricingParameters._T         
        sigma = optionContract._Sigma  
        divi = pricingParameters._div  
          
        #calculate delta T      
        deltaT = float(T) / NTS  
   
        Zmin = -7.5  
        Zmax = 7.5  
        
        dZ = float((Zmax-Zmin)/NAS)     
          
        #time step is constraint by stability to be a function  of dZ  
        #but if the user wants a smaller time step, that is fine too  
        dt = min(dZ**2/(3*sigma*sigma), float(T)/NTS)  
          
        dividends, dividendsStep = PDEPricer.preparedivs(divi, T, deltaT)  
        r, rArray, rateChangeStep, DFArray = PDEPricer.preparerates(r2, T, dt, NTS)  
       
              
        #we get the rate  
        rr = rArray[-1][0]  
        NTS = int(T/dt)+1  
          
        a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
        b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
        c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
             
         
        VC = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        Z = np.array([(Zmin+i*dZ) for i in xrange(int(NAS)+1)])  
          
          
        if regul:  
            VC = payoff(np.exp(Z), rr, sigma, dt)  
            NTS = NTS-1  
        else:  
            VC = payoff(np.exp(Z))  
              
        a1 = (Z[NAS]-Z[NAS-1])  
        b1 = (Z[NAS-1]-Z[NAS-2])  
              
      
        for k in range(NTS, -1, -1):  
              
            #we need to change the a, b, c only when rate changes  
            if (k in rateChangeStep[:]):  
                rr = (r[1][np.nonzero(rateChangeStep == (k))[0]])  
                a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
            if (k in dividendsStep):  
                if stepAdapt:  
                      
                    #we perform a fraction of step for the pde, place the dividend and finish the step  
                    #so no matter the step size, the div is at the right place.  
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                    fraction = -(k*dt-dividends[0][np.nonzero(dividendsStep == (k))[0]])  
                    dtf = dt-fraction  
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                      
                    VC[0] = VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    tckC = si.splrep(Z, VC, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckC, der = 0)              
                      
                                   
                    VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    #check for exercise  
                    #VC = earlyCall(VC, np.exp(Z), Kk)  
                      
                    #we finish the step  
                    dtf = fraction  
                     
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                      
                 
                else:  
                    #otherwise perform the step and then place the dividend  
                     #call  
                    VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                    VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                      
                    tckC = si.splrep(np.exp(Z), VC, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckC, der = 0)              
                      
            else:  
                #there is no dividends  
                #call  
                VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                 
                   
            VC[0] = VC[0]*(1-rr*dt)  
            VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
              
            #check early exercise  
            #VC = earlyCall(VC, np.exp(Z), Kk)  
              
              
        #get the prices and the grid from the vector      
        tck = si.splrep(Z, VC)  
        option.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        option.delta = (p1-p2)/2  
        option.gamma = p1+p2-2*option.price  
        VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
        tck = si.splrep(Z, VC)  
        option.theta = (option.price-si.splev(np.log(S0), tck, der = 0) )/dt  
                 
          
        return VC        
    @staticmethod  
    def PDEUSDuo(pricingParameters, optionContract, NTS, NAS, regul, stepAdapt):  
              
        #parameters  
        S0 = pricingParameters._S0  
        r2 = pricingParameters._r2  
        T = pricingParameters._T         
        sigma = optionContract._Sigma  
        K = optionContract._K  
        divi = pricingParameters._div  
          
        #calculate delta T      
        deltaT = float(T) / NTS  
   
        Zmin = -7.5  
        Zmax = 7.5  
        
        dZ = float((Zmax-Zmin)/NAS)     
          
        #time step is constraint by stability to be a function  of dZ  
        #but if the user wants a smaller time step, that is fine too  
        dt = min(dZ**2/(3*sigma*sigma), float(T)/NTS)  
          
        dividends, dividendsStep = PDEPricer.preparedivs(divi, T, deltaT)  
        r, rArray, rateChangeStep, DFArray = PDEPricer.preparerates(r2, T, dt, NTS)  
       
              
        #we get the rate  
        rr = rArray[-1][0]  
        NTS = int(T/dt)+1  
          
        a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
        b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
        c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
             
         
        VC = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        VP = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        Z = np.array([(Zmin+i*dZ) for i in xrange(int(NAS)+1)])  
          
          
        if regul:  
            VC = optionContract.expiCR(np.exp(Z), rr, sigma, dt)  
            VP = optionContract.expiPR(np.exp(Z), rr, sigma, dt)  
            NTS = NTS-1  
        else:  
            VC = optionContract.expiC(np.exp(Z))  
            VP = optionContract.expiP(np.exp(Z))  
      
        a1 = (Z[NAS]-Z[NAS-1])  
        b1 = (Z[NAS-1]-Z[NAS-2])  
              
      
        for k in range(NTS, -1, -1):  
              
            #we need to change the a, b, c only when rate changes  
            if (k in rateChangeStep[:]):  
                rr = (r[1][np.nonzero(rateChangeStep == (k))[0]])  
                a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
            if (k in dividendsStep):  
                if stepAdapt:  
                      
                    #we perform a fraction of step for the pde, place the dividend and finish the step  
                    #so no matter the step size, the div is at the right place.  
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                    fraction = -(k*dt-dividends[0][np.nonzero(dividendsStep == (k))[0]])  
                    dtf = dt-fraction  
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                    VP[1:-1] = af*VP[:-2]+bf*VP[1:-1]+cf*VP[2:]  
                 
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    tckC = si.splrep(Z, VC, k = 1)  
                    tckP = si.splrep(Z, VP, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckC, der = 0)              
                    VP = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckP, der = 0)             
                      
                                   
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    #check for exercise  
                    VC = optionContract.earlyC(VC, np.exp(Z))  
                    VP = optionContract.earlyP(VP, np.exp(Z))  
                      
                    #we finish the step  
                    dtf = fraction  
                     
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                    VP[1:-1] = af*VP[:-2]+bf*VP[1:-1]+cf*VP[2:]  
                 
                else:  
                    #otherwise perform the step and then place the dividend  
                     #call  
                    VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                    VC[0] = VC[0]*(1-rr*dt)  
                      
                    #put  
                    VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
                      
                    VC[0] = 0  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                    VP[0] = K  
                    VP[NAS] = 0  
                      
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                      
                    tckC = si.splrep(np.exp(Z), VC, k = 1)  
                    tckP = si.splrep(np.exp(Z), VP, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckC, der = 0)              
                    VP = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckP, der = 0)    
                     
                      
            else:  
                #there is no dividends  
                #call  
                VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
               #put  
                VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
                   
            VC[0] = 0  
            VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
            VP[0] = K  
            VP[NAS] = 0  
            #check early exercise  
            VC = optionContract.earlyC(VC, np.exp(Z))  
            VP = optionContract.earlyP(VP, np.exp(Z))  
              
        #get the prices and the grid from the vector      
        tck = si.splrep(Z, VC)  
        optionContract.call.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        optionContract.call.delta = (p1-p2)/2  
        optionContract.call.gamma = p1+p2-2*optionContract.call.price  
        VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
        tck = si.splrep(Z, VC)  
        optionContract.call.theta = (optionContract.call.price-si.splev(np.log(S0), tck, der = 0) )/dt  
          
          
        tck = si.splrep(Z, VP)  
        optionContract.put.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        optionContract.put.delta = (p1-p2)/2  
        optionContract.put.gamma = p1+p2-2*optionContract.put.price  
        VP[1:-1] = a*VP[:-2]+b*VP[1:-1]+c*VP[2:]  
        tck = si.splrep(Z, VP)  
        optionContract.put.theta = (optionContract.put.price-si.splev(np.log(S0), tck, der = 0) )/dt  
          
          
        return VC                
    @staticmethod  
    def PDEUSSolo(pricingParameters, optionContract, NTS, NAS, regul, stepAdapt, option, payoff, early):  
              
        #parameters  
        S0 = pricingParameters._S0  
        r2 = pricingParameters._r2  
        T = pricingParameters._T         
        sigma = optionContract._Sigma  
        divi = pricingParameters._div  
          
        #calculate delta T      
        deltaT = float(T) / NTS  
   
        Zmin = -7.5  
        Zmax = 7.5  
        
        dZ = float((Zmax-Zmin)/NAS)     
          
        #time step is constraint by stability to be a function  of dZ  
        #but if the user wants a smaller time step, that is fine too  
        dt = min(dZ**2/(3*sigma*sigma), float(T)/NTS)  
          
        dividends, dividendsStep = PDEPricer.preparedivs(divi, T, deltaT)  
        r, rArray, rateChangeStep, DFArray = PDEPricer.preparerates(r2, T, dt, NTS)  
       
              
        #we get the rate  
        rr = rArray[-1][0]  
        NTS = int(T/dt)+1  
          
        a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
        b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
        c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
             
         
        VC = np.asarray( [0 for i in xrange(int(NAS)+1)])  
        Z = np.array([(Zmin+i*dZ) for i in xrange(int(NAS)+1)])  
          
          
        if regul:  
            VC = payoff(np.exp(Z), rr, sigma, dt)  
            NTS = NTS-1  
        else:  
            VC = payoff(np.exp(Z))  
              
        a1 = (Z[NAS]-Z[NAS-1])  
        b1 = (Z[NAS-1]-Z[NAS-2])  
              
      
        for k in range(NTS, -1, -1):  
              
            #we need to change the a, b, c only when rate changes  
            if (k in rateChangeStep[:]):  
                rr = (r[1][np.nonzero(rateChangeStep == (k))[0]])  
                a = np.asarray( [(1/(1+rr*dt))*(-(rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                b = np.asarray( [(1/(1+rr*dt))*(1-(sigma**2)*dt/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                c = np.asarray( [(1/(1+rr*dt))*((rr-sigma**2/2)*dt/(2*dZ)+(sigma**2)*dt/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
            if (k in dividendsStep):  
                if stepAdapt:  
                      
                    #we perform a fraction of step for the pde, place the dividend and finish the step  
                    #so no matter the step size, the div is at the right place.  
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                    fraction = -(k*dt-dividends[0][np.nonzero(dividendsStep == (k))[0]])  
                    dtf = dt-fraction  
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                      
                    VC[0] = VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    tckC = si.splrep(Z, VC, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.log(np.maximum(np.exp(Z)-div, np.exp(Zmin))), tckC, der = 0)              
                      
                                   
                    VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    #check for exercise  
                    VC = early(VC, np.exp(Z))  
                      
                    #we finish the step  
                    dtf = fraction  
                     
                    af = np.asarray( [(1/(1+rr*dtf))*(-(rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
                    bf = np.asarray( [(1/(1+rr*dtf))*(1-(sigma**2)*dtf/(dZ*dZ)) for i in xrange(int(NAS-1))])  
                    cf = np.asarray( [(1/(1+rr*dtf))*((rr-sigma**2/2)*dtf/(2*dZ)+(sigma**2)*dtf/(2*dZ*dZ)) for i in xrange(int(NAS-1))])  
              
                    VC[1:-1] = af*VC[:-2]+bf*VC[1:-1]+cf*VC[2:]  
                      
                 
                else:  
                    #otherwise perform the step and then place the dividend  
                     #call  
                    VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                    VC[0] = VC[0]*(1-rr*dt)  
                    VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
                      
                    div = dividends[1][np.nonzero(dividendsStep == (k))[0]]   
                      
                    tckC = si.splrep(np.exp(Z), VC, k = 1)  
                    #option prices are interpolated on the grid  
                    VC = si.splev(np.maximum(np.exp(Z)-div, np.exp(Zmin)), tckC, der = 0)              
                     
                      
            else:  
                #there is no dividends  
                #call  
                VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
                 
                   
            VC[0] = VC[0]*(1-rr*dt)  
            VC[NAS] = (VC[NAS-1]/a1+VC[NAS-1]/b1 -VC[NAS-2]/b1)*a1  
              
            #check early exercise  
            VC = early(VC, np.exp(Z))  
              
              
        #get the prices and the grid from the vector      
        tck = si.splrep(Z, VC)  
        option.price = si.splev(np.log(S0), tck, der = 0)   
        p1 = si.splev(np.log(S0+1), tck, der = 0)  
        p2 = si.splev(np.log(S0-1), tck, der = 0)  
        option.delta = (p1-p2)/2  
        option.gamma = p1+p2-2*option.price  
        VC[1:-1] = a*VC[:-2]+b*VC[1:-1]+c*VC[2:]  
        tck = si.splrep(Z, VC)  
        option.theta = (option.price-si.splev(np.log(S0), tck, der = 0) )/dt  
                     
          
        return VC   
    def __init__(self):  
        self.SetPricerId("PDE")    
        super(PDEPricer, self).__init__()          
          
      
    def SetPricerParameters(self, N, NAS = 1, Regul = False, StepAdapt = False):  
        self._NumberAssetSteps = NAS  
        self._NumberSteps = N  
        self._Regul = Regul  
        self._StepAdapt = StepAdapt  
  
    def PrintSelf(self):  
        print self._pricerId  
        print "Number of time steps: ", self._NumberSteps  
        print "Number of asset steps: ", self._NumberAssetSteps  
        print "StepAdapt on: ", self._StepAdapt  
        print "Regul on: ", self._Regul  
        self._PricingParameters.printSelf()  
          
        for optionContract in self._contractList:  
            optionContract.printSelf()  
  
    def CalculateContract(self, optionContract):  
        # Check input parameters  
         
        if isinstance(optionContract, OptionContractUS):  
              
            if isinstance(optionContract, OptionContractDuo):  
                PDEPricer.PDEUSDuo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt)  
            elif isinstance(optionContract, OptionContractCall):  
                if self._Regul:  
                    PDEPricer.PDEUSSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt, optionContract.call, optionContract.expiCR, optionContract.earlyC)  
                else:  
                    PDEPricer.PDEUSSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, False, optionContract.call, optionContract.expiC, optionContract.earlyC)  
            elif isinstance(optionContract, OptionContractPut):  
                if self._Regul:  
                    PDEPricer.PDEUSSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt, optionContract.put, optionContract.expiPR, optionContract.earlyP)  
                else:  
                    PDEPricer.PDEUSSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, False, optionContract.put, optionContract.expiP, optionContract.earlyP)  
        else:  
              
            if isinstance(optionContract, OptionContractDuo):  
                PDEPricer.PDEEUDuo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt)  
            elif isinstance(optionContract, OptionContractCall):  
                if self._Regul:  
                    PDEPricer.PDEEUSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt, optionContract.call, optionContract.expiCR)  
                else:  
                    PDEPricer.PDEEUSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, False, optionContract.call, optionContract.expiC)  
            elif isinstance(optionContract, OptionContractPut):  
                if self._Regul:  
                    PDEPricer.PDEEUSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, self._StepAdapt, optionContract.put, optionContract.expiPR)  
                else:  
                    PDEPricer.PDEEUSolo(self._PricingParameters, optionContract, self._NumberSteps, self._NumberAssetSteps, self._Regul, False, optionContract.put, optionContract.expiP)  
            else:  
                raise PricerError("PDEPricer", "Contract not yet supported by pricing algorithm")