import numpy as np

def flux_jumps(mesh, u):
    N = len(mesh) - 1   # number of elements
    jumps = np.zeros(N + 1) # N-1 edges
    
    for i in range(1, N):
        upper = (u[i + 1] - u[i])/(mesh[i + 1] - mesh[i])
        lower = (u[i] - u[i - 1])/(mesh[i] - mesh[i - 1])
        jumps[i] = upper - lower

    return jumps


def residual_errors(mesh, u):
    N = len(mesh) - 1   # number of elements
    errors = np.zeros(N)

    jumps = flux_jumps(mesh, u)

    for i in range(N):
        errors[i] = np.sqrt((jumps[i]**2 + jumps[i + 1]**2)/2)  # Braess (8.10)
        
    #print("errors:\n", errors)

    return errors



def adapt_h(mesh, u, alpha):
    N = len(mesh) - 1   # number of elements

    errors = residual_errors(mesh, u)
    threshhold = alpha * abs(max(errors))

    # refine mesh
    refined_mesh = [mesh[0]]
    for i in range(N):
        if abs(errors[i]) <= threshhold:
            refined_mesh.append(mesh[i + 1])
        else:
            refined_mesh.append(mesh[i] + (mesh[i + 1] - mesh[i])/2)
            refined_mesh.append(mesh[i + 1])


    #print("refined mesh:\n", refined_mesh)

    return refined_mesh


def adapt_r(mesh, u):  
    N = len(mesh) - 1   # number of elements

    rho = np.abs(flux_jumps(mesh, u))  # rho ... mesh density function

    p = np.zeros(N)     # piecewise constant function on the mesh elements
    for i in range(N):
        p[i] = (rho[i] + rho[i + 1])/2

    P = np.zeros(N + 1)     # \int_0^{x_j} p(x) dx for j = 1,...,N
    for j in range(1, N + 1):
        h_j = mesh[j] - mesh[j - 1]
        P[j] = P[j - 1] + h_j*p[j - 1]  # add integral over j-th interval

    moved_mesh = np.zeros(N + 1)
    moved_mesh[0] = mesh[0]
    moved_mesh[-1] = mesh[-1]

    for j in range(1, N):  # calculate the new nodes with De Boor's algorithm
        xi_j = j/N
        k = np.searchsorted(P, xi_j*P[-1], side="left")  # searches for index k, such that xi_j*P[-1] > P[k]
        assert(P[k - 1] < xi_j*P[-1])
        assert(xi_j*P[-1] <= P[k])

        moved_mesh[j] = mesh[k - 1] + (xi_j*P[-1] - P[k - 1])/p[k - 1]
        print("orign_mesh[j] =", mesh[j])
        print("moved_mesh[j] =", moved_mesh[j])
    print("done\n")


    #print("moved mesh:\n", moved_mesh)

    return moved_mesh