import sys
from numpy import *
import matplotlib.pyplot as plt

def main(N, U):
    L = 1.0
    rho = 1.0
    gamma = 0.05
    faces = linspace(0, L, N+1)
    nodes = 0.5*(faces[1:] + faces[:-1])
    delta = faces[1]-faces[0]
    phi0 = 1.0
    phi1 = 0.0

    # Set up linear system of equations
    A = zeros((N, N))
    b = zeros(N)

    D = gamma / delta
    F = rho * U

    print("Peclet number = {}".format(F/D))

    # ap*T_P - aw*T_W - ae*T_E = 0
    for i in range(1, N-1):
        aw = D + F / 2.0
        ae = D - F / 2.0
        sp = 0
        su = 0
        A[i, i-1] = -aw
        A[i, i]   = ae + aw - sp
        A[i, i+1] = -ae
        b[i] = su

    # Node 1
    ae = D - F / 2.0
    aw = 0
    sp = -(2*D + F)
    su = (2*D + F)*phi0
    A[0, 0] = aw + ae - sp
    A[0, 1] = -ae
    b[0] = su

    # Node N
    aw = D + F / 2.0
    ae = 0
    sp = -(2*D - F)
    su = (2*D - F)*phi1
    A[N-1, N-1] = aw + ae - sp
    A[N-1, N-2] = -aw
    b[N-1] = su

    # Exact solution
    y_e = linspace(0, L, 1000)
    def exact(y):
        return phi0 + (phi1-phi0)*(exp(rho*U*y/gamma)-1)/(exp(rho*U*L/gamma)-1)

    phi = linalg.solve(A, b)
    plt.plot(nodes, phi)
    plt.plot(y_e, exact(y_e), 'r')
    return delta, sqrt(sum((phi-exact(nodes))**2)/N)
    #return delta, linalg.norm(phi-exact(nodes))

if __name__ == "__main__":

    if len(sys.argv) > 1:
        N = int(sys.argv[-1])
        U = float(sys.argv[-2])
        main(N, U)
        plt.legend(('Numerical', 'Exact'), loc='lower left')
        plt.title("Central difference with {} CVs".format(N))
        plt.show()

    else:
        err = []
        mesh = range(3, 12)
        for i in mesh:
            err.append(main(2**i, 2.5))

        err = array(err)
        for i in range(1, err.shape[0]):
            print(err[i, 0], err[i, 1], log(err[i, 1]/err[i-1, 1]) / log(err[i, 0]/err[i-1, 0]))

        leg = ["N = {}".format(2**i) for i in mesh]
        leg.append("Exact")
        plt.legend(leg, loc='lower left')

        plt.figure()
        plt.loglog(err[:, 0], err[:, 1])
        plt.loglog(err[:, 0], err[:, 0]/err[0,0]*err[0,1], 'b')
        plt.loglog(err[:, 0], err[:, 0]**2/err[0,0]**2*err[0,1], 'r')
        plt.legend(('Central difference', r'$\mathcal{O}(1/h)$', r'$\mathcal{O}(1/h^2)$'))

        plt.xlabel('mesh size')
        plt.ylabel('L2 error norm')

        plt.show()