import numpy as np
import math
#from scipy.interpolate import lagrange
#from scipy.interpolate import interp1d

a = 0.4
k1 = 64
k2 = 64

def u(x,t):
    return np.sin(math.pi * (x - a * t)) + (x - a * t)**3

def getGrids(n_x, n_t, x_0 = 0., x_n = 1., t_0 =0., t_n = 1):
    xGrid = np.linspace(x_0, x_n, n_x)
    tGrid = np.linspace(t_0, t_n, n_t)
    h, tau = xGrid[1], tGrid[1]
    c = a * tau / h
    if c <= 1:
        return xGrid, tGrid, c/2
    else:
        print("Incorrect h and tau")
        return None

def fillEdges(xGrid, tGrid):
    lx, lt = len(xGrid), len(tGrid)
    u1 = np.zeros((lx, lt))
    u1[:,0] = [u(xGrid[i],tGrid[0]) for i in range(lx)]
    u1[0] = [u(xGrid[0],tGrid[i]) for i in range(lt)]
    #u1[-1] = [u(xGrid[-1],tGrid[i]) for i in range(lt)]
    return u1

def fillReal(xGrid, tGrid):
     lx, lt = len(xGrid), len(tGrid)
     u1 = np.zeros((lx, lt))
     u1 = [[u(xGrid[i], tGrid[j]) for j in range(lt)] for i in range(lx)]
     return u1

def Laks(u1, xGrid, tGrid, c):
    u2 = u1.copy()
    lx, lt = len(xGrid) - 1, len(tGrid) - 1
    xp = [xGrid[k] for k in range(lx - 4,lx)]
    for j in range(lt):
        for i in range(1 ,lx):
            u2[i][j + 1] = u2[i + 1][j] * (0.5 - c) + u2[i - 1][j] * (0.5 + c)
        fp = [u2[k][j] for k in range(lx - 4,lx)]
        u2[-1][j + 1] = np.interp(xGrid[-1],xp, fp)
        #print(u2, "\n\n")
        #f = interp1d(xp, fp)
        #u2[-1][j] = f(xGrid[-1])
        #poly = lagrange(xp,fp)
        #u2[-1][j] = poly(xGrid[-1])
        #print(xp, fp)

    return u2


xGrid, tGrid, c = getGrids(k1, k2)
u1 = fillEdges(xGrid, tGrid)
u2 = fillReal(xGrid, tGrid)
u3 = Laks(u1, xGrid, tGrid, c)
#print(u2, '\n')
#print(u3)
error = abs(u2 - u3).max()
print('\nn_h = ', k1,'  n_t = ', k2)
print('\nError = ', error)