# -*- coding: utf-8 -*-
import numpy as np
from colorama import Fore, Back  # Used to write colored text (for printing) .Fore = Font color. Back = Background color.

"""
This module contains the functions LUdecomp3 and LUsolve3 for the decomposition and solution phases of a tridiagonal matrix. 
In LUsolve3, the vector y writes over the constant vector b during forward substitution. Similarly, 
the solution vector x overwrites y in the back substitution process. In other words, b contains the solution upon exit from LUsolve3.
"""
       	
def LUdecomp3(c,d,e):   #Analysis of the unknown as in Gauss Elimination method.
    n = d.shape[0]
    for k in range(1,n):
        lam = c[k-1]/d[k-1]
        d[k] = d[k] - lam*e[k-1]
        c[k-1] = lam
    return c,d,e


def LUsolve3(c,d,e,b):
    n = d.shape[0]
    for k in range(1,n):
        b[k] = b[k] - c[k-1]*b[k-1]

    b[n-1] = b[n-1]/d[n-1]
    for k in range(n-2,-1,-1):
        b[k] = (b[k] - e[k]*b[k+1])/d[k]
    return b 

"""
The first stage of cubic spline interpolation is to set up  and solve them for the unknown k’s (recall that k0 = kn = 0). 
This task is carried out by the function curvatures. The second stage is the computation of the interpolant at x  
This step can be repeated any number of times with different values of x using the function evalSpline. 
The function findSegment embedded in evalSpline finds the segment of the spline that contains x using the method of bisection.
"""       	

def  curvatures(xData,yData):
    n  =  len(xData)  -  1
    c  =  np.zeros(n)
    d  =  np.ones(n+1)
    e  =  np.zeros(n)
    k  =  np.zeros(n+1)
    c[0:n-1]=  xData[0:n-1]  -  xData[1:n]
    d[1:n]  =  2.0*(xData[0:n-1]  -  xData[2:n+1])
    e[1:n]  =  xData[1:n]  -  xData[2:n+1]
    k[1:n]  =  6.0*(yData[0:n-1]  -  yData[1:n]) \
        /(xData[0:n-1]  -  xData[1:n])           \
        -6.0*(yData[1:n]  -  yData[2:n+1])       \
        /(xData[1:n]  -  xData[2:n+1])
    LUdecomp3(c,d,e) #As LUdecomp3 and LUsolve3 can be written in this order, they can be written and called before curvatures(xData,yData).
    LUsolve3(c,d,e,k)
    return  k 

       	

def  evalSpline(xData,yData,k,x):
    
    def findSegment(xData,x):
        iLeft  =  0
        iRight  =  len(xData)-  1
        while 1:
            if (iRight-iLeft)<=1: return iLeft
            i=(iLeft+iRight)//2
            if x < xData[i] :  iRight=i
            else: iLeft=i          
                          
    i=findSegment(xData,x)
    h  =  xData[i]-xData[i+1]
    y  =  ((x-xData[i+1])**3/h-(x-xData[i+1])*h)*k[i]/6.0 \
        - ((x-xData[i])**3/h-(x-xData[i])*h)*k[i+1]/6.0   \
        + (yData[i]*(x-xData[i+1])                        \
        - yData[i+1]*(x-xData[i]))/h
    
    return  y 
 
      	

def  findCoefficients(xData,yData,k,m) :   #finding coefficients.
    h=xData[m+1]-xData[m]
    d=(yData[m+1]-yData[m])/h
    
    s=np.array([yData[m],d-h*(2*k[m]+k[m+1])/6,k[m]/2,(k[m+1]-k[m])/6/h])
    return s 


xData=np.array([1,2,3,4,5], dtype=float)
yData=np.array([0,1,0,1,0], dtype=float)
print (xData, yData)


k= curvatures(xData,yData)
print (Fore.GREEN+Back.BLACK,k)


x=1.5   # x = 1.5 and x = 4.5 (because of symmetry, these values should be equal)
y= evalSpline(xData,yData,k,x)
print (Fore.RED+Back.CYAN,'f(x)=',y )