#!/usr/bin/env python3
# -*- coding: UTF8 -*-
#
# Copyright (C) 2018-2020 Konstantinos Poulios, Yves Renard.
#
# This file is a part of GetFEM++
#
# GetFEM++  is  free software;  you  can  redistribute  it  and/or modify it
# under  the  terms  of the  GNU  Lesser General Public License as published
# by  the  Free Software Foundation;  either version 2.1 of the License,  or
# (at your option) any later version.
# This program  is  distributed  in  the  hope  that it will be useful,  but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or  FITNESS  FOR  A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.
# You  should  have received a copy of the GNU Lesser General Public License
# along  with  this program;  if not, write to the Free Software Foundation,
# Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301, USA.
#
############################################################################
import getfem as gf
import numpy as np
gf.util_trace_level(1)

# Input data
D = 100.            # disc length
H = 2.              # disc thickness

N_D1 = 10           # number of elements in disc core
N_D2 = 3            # number of elements in disc peel

E = 5e3             # Elastic Young's modulus
nu = 0.4            # Poisson's ratio

pmax = 1.5e2        # Maximum applied pressure
steps = 300

disp_fem_order = 2        # displacements finite element order
n_fem_order = 2           # thickness vector finite element order

integration_degree = 5    # 9 gauss points per quad

# mesh generation
mesh = gf.Mesh("import", "structured_ball",
               "GT='GT_QK(2,2)';ORG=[0,0];SIZES=[%f];NSUBDIV=[%i,%i];SYMMETRIES=0"
               % (D/2., N_D1, N_D2))
mesh.transform([[1,0],[0,1],[0,0]])
# mesh regions
DIR_BOUNDARY = 6
mesh.set_region(DIR_BOUNDARY, mesh.outer_faces(2))

# FEM
mfu = gf.MeshFem(mesh, 3)
mfu.set_classical_fem(disp_fem_order)
mfdir = mfu

mfn = gf.MeshFem(mesh, 3)
mfn.set_classical_fem(n_fem_order)

mfout = gf.MeshFem(mesh)
mfout.set_classical_discontinuous_fem(disp_fem_order)

# numerical integration
mim = gf.MeshIm(mesh, integration_degree)

# Model
md = gf.Model("real")
md.add_fem_variable("u", mfu)                     # displacements variable
md.add_fem_variable("n", mfn)                     # deformed normal vector variable
md.set_variable("n", md.interpolation("[0,0,1]", mfn, -1))
md.add_initialized_data("gamma", 0.)              # numerical continuation load multiplier
md.add_initialized_data("pmax", pmax)             # maximum pressure
md.add_initialized_data('kappa', E/(3*(1-2*nu)))  # Bulk modulus
md.add_initialized_data('mu', E/(2*(1+nu)))       # Shear modulus
md.add_initialized_data('H', H)

md.add_macro("F", "[1,0,0;0,1,0;0,0,0]+Grad(u)+n@[0,0,1]")
md.add_nonlinear_term(mim, "H*0.5*kappa*sqr(log(Det(F)))+"
                           "H*0.5*mu*(pow(Det(F),-2/3)*Norm_sqr(F)-3)")
md.add_nonlinear_term(mim, "gamma*pmax*Det(F)*((Inv(F)'*[0,0,-1]).Test_u)")

md.add_Dirichlet_condition_with_multipliers(mim, "u", mfdir, DIR_BOUNDARY)

# Numerical continuation
S = gf.ContStruct(md, "gamma", 1./mfu.nbdof(), "variable_name", "u",
                      "h_init", 1e-1,          "h_max", 1e0,
                      "h_min", 1e-4,           "h_dec",  1/np.sqrt(10.),
                      "h_inc",  np.sqrt(10.),  "lsolver", "mumps",
                      "max_iter", 10,          "thr_iter", 6,
                      "max_res", 1e-5,         "max_diff", 1e-8,
                      "min_cos", 0.999,        "singularities", 0, "very_noisy")

print("Displacement dofs: %i\nTotal model dofs: %i" % (mfu.nbdof(), md.nbdof()))
md.add_macro("sigmaD", "mu*pow(Det(F),-5./3.)*Deviator(Left_Cauchy_Green(F))")
md.add_macro("sigmaH", "kappa*log(Det(F))/Det(F)*Id(3)")
with open("gf_membrane.dat", "w") as f:
   for step in range(steps+1):
      if md.variable("gamma") >= 1.:
         break # simulation completed
      print("SOLVING STEP %i" % step)
      if step == 0:
         md.solve("noisy", "lsolver", "mumps", "max_iter", 60, "max_res", 1e-7, #)[0]
                  "lsearch", "simplest", "alpha max ratio", 10., "alpha min", 0.1,
                  "alpha mult", 0.6, "alpha threshold res", 5e3)[0]
         init_dir = 1.
         gamma = 0.
         x = md.from_variables()
         [tx, tgamma, h] = S.init_Moore_Penrose_continuation(x, gamma, init_dir)
         limit_points = []
      else:
         x = md.from_variables()
         x, gamma, tx, tgamma, h, h0, sing_label = \
           S.Moore_Penrose_continuation(x, gamma, tx, tgamma, h)
         if h == 0:
            break # Continuation has failed
         elif sing_label == "limit point":
            limit_points.append(step)
         md.to_variables(x)

      VM = md.local_projection(mim, "sqrt(1.5)*Norm(sigmaD)", mfout)
      output = (mfout, VM, "Von Mises Stress")
      for ij in ((1,1),(2,2),(3,3),(1,2),(2,3),(3,1)):
        SIGMA = md.local_projection(mim, "sigmaH(%i,%i)+sigmaD(%i,%i)" % (ij + ij), mfout)
        output += (mfout, SIGMA, "Cauchy Stress %i %i" % ij)
      output = (mfu, md.variable("u"), "Displacements",
                mfn, md.variable("n"), "Thickness vector") + output

      mfout.export_to_vtk("gf_membrane_%i.vtk" % step, *output)
      f.write(("step=%i gamma=%e\n") % (step, md.variable("gamma")))
      f.flush()