import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate

# Basis functions
# ---- TODO: order p? ----
def phi(k, mesh, x):
    if x > mesh[k-1] and x < mesh[k]:
        return (x-mesh[k-1]) / (mesh[k] - mesh[k-1])
    elif x > mesh[k] and x < mesh[k+1]:
        return (mesh[k+1]-x) / (mesh[k+1] - mesh[k])
    elif x == mesh[k]:
        return 1
    else:
        return 0
def phi_prime(k, mesh, x):
    if x > mesh[k-1] and x < mesh[k]:
        return 1 / (mesh[k] - mesh[k-1])
    elif x > mesh[k] and x < mesh[k+1]:
        return -1 / (mesh[k+1] - mesh[k])
    elif x == mesh[k]:
        return print("oh no")
    else:
        return 0
    
# vertex flux jump between elements
def flux_jumps(mesh, u):
    m = len(mesh)

    r = (u[2:]-u[1:-1]) / (mesh[2:]-mesh[1:-1])
    l = (u[1:-1]-u[:-2]) / (mesh[1:-1]-mesh[:-2])
    jumps = np.zeros(m)                      # 0 jump at bnd?
    jumps[1:-1] = r - l

    jumps[0] = jumps[1]                      # or same jump at bnd?
    jumps[-1] = jumps[-2]                    # or same jump at bnd?
    return jumps

# h-adaptivity (refine neighboring elments, if flux jump over certain threshold)
def adapt_h(mesh, jumps):
    new_mesh = [mesh[0]]
    jumps = jumps[1:-1]                      # only interior jumps (excluding bnd nodes needed for De Boor)
    
    # define threshold
    threshold = 0.5 * np.max(np.abs(jumps))
    # mark elements
    marked_nodes = np.abs(jumps) > threshold
    marked_el = np.zeros(len(mesh)-1, dtype=bool)
    for i, refine in enumerate(marked_nodes):
        if refine:
            marked_el[i] = True
            marked_el[i+1] = True

    # make new mesh
    for k, refine in enumerate(marked_el):
        l = mesh[k]
        r = mesh[k+1]
        if refine:
            new_mesh.extend(np.linspace(l, r, 3)[1:])
        else:
            new_mesh.append(r)

    return np.asarray(new_mesh, dtype=float)

# r-adaptivity (one iteration of De Boor's algorithm, moving mesh nodes for equidistributing mesh)
def adapt_r(mesh, rho):
    m = len(mesh)
    
    p = 0.5 * (rho[:-1] + rho[1:])                  # piecewise constant function on elements
    
    P = np.zeros(m)                             
    for i in range(1,m):
        P[i] = P[i-1] + (mesh[i]-mesh[i-1])*p[i-1]  # approx integral of p, from node 0 to node i
    Pb = P[-1]                                      # integral of p, over whole mesh 

    new_mesh = mesh.copy()
    for j in range(1, m-1):
        xi_j = (j)/(m-1)
        k = np.searchsorted(P, xi_j*Pb)             # search k s.t.: P[node k-1] < xi_j*Pb < P[node k]
        k = max(k,1)                                # if k=0
        new_mesh[j] = mesh[k-1] + 2*(xi_j*Pb - P[k-1]) / (rho[k-1]+rho[k])    # new node  j
    return new_mesh