'''
Main program for running Galerkin Least Squares method for the finite element method.
All elemental choices are kept in the "Elements" folder and are to be imported modularly and they are class based
Mesh numbering starts at 0


'''


import matplotlib.pyplot as plt
import numpy.linalg as npl
#imports from other files
from elements import *
from assembly import assemble, assemble_ijv
from read_mesh import read_in
from EBC import Apply_EBC


# Mesh info read in


mesh = '../mesh/basic.txt'
connect = '../mesh/basic_con.txt'
sides = '../mesh/basic_sides.txt'

mesh,con_mat,top,left,bottom,right = read_in(mesh,connect,sides)
#print(bottom)
#running of code!


dim = 5
K,F = assemble(con_mat,mesh)
#i_vals,j_vals,v_vals,F = assemble_ijv(con_mat,mesh)
#BC_vals = np.zeros(dim) #set top EBC to zero
#BC_vals = [1,1,1]
#boundary = bottom
BC_vals = [2,4,6,8,10,0,0,0,0,0]

boundary = bottom+top

K_EBC,F_EBC,N_wo = Apply_EBC(K,F,mesh,boundary,BC_vals)

#Basic - very basic - solving
Sol = npl.solve(K_EBC,F_EBC)
#Sol = ssl.spsolve(K_EBC,F_EBC)
#print(Sol)

#plt.plot(N_wo,Sol)
#plt.show()
#Put back
Mat_sol = np.zeros((dim,dim))
count1= 0
count2 = 0
BC = 0
for i in range(dim):
    for j in range(dim):
        if count1 in boundary:
            Mat_sol[i][j] = BC_vals[BC]
            BC += 1
        else:
            Mat_sol[i][j] = Sol[count2]
            count2 += 1
        count1 += 1

fig,ax=plt.subplots()

x = np.linspace(0, 1, dim)
y = np.linspace(0, 1, dim)
X,Y = np.meshgrid(x,y, indexing='xy')
z = plt.contourf(X ,Y , Mat_sol)
fig.colorbar(z,ax=ax)
plt.show()

'''
#print(Sol)
#now to put back IC to solutions
Sol_final = np.array(np.zeros(len(mesh)))
count = 0
for i in range(len(Sol_final)):
    if (i in top):
        pass #leave as zero like boundary...
    else:
        Sol_final[i] = Sol[count] #next number in solution matrix
        count += 1

#print(Sol_final)
#need to turn solution into matrix now

sol_mat = np.zeros((dim,dim))
count  = 0
for i in range(dim):
    for j in range(dim):
        sol_mat[i][j] = Sol_final[count]
        count +=1

fig,ax=plt.subplots()

x = np.linspace(0, 1, dim)
y = np.linspace(0, 1, dim)
X,Y = np.meshgrid(x,y, indexing='xy')
z = plt.contourf(X ,Y , sol_mat)
fig.colorbar(z,ax=ax)
plt.show()
'''