from __future__ import absolute_import
import numpy as nm

import sfepy.linalg as la
from sfepy.discrete.integrals import Integral
from six.moves import range

def prepare_remap(indices, n_full):
    """
    Prepare vector for remapping range `[0, n_full]` to its subset given
    by `indices`.
    """
    remap = nm.empty((n_full,), dtype=nm.int32)
    remap.fill(-1)
    remap[indices] = nm.arange(indices.shape[0], dtype=nm.int32)

    return remap

def invert_remap(remap):
    """
    Return the inverse of `remap`, i.e. a mapping from a sub-range
    indices to a full range, see :func:`prepare_remap()`.
    """
    if remap is not None:
        inverse = nm.where(remap >= 0)[0].astype(nm.int32)

    else:
        inverse = None

    return inverse

def prepare_translate(old_indices, new_indices):
    """
    Prepare vector for translating `old_indices` to `new_indices`.

    Returns
    -------
    translate : array
        The translation vector. Then `new_ar = translate[old_ar]`.
    """
    old_indices = nm.asarray(old_indices)
    new_indices = nm.asarray(new_indices)

    translate = nm.zeros(old_indices.max() + 1, dtype=new_indices.dtype)
    translate[old_indices] = new_indices

    return translate

def compute_nodal_normals(nodes, region, field, return_imap=False):
    """
    Nodal normals are computed by simple averaging of element normals of
    elements every node is contained in.
    """
    dim = region.dim

    field.domain.create_surface_group(region)
    field.setup_surface_data(region)

    # Custom integral with quadrature points very close to facet vertices.
    coors = field.gel.surface_facet.coors
    centre = coors.sum(axis=0) / coors.shape[0]
    qp_coors = coors + 1e-8 * (centre - coors)
    # Unit normals -> weights = ones.
    qp_weights = nm.ones(qp_coors.shape[0], dtype=nm.float64)

    integral = Integral('aux', coors=qp_coors, weights=qp_weights)

    normals = nm.zeros((nodes.shape[0], dim), dtype=nm.float64)
    mask = nm.zeros((nodes.max() + 1,), dtype=nm.int32)
    imap = nm.empty_like(mask)
    imap.fill(nodes.shape[0]) # out-of-range index for normals.
    imap[nodes] = nm.arange(nodes.shape[0], dtype=nm.int32)

    cmap, _ = field.get_mapping(region, integral, 'surface')
    e_normals = cmap.normal[..., 0]

    sd = field.domain.surface_groups[region.name]
    econn = sd.get_connectivity()
    mask[econn] += 1

    # normals[imap[econn]] += e_normals
    im = imap[econn]
    for ii, en in enumerate(e_normals):
        normals[im[ii]] += en

    # All nodes must have a normal.
    if not nm.all(mask[nodes] > 0):
        raise ValueError('region %s has not complete faces!' % region.name)

    norm = la.norm_l2_along_axis(normals)[:, nm.newaxis]
    if (norm < 1e-15).any():
        raise ValueError('zero nodal normal! (a node in volume?)')
    normals /= norm

    if return_imap:
        return normals, imap

    else:
        return normals

def _get_edge_path(graph, seed, mask, cycle=False):
    """
    Get a path in an edge graph starting with seed. The mask is incremented by
    one at positions of the path vertices.
    """
    if mask[seed]:
        return []

    path = [seed]
    mask[seed] = 1

    row = graph[seed].indices
    nv = len(row)
    while nv:
        if nv == 2:
            if mask[row[0]]:
                if mask[row[1]]:
                    if cycle:
                        path.append(seed)
                    break