#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import getfem as gf
gf.util('trace level', 1)

def SolveBasic3DMecaProblem(version):
    N1 = 2; N2 = 4; h = 20.; DX = 1./N1; DY = (1.*h)/N2;
    m = gf.Mesh('cartesian', np.arange(-0.5, 0.5+DX,DX), np.arange(0., h+DY,DY),
                np.arange(-1.5, 1.5+3*DX,3*DX))
    mfu  = gf.MeshFem(m, 3)
    mfu.set_fem(gf.Fem('FEM_QK(3,2)'))
    mim = gf.MeshIm(m, gf.Integ('IM_GAUSS_PARALLELEPIPED(3,4)'))

    ftop = m.outer_faces_with_direction([0.,  1., 0.], 0.5)
    fbot = m.outer_faces_with_direction([0., -1., 0.], 0.5)

    m.set_region(1, ftop)
    m.set_region(2, fbot)
    m.set_region(3, np.append(ftop,fbot,axis=1));

    model = gf.Model('real')
    model.add_fem_variable('u', mfu)

    model.add_initialized_data('paramsIMR', [1,1])
    if version=="Predefined":
        lawname = 'Incompressible Mooney Rivlin'
        model.add_finite_strain_elasticity_brick(mim, lawname, 'u', 'paramsIMR')
    elif version=="AssemblyPredefined":
        model.add_nonlinear_generic_assembly_brick(mim, "Incompressible_Mooney_Rivlin_potential(Grad_u, [paramsIMR(1);paramsIMR(2)])")
    elif version=="GWFL":
        model.add_macro("F", "Grad_u+Id(3)")
        model.add_nonlinear_generic_assembly_brick(mim, 
                          "paramsIMR(1)* ( Matrix_j1(Right_Cauchy_Green(F)) - 3 )\
                          + paramsIMR(2)* ( Matrix_j2(Right_Cauchy_Green(F)) - 3 )")
    else:
        raise Exception("Version not available")

    #Add finite strain incompressibility condition
    mfp = gf.MeshFem(m,1)
    mfp.set_classical_discontinuous_fem(1)
    model.add_fem_variable('p', mfp)
    model.add_finite_strain_incompressibility_brick(mim, 'u', 'p')

    model.add_initialized_data('f', [0,-0.1,0])
    model.add_source_term_brick(mim, 'u', 'f') 

    model.add_initialized_data('u_d', [0,0,0])
    model.add_Dirichlet_condition_with_multipliers(mim, 'u', 1, 3, "u_d")

    model.variable_list()

    maxIter=50
    solverIt=model.solve('very noisy','max_iter',maxIter, 'max_res',1e-7, 'lsearch','simplest')

    if not solverIt[0]<maxIter:
        raise Exception("Failed to converge")

    displacements = model.variable('u')
    return displacements

def CheckConsistency_GWFL_PredefinedBrick():
    print("Checking consitency between Predefined brick and AssemblyPredefined brick")
    solutionPredefined=SolveBasic3DMecaProblem(version="Predefined")

    solutionAssemblyPredefined=SolveBasic3DMecaProblem(version="AssemblyPredefined")
    np.testing.assert_almost_equal(solutionPredefined,solutionAssemblyPredefined)

    print("Checking consitency between Predefined brick and GWFL")
    solutionGWFL=SolveBasic3DMecaProblem(version="GWFL")
    np.testing.assert_almost_equal(solutionPredefined,solutionGWFL)
if __name__ == '__main__':
    CheckConsistency_GWFL_PredefinedBrick()