import numpy as nm
import six

from sfepy import data_dir
from sfepy.base.base import Struct, output
from sfepy.terms.terms_hyperelastic_ul import HyperElasticULFamilyData
from sfepy.homogenization.micmac import get_homog_coefs_nonlinear
import sfepy.linalg as la

hyperelastic_data = {}


def post_process(out, pb, state, extend=False):
    if isinstance(state, dict):
        pass
    else:
        pb.update_materials_flag = 2
        stress = pb.evaluate('ev_integrate_mat.1.Omega(solid.S, u)',
                             mode='el_avg')

        out['cauchy_stress'] = Struct(name='output_data',
                                      mode='cell',
                                      data=stress,
                                      dofs=None)

        strain = pb.evaluate('ev_integrate_mat.1.Omega(solid.E, u)',
                             mode='el_avg')

        out['green_strain'] = Struct(name='output_data',
                                     mode='cell',
                                     data=strain,
                                     dofs=None)

        pb.update_materials_flag = 0

        if pb.conf.options.get('recover_micro', False):
            happ = pb.homogen_app
            if pb.ts.step == 0:
                rname = pb.conf.options.recovery_region
                rcells = pb.domain.regions[rname].get_cells()
                sh = hyperelastic_data['homog_mat_shape']

                happ.app_options.store_micro_idxs = sh[1] * rcells
            else:
                hpb = happ.problem
                recovery_hook = hpb.conf.options.get('recovery_hook', None)
                if recovery_hook is not None:
                    recovery_hook = hpb.conf.get_function(recovery_hook)
                    rname = pb.conf.options.recovery_region
                    rcoors = []
                    for ii in happ.app_options.store_micro_idxs:
                        key = happ.get_micro_cache_key('coors', ii, pb.ts.step)
                        if key in happ.micro_state_cache:
                            rcoors.append(happ.micro_state_cache[key])

                    recovery_hook(hpb, rcoors, pb.domain.regions[rname], pb.ts)

    return out


def get_homog_mat(ts, coors, mode, term=None, problem=None, **kwargs):
    if problem.update_materials_flag == 2 and mode == 'qp':
        out = hyperelastic_data['homog_mat']
        return {k: nm.array(v) for k, v in six.iteritems(out)}
    elif problem.update_materials_flag == 0 or not mode == 'qp':
        return

    output('get_homog_mat')
    dim = problem.domain.mesh.dim

    update_var = problem.conf.options.mesh_update_variables[0]
    state_u = problem.equations.variables[update_var]
    state_u.field.clear_mappings()
    family_data = problem.family_data(state_u, term.region,
                                      term.integral, term.integration)

    mtx_f = family_data.mtx_f.reshape((coors.shape[0],)
                                      + family_data.mtx_f.shape[-2:])
    out = get_homog_coefs_nonlinear(ts, coors, mode, mtx_f,
                                    term=term, problem=problem,
                                    iteration=problem.iiter, **kwargs)

    out['E'] = 0.5 * (la.dot_sequences(mtx_f, mtx_f, 'ATB') - nm.eye(dim))

    hyperelastic_data['time'] = ts.step
    hyperelastic_data['homog_mat_shape'] = family_data.det_f.shape[:2]
    hyperelastic_data['homog_mat'] = \
        {k: nm.array(v) for k, v in six.iteritems(out)}

    return out


def ulf_iteration_hook(pb, nls, vec, it, err, err0):
    vec = pb.equations.make_full_vec(vec)
    pb.equations.set_variables_from_state(vec)

    update_var = pb.conf.options.mesh_update_variables[0]
    state_u = pb.equations.variables[update_var]

    nods = state_u.field.get_dofs_in_region(state_u.field.region, merge=True)
    coors = pb.domain.get_mesh_coors().copy()
    coors[nods, :] += state_u().reshape(len(nods), state_u.n_components)

    if len(state_u.field.mappings0) == 0:
        state_u.field.save_mappings()

    state_u.field.clear_mappings()
    pb.set_mesh_coors(coors, update_fields=False, actual=True,
                      clear_all=False)

    pb.iiter = it
    pb.update_materials_flag = True
    pb.update_materials()
    pb.update_materials_flag = False


class MyEvalResidual(object):
    def __init__(self, problem, matrix_hook=None):
        self.problem = problem
        self.matrix_hook = problem.matrix_hook

    def eval_residual(self, vec, is_full=False):
        if not is_full:
            vec = self.problem.equations.make_full_vec(vec)
        vec_r = self.problem.equations.eval_residuals(vec * 0)

        return vec_r


def ulf_init(pb):
    pb.family_data = 
HyperElasticULFamilyData()