import numpy as np
from adapt import *

# PDE:
#  -(lambda(x)u'(x))' = 0         x in (0,1)
#                u(0) = 0
#                u(1) = 1
#           lambda(x) = |  1     x in (0,1/sqrt(2))
#                       | 10     x in (1/sqrt(2),1)   

# rhs
def f(x):
    return 0

def lam(x):
    if x >= 0 and x <= 1/np.sqrt(2):
        return 1
    elif x <= 1 and x > 1/np.sqrt(2):
        return 10
    else:
        return print("lambda undefined")

# Stiffness and Load
def K_loc(k, mesh):
    K_loc = np.zeros((2,2))
    K_loc[0,0] = integrate.quad(lambda x: lam(x)*phi_prime(k-1, mesh, x)**2, mesh[k-1], mesh[k])[0]
    K_loc[1,0] = integrate.quad(lambda x: lam(x)*phi_prime(k-1, mesh, x)*phi_prime(k, mesh, x), mesh[k-1], mesh[k])[0]
    K_loc[0,1] = integrate.quad(lambda x: lam(x)*phi_prime(k, mesh, x)*phi_prime(k-1, mesh, x), mesh[k-1], mesh[k])[0]
    K_loc[1,1] = integrate.quad(lambda x: lam(x)*phi_prime(k, mesh, x)**2, mesh[k-1], mesh[k])[0]
    return K_loc
def F_loc(k, mesh):
    F_loc = np.zeros(2)
    F_loc[0] = integrate.quad(lambda x: f(x) * phi(k-1, mesh, x), mesh[k-1], mesh[k])[0]
    F_loc[1] = integrate.quad(lambda x: f(x) * phi(k, mesh, x), mesh[k-1], mesh[k])[0]
    return F_loc

# Assembling
def Assemble(mesh):
    m = len(mesh)
    n = m-1

    K = np.zeros((m,m))
    F = np.zeros(m)
    for k in range(1,m):
        K[k-1:k+1,k-1:k+1] += K_loc(k, mesh)
        F[k-1:k+1] += F_loc(k, mesh)

    # 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
    return K,F

def plotting(mesh, u, comment):
    exact_x = [0, 1/np.sqrt(2), 1]
    exact = [0, 10/(np.sqrt(2)+9), 1]
    plt.plot(exact_x, exact, "--", linewidth=1, color="red", label="exact")
    plt.title(f"n = {n} | {comment}")
    plt.xlabel("x")
    plt.ylabel("u(x)")
    plt.plot(mesh, u, "-o", label="u_h")
    plt.xticks(mesh, labels=[])
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("task_b.png", dpi=300)
    plt.show()
    return 0

##############################################################################

mesh = np.linspace(0, 1, 10)
m = len(mesh)
n = m-1

K, F = Assemble(mesh)                                 # assemble
u = np.linalg.solve(K, F)                             # solve
jumps = flux_jumps(mesh, u)                           # flux jumps
plotting(mesh, u, "before adapting")                  # plotting

iterations = 3
for it in range(iterations):
    # mesh = adapt_h(mesh, jumps)                       # h-adaptivity
    mesh = adapt_r(mesh, np.abs(jumps))               # r-adaptivity (positive density (jumps)!)
    m = len(mesh)
    n = m-1

    K, F = Assemble(mesh)                             # assemble
    u = np.linalg.solve(K, F)                         # solve
    jumps = flux_jumps(mesh, u)                       # flux jumps
    # plotting(mesh, u, f"iteration {it+1}")            # plotting each iteration

# print(jumps)
plotting(mesh, u, f"after {iterations} iterations")   # plotting