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

from computeEquationNumber import computeEquationNumber
from assembleGlobalStiffnessMatrix import assembleInitialGlobalMaterialStiffnessMatrix
from plotter import plotDeformedStructure2d, plotForceDisplacement, plotInternalForces2dBeam
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=2

#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='linear' # 'linear' or 'corotational'
geoTransfBeam='linear'


# 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)
KTot_global=KMaterial_global


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

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.1
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)


# Force or displacement control integrator FIXED STEP SIZE
integrator='forceControl'
lambda_max=4803400*0.5
nIncrements=500
DeltaLambdaBar=lambda_max/nIncrements
tol=1e-4
nIterMax=1000

# integrator='displacementControl'
# uMax=1350
# nIncrements=500
# DeltaVBar=uMax/nIncrements
# q_ctrDof=nDof*(nEl_column+1)-3; # Control DOF
# tol=1e-4
# nIterMax=1000

if integrator=='forceControl':
    aIntegrator=np.zeros(nDofTot)
    bIntegrator=1
    cIntegrator=np.array([DeltaLambdaBar,0])
elif integrator=='displacementControl':
    aIntegrator=np.zeros(nDofTot)
    aIntegrator[q_ctrDof]=1
    bIntegrator=0
    cIntegrator=np.array([DeltaVBar,0])


## Solve the system
for n in range(nIncrements):
    converged=False
    i=1

    # Iteration
    while not converged:

        DeltaV_u[freeDofs] = np.linalg.solve(KTot_global[np.ix_(freeDofs, freeDofs)], -F_unb[freeDofs])
        DeltaV_f[freeDofs] = np.linalg.solve(KTot_global[np.ix_(freeDofs, freeDofs)], F_ext[freeDofs])
        if i==1:
            DeltaLambda=(cIntegrator[0]-np.dot(aIntegrator,DeltaV_u))/(np.dot(aIntegrator,DeltaV_f)+bIntegrator)
        else:   
            DeltaLambda=(cIntegrator[1]-np.dot(aIntegrator,DeltaV_u))/(np.dot(aIntegrator,DeltaV_f)+bIntegrator)


        lambdaIntegrator += DeltaLambda
        DeltaV = DeltaV_u + DeltaLambda * DeltaV_f
        v += DeltaV


        #Element state determination
        F_int=np.zeros(nDofTot)
        KTot_global=np.zeros((nDofTot,nDofTot))
        QLocal_elem_step = [] 

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

            QLocal_elem_step.append(QLocal_elem.copy())

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

        # Compute unbalanced force vector
        F_ext_tot+=DeltaLambda*F_ext
        F_unb=F_int-F_ext_tot

        # Check convergence
        test=np.linalg.norm(F_unb[freeDofs])
        if test<tol:
            converged=True
            v_Array.append(v.copy())
            F_int_Array.append(F_int.copy())
            lambda_Vector.append(lambdaIntegrator)
            QLocal_elem_list.append(QLocal_elem_step)

        if i==nIterMax and not converged:
            raise ValueError("Failed to converge")
        
        i+=1


# Convert lists to NumPy arrays before indexing
v_Array = np.array(v_Array)
F_int_Array = np.array(F_int_Array)
lambda_Vector = np.array(lambda_Vector)
QLocal_elem_list = np.array(QLocal_elem_list)


# Plot the load-displacement results
# plotForceDisplacement(v_Array[:,nDof*(nEl_column+1)-3],F_int_Array[:,nDof*(nEl_column+1)-3]/1e3,theTitle='forceDisplacement')


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


# Plot internal forces
plotInternalForces2dBeam(AllElement_data,QLocal_elem_list[-1,:],theTitle='M_linear')


test=1