from __future__ import print_function
from __future__ import absolute_import
import numpy as nm
import sys
from six.moves import range
sys.path.append('.')

from sfepy.base.base import output, assert_
from sfepy.base.ioutils import ensure_path
from sfepy.linalg import cycle
from sfepy.discrete.fem.mesh import Mesh
from sfepy.mesh.mesh_tools import elems_q2t

def get_tensor_product_conn(shape):
    """
    Generate vertex connectivity for cells of a tensor-product mesh of the
    given shape.

    Parameters
    ----------
    shape : array of 2 or 3 ints
        Shape (counts of nodes in x, y, z) of the mesh.

    Returns
    -------
    conn : array
        The vertex connectivity array.
    desc : str
        The cell kind.
    """
    shape = nm.asarray(shape)
    dim = len(shape)
    assert_(1 <= dim <= 3)

    n_nod = nm.prod(shape)
    n_el = nm.prod(shape - 1)

    grid = nm.arange(n_nod, dtype=nm.int32)
    grid.shape = shape

    if dim == 1:
        conn = nm.zeros((n_el, 2), dtype=nm.int32)
        conn[:, 0] = grid[:-1]
        conn[:, 1] = grid[1:]
        desc = '1_2'

    elif dim == 2:
        conn = nm.zeros((n_el, 4), dtype=nm.int32)
        conn[:, 0] = grid[:-1, :-1].flat
        conn[:, 1] = grid[1:, :-1].flat
        conn[:, 2] = grid[1:, 1:].flat
        conn[:, 3] = grid[:-1, 1:].flat
        desc = '2_4'

    else:
        conn = nm.zeros((n_el, 8), dtype=nm.int32)
        conn[:, 0] = grid[:-1, :-1, :-1].flat
        conn[:, 1] = grid[1:, :-1, :-1].flat
        conn[:, 2] = grid[1:, 1:, :-1].flat
        conn[:, 3] = grid[:-1, 1:, :-1].flat
        conn[:, 4] = grid[:-1, :-1, 1:].flat
        conn[:, 5] = grid[1:, :-1, 1:].flat
        conn[:, 6] = grid[1:, 1:, 1:].flat
        conn[:, 7] = grid[:-1, 1:, 1:].flat
        desc = '3_8'

    return conn, desc

def gen_block_mesh(dims, shape, centre, mat_id=0, name='block',
                   coors=None, verbose=True):
    """
    Generate a 2D or 3D block mesh. The dimension is determined by the
    lenght of the shape argument.

    Parameters
    ----------
    dims : array of 2 or 3 floats
        Dimensions of the block.
    shape : array of 2 or 3 ints
        Shape (counts of nodes in x, y, z) of the block mesh.
    centre : array of 2 or 3 floats
        Centre of the block.
    mat_id : int, optional
        The material id of all elements.
    name : string
        Mesh name.
    verbose : bool
        If True, show progress of the mesh generation.

    Returns
    -------
    mesh : Mesh instance
    """
    dims = nm.asarray(dims, dtype=nm.float64)
    shape = nm.asarray(shape, dtype=nm.int32)
    centre = nm.asarray(centre, dtype=nm.float64)

    dim = shape.shape[0]

    centre = centre[:dim]
    dims = dims[:dim]

    n_nod = nm.prod(shape)
    output('generating %d vertices...' % n_nod, verbose=verbose)

    x0 = centre - 0.5 * dims
    dd = dims / (shape - 1)

    ngrid = nm.mgrid[[slice(ii) for ii in shape]]
    ngrid.shape = (dim, n_nod)

    coors = x0 + ngrid.T * dd
    output('...done', verbose=verbose)

    n_el = nm.prod(shape - 1)
    output('generating %d cells...' % n_el, verbose=verbose)

    mat_ids = nm.empty((n_el,), dtype=nm.int32)
    mat_ids.fill(mat_id)

    conn, desc = get_tensor_product_conn(shape)
    output('...done', verbose=verbose)

    mesh = 
Mesh.from_data(name, coors, None, [conn], [mat_ids], [desc])