# This example implements 2nd-level homogenization of Biot-Darcy-Brinkman model of flow in deformable
# double porous media.
# The mathematical model is described in:
#
#<NAME>., <NAME>., <NAME>.
#The Biot-Darcy-Brinkman model of flow in deformable double porous media; homogenization and numerical modelling.
# Computers and Mathematics with applications, 78(9):3044-3066, 2019,
# https://doi.org/10.1016/j.camwa.2019.04.004
#
# Run calculation of homogenized coefficients:
#
#   ./homogen.py example_perfusion_BDB/perf_BDB_mes.py
#
# The results are stored in `example_perfusion_BDB/results/meso` directory.
#

import numpy as nm
from sfepy import data_dir

import os.path as osp
from sfepy.discrete.fem.mesh import Mesh
import sfepy.discrete.fem.periodic as per
from sfepy.mechanics.tensors import dim2sym
from sfepy.homogenization.utils import define_box_regions
import sfepy.homogenization.coefs_base as cb
from sfepy.homogenization.micmac import get_homog_coefs_linear


data_dir = 'example_perfusion_BDB'

def coefs2qp(coefs, nqp):
    out = {}
    for k, v in coefs.items():
        if type(v) not in [nm.ndarray, float]:
            continue

        if type(v) is nm.ndarray:
            if len(v.shape) >= 3:
                out[k] = v

        out[k] = nm.tile(v, (nqp, 1, 1))

    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



# get homogenized coefficients, recalculate them if necessary
def get_homog(coors, mode, pb,micro_filename, **kwargs):
    if not (mode == 'qp'):
        return
    nqp = coors.shape[0]
    coefs_filename = 'coefs_micro'
    coefs_filename = osp.join(pb.conf.options.get('output_dir', '.'),
                              coefs_filename) + '.h5'
    coefs = get_homog_coefs_linear(0, 0, None,
                                   micro_filename=micro_filename,coefs_filename = coefs_filename    )
    coefs['B'] = coefs['B'][:, nm.newaxis]
    for k in coefs.keys():
        v = coefs[k]
        if type(v) is nm.ndarray:
            if len(v.shape) == 0:
                coefs[k] = v.reshape((1, 1))
            elif len(v.shape) == 1:
                coefs[k] = v[:, nm.newaxis]
        elif isinstance(v, float):
            coefs[k] = nm.array([[v]])

    out = coefs2qp(coefs, nqp)

    return out

def define(filename_mesh=None):

    eta = 3.6e-3
    if filename_mesh is None:
        filename_mesh = osp.join(data_dir, 'meso_perf_puc.vtk')

    mesh = Mesh.from_file(filename_mesh)

    poroela_micro_file = osp.join(data_dir, 'perf_BDB_mic.py')
    dim = 3
    sym = (dim + 1) * dim // 2
    sym_eye = 'nm.array([1,1,0])' if 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)

    regions.update({
        'Z': 'all',
        'Gamma_Z': ('vertices of surface', 'facet'),
        # matrix
        'Zm': 'cells of group 1',
        'Zm_left': ('r.Zm *v r.Left', 'vertex'),
        'Zm_right': ('r.Zm *v r.Right', 'vertex'),
        'Zm_bottom': ('r.Zm *v r.Bottom', 'vertex'),
        'Zm_top': ('r.Zm *v r.Top', 'vertex'),
        'Gamma_Zm': ('r.Zm *v r.Zc', 'facet', 'Zm'),
        # canal
        'Zc': 'cells of group 2',
        'Zc0': ('r.Zc -v r.Gamma_Zc', 'vertex'),
        'Zc_left': ('r.Zc0 *v r.Left', 'vertex'),
        'Zc_right': ('r.Zc0 *v r.Right', 'vertex'),
        'Zc_bottom': ('r.Zc0 *v r.Bottom', 'vertex'),
        'Zc_top': ('r.Zc0 *v r.Top', 'vertex'),
        'Gamma_Zc': ('r.Zm *v r.Zc', 'facet', 'Zc'),
        "Surface": ("vertices of surface", "facet"),
        'Center_c': ('vertex 5346', 'vertex'),  # canal center
    })

    if dim == 3:
        regions.update({
            'Zm_far': ('r.Zm *v r.Far', 'vertex'),
            'Zm_near': ('r.Zm *v r.Near', 'vertex'),
            'Zc_far': ('r.Zc0 *v r.Far', 'vertex'),
            'Zc_near': ('r.Zc0 *v r.Near', 'vertex'),
        })


    fields = {
        'one': ('real', 'scalar', 'Z', 1),
        'displacement': ('real', 'vector', 'Zm', 1),
        'pressure_m': ('real', 'scalar', 'Zm', 1),
        'pressure_c': ('real', 'scalar', 'Zc', 1),
        'displacement_c': ('real', 'vector', 'Zc', 1),
        'velocity': ('real', 'vector', 'Zc', 2),
    }

    variables = {
        # displacement
        'u': ('unknown field', 'displacement', 0),
        'v': ('test field', 'displacement', 'u'),
        'Pi_u': ('parameter field', 'displacement', 'u'),
        'U1': ('parameter field', 'displacement', '(set-to-None)'),
        'U2': ('parameter field', 'displacement', '(set-to-None)'),

        'uc': ('unknown field', 'displacement_c', 4),
        'vc': ('test field', 'displacement_c', 'uc'),
        # velocity
        'w': ('unknown field', 'velocity', 1),
        'z': ('test field', 'velocity', 'w'),
        'Pi_w': ('parameter field', 'velocity', 'w'),
        'W1': ('parameter field', 'velocity', '(set-to-None)'),
        'W2': ('parameter field', 'velocity', '(set-to-None)'),
        # pressure
        'pm': ('unknown field', 'pressure_m', 2),
        'qm': ('test field', 'pressure_m', 'pm'),
        'Pm1': ('parameter field', 'pressure_m', '(set-to-None)'),
        'Pm2': ('parameter field', 'pressure_m', '(set-to-None)'),
        'Pi_pm': ('parameter field', 'pressure_m', 'pm'),
        'pc': ('unknown field', 'pressure_c', 3),
        'qc': ('test field', 'pressure_c', 'pc'),
        'Pc1': ('parameter field', 'pressure_c', '(set-to-None)'),
        'Pc2': ('parameter field', 'pressure_c', '(set-to-None)'),
        # one
        'one': ('parameter field', 'one', '(set-to-None)'),
    }

    functions = {
        'match_x_plane': (per.match_x_plane,),
        'match_y_plane': (per.match_y_plane,),
        'match_z_plane': (per.match_z_plane,),
        'get_homog': (lambda ts, coors, mode=None, problem=None, **kwargs:\
            get_homog(coors, mode, problem, poroela_micro_file, **kwargs),),
    }
    materials = {

        'hmatrix': 'get_homog',
        'fluid': ({
                    'eta_c': eta* nm.eye(
dim2sym(dim)