import numpy as np
from computeEquationNumber import computeEquationNumber
from assembleGlobalStiffnessMatrix import assembleGlobalStiffnessMatrix
from plotter import plotDeformedStructure2d, plotInternalForces2dTruss


#Dimension and degrees of freedom of the problem
nDim=2
nDof=2 #nDof per node
nNodesPerElement=2


#Matrerial properties
E=200000 #Young's modulus


#Node coordinates
ptA=np.array([0,0])*1e3
ptB=np.array([2,2])*1e3
ptC=np.array([-4,2])*1e3
coordinates=np.array([ptA,ptB,ptC])
nNodes=coordinates.shape[0]
nDofTot=nNodes*nDof 


# Connectivity 
element1=np.array([0,1])    
element2=np.array([0,2])
connectivity=np.array([element1,element2])
nElements=connectivity.shape[0]


# Elements properties
areas=np.array([1000,4000])

element_data = [] #contains the properties of each element
#2dTruss: '2dTruss', nodes, coordNode1, coordNode2, E A
elem1=['2dTruss',connectivity[0],coordinates[connectivity[0,0]],coordinates[connectivity[0,1]],E,areas[0]]
elem2=['2dTruss',connectivity[1],coordinates[connectivity[1,0]],coordinates[connectivity[1,1]],E,areas[1]]
element_data.append(elem1)  
element_data.append(elem2)


# Equation number matrix
numEquations=computeEquationNumber(nDof,element_data)


# Assemble global stiffness matrix
K_global=assembleGlobalStiffnessMatrix(nDofTot,numEquations,element_data)


# Assemble load and displacemement vectors
f_global=np.zeros((nDofTot,1))
u_global=np.zeros((nDofTot,1))


# Apply boundary conditions
allDof=np.arange(nDofTot)
fixedDof=np.array([2,3,4,5])
freeDof=np.setdiff1d(allDof, fixedDof)
u_global[1]=-5


# Solve the system to get forces and displacements
f_global=np.dot(K_global,u_global)


# Print results
print('f_x:',f_global[0]/1e3)
print('f_y:',f_global[1]/1e3)
print('angle alpha:',np.arctan(f_global[1]/f_global[0])*180/np.pi)


# Plot deformed structure
plotDeformedStructure2d(coordinates,connectivity,u_global,scaleFactor=1e2,theTitle='deformedStructure')


# Plot internal forces
plotInternalForces2dTruss(coordinates,connectivity, E, areas, u_global, theTitle='axialLoad')


test=1