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