import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as sl
import sys


def Main_Program():
    # Getting input data from the terminal.
    p = int(sys.argv[1])        # Polynomial degree.
    N = int(sys.argv[2])        # Number of intervals.
    sig = float(sys.argv[3])    # Sigma-parameter.
    a = eval(sys.argv[4])       # Left interval limit.
    b = eval(sys.argv[5])       # Right interval limit.
    u1 = eval(sys.argv[6])      # Left Dirichlet condition.
    u2 = eval(sys.argv[7])      # Right Dirichlet condition.

    # Assembly of the linear system.
    A,f = Assemble_Linear_System(p,N,sig,a,b)

    # Solving the linear system.
    u = Solve_Linear_System(A,f,p,N,u1,u2)

    # Calculating the global error.
    ERR = Compute_Global_Error(u,p,N,a,b)
    print(ERR)


def Assemble_Linear_System(p,N,sig,a,b):
    # Initializing assembly parameters.
    T = np.linspace(a,b,N+1)    # Uniform grid on the domain.
    h = (b-a)/N                 # Uniform element length.
    dim = N*p+1                 # Degrees of freedom.
    A = np.zeros((dim,dim))     # Linear matrix.
    f = np.zeros(dim)           # Load vector.
    X,W = Gauss_Data()          # Gaussian quadrature data.

    # Starting the assembly.
    for i in range(0,N):
        A_e = (1/h)*Element_Matrix(p,"Stiffness")
        if sig != 0:
            A_e += h*sig*Element_Matrix(p,"Mass")
        A[p*i:p*(i+1)+1,p*i:p*(i+1)+1] += A_e
        f_e = Element_Vector(p,h,T[i],X,W)
        f[p*i:p*(i+1)+1] += f_e

    return A,f


def Element_Matrix(p,type):
    if type == "Mass":
        if p == 1:
            return (1/6)*np.matrix([[2,1],[1,2]],dtype="float")
        elif p == 2:
            return (1/30)*np.matrix([[4,2,-1],[2,16,2],[-1,2,4]],dtype="float")
    elif type == "Stiffness":
        if p == 1:
            return np.matrix([[1,-1],[-1,1]],dtype="float")
        elif p == 2:
            return (1/3)*np.matrix([[7,-8,1],[-8,16,-8],[1,-8,7]],dtype="float")


def Element_Vector(p,h,c,X,W):
    f_e = np.zeros(p+1)
    for i in range(0,len(X)):
        f_e += W[i]*Force(h*X[i]+c)*Basis_Functions(p,h*X[i]+c)
    f_e = h*f_e
    return f_e


def Gauss_Data():
    X = (np.array([-0.9061798459386641,-0.5384693101056831,0,0.5384693101056831,
    0.9061798459386641])+1)/2
    W = 0.5*np.array([0.23692688505618908,0.4786286704993664,128/225,
    0.4786286704993664,0.23692688505618908])
    return X,W


def Solve_Linear_System(A,f,p,N,u1,u2):
    # Condensating the equation system.
    A = A[1:N*p,1:N*p]
    A = sp.csr_matrix(A)
    f = f[1:N*p]

    # Solving the system.
    u0 = sl.spsolve(A,f)

    # Enforcing Dirichlet conditions.
    u = np.zeros(N*p+1)
    u[1:N*p] = u0
    u[0] = u1
    u[N*p] = u2
    return u


def Compute_Global_Error(u,p,N,a,b):
    ERR = 0
    h = (b-a)/N
    X,W = Gauss_Data()
    T = np.linspace(a,b,N+1)
    for i in range(0,N):
        e = 0
        c = T[i]
        co = u[p*i:p*(i+1)+1]
        for j in range(0,len(X)):
            b = Basis_Functions(p,h*X[j]+c)
            u_ex = Exact(h*X[j]+c)
            e += W[j]*(np.dot(co,b)-u_ex)**2
        e = h*e
        ERR += e
    ERR = np.sqrt(ERR)

    return ERR


def Basis_Functions(p,x):
    if p == 1:
        return np.array([1-x,x])
    elif p == 2:
        return np.array([(2*x-1)*(x-1),4*x*(1-x),x*(2*x-1)])


def Exact(x):
    # u = np.sin(np.pi*x)
    u = np.exp(-x**2)
    return u


def Force(x):
    # f = np.pi**2*np.sin(np.pi*x)
    f = 4*(1-x**2)*np.exp(-x**2)
    return f


Main_Program()