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=70000 #Young's modulus


#Node coordinates
ptA=np.array([0,0])*1e3
ptB=np.array([8,0])*1e3
ptC=np.array([4,3])*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])
element3=np.array([1,2])
connectivity=np.array([element1,element2,element3])
nElements=connectivity.shape[0]


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

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]]
elem3=['2dTruss',connectivity[2],coordinates[connectivity[2,0]],coordinates[connectivity[2,1]],E,areas[2]]
element_data.append(elem1)  
element_data.append(elem2)
element_data.append(elem3)


# 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([0,1,3])
freeDof=np.setdiff1d(allDof, fixedDof)
f_global[5]=-2000e3


# Solve the system to get forces and displacements
u_global[freeDof]=np.linalg.solve(K_global[np.ix_(freeDof,freeDof)],f_global[freeDof])
f_global[fixedDof]=np.dot(K_global[np.ix_(fixedDof,allDof)],u_global)


# Print results
print('R1_x:',f_global[0]/1e3)
print('R1_y:',f_global[1]/1e3)
print('R2_y:',f_global[3]/1e3)

# 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