import numpy as np
import matplotlib.pyplot as plt
#take in Data in the form [[x1,y1],[x2,y2]...]
a = np.array([[1, 4],[2,2],[3,0],[5,3]], dtype=float)

#Define the function S.
b = len(a) - 1
s = np.zeros(((4*b),(4*b)), dtype=float)
B = np.zeros(4*b, dtype=float)
def S1(x,n,ns):
    s[2*ns, 4*ns] = 1.0
    B[2*ns] = (a[ns, 1])
    pass
def S(x,n, ns):
    s[2*ns+1, 4*ns] = 1.0
    s[2*ns+1, 4*ns + 1] = (x - n)
    s[2*ns+1, 4*ns + 2] = ((x - n)**2)
    s[2*ns+1, 4*ns + 3] = ((x - n)**3)
    B[2*ns+1] = (a[ns+1, 1])
    pass

#Populate the matrix
c = 0
while c < b:
    S1(a[c, 0], a[c, 0], c)
    S(a[c+1, 0], a[c, 0], c)
    c = c +1
    pass
print(s)

#Set fprime equivalencies
c = c + b
e = 0
while e < b-1:
    n1 = (a[e,0])
    x = (a[e+1,0])
    n2 = (a[e+1,0]) 
    s[c, 4*e+1] = 1
    s[c, 4*e+2] = (2*x - 2*n1)
    s[c, 4*e+3] = (3*(x**2) - 6*n1*x + 3*n1**2)
    s[c, 4*e+5] = -1
    s[c, 4*e+6] = -(2*x - 2*n2)
    s[c, 4*e+7] = -(3*(x**2) - 6*n2*x + 3*n2**2)
    c = c + 1
    e = e + 1
    pass
print (s)

#Set fsecondprime equivalencies
e = 0
while e < b-1:
    n1 = a[e,0] 
    x = a[e+1,0] 
    n2 = a[e+1,0] 
    s[c, 4*e+2] = 2
    s[c, 4*e+3] = (6*x - 6*n1)
    s[c, 4*e+6] = -2
    s[c, 4*e+7] = -(6*x - 6*n2)
    c = c + 1
    e = e + 1
    pass
print (s)

# set boundary condition
n1 = (a[0,0])
x1 = (a[0,0]) 
n2 = (a[b-1,0])
x2 = (a[b,0]) 
s[c, 2] = 2
s[c, 3] = (6*x1 - 6*n1)
B[c] = 0.0
c = c + 1
s[c, 4*b-2] = 2
s[c, 4*b-1] = (6*x2 - 6*n2)
B[c] = 0.0

#Solve and Rearrange
s = np.array(s, dtype=float)
B = np.reshape(B,(len(B),1))
print(s)
print(B)
print(np.linalg.solve(s,B))
SolutionVector = np.linalg.solve(s,B)
SolutionMatrix = np.reshape(SolutionVector, (b,4))
print(SolutionMatrix)

#Define S(x) to use when making the data
def Su(a, b, c, d, n, x):
    return a + b*(x-n) + c*((x-n)**2) + d*((x-n)**3)

#Create list to plot
e = 0
k = float(0)
h = .1
plotVector = []
kVector = []
while e < b:
    while k <= a[e+1,0]:
        ai = float(SolutionMatrix[e,0])
        bi = float(SolutionMatrix[e,1])
        ci = float(SolutionMatrix[e,2])
        di = float(SolutionMatrix[e,3])
        ni = float(a[e,0])
        plotVector.append(round(Su(ai,bi,ci,di,ni,k,), 6))
        k = round(k + h, 6)
        kVector.append(k)
        pass
    e = round(e + 1, 6)
    pass
print(plotVector)
print(len(plotVector))

#Plot the data
plt.plot(kVector,plotVector, '.')
plt.show()