# This example implements homogenization of piezoeletric porous media.
# The mathematical model and numerical results are described in: 
#
# <NAME>., <NAME>.
# Homogenization of the fluid-saturated piezoelectric porous media.
# International Journal of Solids and Structures
# Volume 147, 15 August 2018, Pages 110-125
# https://doi.org/10.1016/j.ijsolstr.2018.05.017
#
# Run calculation of homogeized coefficients:
#
#   ./homogen.py example_poropiezo-1/poropiezo_micro_dfc.py
#
# The results are stored in `example_poropiezo-1/results` directory.
#

import sys
import numpy as nm
import os.path as osp
from sfepy.mechanics.matcoefs import stiffness_from_youngpoisson
from sfepy.homogenization.utils import coor_to_sym, define_box_regions
import sfepy.discrete.fem.periodic as per
from sfepy.discrete.fem.mesh import Mesh
import sfepy.homogenization.coefs_base as cb
from sfepy.base.base import Struct

data_dir = 'example_poropiezo-1'


def data_to_struct(data):
    out = {}
    for k, v in data.items():
        out[k] = Struct(name='output_data',
                        mode='cell' if v[2] == 'c' else 'vertex',
                        data=v[0],
                        var_name=v[1],
                        dofs=None)

    return out


def get_periodic_bc(var_tab, dim=3, dim_tab=None):
    if dim_tab is None:
        dim_tab = {'x': ['left', 'right'],
                   'z': ['bottom', 'top'],
                   'y': ['near', 'far']}

    periodic = {}
    epbcs = {}

    for ivar, reg in var_tab:
        periodic['per_%s' % ivar] = pers = []
        for idim in 'xyz'[0:dim]:
            key = 'per_%s_%s' % (ivar, idim)
            regs = ['%s_%s' % (reg, ii) for ii in dim_tab[idim]]
            epbcs[key] = (regs, {'%s.all' % ivar: '%s.all' % ivar},
                          'match_%s_plane' % idim)
            pers.append(key)

    return epbcs, periodic


# reconstruct displacement, and electric fields at the microscopic level,
# see Section 6.2
def recovery_micro_dfc(pb, corrs, macro):
    eps0 = macro['eps0']
    mesh = pb.domain.mesh
    regions = pb.domain.regions
    dim = mesh.dim
    Yms_map = regions['Yms'].get_entities(0)
    Ym_map = regions['Ym'].get_entities(0)

    gl = '_' + list(corrs.keys())[0].split('_')[-1]
    u1 = -corrs['corrs_p' + gl]['u'] * macro['press'][Yms_map, :]
    phi = -corrs['corrs_p' + gl]['r'] * macro['press'][Ym_map, :]

    for ii in range(2):
        u1 += corrs['corrs_k%d' % ii + gl]['u'] * macro['phi'][ii]
        phi += corrs['corrs_k%d' % ii + gl]['r'] * macro['phi'][ii]

    for ii in range(dim):
        for jj in range(dim):
            kk = coor_to_sym(ii, jj, dim)
            phi += corrs['corrs_rs' + gl]['r_%d%d' % (ii, jj)]\
                * nm.expand_dims(macro['strain'][Ym_map, kk], axis=1)
            u1 += corrs['corrs_rs' + gl]['u_%d%d' % (ii, jj)]\
                * nm.expand_dims(macro['strain'][Yms_map, kk], axis=1)

    u = macro['u'][Yms_map, :] + eps0 * u1

    mvar = pb.create_variables(['u', 'r', 'svar'])

    e_mac_Yms = [None] * macro['strain'].shape[1]

    for ii in range(dim):
        for jj in range(dim):
            kk = coor_to_sym(ii, jj, dim)
            mvar['svar'].set_data(macro['strain'][:, kk])
            mac_e_Yms = pb.evaluate('ev_volume_integrate.i2.Yms(svar)',
                                    mode='el_avg',
                                    var_dict={'svar': mvar['svar']})

            e_mac_Yms[kk] = mac_e_Yms.squeeze()

    e_mac_Yms = nm.vstack(e_mac_Yms).T[:, nm.newaxis, :, nm.newaxis]

    mvar['r'].set_data(phi)
    E_mic = pb.evaluate('ev_grad.i2.Ym(r)',
                        mode='el_avg',
                        var_dict={'r': mvar['r']}) / eps0

    mvar['u'].set_data(u1)
    e_mic = pb.evaluate('ev_cauchy_strain.i2.Yms(u)',
                        mode='el_avg',
                        var_dict={'u': mvar['u']})
    e_mic += e_mac_Yms

    out = {
        'u0': (macro['u'][Yms_map, :], 'u', 'p'),  # macro displacement
        'u1': (u1, 'u', 'p'),  # local displacement corrections, see eq. (58)
        'u': (u, 'u', 'p'),  # total displacement
        'e_mic': (e_mic, 'u', 'c'),  # micro strain field, see eq. (58)
        'phi': (phi, 'r', 'p'),  # electric potential, see eq. (57)
        'E_mic': (E_mic, 'r', 'c'),  # electric field, see eq. (58)
    }

    return data_to_struct(out)


# define homogenized coefficients and subproblems for correctors
def define(grid0=100, filename_mesh=None):
    eps0 = 0.01 / grid0

    if filename_mesh is None:
        filename_mesh = osp.join(data_dir, 'piezo_mesh_micro_dfc.vtk')

    mesh = Mesh.from_file(filename_mesh)
    n_conduct = len(nm.unique(mesh.cmesh.cell_groups)) - 2

    sym_eye = 'nm.array([1,1,0])' if mesh.dim == 2 else\
        'nm.array([1,1,1,0,0,0])'

    bbox = mesh.get_bounding_box()
    regions = 
define_box_regions(mesh.dim, bbox[0], bbox[1], eps=1e-3)