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])
    return mesh

def gen_cylinder_mesh(dims, shape, centre, axis='x', force_hollow=False,
                      is_open=False, open_angle=0.0, non_uniform=False,
                      name='cylinder', verbose=True):
    """
    Generate a cylindrical mesh along an axis. Its cross-section can be
    ellipsoidal.

    Parameters
    ----------
    dims : array of 5 floats
        Dimensions of the cylinder: inner surface semi-axes a1, b1, outer
        surface semi-axes a2, b2, length.
    shape : array of 3 ints
        Shape (counts of nodes in radial, circumferential and longitudinal
        directions) of the cylinder mesh.
    centre : array of 3 floats
        Centre of the cylinder.
    axis: one of 'x', 'y', 'z'
        The axis of the cylinder.
    force_hollow : boolean
        Force hollow mesh even if inner radii a1 = b1 = 0.
    is_open : boolean
        Generate an open cylinder segment.
    open_angle : float
        Opening angle in radians.
    non_uniform : boolean
        If True, space the mesh nodes in radial direction so that the element
        volumes are (approximately) the same, making thus the elements towards
        the outer surface thinner.
    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)

    a1, b1, a2, b2, length = dims
    nr, nfi, nl = shape
    origin = centre - nm.array([0.5 * length, 0.0, 0.0])

    dfi = 2.0 * (nm.pi - open_angle) / nfi
    if is_open:
        nnfi = nfi + 1
    else:
        nnfi = nfi

    is_hollow = force_hollow or not (max(abs(a1), abs(b1)) < 1e-15)

    if is_hollow:
        mr = 0
    else:
        mr = (nnfi - 1) * nl

    grid = nm.zeros((nr, nnfi, nl), dtype=nm.int32)

    n_nod = nr * nnfi * nl - mr
    coors = nm.zeros((n_nod, 3), dtype=nm.float64)

    angles = nm.linspace(open_angle, open_angle+(nfi)*dfi, nfi+1)
    xs = nm.linspace(0.0, length, nl)
    if non_uniform:
        ras = nm.zeros((nr,), dtype=nm.float64)
        rbs = nm.zeros_like(ras)
        advol = (a2**2 - a1**2) / (nr - 1)
        bdvol = (b2**2 - b1**2) / (nr - 1)
        ras[0], rbs[0] = a1, b1
        for ii in range(1, nr):
            ras[ii] = nm.sqrt(advol + ras[ii-1]**2)
            rbs[ii] = nm.sqrt(bdvol + rbs[ii-1]**2)
    else:
        ras = nm.linspace(a1, a2, nr)
        rbs = nm.linspace(b1, b2, nr)

    # This is 3D only...
    output('generating %d vertices...' % n_nod, verbose=verbose)
    ii = 0
    for ix in range(nr):
        a, b = ras[ix], rbs[ix]
        for iy, fi in enumerate(angles[:nnfi]):
            for iz, x in enumerate(xs):
                grid[ix,iy,iz] = ii
                coors[ii] = origin + [x, a * nm.cos(fi), b * nm.sin(fi)]
                ii += 1

                if not is_hollow and (ix == 0):
                    if iy > 0:
                        grid[ix,iy,iz] = grid[ix,0,iz]
                        ii -= 1
    assert_(ii == n_nod)
    output('...done', verbose=verbose)

    n_el = (nr - 1) * nfi * (nl - 1)
    conn = nm.zeros((n_el, 8), dtype=nm.int32)

    output('generating %d cells...' % n_el, verbose=verbose)
    ii = 0
    for (ix, iy, iz) in cycle([nr-1, nnfi, nl-1]):
        if iy < (nnfi - 1):
            conn[ii,:] = [grid[ix  ,iy  ,iz  ], grid[ix+1,iy  ,iz  ],
                          grid[ix+1,iy+1,iz  ], grid[ix  ,iy+1,iz  ],
                          grid[ix  ,iy  ,iz+1], grid[ix+1,iy  ,iz+1],
                          grid[ix+1,iy+1,iz+1], grid[ix  ,iy+1,iz+1]]
            ii += 1
        elif not is_open:
            conn[ii,:] = [grid[ix  ,iy  ,iz  ], grid[ix+1,iy  ,iz  ],
                          grid[ix+1,0,iz  ], grid[ix  ,0,iz  ],
                          grid[ix  ,iy  ,iz+1], grid[ix+1,iy  ,iz+1],
                          grid[ix+1,0,iz+1], grid[ix  ,0,iz+1]]
            ii += 1

    mat_id = nm.zeros((n_el,), dtype = nm.int32)
    desc = '3_8'

    assert_(n_nod == (conn.max() + 1))
    output('...done', verbose=verbose)

    if axis == 'z':
        coors = coors[:,[1,2,0]]
    elif axis == 'y':
        coors = coors[:,[2,0,1]]

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

def _spread_along_axis(axis, coors, tangents, grading_fun):
    """
    Spread the coordinates along the given axis using the grading function, and
    the tangents in the other two directions.
    """
    oo = list(set([0, 1, 2]).difference([axis]))
    c0, c1, c2 = coors[:, axis], coors[:, oo[0]], coors[:, oo[1]]

    out = nm.empty_like(coors)

    mi, ma = c0.min(), c0.max()
    nc0 = (c0 - mi) / (ma - mi)
    out[:, axis] = oc0 = grading_fun(nc0) * (ma - mi) + mi

    nc = oc0 - oc0.min()

    mi, ma = c1.min(), c1.max()
    n1 = 2 * (c1 - mi) / (ma - mi) - 1
    out[:, oo[0]] = c1 + n1 * nc * tangents[oo[0]]

    mi, ma = c2.min(), c2.max()
    n2 = 2 * (c2 - mi) / (ma - mi) - 1
    out[:, oo[1]] = c2 + n2 * nc * tangents[oo[1]]

    return out

def _get_extension_side(side, grading_fun, mat_id,
                        b_dims, b_shape, e_dims, e_shape, centre):
    """
    Get a mesh extending the given side of a block mesh.
    """
    # Pure extension dimensions.
    pe_dims = 0.5 * (e_dims - b_dims)
    coff = 0.5 * (b_dims + pe_dims)
    cc = centre + coff * nm.eye(3)[side]

    if side == 0: # x axis.
        dims = [pe_dims[0], b_dims[1], b_dims[2]]
        shape = [e_shape, b_shape[1], b_shape[2]]
        tangents = [0, pe_dims[1] / pe_dims[0], pe_dims[2] / pe_dims[0]]

    elif side == 1: # y axis.
        dims = [b_dims[0], pe_dims[1], b_dims[2]]
        shape = [b_shape[0], e_shape, b_shape[2]]
        tangents = [pe_dims[0] / pe_dims[1], 0, pe_dims[2] / pe_dims[1]]

    elif side == 2: # z axis.
        dims = [b_dims[0], b_dims[1], pe_dims[2]]
        shape = [b_shape[0], b_shape[1], e_shape]
        tangents = [pe_dims[0] / pe_dims[2], pe_dims[1] / pe_dims[2], 0]

    e_mesh = gen_block_mesh(dims, shape, cc, mat_id=mat_id, verbose=False)
    e_mesh.coors[:] = _spread_along_axis(side, e_mesh.coors, tangents,
                                         grading_fun)

    return e_mesh, shape

def gen_extended_block_mesh(b_dims, b_shape, e_dims, e_shape, centre,
                            grading_fun=None, name=None):
    """
    Generate a 3D mesh with a central block and (coarse) extending side meshes.

    The resulting mesh is again a block. Each of the components has a different
    material id.

    Parameters
    ----------
    b_dims : array of 3 floats
        The dimensions of the central block.
    b_shape : array of 3 ints
        The shape (counts of nodes in x, y, z) of the central block mesh.
    e_dims : array of 3 floats
        The dimensions of the complete block (central block + extensions).
    e_shape : int
        The count of nodes of extending blocks in the direction from the
        central block.
    centre : array of 3 floats
        The centre of the mesh.
    grading_fun : callable, optional
        A function of :math:`x \in [0, 1]` that can be used to shift nodes in
        the extension axis directions to allow smooth grading of element sizes
        from the centre. The default function is :math:`x**p` with :math:`p`
        determined so that the element sizes next to the central block have the
        size of the shortest edge of the central block.
    name : string, optional
        The mesh name.

    Returns
    -------
    mesh : Mesh instance
    """
    b_dims = nm.asarray(b_dims, dtype=nm.float64)
    b_shape = nm.asarray(b_shape, dtype=nm.int32)
    e_dims = nm.asarray(e_dims, dtype=nm.float64)
    centre = nm.asarray(centre, dtype=nm.float64)

    # Pure extension dimensions.
    pe_dims = 0.5 * (e_dims - b_dims)
    # Central block element sizes.
    dd = (b_dims / (b_shape - 1))
    # The "first x" going to grading_fun.
    nc = 1.0 / (e_shape - 1)
    # Grading power and function.
    power = nm.log(dd.min() / pe_dims.min()) / nm.log(nc)
    grading_fun = (lambda x: x**power) if grading_fun is None else grading_fun

    # Central block mesh.
    b_mesh = gen_block_mesh(b_dims, b_shape, centre, mat_id=0, verbose=False)

    # 'x' extension.
    e_mesh, xs = _get_extension_side(0, grading_fun, 10,
                                     b_dims, b_shape, e_dims, e_shape, centre)
    mesh = b_mesh + e_mesh

    # Mirror by 'x'.
    e_mesh.coors[:, 0] = (2 * centre[0]) - e_mesh.coors[:, 0]
    e_mesh.cmesh.cell_groups.fill(11)
    mesh = mesh + e_mesh

    # 'y' extension.
    e_mesh, ys = _get_extension_side(1, grading_fun, 20,
                                     b_dims, b_shape, e_dims, e_shape, centre)
    mesh = mesh + e_mesh

    # Mirror by 'y'.
    e_mesh.coors[:, 1] = (2 * centre[1]) - e_mesh.coors[:, 1]
    e_mesh.cmesh.cell_groups.fill(21)
    mesh = mesh + e_mesh

    # 'z' extension.
    e_mesh, zs = _get_extension_side(2, grading_fun, 30,
                                     b_dims, b_shape, e_dims, e_shape, centre)
    mesh = mesh + e_mesh

    # Mirror by 'z'.
    e_mesh.coors[:, 2] = (2 * centre[2]) - e_mesh.coors[:, 2]
    e_mesh.cmesh.cell_groups.fill(31)
    mesh = mesh + e_mesh

    if name is not None:
        mesh.name = name

    # Verify merging by checking the number of nodes.
    n_nod = (nm.prod(nm.maximum(b_shape - 2, 0)) + 2 * nm.prod(xs)
             + 2 * (max(ys[0] - 2, 0) * ys[1] * ys[2])
             + 2 * (max(zs[0] - 2, 0) * max(zs[1] - 2, 0) * zs[2]))
    if n_nod != mesh.n_nod:
        raise ValueError('Merge of meshes failed! (%d == %d)'
                         % (n_nod, mesh.n_nod))

    return mesh

def tiled_mesh1d(conn, coors, ngrps, idim, n_rep, bb, eps=1e-6, ndmap=False):
    from sfepy.discrete.fem.periodic import match_grid_plane

    s1 = nm.nonzero(coors[:,idim] < (bb[0] + eps))[0]
    s2 = nm.nonzero(coors[:,idim] > (bb[1] - eps))[0]

    if s1.shape != s2.shape:
        raise ValueError('incompatible shapes: %s == %s'\
              % (s1.shape, s2.shape))

    (nnod0, dim) = coors.shape
    nnod = nnod0 * n_rep - s1.shape[0] * (n_rep - 1)
    (nel0, nnel) = conn.shape
    nel = nel0 * n_rep

    dd = nm.zeros((dim,), dtype=nm.float64)
    dd[idim] = bb[1] - bb[0]

    m1, m2 = 
match_grid_plane(coors[s1], coors[s2], idim)