import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=2)


p = 70
N_vec = [10, 20, 30, 40, 70]
for N in N_vec:

    x = np.linspace(0, 1, N + 1)

    A = np.zeros((N + 1, N + 1))

    for i in range(1, N + 1):
        h = x[i] - x[i - 1]


        a_11 =  1./h - p/2.
        a_12 = -1./h + p/2.
        a_21 = -1./h - p/2.
        a_22 =  1./h + p/2.
                                            
        A[i - 1, i - 1] += a_11
        A[i - 1, i] += a_12
        A[i, i - 1] += a_21
        A[i, i] += a_22

    print("A =\n", A)



    # take dirichlet data into account
    u_g = np.zeros(N + 1)
    u_g[0] = 0
    u_g[N] = 1
    print("u_g =\n", u_g)

    # remove first and last row of A
    A_g = A[1:N, :]
    #print("A_g =\n", A_g)

    # assemble RHS with dirichlet data
    f = -A_g.dot(u_g)
    #print(f)

    # matrix for the inner nodes (excluding nodes with dirichlet bcs)
    A_0 = A[1:N, 1:N]   
    #print(A_0)

    # solve for u_0 (free dofs)
    u_0 = np.linalg.solve(A_0, f)

    # assemble "u = u_0 + u_g"
    u = np.concatenate([[0], u_0, [1]])


    print("u =\n", u)

    plt.plot(x, u, '-')


# plotting the exact solution
plt.plot(x, (np.exp(p*x) - 1.)/(np.exp(p) - 1.))

plt.xlabel('x')
plt.ylabel('u_h(x)')
plt.legend(N_vec + ['exact'])
plt.title("Comparing discrete solution for increasing number of elements")
plt.grid()
plt.show()