from __future__ import absolute_import
from sfepy.discrete.fem import Mesh, FEDomain
import scipy.sparse as sps
import numpy as nm
from sfepy.base.compat import factorial
from sfepy.base.base import output
from six.moves import range

def elems_q2t(el):

    nel, nnd = el.shape
    if nnd > 4:
        q2t = nm.array([[0, 2, 3, 6],
                        [0, 3, 7, 6],
                        [0, 7, 4, 6],
                        [0, 5, 6, 4],
                        [1, 5, 6, 0],
                        [1, 6, 2, 0]])

    else:
        q2t = nm.array([[0, 1, 2],
                        [0, 2, 3]])

    ns, nn = q2t.shape
    nel *= ns

    out = nm.zeros((nel, nn), dtype=nm.int32);

    for ii in range(ns):
        idxs = nm.arange(ii, nel, ns)

        out[idxs,:] = el[:, q2t[ii,:]]

    return nm.ascontiguousarray(out)

def triangulate(mesh, verbose=False):
    """
    Triangulate a 2D or 3D tensor product mesh: quadrilaterals->triangles,
    hexahedrons->tetrahedrons.

    Parameters
    ----------
    mesh : Mesh
        The input mesh.

    Returns
    -------
    mesh : Mesh
        The triangulated mesh.
    """
    conns = None
    for k, new_desc in [('3_8', '3_4'), ('2_4', '2_3')]:
        if k in mesh.descs:
            conns = mesh.get_conn(k)
            break

    if conns is not None:
        nelo = conns.shape[0]
        output('initial mesh: %d elements' % nelo, verbose=verbose)

        new_conns = elems_q2t(conns)
        nn = new_conns.shape[0] // nelo
        new_cgroups = nm.repeat(mesh.cmesh.cell_groups, nn)

        output('new mesh: %d elements' % new_conns.shape[0], verbose=verbose)
        mesh = Mesh.from_data(mesh.name, mesh.coors,
                              mesh.cmesh.vertex_groups,
                              [new_conns], [new_cgroups], [new_desc])

    return mesh

def smooth_mesh(mesh, n_iter=4, lam=0.6307, mu=-0.6347,
                weights=None, bconstr=True,
                volume_corr=False):
    """
    FE mesh smoothing.

    Based on:

    [1] <NAME>, <NAME>, Smooth surface meshing for automated
    finite element model generation from 3D image data, Journal of
    Biomechanics, Volume 39, Issue 7, 2006, Pages 1287-1295,
    ISSN 0021-9290, 10.1016/j.jbiomech.2005.03.006.
    (http://www.sciencedirect.com/science/article/pii/S0021929005001442)

    Parameters
    ----------
    mesh : mesh
        FE mesh.
    n_iter : integer, optional
        Number of iteration steps.
    lam : float, optional
        Smoothing factor, see [1].
    mu : float, optional
        Unshrinking factor, see [1].
    weights : array, optional
        Edge weights, see [1].
    bconstr: logical, optional
        Boundary constraints, if True only surface smoothing performed.
    volume_corr: logical, optional
        Correct volume after smoothing process.

    Returns
    -------
    coors : array
        Coordinates of mesh nodes.
    """

    def laplacian(coors, weights):

        n_nod = coors.shape[0]
        displ = (weights - sps.identity(n_nod)) * coors

        return displ

    def taubin(coors0, weights, lam, mu, n_iter):

        coors = coors0.copy()

        for ii in range(n_iter):
            displ = laplacian(coors, weights)
            if nm.mod(ii, 2) == 0:
                coors += lam * displ
            else:
                coors += mu * displ

        return coors

    def get_volume(el, nd):
        from sfepy.linalg.utils import dets_fast

        dim = nd.shape[1]
        nnd = el.shape[1]

        etype = '%d_%d' % (dim, nnd)
        if etype == '2_4' or etype == '3_8':
            el = elems_q2t(el)

        nel = el.shape[0]

        #bc = nm.zeros((dim, ), dtype=nm.double)
        mul = 1.0 / factorial(dim)
        if dim == 3:
            mul *= -1.0

        mtx = nm.ones((nel, dim + 1, dim + 1), dtype=nm.double)
        mtx[:,:,:-1] = nd[el,:]
        vols = mul * 
dets_fast(mtx)