import numpy as np
from scipy.sparse.linalg import eigs

from computeEquationNumber import computeEquationNumber
from assembleGlobalStiffnessMatrix import assembleInitialGlobalMaterialStiffnessMatrix
from plotter import plotDeformedStructure2d, plotInternalForces2dTruss
from elemStateDetermin import elemStateDetermin


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


#Matrerial properties
E=200000 #Young's modulus


#Frame properties
h=7*1e3
l1=6*1e3
l2=12*1e3
Ab=16000
Ib=1.186*10**9
Ac=8600
Ic=1.032*10**8
nEl_column=50

#Node coordinates
pts_leftCol = []
pts_middleCol = []
pts_rightCol = []
for i in range(nEl_column+1):
    y = i * (h / nEl_column )
    pts_leftCol.append((0, y))
    pts_middleCol.append((l1, y))
    pts_rightCol.append((l1+l2, y))
coordinates=np.array(pts_leftCol+pts_middleCol+pts_rightCol)
nNodes=coordinates.shape[0]
nDofTot=nNodes*nDof 


# Connectivity 
elems_leftCol = []
elems_middleCol = []
elems_rightCol = []
for i in range(nEl_column):
    elems_leftCol.append((i+1, i+2))
    elems_middleCol.append((i+nEl_column+2, i+nEl_column+3))
    elems_rightCol.append((i+2*nEl_column+3, i+2*nEl_column+4))
elems_beam=[(nEl_column+1, 2*(nEl_column+1)),(2*(nEl_column+1), 3*(nEl_column+1))]
connectivity=np.array(elems_leftCol+elems_middleCol+elems_rightCol+elems_beam)-1
# print(coordinates)
# print(connectivity)
nElements=connectivity.shape[0]


# Geometric transformation local <-> basic reference frame
geoTransfCols='corotational'
geoTransfBeam='corotational'


# Elements properties
AllElement_data = [] #contains the properties of each element
#2dTruss: '2dTruss', nodes, coordNode1, coordNode2, E, A, geoTransf
#2dBeam: '2dBeam', nodes, coordNode1, coordNode2, E, A, I, geoTransf
for i in range(3*nEl_column):
    elem=['2dElasticBeam',connectivity[i],coordinates[connectivity[i,0]],coordinates[connectivity[i,1]],E,Ac,Ic,geoTransfCols]
    AllElement_data.append(elem)
for i in range(2):
    elem=['2dElasticBeam',connectivity[3*nEl_column+i],coordinates[connectivity[3*nEl_column+i,0]],coordinates[connectivity[3*nEl_column+i,1]],E,Ab,Ib,geoTransfBeam]
    AllElement_data.append(elem)


# Equation number matrix
numEquations=computeEquationNumber(nDof,AllElement_data)
#print(numEquations)


# Assemble initial global material stiffness matrix
KMaterial_global=assembleInitialGlobalMaterialStiffnessMatrix(nDofTot,numEquations,AllElement_data)


# Initialize variables
v_Array=[]
F_int_Array=[]
lambda_Vector=[]

F_int=np.zeros(nDofTot)
lambdaIntegrator=0.
v=np.zeros(nDofTot)
F_unb=np.zeros(nDofTot)
F_ext_tot=np.zeros(nDofTot)

DeltaV_u=np.zeros(nDofTot)
DeltaV_f=np.zeros(nDofTot)


# Definition of the boundary conditions
alpha=0.
F_ext=np.zeros(nDofTot)
F_ext[nDof*(nEl_column+1)-3]=alpha
F_ext[nDof*(nEl_column+1)-2]=-0.5
F_ext[2*nDof*(nEl_column+1)-2]=-1.0
F_ext[3*nDof*(nEl_column+1)-2]=-1.0
allDofs=np.arange(nDofTot)
fixedDofs=np.array([0,1,2,nDof*(nEl_column+1),nDof*(nEl_column+1)+1,nDof*(nEl_column+1)+2,2*nDof*(nEl_column+1),2*nDof*(nEl_column+1)+1,2*nDof*(nEl_column+1)+2])
freeDofs=np.setdiff1d(allDofs, fixedDofs)


# Apply reference load in one step
DeltaLambdaBar=1





## Determine the critical buckling load
DeltaLambda=DeltaLambdaBar

# Update of load and displacements
DeltaV_u[freeDofs] = np.linalg.solve(KMaterial_global[np.ix_(freeDofs, freeDofs)], -F_unb[freeDofs])
DeltaV_f[freeDofs] = np.linalg.solve(KMaterial_global[np.ix_(freeDofs, freeDofs)], F_ext[freeDofs])
lambdaIntegrator += DeltaLambda
DeltaV = DeltaV_u + DeltaLambda * DeltaV_f
v += DeltaV


#Element state determination
F_int=np.zeros(nDofTot)
KMaterial_global=np.zeros((nDofTot,nDofTot))
KGeom_global=np.zeros((nDofTot,nDofTot))

for elemIdx, elem in enumerate(AllElement_data):
    #Perform element state determination
    QGlobal_elem, KMaterialGlobal_elem, KGeomGlobal_elem, QLocal_elem = elemStateDetermin(elemIdx,elem,numEquations, v)

    # Assemble element state determination results
    dof_elem=np.array(numEquations[elemIdx]).astype(int)
    F_int[dof_elem]+=QGlobal_elem
    KMaterial_global[np.ix_(dof_elem,dof_elem)]+=KMaterialGlobal_elem
    KGeom_global[np.ix_(dof_elem,dof_elem)]+=KGeomGlobal_elem

# Resolution of the eigenvalue problem
num_eigenvalues = min(7, freeDofs.size)

Kg_sub = KGeom_global[np.ix_(freeDofs, freeDofs)]
Ke_sub = KMaterial_global[np.ix_(freeDofs, freeDofs)]

eigenvalues, eigenvectors = eigs(Kg_sub, M=Ke_sub, k=num_eigenvalues, which='LM')
eigenvalues= np.reciprocal(eigenvalues)
ind1 = np.isclose(np.imag(eigenvalues), 0)
B = eigenvalues[ind1]
V2 = eigenvectors[:, ind1]

Pcr_index = np.argmin(np.abs(B))
Pcr = np.abs(B[Pcr_index])
V2 = V2[:, Pcr_index]


# Print results
print('Pcr=',Pcr/1e3, 'kN')


# Plot deformed structure
u_global=np.zeros(nDofTot)
u_global[freeDofs]=V2
plotDeformedStructure2d(nDof,coordinates,connectivity,u_global,scaleFactor=1e5,theTitle='buckledShape')


test=1