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=3 #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])
element3=np.array([0,0])
element4=np.array([1,1])
element5=np.array([2,2])
connectivity=np.array([element1,element2,element3,element4,element5])
nElements=connectivity.shape[0]


# Elements properties
areas=np.array([1000,4000])
inertias=np.array([3000,12000])
kSprings=np.array([50000.0,1000000.0,1000000.0])

element_data = [] #contains the properties of each element
#2dTruss: '2dTruss', nodes, coordNode1, coordNode2, E A
#2dBeam: '2dBeam', nodes, coordNode1, coordNode2, E A I
#spring: 'spring', node, globalDof, k       where globalDof=0 for X, 1 for Y, 2 for rotation
elem1=['2dBeam',connectivity[0],coordinates[connectivity[0,0]],coordinates[connectivity[0,1]],E,areas[0],inertias[0]]
elem2=['2dBeam',connectivity[1],coordinates[connectivity[1,0]],coordinates[connectivity[1,1]],E,areas[1],inertias[1]]
elem3=['spring',connectivity[2],1,kSprings[0]]
elem4=['spring',connectivity[3],2,kSprings[1]]
elem5=['spring',connectivity[4],2,kSprings[2]]
element_data.append(elem1)  
element_data.append(elem2)
element_data.append(elem3)
element_data.append(elem4)
element_data.append(elem5)


# 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([3,4,6,7])
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


test=1