import numpy as np
import matplotlib.pyplot as plt
import adaptivity_schemes

np.set_printoptions(precision=2)


def lam_func(x):
    n = len(x)
    lam_vec = np.zeros(n)
    for i in range(n):    
        if (x[i] > 1/np.sqrt(2)):
            lam_vec[i] = 10
        else:
            lam_vec[i] = 1
    return lam_vec
    


def Solve_6B(mesh):
    N = len(mesh) - 1   # number of elements

    A = np.zeros((N + 1, N + 1))

    lam_vec = lam_func(mesh)

    for i in range(1, N + 1):
        h = mesh[i] - mesh[i - 1]

        a_11 = lam_vec[i]/h
        a_12 = -lam_vec[i]/h
        a_21 = -lam_vec[i]/h
        a_22 = lam_vec[i]/h
                                        
        A[i - 1, i - 1] += a_11
        A[i - 1, i] += a_12
        A[i, i - 1] += a_21
        A[i, i] += a_22

    #print("A =\n", A)
    
    # take dirichlet data into account
    u_g = np.zeros(N + 1)
    u_g[0] = 0
    u_g[N] = 1
    #print("u_g =\n", u_g)

    # remove first and last row of A
    A_g = A[1:N, :]
    #print("A_g =\n", A_g)

    # assemble RHS with dirichlet data
    f = -A_g.dot(u_g)
    #print(f)

    # matrix for the inner nodes (excluding nodes with dirichlet bcs)
    A_0 = A[1:N, 1:N]   
    #print(A_0)

    # solve for u_0 (free dofs)
    u_0 = np.linalg.solve(A_0, f)

    # assemble "u = u_0 + u_g"
    u = np.concatenate([[0], u_0, [1]])
    #print("u =\n", u)

    return u


########## h-adaptivity ##########
N = 2   # number of elements
mesh = np.linspace(0, 1, N + 1)
u = Solve_6B(mesh)

plt.plot(mesh, u, '-o')
plt.grid()
plt.xlabel('x')
plt.ylabel('u_h(x)')
plt.title("h-adaptivity")


N_vec = ["0 refinements, " + str(N) + " elements"]
refinements = 5   # number of refinements
for i in range(refinements):
    mesh = adaptivity_schemes.adapt_h(mesh, lam_func(mesh)*u, 0.9)
    u = Solve_6B(mesh)
    plt.plot(mesh, u, '-o')

    N_vec.append(str(i + 1) + " refinements, " + str(len(mesh) - 1) + " elements")

plt.legend(N_vec)
plt.show()



########## r-adaptivity ##########
N = 5
mesh = np.linspace(0, 1, N + 1)
u = Solve_6B(mesh)
plt.plot(mesh, u, '-o')
title = "r-adaptivity with " + str(N) + " elements"
plt.title(title)

adaptations_vec = ["0 adaptations"]
adaptations = 4   # number of iterations
for i in range(adaptations):
    mesh = adaptivity_schemes.adapt_r(mesh, lam_func(mesh)*u)
    u = Solve_6B(mesh)
    plt.plot(mesh, u, '-o')

    adaptations_vec.append(str(i + 1) + " adaptations")


plt.legend(adaptations_vec)
plt.xlabel('x')
plt.ylabel('u_h(x)')
plt.grid()
plt.show()