import numpy as np

def differentiate(u, dt):
    dudt = np.zeros(len(u))
    for i in range(1, len(dudt)-1, 1):
        dudt[i] = (u[i+1] - u[i-1])/(2*dt)
    i = 0
    dudt[i] = (u[i+1] - u[i])/dt
    i = len(dudt)-1
    dudt[i] = (u[i] - u[i-1])/dt
    return dudt

def test_differentiate():
    """Test differentiate with a quadratic u."""
    t = np.linspace(0, 4, 9)
    u = 3*t**2 + 2*t + 7
    dudt = differentiate(u, dt=t[1]-t[0])
    error = dudt - 6*t - 2
    # print the error at every mesh point
    print ('error: ', error)

def differentiate_vec(u, dt):
    dudt = np.zeros(len(u))
    dudt[1:-1] = (u[2:] - u[0:-2])/(2*dt)
    dudt[0] = (u[1] - u[0])/dt
    dudt[-1] = (u[-1] - u[-2])/dt
    #dudt[0] = (-3*u[0] + 4*u[1] - u[2])/(2*dt)
    #dudt[-1] = (3*u[-1] - 4*u[-2] + u[-3])/(2*dt)
    return dudt

def test_differentiate_vec():
    """Test differentiate_vec by comparing with differentiate."""
    t = np.linspace(0, 4, 9)
    u = 3*t**2 + 2*t + 7
    dudt_expected = differentiate(u, dt=t[1]-t[0])
    dudt_computed = differentiate_vec(u, dt=t[1]-t[0])
    diff = abs(dudt_expected - dudt_computed).max()
    tol = 1E-15
    assert diff < tol


def test_convergence():
    import sympy as sym
    t = sym.symbols('t')
    u = sym.exp(-0.1*t)
    dudt = sym.diff(u, t)
    print ('u:', u)
    print ('du/dt:', dudt, '\n')
    # Turn sympy expressions to numerical Python functions
    u_func = sym.lambdify([t], u, modules='numpy')
    dudt_e = sym.lambdify([t], dudt, modules='numpy')

    print (' k      Nt   E/dt**2      E      Ei/dt**2    Ei')
    k_values = range(1, 20)
    for k in k_values:
        # Compute mesh
        a=0; b=10; Nt = 2**k;  dt = float(b-a)/Nt;  T = Nt*dt
        t = np.linspace(a, T, Nt+1)
        dudt = differentiate_vec(u_func(t), dt)
        # Approximate L2 integral of the error
        E  = np.sqrt(dt*np.sum((dudt_e(t) - dudt)**2))
        Ei = np.sqrt(dt*np.sum((dudt_e(t)[1:-1] - dudt[1:-1])**2))
        print ('%2d  %6d  %.2E  %.3E  %.2E  %.3E' %(k, Nt, E/dt**2, E, Ei/dt**2, Ei))


test_differentiate()
test_differentiate_vec()
test_convergence()