#!/usr/bin/env python
# -*- coding: utf-8 -*-
import getfem as gf
import numpy as np

if __name__ =="__main__":
    #Parameters
    NX=20.0
    degree = 1
    E=5e5 #Young modulus
    Nu=0.3 #Poisson ratio
    Lambda = E*Nu/((1+Nu)*(1-2*Nu))
    Mu =E/(2*(1+Nu))
    clambdastar = 2*Lambda*Mu/(Lambda+2*Mu)
    traction_force = 25 #Surfacic load on the Neumann boundary of the body in the vertical direction
    d=2                   # Dimension of the problem
    niter = 100  # Maximum number of iterations for the solver.

    mesh=gf.Mesh('regular simplices', np.arange(0,4+1/NX,1/NX), np.arange(0,1+1/NX,1/NX))
    mfu=gf.MeshFem(mesh,2) # Displacement
    mff=gf.MeshFem(mesh,1) # plot Von-mises

    #Assign mesh
    mfu.set_fem(gf.Fem('FEM_PK(2,%d)' % (degree,)));
    P=mesh.pts()

    #Boundaries definition
    ctop=(abs(P[1,:]-1) < 1e-6)
    cbot=(abs(P[1,:]) < 1e-6)
    cright=(abs(P[0,:]-4) < 1e-6)
    cleft=(abs(P[0,:]) < 1e-6)

    pidtop=gf.compress(ctop, range(0, mesh.nbpts()))
    pidbot=gf.compress(cbot, range(0, mesh.nbpts()))
    pidright=gf.compress(cright, range(0, mesh.nbpts()))
    pidleft=gf.compress(cleft, range(0, mesh.nbpts()))

    ftop=mesh.faces_from_pid(pidtop)
    fbot=mesh.faces_from_pid(pidbot)
    fright=mesh.faces_from_pid(pidright)
    fleft=mesh.faces_from_pid(pidleft)

    #Boundaries creation
    NEUMANN_BOUNDARY = 1
    DIRICHLET_BOUNDARY = 2
    FREE_BOUNDARY = 3
    DIRICHLET_BOUNDARY_2 = 4

    mesh.set_region(NEUMANN_BOUNDARY,ftop)
    mesh.set_region(DIRICHLET_BOUNDARY,fleft)
    mesh.set_region(FREE_BOUNDARY,fbot)
    mesh.set_region(DIRICHLET_BOUNDARY_2,fright)

    #Element used
    mim=gf.MeshIm(mesh, gf.Integ('IM_TRIANGLE(5)'))

    md = gf.Model('real')
    md.add_fem_variable('u', mfu)
        
    # Elasticity model 
    md.add_initialized_data('cmu', [Mu])
    md.add_initialized_data('clambda', [clambdastar])
    md.add_isotropic_linearized_elasticity_brick(mim, 'u', 'clambda', 'cmu')

    #Add neumann condition
    md.add_initialized_data('traction', [0, -traction_force])
    md.add_source_term_brick(mim, 'u', 'traction',NEUMANN_BOUNDARY)

    md.assembly()
    tangentMatrix=md.tangent_matrix()
    #print(tangentMatrix.size())
    rhs=md.rhs()

    #Boundary conditions terms
    md.add_Dirichlet_condition_with_multipliers(mim, 'u', 1, DIRICHLET_BOUNDARY)
    
    md.solve('max_res', 1E-6, 'max_iter', niter, 'very noisy')

    refIndices=mfu.basic_dof_on_region(DIRICHLET_BOUNDARY)

    #Compute nodal force from solution
    U=md.variable('u')
    residualMult=tangentMatrix.mult(U)-rhs
    #print(residualMult[refIndices])

    #Compute nodal force with multipliers
    nodalForces=md.variable('mult_on_u')
    #print(nodalForces)

    extendedLambda=np.zeros(mfu.nbdof())
    extendedLambda[refIndices]=nodalForces

    #Build New model and use nodal forces for boundary condition
    print("Build new model")
    mdNew = gf.Model('real')
    mdNew.add_fem_variable('unew', mfu)

    # Elasticity model 
    mdNew.add_initialized_data('cmu', [Mu])
    mdNew.add_initialized_data('clambda', [clambdastar])
    mdNew.add_isotropic_linearized_elasticity_brick(mim, 'unew', 'clambda', 'cmu')

    #Add neumann condition
    mdNew.add_initialized_data('traction', [0, -traction_force])
    mdNew.add_source_term_brick(mim, 'unew', 'traction',NEUMANN_BOUNDARY)

    #Boundary conditions terms
    mdNew.add_explicit_rhs('unew', residualMult)
    
    mdNew.solve('max_res', 1E-6, 'max_iter', niter, 'very noisy')
    Unew=mdNew.variable('unew')

    relaL2Error=gf.compute_L2_norm(mfu,Unew-U,mim)/gf.compute_L2_norm(mfu,U,mim)
    print("Relative L2 error: ",relaL2Error)
    mfu.export_to_pos("SimpleElasticitySolution.pos", mfu, U, 'Displacement',mfu, Unew, 'Displacement_test',mfu, extendedLambda, 'Multipliers',mfu, residualMult, 'MultipliersResidual')