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

np.set_printoptions(precision=2)


def Solve_6A(mesh, p):
    N = len(mesh) - 1   # number of elements

    f = lambda x : 2*p**3*x/((p**2*x**2 + 1)**2)
    g_b = p/(p**2 + 1)



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

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

        a_11 =  1./h
        a_12 = -1./h
        a_21 = -1./h
        a_22 =  1./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


        phi_lower = lambda x : (mesh[i] - x)/h
        f_vec[i-1] += integrate.quad(lambda x : f(x)*phi_lower(x), mesh[i - 1], mesh[i])[0]

        phi_upper = lambda x : (x - mesh[i - 1])/h
        f_vec[i] += integrate.quad(lambda x : f(x)*phi_upper(x), mesh[i - 1], mesh[i])[0]



    # take Neumann data into account
    A[N, N] += 0
    f_vec[N] += g_b


    # take Dirichlet data into account
    u_g = np.zeros(N + 1)
    u_g[0] = -np.arctan(p)
    #print("u_g =\n", u_g)

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

    # remove first row of f_vec
    f_g = f_vec[1:N+1]

    # assemble RHS with dirichlet data
    f_g -= A_g.dot(u_g)
    #print("f_g =\n", f_g)

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

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

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

    return u



p = 100
########## h-adaptivity ##########
N = 5   # number of elements
mesh = np.linspace(-1, 1, N + 1)
u = Solve_6A(mesh, p)

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 = 4   # number of refinements
for i in range(refinements):
    mesh = adaptivity_schemes.adapt_h(mesh, u, 0.7)
    u = Solve_6A(mesh, p)
    plt.plot(mesh, u, '-o')

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


# plot exact solution
x = np.linspace(-1, 1, 50)
plt.plot(x, np.arctan(p*x))
N_vec.append("exact")

plt.legend(N_vec)
plt.show()



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

adaptations_vec = ["0 adaptations"]
adaptations = 5   # number of iterations
for i in range(adaptations):
    mesh = adaptivity_schemes.adapt_r(mesh, u)

    u = Solve_6A(mesh, p)
    plt.plot(mesh, u, '-o')

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


# plot exact solution
x = np.linspace(-1, 1, 50)
plt.plot(x, np.arctan(p*x))
adaptations_vec.append("exact")


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