#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import getfem as gf

# Import the mesh : disc
m = gf.Mesh('load', 'disc_P2_h2.mesh')

d = m.dim() # Mesh dimension

# Parameters of the model
enforceContactWith="Penalization"
E,nu=1e6,0.3
friction_coeff = 0.4  # coefficient of friction
vertical_force = 0.05 # Volumic load in the vertical direction
r = 10.0               # Augmentation parameter

niter = 100  # Maximum number of iterations for Newton's algorithm.

# Signed distance representing the obstacle
obstacle = 'y'

# Selection of the contact and Dirichlet boundaries
GAMMAC = 1
GAMMAD = 2

border = m.outer_faces()
normals = m.normal_of_faces(border)
contact_boundary = border[:,np.nonzero(normals[d-1] < -0.01)[0]]
m.set_region(GAMMAC, contact_boundary)
contact_boundary = border[:,np.nonzero(normals[d-1] > 0.01)[0]]
m.set_region(GAMMAD, contact_boundary)

# Finite element methods
u_degree = 2

mfu = gf.MeshFem(m, d)
mfu.set_classical_fem(u_degree)

mfd = gf.MeshFem(m, 1)
mfd.set_classical_fem(u_degree)

# Integration method
mim = gf.MeshIm(m, 4)
mim_friction = gf.MeshIm(m,gf.Integ('IM_STRUCTURED_COMPOSITE(IM_TRIANGLE(4),4)'))

# Volumic density of force
nbdofd = mfd.nbdof()
nbdofu = mfu.nbdof()
F = np.zeros(nbdofd*d)
F[d-1:nbdofd*d:d] = -vertical_force;

# Elasticity model
md = gf.Model('real')
md.add_fem_variable('u', mfu)
cmu = E/(2*(1+nu))  
clambda = E*nu/((1+nu)*(1-2*nu))              
clambda = 2*clambda*cmu/(clambda+2*cmu)
md.add_initialized_data('cmu', [cmu])
md.add_initialized_data('clambda', [clambda])
md.add_isotropic_linearized_elasticity_brick(mim, 'u', 'clambda', 'cmu')
md.add_initialized_fem_data('volumicload', mfd, F)
md.add_source_term_brick(mim, 'u', 'volumicload')

Ddata = np.zeros(d)
Ddata[d-1] = -5
md.add_initialized_data('Ddata', Ddata)
md.add_Dirichlet_condition_with_multipliers(mim, 'u', u_degree, GAMMAD, 'Ddata')

if enforceContactWith =="Penalization":
   md.add_initialized_data('r', [r])
   md.add_initialized_data('friction_coeff', [friction_coeff])
   OBS = mfd.eval(obstacle)
   md.add_initialized_fem_data('obstacle', mfd, OBS);
   md.add_penalized_contact_with_rigid_obstacle_brick(mim_friction, 'u', 'obstacle', 'r', 'friction_coeff', GAMMAC)
elif enforceContactWith=="Multipliers":
   mflambda= gf.MeshFem(m, 2)
   mflambda.set_classical_fem(1)
   md.add_filtered_fem_variable("lambda", mflambda,GAMMAC)
   md.add_initialized_data('r', [r])
   md.add_initialized_data('N1', [0.,-1.0])
   md.add_initialized_data('friction_coeff', [friction_coeff])
   md.add_linear_generic_assembly_brick(mim, '-lambda.(Test_u)', GAMMAC)
   md.add_linear_generic_assembly_brick(mim, '-(1/r)*lambda.Test_lambda', GAMMAC)
   md.add_nonlinear_generic_assembly_brick(mim, '(1/r)*Coulomb_friction_coupled_projection(lambda, N1, u,X(2)-u.N1, friction_coeff, r).Test_lambda', GAMMAC)
else:
   raise Exception("method not available")


md.variable_list()
# Solve the problem
md.solve('max_res', 1E-8, 'very noisy', 'max_iter', niter, 'lsearch', 'default') #, 'with pseudo potential')

U = md.get('variable', 'u')

mfd.export_to_pos('static_contact.pos', mfu, U, 'Displacement')