import numpy as np
import matplotlib.pyplot as plt

# Peclet problem:
#  -u''(x) + pu'(x) = 0    x in (0,1)
#              u(0) = 0
#              u(1) = 1

# parameters
p = 70
n_values = [10,20,30,40,70]       # elements

# exact solution
x_exact = np.linspace(0,1,1000)
u_exact = (np.exp(p*x_exact)-1)/(np.exp(p)-1)

for n in n_values:
    # 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 = (p/2) * np.array([[ -1, 1], [ -1, 1]])
    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

    # Boundary conditions
    # Dirichlet: u(0) = 0
    K[0,:] = 0
    K[0,0] = 1
    F[0]   = 0   
    # Dirichlet: u(1) = 1
    K[-1,:] = 0
    K[-1,-1] = 1
    F[-1]   = 1

    u = np.linalg.solve(K, F)
    plt.plot(x, u, "-o", markersize=2, label=f"n = {n}")

plt.plot(x_exact, u_exact, "black", label="exact")
plt.title(f"p = {p}")
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_c.png", dpi=300)
plt.show()