import numpy as np
from adapt import *

# PDE:
#  -u''(x) = f(x)       x in (-1,1)
#    u(-1) = -arctan(p)
#    u'(1) = p / (p^2 + 1)
# 
# weak form
# int u'v' dx = int f(x) * v(x) dx + p/(p^2+1) * v(1)
# 

# rhs
def f(x):
    return 2 * p**3 * x / (p**2 * x**2 + 1)**2

# Stiffness and Load
def K_loc(k, mesh):
    K_loc = np.zeros((2,2))
    K_loc[0,0] = integrate.quad(lambda x: phi_prime(k-1, mesh, x)**2, mesh[k-1], mesh[k])[0]
    K_loc[1,0] = integrate.quad(lambda 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: 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: 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(-1) = -arctan(p)
    K[0,:] = 0
    K[0,0] = 1
    F[0]   = -np.arctan(p)
    # Neumann: u'(1) = p / (p^2 + 1)
    F[-1] += p/(p**2+1) * phi(n, mesh, 1)
    return K,F

def plotting(mesh, u, comment):
    exact_x = np.linspace(-1,1,1000)
    exact = np.arctan(p*exact_x)
    plt.plot(exact_x, exact, "--", linewidth=1, color="red", label="exact")
    plt.title(f"p = {p} | 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_a.png", dpi=300)
    plt.show()
    return 0

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

# parameters
p_list = [5,10,20,100]
for p in p_list:

    # h-adaptivity
    mesh = np.array([-1.0, -0.2, 0, 0.7, 1.0])            # with 0 as node
    # mesh = np.array([-1.0,-0.131,0.372,1.0])              # without 0 as node

    # r-adaptivity
    # mesh = np.linspace(-1, 1, 11)                         # with 0 as node
    # mesh = np.linspace(-1, 1, 10)                         # without 0 as node

    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 = 6
    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