#!/usr/bin/env python

# This code was adapted from http://sfepy.org/doc-devel/mat_optim.html.
"""
Compare various elastic materials w.r.t. uniaxial tension/compression test.

Requires Matplotlib.
"""
from __future__ import absolute_import
from argparse import ArgumentParser, RawDescriptionHelpFormatter
import sys
import six
sys.path.append('.')
from sfepy.base.base import output
from sfepy.base.conf import ProblemConf, get_standard_keywords
from sfepy.discrete import Problem
from sfepy.base.plotutils import plt
from matplotlib.collections import PolyCollection
from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection

import numpy as np
from functools import partial

def define( K=8.333, mu_nh=3.846, mu_mr=1.923, kappa=1.923, lam=5.769, mu=3.846 ):
    """Define the problem to solve."""
    from sfepy.discrete.fem.meshio import UserMeshIO
    from sfepy.mesh.mesh_generators import gen_block_mesh
    from sfepy.mechanics.matcoefs import stiffness_from_lame

    def mesh_hook(mesh, mode):
        """
        Generate the block mesh.
        """
        if mode == 'read':
            mesh = gen_block_mesh([2, 2, 3], [2, 2, 4], [0, 0, 1.5], name='el3',
                                  verbose=False)
            return mesh

        elif mode == 'write':
            pass

    filename_mesh = UserMeshIO(mesh_hook)

    options = {
        'nls' : 'newton',
        'ls' : 'ls',
        'ts' : 'ts',
        'save_times' : 'all',
    }

    functions = {
        'linear_pressure' : (linear_pressure,),
        'empty' : (lambda ts, coor, mode, region, ig: None,),
    }

    fields = {
        'displacement' : ('real', 3, 'Omega', 1),
    }

    # Coefficients are chosen so that the tangent stiffness is the same for all
    # material for zero strains.
    materials = {
        'solid' : ({
            'K'  : K, # bulk modulus
            'mu_nh' : mu_nh, # shear modulus of neoHookean term
            'mu_mr' : mu_mr, # shear modulus of Mooney-Rivlin term
            'kappa' : kappa, # second modulus of Mooney-Rivlin term
            # elasticity for LE term
            'D' : stiffness_from_lame(dim=3, lam=lam, mu=mu),
        },),
        'load' : 'empty',
    }

    variables = {
        'u' : ('unknown field', 'displacement', 0),
        'v' : ('test field', 'displacement', 'u'),
    }

    regions = {
        'Omega' : 'all',
        'Bottom' : ('vertices in (z < 0.1)', 'facet'),
        'Top' : ('vertices in (z > 2.9)', 'facet'),
    }

    ebcs = {
        'fixb' : ('Bottom', {'u.all' : 0.0}),
        'fixt' : ('Top', {'u.[0,1]' : 0.0}),
    }

    integrals = {
        'i' : 1,
        'isurf' : 2,
    }
    equations = {
        'linear' : """dw_lin_elastic.i.Omega(solid.D, v, u)
                    = dw_surface_ltr.isurf.Top(load.val, v)""",
        'neo-Hookean' : """dw_tl_he_neohook.i.Omega(solid.mu_nh, v, u)
                         + dw_tl_bulk_penalty.i.Omega(solid.K, v, u)
                         = dw_surface_ltr.isurf.Top(load.val, v)""",
        'Mooney-Rivlin' : """dw_tl_he_neohook.i.Omega(solid.mu_mr, v, u)
                           + dw_tl_he_mooney_rivlin.i.Omega(solid.kappa, v, u)
                           + dw_tl_bulk_penalty.i.Omega(solid.K, v, u)
                           = dw_surface_ltr.isurf.Top(load.val, v)""",
    }

    solvers = {
        'ls' : ('ls.scipy_direct', {}),
        'newton' : ('nls.newton', {
            'i_max'      : 5,
            'eps_a'      : 1e-10,
            'eps_r'      : 1.0,
        }),
        'ts' : ('ts.simple', {
            't0'    : 0,
            't1'    : 1,
            'dt'    : None,
            'n_step' : 26, # has precedence over dt!
            'verbose' : 1,
        }),
    }

    return locals()

##
# Pressure tractions.
def linear_pressure(ts, coor, mode=None, coef=1, **kwargs):
    if mode == 'qp':
        val = np.tile(coef * ts.step, (coor.shape[0], 1, 1))
        return {'val' : val}

def store_top_u(displacements):
    """Function _store() will be called at the end of each loading step. Top
    displacements will be stored into `displacements`."""
    def _store(problem, ts, state):

        top = problem.domain.regions['Top']
        top_u = problem.get_variables()['u'].get_state_in_region(top)
        displacements.append(np.mean(top_u[:,-1]))

    return _store

def solve_branch(problem, branch_function, material_type):

    eq = problem.conf.equations[material_type]
    problem.set_equations({material_type : eq})

    load = problem.get_materials()['load']
    load.set_function(branch_function)

    out = []
    problem.solve(save_results=False, step_hook=store_top_u(out))
    displacements = np.array(out, dtype=np.float64)

    return displacements

helps = {
    'no_plot' : 'do not show plot window',
}

def plot_mesh(pb):
        # plot mesh for macro problem
        coors = pb.domain.mesh.coors
        graph = pb.domain.mesh.get_conn(pb.domain.mesh.descs[0])
        fig2 = 
plt.figure(figsize=(5,6))