import sys
import numpy as np
import math
from matplotlib import pyplot as plt
## Input file syntax:
##    *Node
##    1, 0.0, 0.0
##    2, 0.0, 1.0
##    3, 1.0, 1.0
##    4, 1.0, 0.0
##    *Element
##    1, 1, 2, 3, 4
##    *Step
##    *Boundary
##    1, 1, 2, 0.0          # nodeId, dof1, dof2, value
##    2, 1, 1, 0.0
##    3, 1, 1, 0.01
##    4, 1, 1, 0.01
##    4, 2, 2, 0.0

class FEM:
    
    def shape(self,xi):
        """Shape functions for a 4-node, isoparametric element
            N_i(xi,eta) where i=[1,2,3,4]
            Input: 1x2,  Output: 1x4"""
        xi,eta = tuple(xi)
        N = [(1.0-xi)*(1.0-eta), (1.0+xi)*(1.0-eta), (1.0+xi)*(1.0+eta), (1.0-xi)*(1.0+eta)]
        return 0.25 * np.array(N)
    def gradshape(self,xi):
        """Gradient of the shape functions for a 4-node, isoparametric element.
            dN_i(xi,eta)/dxi and dN_i(xi,eta)/deta
            Input: 1x2,  Output: 2x4"""
        xi,eta = tuple(xi)
        dN = [[-(1.0-eta),  (1.0-eta), (1.0+eta), -(1.0+eta)],
              [-(1.0-xi), -(1.0+xi), (1.0+xi),  (1.0-xi)]]
        return 0.25 * np.array(dN)
    def local_error(self,str):
        print("*** ERROR ***")
        print(str)
        sys.exit(3)
    def read_inp_file(self,inpFileName, nodes, conn, boundary):
        print('\n** Reading input file')
        inpFile = open(inpFileName, 'r')
        lines = inpFile.readlines()
        inpFile.close()
        state = 0
        for line in lines:
            line = line.strip()
            if len(line) <= 0: continue
            if line[0] == '*':
                state = 0
            if line.lower() == "*node":
                state = 1
                continue
            if line.lower() == "*element":
                state = 2
                continue
            if line.lower() == "*boundary":
                state = 3
                continue
            if state == 0:
                continue
            if state == 1:
                # read nodes
                values = line.split(",")
                if len(values) != 3:
                    self.local_error("A node definition needs 3 values")
                nodeNr = int(values[0]) - 1  # zero indexed
                xx = float(values[1])
                yy = float(values[2])
                nodes.append([xx,yy])   # assume the nodes are ordered 1, 2, 3...
                continue
            if state == 2:
                # read elements
                values = line.split(",")
                if len(values) != 5:
                    self.local_error("An element definition needs 5 values")
                elemNr = int(values[0])
                n1 = int(values[1]) - 1  # zero indexed
                n2 = int(values[2]) - 1
                n3 = int(values[3]) - 1
                n4 = int(values[4]) - 1
                #conn.append([n1, n2, n3, n4]) # assume elements ordered 1, 2, 3
                conn.append([n1, n4, n3, n2]) # assume elements ordered 1, 2, 3
                continue
            if state == 3:
                # read displacement boundary conditions
                values = line.split(",")
                if len(values) != 4:
                    self.local_error("A displacement boundary condition needs 4 values")
                nodeNr = int(values[0]) - 1  # zero indexed
                dof1 = int(values[1])
                dof2 = int(values[2])
                val = float(values[3])
                if dof1 == 1:
                    boundary.append([nodeNr,1,val])
                if dof2 == 2:
                    boundary.append([nodeNr,2,val])
                continue
    
    def fem2d(self,FILE_PATH,plot_type):
        ##
        ## Main Program
        ##
        nodes = []
        conn = []
        boundary = []
        if FILE_PATH == None: self.local_error('No input file provided.')
        print('Input file:', FILE_PATH)
        avail=['u1','u2','s11','s22','s12','e11','e22','e12','all']
        if plot_type not in avail : self.local_error('Invalid Request')
        self.read_inp_file(FILE_PATH, nodes, conn, boundary)
        nodes = np.array(nodes)
        num_nodes = len(nodes)
        
        print('   number of nodes:', len(nodes))
        print('   number of elements:', len(conn))
        print('   number of displacement boundary conditions:', len(boundary))
    
        ###############################
        # Plane-strain material tangent (see Bathe p. 194)
        # C is 3x3
        E = 100.0
        v = 0.3
        C = E/(1.0+v)/(1.0-2.0*v) * np.array([[1.0-v, v, 0.0], [v, 1.0-v, 0.0], [0.0, 0.0, 0.5-v]])
        ###############################
        # Make stiffness matrix
        # if N is the number of DOF, then K is NxN
        K = np.zeros((2*num_nodes, 2*num_nodes))    # square zero matrix
        # 2x2 Gauss Quadrature (4 Gauss points)
        # q4 is 4x2
        q4 = np.array([[-1,-1],[1,-1],[-1,1],[1,1]]) / math.sqrt(3.0)
        print('\n** Assembling stiffness matrix')
        # strain in an element: [strain] = B    U
        #                        3x1     = 3x8  8x1
        #
        # strain11 = B11 U1 + B12 U2 + B13 U3 + B14 U4 + B15 U5 + B16 U6 + B17 U7 + B18 U8
        #          = B11 u1          + B13 u1          + B15 u1          + B17 u1
        #          = dN1/dx u1       + dN2/dx u1       + dN3/dx u1       + dN4/dx u1
        B = np.zeros((3,8))
        # conn[0] is node numbers of the element
        for c in conn:     # loop through each element
            # coordinates of each node in the element
            # shape = 4x2
            # for example:
            #    nodePts = [[0.0,   0.0],
            #               [0.033, 0.0],
            #               [0.033, 0.066],
            #               [0.0,   0.066]]
            nodePts = nodes[c,:]
            Ke = np.zeros((8,8))    # element stiffness matrix is 8x8
            for q in q4:            # for each Gauss point
                # q is 1x2, N(xi,eta)
                dN = self.gradshape(q)       # partial derivative of N wrt (xi,eta): 2x4
                J  = np.dot(dN, nodePts).T # J is 2x2
                dN = np.dot(np.linalg.inv(J), dN)    # partial derivative of N wrt (x,y): 2x4
                # assemble B matrix  [3x8]
                B[0,0::2] = dN[0,:]
                B[1,1::2] = dN[1,:]
                B[2,0::2] = dN[1,:]
                B[2,1::2] = dN[0,:]
                # element stiffness matrix
                Ke += np.dot(np.dot(B.T,C),B) * np.linalg.det(J)
            # Scatter operation
            for i,I in enumerate(c):
                for j,J in enumerate(c):
                    K[2*I,2*J]     += Ke[2*i,2*j]
                    K[2*I+1,2*J]   += Ke[2*i+1,2*j]
                    K[2*I+1,2*J+1] += Ke[2*i+1,2*j+1]
                    K[2*I,2*J+1]   += Ke[2*i,2*j+1]
        ###############################
        # Assign nodal forces and boundary conditions
        #    if N is the number of nodes, then f is 2xN
        f = np.zeros((2*num_nodes))          # initialize to 0 forces
        # How about displacement boundary conditions:
        #    [k11 k12 k13] [u1] = [f1]
        #    [k21 k22 k23] [u2]   [f2]
        #    [k31 k32 k33] [u3]   [f3]
        #
        #    if u3=x then
        #       [k11 k12 k13] [u1] = [f1]
        #       [k21 k22 k23] [u2]   [f2]
        #       [k31 k32 k33] [ x]   [f3]
        #   =>
        #       [k11 k12 k13] [u1] = [f1]
        #       [k21 k22 k23] [u2]   [f2]
        #       [  0   0   1] [u3]   [ x]
        #   the reaction force is
        #       f3 = [k31 k32 k33] * [u1 u2 u3]
        for i in range(len(boundary)):  # apply all boundary displacements
            nn  = boundary[i][0]
            dof = boundary[i][1]
            val = boundary[i][2]
            j = 2*nn
            if dof == 2: j = j + 1
            K[j,:] = 0.0
            K[j,j] = 1.0
            f[j] = val
        ###############################
        #print('\n** Solve linear system: Ku = f')    # [K] = 2N x 2N, [f] = 2N x 1, [u] = 2N x 1
        u = np.linalg.solve(K, f)
        ###############################
        #print('\n** Post process the data')
        # (pre-allocate space for nodal stress and strain)
        node_strain = []
        node_stress = []
        for ni in range(len(nodes)):
            node_strain.append([0.0, 0.0, 0.0])
            node_stress.append([0.0, 0.0, 0.0])
        node_strain = np.array(node_strain)
        node_stress = np.array(node_stress)
    
        print(f'min displacements: u1={min(u[0::2]):.4g}, u2={min(u[1::2]):.4g}')
        print(f'max displacements: u1={max(u[0::2]):.4g}, u2={max(u[1::2]):.4g}')
    
        emin = np.array([ 9.0e9,  9.0e9,  9.0e9])
        emax = np.array([-9.0e9, -9.0e9, -9.0e9])
        smin = np.array([ 9.0e9,  9.0e9,  9.0e9])
        smax = np.array([-9.0e9, -9.0e9, -9.0e9])
        for c in conn:    # for each element (conn is Nx4)
                                            # c is like [2,5,22,53]
            nodePts = nodes[c,:]            # 4x2, eg: [[1.1,0.2], [1.2,0.3], [1.3,0.4], [1.4, 0.5]]
            for q in q4:                    # for each integration pt, eg: [-0.7,-0.7]
                dN = self.gradshape(q)                    # 2x4
                J  = np.dot(dN, nodePts).T            # 2x2
                dN = np.dot(np.linalg.inv(J), dN)    # 2x4
                B[0,0::2] = dN[0,:]                    # 3x8
                B[1,1::2] = dN[1,:]
                B[2,0::2] = dN[1,:]
                B[2,1::2] = dN[0,:]
    
                UU = np.zeros((8,1))                # 8x1
                UU[0] = u[2*c[0]]
                UU[1] = u[2*c[0] + 1]
                UU[2] = u[2*c[1]]
                UU[3] = u[2*c[1] + 1]
                UU[4] = u[2*c[2]]
                UU[5] = u[2*c[2] + 1]
                UU[6] = u[2*c[3]]
                UU[7] = u[2*c[3] + 1]
                # get the strain and stress at the integration point
                strain = np.matmul(B,UU) #B @ UU        # (B is 3x8) (UU is 8x1)         => (strain is 3x1)
                stress = np.matmul(C,strain) #C @ strain    # (C is 3x3) (strain is 3x1)     => (stress is 3x1)
                emin[0] = min(emin[0], strain[0][0])
                emin[1] = min(emin[1], strain[1][0])
                emin[2] = min(emin[2], strain[2][0])
                emax[0] = max(emax[0], strain[0][0])
                emax[1] = max(emax[1], strain[1][0])
                emax[2] = max(emax[2], strain[2][0])
    
                node_strain[c[0]][:] = strain.T[0]
                node_strain[c[1]][:] = strain.T[0]
                node_strain[c[2]][:] = strain.T[0]
                node_strain[c[3]][:] = strain.T[0]
                node_stress[c[0]][:] = stress.T[0]
                node_stress[c[1]][:] = stress.T[0]
                node_stress[c[2]][:] = stress.T[0]
                node_stress[c[3]][:] = stress.T[0]
                smax[0] = max(smax[0], stress[0][0])
                smax[1] = max(smax[1], stress[1][0])
                smax[2] = max(smax[2], stress[2][0])
                smin[0] = min(smin[0], stress[0][0])
                smin[1] = min(smin[1], stress[1][0])
                smin[2] = min(smin[2], stress[2][0])
        print(f'   min strains: e11={emin[0]:.4g}, e22={emin[1]:.4g}, e12={emin[2]:.4g}')
        print(f'   max strains: e11={emax[0]:.4g}, e22={emax[1]:.4g}, e12={emax[2]:.4g}')
        print(f'   min stresses  s11={smin[0]:.4g}, s22={smin[1]:.4g}, s12={smin[2]:.4g}')
        print(f'   max stresses:  s11={smax[0]:.4g}, s22={smax[1]:.4g}, s12={smax[2]:.4g}')
        ###############################
        print('\n** Plotting displacement')
        xvec = []
        yvec = []
        res=[]
        resu1  = []
        resu2  = []
        ress11 = []
        ress22 = []
        ress12 = []
        rese11 = []
        rese22 = []
        rese12 = []
        
        for ni,pt in enumerate(nodes):
            xvec.append(pt[0] + u[2*ni])
            yvec.append(pt[1] + u[2*ni+1])
            resu1.append(u[2*ni]) 
            resu2.append(u[2*ni+1])
            ress11.append(node_stress[ni][0]) 
            ress22.append(node_stress[ni][1]) 
            ress12.append(node_stress[ni][2]) 
            rese11.append(node_strain[ni][0]) 
            rese22.append(node_strain[ni][1]) 
            rese12.append(node_strain[ni][2]) 
            
        if plot_type=='u1':  res=resu1               # x-disp
        elif plot_type=='u2':  res=resu2                # y-disp
        elif plot_type=='s11': res=ress11       # s11
        elif plot_type=='s22': res=ress22       # s22
        elif plot_type=='s12': res=ress12       # s12
        elif plot_type=='e11': res=rese11       # e11
        elif plot_type=='e22': res=rese22        # e22
        elif plot_type=='e12': res=rese12       # e12
        
        if plot_type !='all':
            tri = []
            for c in conn:
                tri.append( [c[0], c[1], c[2]] )
                tri.append( [c[0], c[2], c[3]] )
            #print(f"LENGTH OF XVEC: {len(xvec)}")
            #print(f"LENGTH OF YVEC: {len(yvec)}")
            #print(f"LENGTH OF RES: {len(res)}")
            t = plt.tricontourf(xvec, yvec, res, triangles=tri, levels=14, cmap=plt.cm.jet)
            #plt.scatter(xvec, yvec, marker='o', c='b', s=0.5) # (plot the nodes)
            plt.grid()
            plt.colorbar(t)
            plt.title(plot_type)
            plt.axis('equal')
            plt.show()
            print('Done.')
        else:
            output=[resu1,resu2,ress11,ress22,ress12,rese11,rese22,rese12]
            id=0
            for res in output:
                tri = []
                for c in conn:
                    tri.append( [c[0], c[1], c[2]] )
                    tri.append( [c[0], c[2], c[3]] )
                #print(f"LENGTH OF XVEC: {len(xvec)}")
                #print(f"LENGTH OF YVEC: {len(yvec)}")
                #print(f"LENGTH OF RES: {len(res)}")
                t = plt.tricontourf(xvec, yvec, res, triangles=tri, levels=14, cmap=plt.cm.jet)
                #plt.scatter(xvec, yvec, marker='o', c='b', s=0.5) # (plot the nodes)
                plt.grid()
                plt.colorbar(t)
                plt.title(avail[id])
                plt.axis('equal')
                plt.show()
                id+=1
            print('Done.')
                
        