import numpy as np
import matplotlib.pyplot as plt

# PDE:
#  -u''(x) + a*u(x) = f(x)         x in (0,1)
#              u(0) = 0
#             u'(1) = \alpha*(g_b - u(1))

# parameters
a = 1
f = 3
alpha = 1
gb = 1
n = 10       # elements

# mesh
m = n+1      # nodes
h = 1.0/n
x = np.linspace(0,1,m)

# local stiffness matrix
K_loc = np.zeros((2,2))
A = (1.0/h) * np.array([[ 1,-1], [-1, 1]])
B = (a*h/6) * np.array([[ 2, 1], [ 1, 2]])
K_loc = A+B

# Assembling
K = np.zeros((m,m))
F = np.zeros(m)
for i in range(n):
    K[i:i+2,i:i+2] += K_loc
    F[i:i+2] += np.full(2, f*h/2)

# Boundary conditions
# Dirichlet: u(0) = 0
K[0,:] = 0
K[0,0] = 1
F[0]   = 0   
# Neumann: u'(1) = \alpha*(g_b - u(1))
A[-1,-1] += alpha
F[-1] += alpha*gb

u = np.linalg.solve(K, F)
plt.plot(x, u, "-o", label="u_h")
plt.title(f"n = {n} | h = {h} | a = {a} | f(x) = {f} | alpha = {alpha} | gb = {gb}")
plt.xlabel("x")
plt.ylabel("u(x)")
plt.legend()
plt.grid(True)

print("K = ", K)
print("f = ", F)
print("u = ", u)

plt.tight_layout()
plt.savefig("../task_a.png", dpi=300)
plt.show()