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')