import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=2)


N = 10   # number of elements
a = 1
alpha = 1
g_b = 1
f_const = 1

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


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

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


    a_11 =  1./h + a*h/3.
    a_12 = -1./h + a*h/6.
    a_21 = -1./h + a*h/6.
    a_22 =  1./h + a*h/3
                                      
    A[i - 1, i - 1] += a_11
    A[i - 1, i] += a_12
    A[i, i - 1] += a_21
    A[i, i] += a_22


    f_vec[i] = f_const*h

print("A =\n", A)

# take Neumann data into account
A[N, N] += alpha
f_vec[N] += alpha*g_b


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

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

# remove first row of f_vec
f_g = f_vec[1:N+1]

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

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

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

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


print("u =\n", u)

plt.plot(x, u, '-')
plt.xlabel('x')
plt.ylabel('u_h(x)')
plt.grid()
plt.show()