"""Approximate nearest neighbor search"""
# Author: Maheshakya Wijewardena <maheshakya.10@cse.mrt.ac.lk>
#         Joel Nothman <joel.nothman@gmail.com>

import numpy as np
import warnings

from scipy import sparse

from .base import KNeighborsMixin, RadiusNeighborsMixin
from ..base import BaseEstimator
from ..utils.validation import check_array
from ..utils import check_random_state
from ..metrics.pairwise import pairwise_distances

from ..random_projection import GaussianRandomProjection

__all__ = ["LSHForest"]

HASH_DTYPE = '>u4'
MAX_HASH_SIZE = np.dtype(HASH_DTYPE).itemsize * 8

def _find_matching_indices(tree, bin_X, left_mask, right_mask):
    """Finds indices in sorted array of integers.

    Most significant h bits in the binary representations of the
    integers are matched with the items' most significant h bits.
    """
    left_index = np.searchsorted(tree, bin_X & left_mask)
    right_index = np.searchsorted(tree, bin_X | right_mask,
                                  side='right')
    return left_index, right_index

def _find_longest_prefix_match(tree, bin_X, hash_size,
                               left_masks, right_masks):
    """Find the longest prefix match in tree for each query in bin_X

    Most significant bits are considered as the prefix.
    """
    hi = np.empty_like(bin_X, dtype=np.intp)
    hi.fill(hash_size)
    lo = np.zeros_like(bin_X, dtype=np.intp)
    res = np.empty_like(bin_X, dtype=np.intp)

    left_idx, right_idx = _find_matching_indices(tree, bin_X,
                                                 left_masks[hi],
                                                 right_masks[hi])
    found = right_idx > left_idx
    res[found] = lo[found] = hash_size

    r = np.arange(bin_X.shape[0])
    kept = r[lo < hi]  # indices remaining in bin_X mask
    while kept.shape[0]:
        mid = (lo.take(kept) + hi.take(kept)) // 2

        left_idx, right_idx = _find_matching_indices(tree,
                                                     bin_X.take(kept),
                                                     left_masks[mid],
                                                     right_masks[mid])
        found = right_idx > left_idx
        mid_found = mid[found]
        lo[kept[found]] = mid_found + 1
        res[kept[found]] = mid_found
        hi[kept[~found]] = mid[~found]

        kept = r[lo < hi]

    return res

class ProjectionToHashMixin(object):
    """Turn a transformed real-valued array into a hash"""
    @staticmethod
    def _to_hash(projected):
        if projected.shape[1] % 8 != 0:
            raise ValueError('Require reduced dimensionality to be a multiple '
                             'of 8 for hashing')
        # XXX: perhaps non-copying operation better
        out = np.packbits((projected > 0).astype(int)).view(dtype=HASH_DTYPE)
        return out.reshape(projected.shape[0], -1)

    def fit_transform(self, X, y=None):
        self.fit(X)
        return self.transform(X)

    def transform(self, X, y=None):
        return self._to_hash(super(ProjectionToHashMixin, self).transform(X))

class GaussianRandomProjectionHash(ProjectionToHashMixin,
                                   GaussianRandomProjection):
    """Use GaussianRandomProjection to produce a cosine LSH fingerprint"""
    def __init__(self,
                 n_components=8,
                 random_state=None):
        super(GaussianRandomProjectionHash, self).__init__(
            n_components=n_components,
            random_state=random_state)

def _array_of_arrays(list_of_arrays):
    """Creates an array of array from list of arrays."""
    out = np.empty(len(list_of_arrays), dtype=object)
    out[:] = list_of_arrays
    return out

class LSHForest(BaseEstimator, KNeighborsMixin, RadiusNeighborsMixin):
    """Performs approximate nearest neighbor search using LSH forest.

    LSH Forest: Locality Sensitive Hashing forest [1] is an alternative
    method for vanilla approximate nearest neighbor search methods.
    LSH forest data structure has been implemented using sorted
    arrays and binary search and 32 bit fixed-length hashes.
    Random projection is used as the hash family which approximates
    cosine distance.

    The cosine distance is defined as ``1 - cosine_similarity``: the lowest
    value is 0 (identical point) but it is bounded above by 2 for the farthest
    points. Its value does not depend on the norm of the vector points but
    only on their relative angles.

    Read more in the :ref:`User Guide <approximate_nearest_neighbors>`.

    Parameters
    ----------

    n_estimators : int (default = 10)
        Number of trees in the LSH Forest.

    min_hash_match : int (default = 4)
        lowest hash length to be searched when candidate selection is
        performed for nearest neighbors.

    n_candidates : int (default = 10)
        Minimum number of candidates evaluated per estimator, assuming enough
        items meet the `min_hash_match` constraint.

    n_neighbors : int (default = 5)
        Number of neighbors to be returned from query function when
        it is not provided to the :meth:`kneighbors` method.

    radius : float, optinal (default = 1.0)
        Radius from the data point to its neighbors. This is the parameter
        space to use by default for the :meth`radius_neighbors` queries.

    radius_cutoff_ratio : float, optional (default = 0.9)
        A value ranges from 0 to 1. Radius neighbors will be searched until
        the ratio between total neighbors within the radius and the total
        candidates becomes less than this value unless it is terminated by
        hash length reaching `min_hash_match`.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    Attributes
    ----------

    hash_functions_ : list of GaussianRandomProjectionHash objects
        Hash function g(p,x) for a tree is an array of 32 randomly generated
        float arrays with the same dimenstion as the data set. This array is
        stored in GaussianRandomProjectionHash object and can be obtained
        from ``components_`` attribute.

    trees_ : array, shape (n_estimators, n_samples)
        Each tree (corresponding to a hash function) contains an array of
        sorted hashed values. The array representation may change in future
        versions.

    original_indices_ : array, shape (n_estimators, n_samples)
        Original indices of sorted hashed values in the fitted index.

    References
    ----------

    .. [1] M. Bawa, T. Condie and P. Ganesan, "LSH Forest: Self-Tuning
           Indexes for Similarity Search", WWW '05 Proceedings of the
           14th international conference on World Wide Web,  651-660,
           2005.

    Examples
    --------
      >>> from sklearn.neighbors import LSHForest

      >>> X_train = [[5, 5, 2], [21, 5, 5], [1, 1, 1], [8, 9, 1], [6, 10, 2]]
      >>> X_test = [[9, 1, 6], [3, 1, 10], [7, 10, 3]]
      >>> lshf = LSHForest()
      >>> lshf.fit(X_train)  # doctest: +NORMALIZE_WHITESPACE
      LSHForest(min_hash_match=4, n_candidates=50, n_estimators=10,
                n_neighbors=5, radius=1.0, radius_cutoff_ratio=0.9,
                random_state=None)
      >>> distances, indices = lshf.kneighbors(X_test, n_neighbors=2)
      >>> distances                                        # doctest: +ELLIPSIS
      array([[ 0.069...,  0.149...],
             [ 0.229...,  0.481...],
             [ 0.004...,  0.014...]])
      >>> indices
      array([[1, 2],
             [2, 0],
             [4, 0]])

    """

    def __init__(self, n_estimators=10, radius=1.0, n_candidates=50,
                 n_neighbors=5, min_hash_match=4, radius_cutoff_ratio=.9,
                 random_state=None):
        self.n_estimators = n_estimators
        self.radius = radius
        self.random_state = random_state
        self.n_candidates = n_candidates
        self.n_neighbors = n_neighbors
        self.min_hash_match = min_hash_match
        self.radius_cutoff_ratio = radius_cutoff_ratio

    def _compute_distances(self, query, candidates):
        """Computes the cosine distance.

        Distance is from the query to points in the candidates array.
        Returns argsort of distances in the candidates
        array and sorted distances.
        """
        if candidates.shape == (0,):
            # needed since _fit_X[np.array([])] doesn't work if _fit_X sparse
            return np.empty(0, dtype=np.int), np.empty(0, dtype=float)

        if sparse.issparse(self._fit_X):
            candidate_X = self._fit_X[candidates]
        else:
            candidate_X = self._fit_X.take(candidates, axis=0, mode='clip')
        distances = pairwise_distances(query, candidate_X,
                                       metric='cosine')[0]
        distance_positions = np.argsort(distances)
        distances = distances.take(distance_positions, mode='clip', axis=0)
        return distance_positions, distances

    def _generate_masks(self):
        """Creates left and right masks for all hash lengths."""
        tri_size = MAX_HASH_SIZE + 1
        # Called once on fitting, output is independent of hashes
        left_mask = np.tril(np.ones((tri_size, tri_size), dtype=int))[:, 1:]
        right_mask = left_mask[::-1, ::-1]

        self._left_mask = np.packbits(left_mask).view(dtype=HASH_DTYPE)
        self._right_mask = np.packbits(right_mask).view(dtype=HASH_DTYPE)

    def _get_candidates(self, query, max_depth, bin_queries, n_neighbors):
        """Performs the Synchronous ascending phase.

        Returns an array of candidates, their distance ranks and
        distances.
        """
        index_size = self._fit_X.shape[0]
        # Number of candidates considered including duplicates
        # XXX: not sure whether this is being calculated correctly wrt
        #      duplicates from different iterations through a single tree
        n_candidates = 0
        candidate_set = set()
        min_candidates = self.n_candidates * self.n_estimators
        while (max_depth > self.min_hash_match and
               (n_candidates < min_candidates or
                len(candidate_set) < n_neighbors)):

            left_mask = self._left_mask[max_depth]
            right_mask = self._right_mask[max_depth]
            for i in range(self.n_estimators):
                start, stop = _find_matching_indices(self.trees_[i],
                                                     bin_queries[i],
                                                     left_mask, right_mask)
                n_candidates += stop - start
                candidate_set.update(
                    self.original_indices_[i][start:stop].tolist())
            max_depth -= 1

        candidates = np.fromiter(candidate_set, count=len(candidate_set),
                                 dtype=np.intp)
        # For insufficient candidates, candidates are filled.
        # Candidates are filled from unselected indices uniformly.
        if candidates.shape[0] < n_neighbors:
            warnings.warn(
                "Number of candidates is not sufficient to retrieve"
                " %i neighbors with"
                " min_hash_match = %i. Candidates are filled up"
                " uniformly from unselected"
                " indices." % (n_neighbors, self.min_hash_match))
            remaining = np.setdiff1d(np.arange(0, index_size), candidates)
            to_fill = n_neighbors - candidates.shape[0]
            candidates = np.concatenate((candidates, remaining[:to_fill]))

        ranks, distances = self._compute_distances(query,
                                                   candidates.astype(int))

        return (candidates[ranks[:n_neighbors]],
                distances[:n_neighbors])

    def _get_radius_neighbors(self, query, max_depth, bin_queries, radius):
        """Finds radius neighbors from the candidates obtained.

        Their distances from query are smaller than radius.
        Returns radius neighbors and distances.
        """
        ratio_within_radius = 1
        threshold = 1 - self.radius_cutoff_ratio
        total_candidates = np.array([], dtype=int)
        total_neighbors = np.array([], dtype=int)
        total_distances = np.array([], dtype=float)

        while (max_depth > self.min_hash_match and
               ratio_within_radius > threshold):
            left_mask = self._left_mask[max_depth]
            right_mask = self._right_mask[max_depth]
            candidates = []
            for i in range(self.n_estimators):
                start, stop = _find_matching_indices(self.trees_[i],
                                                     bin_queries[i],
                                                     left_mask, right_mask)
                candidates.extend(
                    self.original_indices_[i][start:stop].tolist())
            candidates = np.setdiff1d(candidates, total_candidates)
            total_candidates = np.append(total_candidates, candidates)
            ranks, distances = self._compute_distances(query, candidates)
            m = np.searchsorted(distances, radius, side='right')
            positions = np.searchsorted(total_distances, distances[:m])
            total_neighbors = np.insert(total_neighbors, positions,
                                        candidates[ranks[:m]])
            total_distances = np.insert(total_distances, positions,
                                        distances[:m])
            ratio_within_radius = (total_neighbors.shape[0] /
                                   float(total_candidates.shape[0]))
            max_depth = max_depth - 1
        return total_neighbors, total_distances

    def fit(self, X, y=None):
        """Fit the LSH forest on the data.

        This creates binary hashes of input data points by getting the
        dot product of input points and hash_function then
        transforming the projection into a binary string array based
        on the sign (positive/negative) of the projection.
        A sorted array of binary hashes is created.

        Parameters
        ----------
        X : array_like or sparse (CSR) matrix, shape (n_samples, n_features)
            List of n_features-dimensional data points. Each row
            corresponds to a single data point.

        Returns
        -------
        self : object
            Returns self.
        """

        self._fit_X = check_array(X, accept_sparse='csr')

        # Creates a g(p,x) for each tree
        self.hash_functions_ = []
        self.trees_ = []
        self.original_indices_ = []

        rng = check_random_state(self.random_state)
        int_max = np.iinfo(np.int32).max

        for i in range(self.n_estimators):
            # This is g(p,x) for a particular tree.
            # Builds a single tree. Hashing is done on an array of data points.
            # `GaussianRandomProjection` is used for hashing.
            # `n_components=hash size and n_features=n_dim.
            hasher = GaussianRandomProjectionHash(MAX_HASH_SIZE,
                                                  rng.randint(0, int_max))
            hashes = hasher.fit_transform(self._fit_X)[:, 0]
            original_index = np.argsort(hashes)
            bin_hashes = hashes[original_index]
            self.original_indices_.append(original_index)
            self.trees_.append(bin_hashes)
            self.hash_functions_.append(hasher)

        self._generate_masks()

        return self

    def _query(self, X):
        """Performs descending phase to find maximum depth."""
        # Calculate hashes of shape (n_samples, n_estimators, [hash_size])
        bin_queries = np.asarray([hasher.transform(X)[:, 0]
                                  for hasher in self.hash_functions_])
        bin_queries = np.rollaxis(bin_queries, 1)

        # descend phase
        depths = [_find_longest_prefix_match(tree, tree_queries, MAX_HASH_SIZE,
                                             self._left_mask, self._right_mask)
                  for tree, tree_queries in zip(self.trees_,
                                                np.rollaxis(bin_queries, 1))]

        return bin_queries, np.max(depths, axis=0)

    def kneighbors(self, X, n_neighbors=None, return_distance=True):
        """Returns n_neighbors of approximate nearest neighbors.

        Parameters
        ----------
        X : array_like or sparse (CSR) matrix, shape (n_samples, n_features)
            List of n_features-dimensional data points.  Each row
            corresponds to a single query.

        n_neighbors : int, opitonal (default = None)
            Number of neighbors required. If not provided, this will
            return the number specified at the initialization.

        return_distance : boolean, optional (default = False)
            Returns the distances of neighbors if set to True.

        Returns
        -------
        dist : array, shape (n_samples, n_neighbors)
            Array representing the cosine distances to each point,
            only present if return_distance=True.

        ind : array, shape (n_samples, n_neighbors)
            Indices of the approximate nearest points in the population
            matrix.
        """
        if not hasattr(self, 'hash_functions_'):
            raise ValueError("estimator should be fitted.")

        if n_neighbors is None:
            n_neighbors = self.n_neighbors

        X = check_array(X, accept_sparse='csr')

        neighbors, distances = [], []
        bin_queries, max_depth = self._query(X)
        for i in range(X.shape[0]):

            neighs, dists = self._get_candidates(X[[i]], max_depth[i],
                                                 bin_queries[i],
                                                 n_neighbors)
            neighbors.append(neighs)
            distances.append(dists)

        if return_distance:
            return np.array(distances), np.array(neighbors)
        else:
            return np.array(neighbors)

    def radius_neighbors(self, X, radius=None, return_distance=True):
        """Finds the neighbors within a given radius of a point or points.

        Return the indices and distances of some points from the dataset
        lying in a ball with size ``radius`` around the points of the query
        array. Points lying on the boundary are included in the results.

        The result points are *not* necessarily sorted by distance to their
        query point.

        LSH Forest being an approximate method, some true neighbors from the
        indexed dataset might be missing from the results.

        Parameters
        ----------
        X : array_like or sparse (CSR) matrix, shape (n_samples, n_features)
            List of n_features-dimensional data points. Each row
            corresponds to a single query.

        radius : float
            Limiting distance of neighbors to return.
            (default is the value passed to the constructor).

        return_distance : boolean, optional (default = False)
            Returns the distances of neighbors if set to True.

        Returns
        -------
        dist : array, shape (n_samples,) of arrays
            Each element is an array representing the cosine distances
            to some points found within ``radius`` of the respective query.
            Only present if ``return_distance=True``.

        ind : array, shape (n_samples,) of arrays
            Each element is an array of indices for neighbors within ``radius``
            of the respective query.
        """
        if not hasattr(self, 'hash_functions_'):
            raise ValueError("estimator should be fitted.")

        if radius is None:
            radius = self.radius

        X = check_array(X, accept_sparse='csr')

        neighbors, distances = [], []
        bin_queries, max_depth = self._query(X)
        for i in range(X.shape[0]):

            neighs, dists = self._get_radius_neighbors(X[[i]], max_depth[i],
                                                       bin_queries[i], radius)
            neighbors.append(neighs)
            distances.append(dists)

        if return_distance:
            return _array_of_arrays(distances), _array_of_arrays(neighbors)
        else:
            return _array_of_arrays(neighbors)

    def partial_fit(self, X, y=None):
        """
        Inserts new data into the already fitted LSH Forest.
        Cost is proportional to new total size, so additions
        should be batched.

        Parameters
        ----------
        X : array_like or sparse (CSR) matrix, shape (n_samples, n_features)
            New data point to be inserted into the LSH Forest.
        """
        X = check_array(X, accept_sparse='csr')
        if not hasattr(self, 'hash_functions_'):
            return self.fit(X)

        if X.shape[1] != self._fit_X.shape[1]:
            raise ValueError("Number of features in X and"
                             " fitted array does not match.")
        n_samples = X.shape[0]
        n_indexed = self._fit_X.shape[0]

        for i in range(self.n_estimators):
            bin_X = self.hash_functions_[i].transform(X)[:, 0]
            # gets the position to be added in the tree.
            positions = self.trees_[i].searchsorted(bin_X)
            # adds the hashed value into the tree.
            self.trees_[i] = np.insert(self.trees_[i],
                                       positions, bin_X)
            # add the entry into the original_indices_.
            self.original_indices_[i] = np.insert(self.original_indices_[i],
                                                  positions,
                                                  np.arange(n_indexed,
                                                            n_indexed +
                                                            n_samples))

        # adds the entry into the input_array.
        if sparse.issparse(X) or sparse.issparse(self._fit_X):
            self._fit_X = sparse.vstack((self._fit_X, X))
        else:
            self._fit_X = np.row_stack((self._fit_X, X))

        return self


"""
===========================
Medial axis skeletonization
===========================

The medial axis of an object is the set of all points having more than one
closest point on the object's boundary. It is often called the **topological
skeleton**, because it is a 1-pixel wide skeleton of the object, with the same
connectivity as the original object.

Here, we use the medial axis transform to compute the width of the foreground
objects. As the function ``medial_axis`` (``skimage.morphology.medial_axis``)
returns the distance transform in addition to the medial axis (with the keyword
argument ``return_distance=True``), it is possible to compute the distance to
the background for all points of the medial axis with this function. This gives
an estimate of the local width of the objects.

For a skeleton with fewer branches, there exists another skeletonization
algorithm in ``skimage``: ``skimage.morphology.skeletonize``, that computes
a skeleton by iterative morphological thinnings.

"""
import numpy as np
from scipy import ndimage as ndi
from skimage.morphology import medial_axis
import matplotlib.pyplot as plt

def microstructure(l=256):
    """
    Synthetic binary data: binary microstructure with blobs.

    Parameters
    ----------

    l: int, optional
        linear size of the returned image

    """
    n = 5
    x, y = np.ogrid[0:l, 0:l]
    mask = np.zeros((l, l))
    generator = np.random.RandomState(1)
    points = l * generator.rand(2, n**2)
    mask[(points[0]).astype(np.int), (points[1]).astype(np.int)] = 1
    mask = ndi.gaussian_filter(mask, sigma=l/(4.*n))
    return mask > mask.mean()

data = microstructure(l=64)

# Compute the medial axis (skeleton) and the distance transform
skel, distance = medial_axis(data, return_distance=True)

# Distance to the background for pixels of the skeleton
dist_on_skel = distance * skel

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))
ax1.imshow(data, cmap=plt.cm.gray, interpolation='nearest')
ax1.axis('off')
ax2.imshow(dist_on_skel, cmap=plt.cm.spectral, interpolation='nearest')
ax2.contour(data, [0.5], colors='w')
ax2.axis('off')

fig.subplots_adjust(hspace=0.01, wspace=0.01, top=1, bottom=0, left=0, right=1)
plt.show()


"""Extrudes a 2D mesh to generate an ExodusII 3D mesh.

Works with and assumes all polyhedra cells (and polygon faces).

To see usage, run:
------------------------------------------------------------

python meshing_ats.py -h

Example distributed with this source, to run:
------------------------------------------------------------

$> cd four-polygon-test
$> python ../meshing_ats.py -n 10 -d 1 ./four_polygon.vtk
$> mkdir run0
$> cd run0
$> ats --xml_file=../test1-fv-four-polygon.xml

Requires building the latest version of Exodus
------------------------------------------------------------

Note that this is typically done in your standard ATS installation,
assuming you have built your Amanzi TPLs with shared libraries (the
default through bootstrap).

In that case, simply ensure that ${AMANZI_TPLS_DIR}/SEACAS/lib is in
your PYTHONPATH.

"""

from __future__ import print_function

import sys,os
import numpy as np
import collections
import argparse

try:
    import exodus
except ImportError:
    sys.path.append(os.path.join(os.environ["SEACAS_DIR"],"lib"))
    import exodus

class SideSet(object):
    def __init__(self, name, setid, elem_list, side_list):
        assert(type(setid) == int)
        assert(type(elem_list) == list or type(elem_list) == np.ndarray)
        assert(type(side_list) == list or type(side_list) == np.ndarray)

        self.name = name
        self.setid = setid
        self.elem_list = elem_list
        self.side_list = side_list

class LabeledSet(object):
    def __init__(self, name, setid, entity, ent_ids):
        assert entity in ['CELL', 'FACE', 'NODE']
        assert(type(setid) == int)
        assert(type(ent_ids) == list or type(ent_ids) == np.ndarray)

        self.name = name
        self.setid = setid
        self.entity = entity
        self.ent_ids = np.array(ent_ids)

class Mesh2D(object):
    def __init__(self, coords, connectivity, labeled_sets=None, check_handedness=True):
        """
        Creates a 2D mesh from coordinates and a list cell-to-node connectivity lists.

        coords          : numpy array of shape (NCOORDS, NDIMS)
        connectivity    : list of lists of integer indices into coords specifying a
                          (clockwise OR counterclockwise) ordering of the nodes around
                          the 2D cell
        labeled_sets    : list of LabeledSet objects
        """
        assert type(coords) == np.ndarray
        assert len(coords.shape) == 2

        self.dim = coords.shape[1]

        self.coords = coords
        self.conn = connectivity
        if labeled_sets is not None:
            self.labeled_sets = labeled_sets
        else:
            self.labeled_sets = []

        self.validate()
        self.edge_counts()
        if check_handedness:
            self.check_handedness()

    def validate(self):
        assert self.coords.shape[1] == 2 or self.coords.shape[1] == 3
        assert type(self.conn) is list
        for f in self.conn:
            assert type(f) is list
            assert len(set(f)) == len(f)
            for i in f:
                assert i < self.coords.shape[0]

        for ls in self.labeled_sets:
            if ls.entity == "NODE":
                size = len(self.coords)
            elif ls.entity == "CELL":
                size = len(self.conn)

            for i in ls.ent_ids:
                assert i < size
        return True

    def num_cells(self):
        return len(self.conn)

    def num_nodes(self):
        return self.coords.shape[0]

    def num_edges(self):
        return len(self.edges())

    @staticmethod
    def edge_hash(i,j):
        return tuple(sorted((i,j)))

    def edges(self):
        return self.edge_counts().keys()

    def edge_counts(self):
        try:
            return self._edges
        except AttributeError:
            self._edges = collections.Counter(self.edge_hash(f[i], f[(i+1)%len(f)]) for f in self.conn for i in range(len(f)))
        return self._edges

    def check_handedness(self):
        for conn in self.conn:
            points = np.array([self.coords[c] for c in conn])
            cross = 0
            for i in range(len(points)):
                im = i - 1
                ip = i + 1
                if ip == len(points):
                    ip = 0

                p = points[ip] - points[i]
                m = points[i] - points[im]
                cross = cross + p[1] * m[0] - p[0] * m[1]
            if cross < 0:
                conn.reverse()

    def plot(self, color=None, ax=None):
        if color is None:
            import colors
            cm = colors.cm_mapper(0,self.num_cells()-1)
            colors = [cm(i) for i in range(self.num_cells())]
        else:
            colors = color

        verts = [[self.coords[i,0:2] for i in f] for f in self.conn]
        from matplotlib import collections
        gons = collections.PolyCollection(verts, facecolors=colors)
        from matplotlib import pyplot as plt
        if ax is None:
            fig,ax = plt.subplots(1,1)
        ax.add_collection(gons)
        ax.autoscale_view()

    @classmethod
    def read_VTK(cls, filename):
        try:
            return cls.read_VTK_Simplices(filename)
        except AssertionError:
            return cls.read_VTK_Unstructured(filename)

    @classmethod
    def read_VTK_Unstructured(cls, filename):
        with open(filename,'r') as fid:
            points_found = False
            polygons_found = False
            while True:
                line = fid.readline().decode('utf-8')
                if not line:
                    # EOF
                    break

                line = line.strip()
                if len(line) == 0:
                    continue

                split = line.split()
                section = split[0]

                if section == 'POINTS':
                    ncoords = int(split[1])
                    points = np.fromfile(fid, count=ncoords*3, sep=' ', dtype='d')
                    points = points.reshape(ncoords, 3)
                    points_found = True

                elif section == 'POLYGONS':
                    ncells = int(split[1])
                    n_to_read = int(split[2])

                    gons = []

                    data = np.fromfile(fid, count=n_to_read, sep=' ', dtype='i')
                    idx = 0
                    for i in range(ncells):
                        n_in_gon = data[idx]
                        gon = list(data[idx+1:idx+1+n_in_gon])

                        # check handedness -- need normals to point up!
                        cross = []
                        for i in range(len(gon)):
                            if i == len(gon)-1:
                                ip = 0
                                ipp = 1
                            elif i == len(gon)-2:
                                ip = i+1
                                ipp = 0
                            else:
                                ip = i+1
                                ipp = i+2
                            d2 = points[gon[ipp]] - points[gon[ip]]
                            d1 = points[gon[i]] - points[gon[ip]]
                            cross.append(np.cross(d2, d1))
                        if (np.array([c[2] for c in cross]).mean() < 0):
                            gon.reverse()

                        gons.append(gon)

                        idx += n_in_gon + 1
                    assert(idx == n_to_read)
                    polygons_found = True

        if not points_found:
            raise RuntimeError("Unstructured VTK must contain sections 'POINTS'")
        if not polygons_found:
            raise RuntimeError("Unstructured VTK must contain sections 'POLYGONS'")
        return cls(points, gons)

    @classmethod
    def read_VTK_Simplices(cls, filename):
        """Stolen from meshio, https://github.com/nschloe/meshio/blob/master/meshio/vtk_io.py"""
        import vtk_io
        with open(filename,'r') as fid:
            data = vtk_io.read_buffer(fid)

        points = data[0]
        if len(data[1]) != 1:
            raise RuntimeError("Simplex VTK file is readable by vtk_io but not by meshing_ats.  Includes: %r"%data[1].keys())

        gons = [v for v in data[1].itervalues()][0]
        gons = gons.tolist()

        # check handedness
        for gon in gons:
            cross = []
            for i in range(len(gon)):
                if i == len(gon)-1:
                    ip = 0
                    ipp = 1
                elif i == len(gon)-2:
                    ip = i+1
                    ipp = 0
                else:
                    ip = i+1
                    ipp = i+2
                d2 = points[gon[ipp]] - points[gon[ip]]
                d1 = points[gon[i]] - points[gon[ip]]
                cross.append(np.cross(d2, d1))
            if (np.array([c[2] for c in cross]).mean() < 0):
                gon.reverse()

        return cls(points, gons)

    @classmethod
    def from_Transect(cls, x, z, width=1):
        """Creates a 2D surface strip mesh from transect data"""
        # coordinates
        if (type(width) is list or type(width) is np.ndarray):
            variable_width = True
            y = np.array([0,1])
        else:
            variable_width = False
            y = np.array([0,width])

        Xc, Yc = np.meshgrid(x, y)
        if variable_width:
            assert(Yc.shape[1] == 2)
            assert(len(width) == Yc.shape[0])
            assert(min(width) > 0.)
            Yc[:,0] = -width/2.
            Yc[:,1] = width/2.

        Xc = Xc.flatten()
        Yc = Yc.flatten()
        Zc = np.concatenate([z,z])

        # connectivity
        nsurf_cells = len(x)-1
        conn = []
        for i in range(nsurf_cells):
            conn.append([i, i+1, nsurf_cells + i + 2, nsurf_cells + i + 1])

        coords = np.array([Xc, Yc, Zc])
        return cls(coords.transpose(), conn)

    @classmethod
    def from_Transect_Guide(cls, x, z, guide):
        """Creates a 2D surface strip mesh from transect data"""
        assert type(guide) == np.ndarray
        assert guide.shape[1] == 3

        # coordinates
        Xc = x
        Yc = np.zeros_like(x)
        Zc = z

        nsteps = guide.shape[0]
        xnew = Xc
        ynew = Yc
        znew = Zc
        for i in range(nsteps):
            xnew = xnew + guide[i][0]
            ynew = ynew + guide[i][1]
            znew = znew + guide[i][2]
            Xc =  np.concatenate([Xc, xnew])
            Yc =  np.concatenate([Yc, ynew])
            Zc =  np.concatenate([Zc, znew])

        # y = np.array([0,1,2])
        # Xc, Yc = np.meshgrid(x, y)
        # Xc = Xc.flatten()
        # Yc = Yc.flatten()
        # Zc = np.concatenate([z,z,z])

        # connectivity
        ns = len(x)
        conn = []
        for j in range(nsteps):
            for i in range(ns-1):
                conn.append([j*ns + i, j*ns + i + 1, (j+1)*ns + i + 1, (j+1)*ns + i ])

        coords = np.array([Xc, Yc, Zc])
        return cls(coords.transpose(), conn)

    @classmethod
    def from_Transect_GuideX(cls, x, z, guide, nsteps):
        """Creates a 2D surface strip mesh from transect data"""
        assert type(guide) == np.ndarray
        assert guide.shape[1] == 3

        # coordinates
        Xc = x
        Yc = np.zeros_like(x)
        Zc = z

        nsteps = guide.shape[0]

        xnew = np.zeros_like(x)
        ynew = np.zeros(len(x))
        znew = np.zeros_like(x)

        xnew[:] = Xc[:]
        ynew[:] = Yc[:]
        znew[:] = Zc[:]

        for i in range(nsteps):
            print(Yc)
            for j in range(len(x)):
                xnew[j] = xnew[j] + guide[j][0]
                ynew[j] = ynew[j] + guide[j][1]
                znew[j] = znew[j] + guide[j][2]

            Xc =  np.concatenate([Xc, xnew])
            Yc =  np.concatenate([Yc, ynew])
            Zc =  np.concatenate([Zc, znew])

        # y = np.array([0,1,2])
        # Xc, Yc = np.meshgrid(x, y)
        # Xc = Xc.flatten()
        # Yc = Yc.flatten()
        # Zc = np.concatenate([z,z,z])

        # connectivity
        ns = len(x)
        conn = []
        for j in range(nsteps):
            for i in range(ns-1):
                conn.append([j*ns + i, j*ns + i + 1, (j+1)*ns + i + 1, (j+1)*ns + i ])

        coords = np.array([Xc, Yc, Zc])
        return cls(coords.transpose(), conn)

class Mesh3D(object):
    def __init__(self, coords, face_to_node_conn, elem_to_face_conn,
                 side_sets=None, labeled_sets=None, material_ids=None):
        """
        Creates a 3D mesh from coordinates and connectivity lists.

        coords            : numpy array of shape (NCOORDS, 3)
        face_to_node_conn : list of lists of integer indices into coords specifying an
                            (clockwise OR counterclockwise) ordering of the nodes around
                            the face
        elem_to_face_conn : list of lists of integer indices into face_to_node_conn
                            specifying a list of faces that make up the elem
        """
        assert type(coords) == np.ndarray
        assert len(coords.shape) == 2
        assert coords.shape[1] == 3

        self.dim = coords.shape[1]

        self.coords = coords
        self.face_to_node_conn = face_to_node_conn
        self.elem_to_face_conn = elem_to_face_conn

        if labeled_sets is not None:
            self.labeled_sets = labeled_sets
        else:
            self.labeled_sets = []

        if side_sets is not None:
            self.side_sets = side_sets
        else:
            self.side_sets = []

        if material_ids is not None:
            self.material_id_list = collections.Counter(material_ids).keys()
            self.material_ids = material_ids
        else:
            self.material_id_list = [10000,]
            self.material_ids = [10000,]*len(self.elem_to_face_conn)

        self.validate()

    def validate(self):
        assert self.coords.shape[1] == 3
        assert type(self.face_to_node_conn) is list
        for f in self.face_to_node_conn:
            assert type(f) is list
            assert len(set(f)) == len(f)
            for i in f:
                assert i < self.coords.shape[0]

        assert type(self.elem_to_face_conn) is list
        for e in self.elem_to_face_conn:
            assert type(e) is list
            assert len(set(e)) == len(e)
            for i in e:
                assert i < len(self.face_to_node_conn)

        for ls in self.labeled_sets:
            if ls.entity == "NODE":
                size = self.num_nodes()
            if ls.entity == "FACE":
                size = self.num_faces()
            elif ls.entity == "CELL":
                size = self.num_cells()

            for i in ls.ent_ids:
                assert i < size

        for ss in self.side_sets:
            for j,i in zip(ss.elem_list, ss.side_list):
                assert j < self.num_cells()
                assert i < len(self.elem_to_face_conn[j])

    def num_cells(self):
        return len(self.elem_to_face_conn)

    def num_faces(self):
        return len(self.face_to_node_conn)

    def num_nodes(self):
        return self.coords.shape[0]

    def write_exodus(self, filename, face_block_mode="one block"):
        """Write the 3D mesh to ExodusII using arbitrary polyhedra spec"""

        # put cells in with blocks, which renumbers the cells, so we have to track sidesets.
        # Therefore we keep a map of old cell to new cell ordering
        #
        # also, though not required by the spec, paraview and visit
        # seem to crash if num_face_blocks != num_elem_blocks.  So
        # make face blocks here too, which requires renumbering the faces.

        # -- first pass, form all elem blocks and make the map from old-to-new
        new_to_old_elems = []
        elem_blks = []
        for i_m,m_id in enumerate(self.material_id_list):
            # split out elems of this material, save new_to_old map
            elems_tuple = [(i,c) for (i,c) in enumerate(self.elem_to_face_conn) if self.material_ids[i] == m_id]
            new_to_old_elems.extend([i for (i,c) in elems_tuple])
            elems = [c for (i,c) in elems_tuple]
            elem_blks.append(elems)

        old_to_new_elems = sorted([(old,i) for (i,old) in enumerate(new_to_old_elems)], lambda a,b: int.__cmp__(a[0],b[0]))

        # -- deal with faces, form all face blocks and make the map from old-to-new
        face_blks = []
        if face_block_mode == "one block":
            # no reordering of faces needed
            face_blks.append(self.face_to_node_conn)

        elif face_block_mode == "n blocks, not duplicated":
            used_faces = np.zeros((len(self.face_to_node_conn),),'bool')
            new_to_old_faces = []
            for i_m,m_id in enumerate(self.material_id_list):
                # split out faces of this material, save new_to_old map
                def used(f):
                    result = used_faces[f]
                    used_faces[f] = True
                    return result

                elem_blk = elem_blks[i_m]
                faces_tuple = [(f,self.face_to_node_conn[f]) for c in elem_blk for (j,f) in enumerate(c) if not used(f)]
                new_to_old_faces.extend([j for (j,f) in faces_tuple])
                faces = [f for (j,f) in faces_tuple]
                face_blks.append(faces)

            # get the renumbering in the elems
            old_to_new_faces = sorted([(old,j) for (j,old) in enumerate(new_to_old_faces)], lambda a,b: int.__cmp__(a[0],b[0]))
            elem_blks = [[[old_to_new_faces[f][1] for f in c] for c in elem_blk] for elem_blk in elem_blks]

        elif face_block_mode == "n blocks, duplicated":
            elem_blks_new = []
            offset = 0
            for i_m, m_id in enumerate(self.material_id_list):
                used_faces = np.zeros((len(self.face_to_node_conn),),'bool')
                def used(f):
                    result = used_faces[f]
                    used_faces[f] = True
                    return result

                elem_blk = elem_blks[i_m]

                tuple_old_f = [(f,self.face_to_node_conn[f]) for c in elem_blk for f in c if not used(f)]
                tuple_new_old_f = [(new,old,f) for (new,(old,f)) in enumerate(tuple_old_f)]

                old_to_new_blk = np.zeros((len(self.face_to_node_conn),),'i')-1
                for new,old,f in tuple_new_old_f:
                    old_to_new_blk[old] = new + offset

                elem_blk_new = [[old_to_new_blk[f] for f in c] for c in elem_blk]
                #offset = offset + len(ftuple_new)

                elem_blks_new.append(elem_blk_new)
                face_blks.append([f for i,j,f in tuple_new_old_f])
            elem_blks = elem_blks_new
        elif face_block_mode == "one block, repeated":
            # no reordering of faces needed, just repeat
            for eblock in elem_blks:
                face_blks.append(self.face_to_node_conn)
        else:
            raise RuntimeError("Invalid face_block_mode: '%s', valid='one block', 'n blocks, duplicated', 'n blocks, not duplicated'"%face_block_mode)

        # open the mesh file
        num_elems = sum(len(elem_blk) for elem_blk in elem_blks)
        num_faces = sum(len(face_blk) for face_blk in face_blks)
        ep = exodus.ex_init_params(title=filename,
                                   num_dim=3,
                                   num_nodes=self.num_nodes(),
                                   num_face=num_faces,
                                   num_face_blk=len(face_blks),
                                   num_elem=num_elems,
                                   num_elem_blk=len(elem_blks),
                                   num_side_sets=len(self.side_sets))
        e = exodus.exodus(filename, mode='w', array_type='numpy', init_params=ep)

        # put the coordinates
        e.put_coord_names(['coordX', 'coordY', 'coordZ'])
        e.put_coords(self.coords[:,0], self.coords[:,1], self.coords[:,2])

        # put the face blocks
        for i_blk, face_blk in enumerate(face_blks):
            face_raveled = [n for f in face_blk for n in f]
            e.put_polyhedra_face_blk(i_blk+1, len(face_blk), len(face_raveled), 0)
            e.put_node_count_per_face(i_blk+1, np.array([len(f) for f in face_blk]))
            e.put_face_node_conn(i_blk+1, np.array(face_raveled)+1)

        # put the elem blocks
        assert len(elem_blks) == len(self.material_id_list)
        for i_blk, (m_id, elem_blk) in enumerate(zip(self.material_id_list, elem_blks)):
            elems_raveled = [f for c in elem_blk for f in c]

            e.put_polyhedra_elem_blk(m_id, len(elem_blk), len(elems_raveled), 0)
            e.put_elem_blk_name(m_id, "MATERIAL_ID_%d"%m_id)
            e.put_face_count_per_polyhedra(m_id, np.array([len(c) for c in elem_blk]))
            e.put_elem_face_conn(m_id, np.array(elems_raveled)+1)

        # add sidesets
        e.put_side_set_names([ss.name for ss in self.side_sets])
        for ss in self.side_sets:
            for elem in ss.elem_list:
                assert old_to_new_elems[elem][0] == elem
            new_elem_list = [old_to_new_elems[elem][1] for elem in ss.elem_list]
            e.put_side_set_params(ss.setid, len(ss.elem_list), 0)
            e.put_side_set(ss.setid, np.array(new_elem_list)+1, np.array(ss.side_list)+1)

        # finish and close
        e.close()

    @classmethod
    def extruded_Mesh2D(cls, mesh2D, layer_types, layer_data, ncells_per_layer, mat_ids):
        """
        Regularly extrude a 2D mesh to make a 3D mesh.

        mesh2D              : a Mesh2D object
        layer_types         : either a string (type) or list of strings (types)
        layer_data          : array of data needed (specific to the type)
        ncells_per_layer    : either a single integer (same number of cells in all
                            : layers) or a list of number of cells in the layer
        mat_ids             : either a single integer (one mat_id for all layers)
                            : or a list of integers (mat_id for each layer)
                            : or a 2D numpy array of integers (mat_id for each layer and
                              each column: [layer_id, surface_cell_id])

        types:
          - 'constant'      : (data=float thickness) uniform thickness
          - 'function'      : (data=function or functor) thickness as a function
                            : of (x,y)
          - 'snapped'       : (data=float z coordinate) snap the layer to
                            : provided z coordinate, telescoping as needed
          - 'node'          : thickness provided on each node of the surface domain
          - 'cell'          : thickness provided on each cell of the surface domain,
                            : interpolate to nodes

        NOTE: dz is uniform through the layer in all but the 'snapped' case
        NOTE: 2D mesh is always labeled 'surface', extrusion is always downwards
        """

        # make the data all lists
        # ---------------------------------
        def is_list(data):
            if type(data) is str:
                return False
            try:
                len(data)
            except TypeError:
                return False
            else:
                return True

        if is_list(layer_types):
            if not is_list(layer_data):
                layer_data = [layer_data,]*len(layer_types)
            else:
                assert len(layer_data) == len(layer_types)

            if not is_list(ncells_per_layer):
                ncells_per_layer = [ncells_per_layer,]*len(layer_types)
            else:
                assert len(ncells_per_layer) == len(layer_types)

        elif is_list(layer_data):
            layer_types = [layer_types,]*len(layer_data)

            if not is_list(ncells_per_layer):
                ncells_per_layer = [ncells_per_layer,]*len(layer_data)
            else:
                assert len(ncells_per_layer) == len(layer_data)

        elif is_list(ncells_per_layer):
            layer_type = [layer_type,]*len(ncells_per_layer)
            layer_data = [layer_data,]*len(ncells_per_layer)
        else:
            layer_type = [layer_type,]
            layer_data = [layer_data,]
            ncells_per_layer = [ncells_per_layer,]

        # helper data and functions for mapping indices from 2D to 3D
        # ------------------------------------------------------------------
        if min(ncells_per_layer) < 0:
            raise RuntimeError("Invalid number of cells, negative value provided.")
        ncells_tall = sum(ncells_per_layer)
        ncells_total = ncells_tall * mesh2D.num_cells()
        nfaces_total = (ncells_tall+1) * mesh2D.num_cells() + ncells_tall * mesh2D.num_edges()
        nnodes_total = (ncells_tall+1) * mesh2D.num_nodes()

        np_mat_ids = np.array(mat_ids, dtype=int)
        if np_mat_ids.size == np.size(np_mat_ids, 0):
            if np_mat_ids.size == 1:
                np_mat_ids = np.full((len(ncells_per_layer), mesh2D.num_cells()), mat_ids[0], dtype=int)
            else:
                np_mat_ids = np.empty((len(ncells_per_layer), mesh2D.num_cells()), dtype=int)
                for ilay in range(len(ncells_per_layer)):
                    np_mat_ids[ilay, :] = np.full(mesh2D.num_cells(), mat_ids[ilay], dtype=int)

        def col_to_id(column, z_cell):
            """Maps 2D cell ID and index in the vertical to a 3D cell ID"""
            return z_cell + column * ncells_tall

        def node_to_id(node, z_node):
            """Maps 2D node ID and index in the vertical to a 3D node ID"""
            return z_node + node * (ncells_tall+1)

        def edge_to_id(edge, z_cell):
            """Maps 2D edge hash and index in the vertical to a 3D face ID of a vertical face"""
            return (ncells_tall + 1) * mesh2D.num_cells() + z_cell + edge * ncells_tall

        # create coordinates
        # ---------------------------------
        coords = np.zeros((mesh2D.coords.shape[0],ncells_tall+1, 3),'d')
        coords[:,:,0:2] = np.expand_dims(mesh2D.coords[:,0:2],1)

        if mesh2D.dim == 3:
            coords[:,0,2] = mesh2D.coords[:,2]
        # else the surface is at 0 depth

        cell_layer_start = 0
        for layer_type, layer_datum, ncells in zip(layer_types, layer_data, ncells_per_layer):
            if layer_type.lower() == 'constant':
                dz = float(layer_datum) / ncells
                for i in range(1,ncells+1):
                    coords[:,cell_layer_start+i,2] = coords[:,cell_layer_start,2] - i * dz

            else:
                # allocate an array of coordinates for the bottom of the layer
                layer_bottom = np.zeros((mesh2D.coords.shape[0],),'d')

                if layer_type.lower() == 'snapped':
                    # layer bottom is uniform
                    layer_bottom[:] = layer_datum

                elif layer_type.lower() == 'function':
                    # layer thickness is given by a function evaluation of x,y
                    for node_col in range(mesh2D.coords.shape[0]):
                        layer_bottom[node_col] = coords[node_col,cell_layer_start,2] - layer_datum(coords[node_col,0,0], coords[node_col,0,1])

                elif layer_type.lower() == 'node':
                    # layer bottom specifically provided through thickness
                    layer_bottom[:] = coords[:,cell_layer_start,2] - layer_datum

                elif layer_type.lower() == 'cell':
                    # interpolate cell thicknesses to node thicknesses
                    import scipy.interpolate
                    centroids = mesh2D.cell_centroids()
                    interp = scipy.interpolate.interp2d(centroids[:,0], centroids[:,1], layer_datum, kind='linear')
                    layer_bottom[:] = coords[:,cell_layer_start,2] - interp(mesh2D.coords[:,0], mesh2D.coords[:,1])

                else:
                    raise RuntimeError("Unrecognized layer_type '%s'"%layer_type)

                # linspace from bottom of previous layer to bottom of this layer
                for node_col in range(mesh2D.coords.shape[0]):
                    coords[node_col,cell_layer_start:cell_layer_start+ncells+1,2] = np.linspace(coords[node_col,cell_layer_start,2], layer_bottom[node_col], ncells+1)

            cell_layer_start = cell_layer_start + ncells

        # create faces, face sets, cells
        bottom = []
        surface = []
        faces = []
        cells = [list() for c in range(ncells_total)]

        # -- loop over the columns, adding the horizontal faces
        for col in range(mesh2D.num_cells()):
            nodes_2 = mesh2D.conn[col]
            surface.append(col_to_id(col,0))
            for z_face in range(ncells_tall + 1):
                i_f = len(faces)
                f = [node_to_id(n, z_face) for n in nodes_2]

                if z_face != ncells_tall:
                    cells[col_to_id(col, z_face)].append(i_f)
                if z_face != 0:
                    cells[col_to_id(col, z_face-1)].append(i_f)

                faces.append(f)
            bottom.append(col_to_id(col,ncells_tall-1))

        # -- loop over the columns, adding the vertical faces
        added = dict()
        vertical_side_cells = []
        vertical_side_indices = []
        for col in range(mesh2D.num_cells()):
            nodes_2 = mesh2D.conn[col]
            for i in range(len(nodes_2)):
                edge = mesh2D.edge_hash(nodes_2[i], nodes_2[(i+1)%len(nodes_2)])
                try:
                    i_e = added[edge]
                except KeyError:
                    # faces not yet added to facelist
                    i_e = len(added.keys())
                    added[edge] = i_e

                    for z_face in range(ncells_tall):
                        i_f = len(faces)
                        assert i_f == edge_to_id(i_e, z_face)
                        f = [node_to_id(edge[0], z_face),
                             node_to_id(edge[1], z_face),
                             node_to_id(edge[1], z_face+1),
                             node_to_id(edge[0], z_face+1)]
                        faces.append(f)
                        face_cell = col_to_id(col, z_face)
                        cells[face_cell].append(i_f)

                        # check if this is an external
                        if mesh2D._edges[edge] == 1:
                            vertical_side_cells.append(face_cell)
                            vertical_side_indices.append(len(cells[face_cell])-1)

                else:
                    # faces already added from previous column
                    for z_face in range(ncells_tall):
                        i_f = edge_to_id(i_e, z_face)
                        cells[col_to_id(col, z_face)].append(i_f)

        # Do some idiot checking
        # -- check we got the expected number of faces
        assert len(faces) == nfaces_total
        # -- check every cell is at least a tet
        for c in cells:
            assert len(c) > 4
        # -- check surface sideset has the right number of entries
        assert len(surface) == mesh2D.num_cells()
        # -- check bottom sideset has the right number of entries
        assert len(bottom) == mesh2D.num_cells()

        # -- len of vertical sides sideset is number of external edges * number of cells, no pinchouts here
        num_sides = ncells_tall * sum(1 for e,c in mesh2D.edge_counts().iteritems() if c == 1)
        assert num_sides == len(vertical_side_cells)
        assert num_sides == len(vertical_side_indices)

        # make the material ids
        material_ids = np.zeros((len(cells),),'i')
        for col in range(mesh2D.num_cells()):
            z_cell = 0
            for ilay in range(len(ncells_per_layer)):
                ncells = ncells_per_layer[ilay]
                for i in range(z_cell, z_cell+ncells):
                    material_ids[col_to_id(col, i)] = np_mat_ids[ilay, col]
                z_cell = z_cell + ncells

        # make the side sets
        side_sets = []
        side_sets.append(SideSet("bottom", 1, bottom, [1,]*len(bottom)))
        side_sets.append(SideSet("surface", 2, surface, [0,]*len(surface)))
        side_sets.append(SideSet("external_sides", 3, vertical_side_cells, vertical_side_indices))

        # reshape coords
        coords = coords.reshape(nnodes_total, 3)

        for e,s in zip(side_sets[0].elem_list, side_sets[0].side_list):
            face = cells[e][s]
            fz_coords = np.array([coords[n] for n in faces[face]])
            #print "bottom centroid = ", np.mean(fz_coords, axis=0)

        for e,s in zip(side_sets[1].elem_list, side_sets[1].side_list):
            face = cells[e][s]
            fz_coords = np.array([coords[n] for n in faces[face]])
            #print "surface centroid = ", np.mean(fz_coords, axis=0)

        # instantiate the mesh
        return cls(coords, faces, cells, side_sets=side_sets, material_ids=material_ids)

def commandline_options():
    parser = argparse.ArgumentParser(description='Extrude a 2D mesh to make a 3D mesh')
    parser.add_argument("-n", "--num-cells", default=10, type=int,
                        help="number of cells to extrude")
    parser.add_argument("-d", "--depth", default=40.0, type=float,
                        help="depth to extrude")
    parser.add_argument("-o", "--outfile", default=None, type=str,
                        help="output filename")
    parser.add_argument("-p", "--plot", default=False, action="store_true",
                        help="plot the 2D mesh")
    parser.add_argument("infile",metavar="INFILE", type=str,
                        help="input filename of surface mesh")

    options = parser.parse_args()

    if options.outfile is None:
        options.outfile = ".".join(options.infile.split(".")[:-1])+".exo"

    if os.path.isfile(options.outfile):
        print('Output file "%s" exists, cowardly not overwriting.'%options.outfile)
        sys.exit(1)

    if not os.path.isfile(options.infile):
        print('No input file provided')
        parser.print_usage()
        sys.exit(1)

    return options

if __name__ == "__main__":
    options = commandline_options()

    m2 = Mesh2D.read_VTK(options.infile)
    if options.plot:
        m2.plot()
    m3 = Mesh3D.extruded_Mesh2D(m2, [options.depth,], [options.num_cells,], [10000,])
    m3.write_exodus(options.outfile)


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""TensorFlowDataFrame implements convenience functions using TensorFlow."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import csv

import numpy as np

from tensorflow.contrib.learn.python.learn.dataframe import dataframe as df
from tensorflow.contrib.learn.python.learn.dataframe.transforms import batch
from tensorflow.contrib.learn.python.learn.dataframe.transforms import csv_parser
from tensorflow.contrib.learn.python.learn.dataframe.transforms import example_parser
from tensorflow.contrib.learn.python.learn.dataframe.transforms import in_memory_source
from tensorflow.contrib.learn.python.learn.dataframe.transforms import reader_source
from tensorflow.contrib.learn.python.learn.dataframe.transforms import sparsify
from tensorflow.contrib.learn.python.learn.dataframe.transforms import split_mask
from tensorflow.python.client import session as sess
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.training import coordinator
from tensorflow.python.training import queue_runner as qr

def _expand_file_names(filepatterns):
  """Takes a list of file patterns and returns a list of resolved file names."""
  if not isinstance(filepatterns, (list, tuple, set)):
    filepatterns = [filepatterns]
  filenames = set()
  for filepattern in filepatterns:
    names = set(gfile.Glob(filepattern))
    filenames |= names
  return list(filenames)

def _dtype_to_nan(dtype):
  if dtype is dtypes.string:
    return b""
  elif dtype.is_integer:
    return np.nan
  elif dtype.is_floating:
    return np.nan
  elif dtype is dtypes.bool:
    return np.nan
  else:
    raise ValueError("Can't parse type without NaN into sparse tensor: %s" %
                     dtype)

def _get_default_value(feature_spec):
  if isinstance(feature_spec, parsing_ops.FixedLenFeature):
    return feature_spec.default_value
  else:
    return _dtype_to_nan(feature_spec.dtype)

class TensorFlowDataFrame(df.DataFrame):
  """TensorFlowDataFrame implements convenience functions using TensorFlow."""

  def run(self,
          num_batches=None,
          graph=None,
          session=None,
          start_queues=True,
          initialize_variables=True,
          **kwargs):
    """Builds and runs the columns of the `DataFrame` and yields batches.

    This is a generator that yields a dictionary mapping column names to
    evaluated columns.

    Args:
      num_batches: the maximum number of batches to produce. If none specified,
        the returned value will iterate through infinite batches.
      graph: the `Graph` in which the `DataFrame` should be built.
      session: the `Session` in which to run the columns of the `DataFrame`.
      start_queues: if true, queues will be started before running and halted
        after producting `n` batches.
      initialize_variables: if true, variables will be initialized.
      **kwargs: Additional keyword arguments e.g. `num_epochs`.

    Yields:
      A dictionary, mapping column names to the values resulting from running
      each column for a single batch.
    """
    if graph is None:
      graph = ops.get_default_graph()
    with graph.as_default():
      if session is None:
        session = sess.Session()
      self_built = self.build(**kwargs)
      keys = list(self_built.keys())
      cols = list(self_built.values())
      if initialize_variables:
        if variables.local_variables():
          session.run(variables.local_variables_initializer())
        if variables.global_variables():
          session.run(variables.global_variables_initializer())
      if start_queues:
        coord = coordinator.Coordinator()
        threads = qr.start_queue_runners(sess=session, coord=coord)
      i = 0
      while num_batches is None or i < num_batches:
        i += 1
        try:
          values = session.run(cols)
          yield collections.OrderedDict(zip(keys, values))
        except errors.OutOfRangeError:
          break
      if start_queues:
        coord.request_stop()
        coord.join(threads)

  def select_rows(self, boolean_series):
    """Returns a `DataFrame` with only the rows indicated by `boolean_series`.

    Note that batches may no longer have consistent size after calling
    `select_rows`, so the new `DataFrame` may need to be rebatched.
    For example:
    '''
    filtered_df = df.select_rows(df["country"] == "jp").batch(64)
    '''

    Args:
      boolean_series: a `Series` that evaluates to a boolean `Tensor`.

    Returns:
      A new `DataFrame` with the same columns as `self`, but selecting only the
      rows where `boolean_series` evaluated to `True`.
    """
    result = type(self)()
    for key, col in self._columns.items():
      try:
        result[key] = col.select_rows(boolean_series)
      except AttributeError as e:
        raise NotImplementedError((
            "The select_rows method is not implemented for Series type {}. "
            "Original error: {}").format(type(col), e))
    return result

  def split(self, index_series, proportion, batch_size=None):
    """Deterministically split a `DataFrame` into two `DataFrame`s.

    Note this split is only as deterministic as the underlying hash function;
    see `tf.string_to_hash_bucket_fast`.  The hash function is deterministic
    for a given binary, but may change occasionally.  The only way to achieve
    an absolute guarantee that the split `DataFrame`s do not change across runs
    is to materialize them.

    Note too that the allocation of a row to one partition or the
    other is evaluated independently for each row, so the exact number of rows
    in each partition is binomially distributed.

    Args:
      index_series: a `Series` of unique strings, whose hash will determine the
        partitioning; or the name in this `DataFrame` of such a `Series`.
        (This `Series` must contain strings because TensorFlow provides hash
        ops only for strings, and there are no number-to-string converter ops.)
      proportion: The proportion of the rows to select for the 'left'
        partition; the remaining (1 - proportion) rows form the 'right'
        partition.
      batch_size: the batch size to use when rebatching the left and right
        `DataFrame`s.  If None (default), the `DataFrame`s are not rebatched;
        thus their batches will have variable sizes, according to which rows
        are selected from each batch of the original `DataFrame`.

    Returns:
      Two `DataFrame`s containing the partitioned rows.
    """
    if isinstance(index_series, str):
      index_series = self[index_series]
    left_mask, = split_mask.SplitMask(proportion)(index_series)
    right_mask = ~left_mask
    left_rows = self.select_rows(left_mask)
    right_rows = self.select_rows(right_mask)

    if batch_size:
      left_rows = left_rows.batch(batch_size=batch_size, shuffle=False)
      right_rows = right_rows.batch(batch_size=batch_size, shuffle=False)

    return left_rows, right_rows

  def split_fast(self, index_series, proportion, batch_size,
                 base_batch_size=1000):
    """Deterministically split a `DataFrame` into two `DataFrame`s.

    Note this split is only as deterministic as the underlying hash function;
    see `tf.string_to_hash_bucket_fast`.  The hash function is deterministic
    for a given binary, but may change occasionally.  The only way to achieve
    an absolute guarantee that the split `DataFrame`s do not change across runs
    is to materialize them.

    Note too that the allocation of a row to one partition or the
    other is evaluated independently for each row, so the exact number of rows
    in each partition is binomially distributed.

    Args:
      index_series: a `Series` of unique strings, whose hash will determine the
        partitioning; or the name in this `DataFrame` of such a `Series`.
        (This `Series` must contain strings because TensorFlow provides hash
        ops only for strings, and there are no number-to-string converter ops.)
      proportion: The proportion of the rows to select for the 'left'
        partition; the remaining (1 - proportion) rows form the 'right'
        partition.
      batch_size: the batch size to use when rebatching the left and right
        `DataFrame`s.  If None (default), the `DataFrame`s are not rebatched;
        thus their batches will have variable sizes, according to which rows
        are selected from each batch of the original `DataFrame`.
      base_batch_size: the batch size to use for materialized data, prior to the
        split.

    Returns:
      Two `DataFrame`s containing the partitioned rows.
    """
    if isinstance(index_series, str):
      index_series = self[index_series]
    left_mask, = split_mask.SplitMask(proportion)(index_series)
    right_mask = ~left_mask
    self["left_mask__"] = left_mask
    self["right_mask__"] = right_mask
    # TODO(soergel): instead of base_batch_size can we just do one big batch?
    # avoid computing the hashes twice
    m = self.materialize_to_memory(batch_size=base_batch_size)
    left_rows_df = m.select_rows(m["left_mask__"])
    right_rows_df = m.select_rows(m["right_mask__"])
    del left_rows_df[["left_mask__", "right_mask__"]]
    del right_rows_df[["left_mask__", "right_mask__"]]

    # avoid recomputing the split repeatedly
    left_rows_df = left_rows_df.materialize_to_memory(batch_size=batch_size)
    right_rows_df = right_rows_df.materialize_to_memory(batch_size=batch_size)
    return left_rows_df, right_rows_df

  def run_one_batch(self):
    """Creates a new 'Graph` and `Session` and runs a single batch.

    Returns:
      A dictionary mapping column names to numpy arrays that contain a single
      batch of the `DataFrame`.
    """
    return list(self.run(num_batches=1))[0]

  def run_one_epoch(self):
    """Creates a new 'Graph` and `Session` and runs a single epoch.

    Naturally this makes sense only for DataFrames that fit in memory.

    Returns:
      A dictionary mapping column names to numpy arrays that contain a single
      epoch of the `DataFrame`.
    """
    # batches is a list of dicts of numpy arrays
    batches = [b for b in self.run(num_epochs=1)]

    # first invert that to make a dict of lists of numpy arrays
    pivoted_batches = {}
    for k in batches[0].keys():
      pivoted_batches[k] = []
    for b in batches:
      for k, v in b.items():
        pivoted_batches[k].append(v)

    # then concat the arrays in each column
    result = {k: np.concatenate(column_batches)
              for k, column_batches in pivoted_batches.items()}
    return result

  def materialize_to_memory(self, batch_size):
    unordered_dict_of_arrays = self.run_one_epoch()

    # there may already be an 'index' column, in which case from_ordereddict)
    # below will complain because it wants to generate a new one.
    # for now, just remove it.
    # TODO(soergel): preserve index history, potentially many levels deep
    del unordered_dict_of_arrays["index"]

    # the order of the columns in this dict is arbitrary; we just need it to
    # remain consistent.
    ordered_dict_of_arrays = collections.OrderedDict(unordered_dict_of_arrays)
    return TensorFlowDataFrame.from_ordereddict(ordered_dict_of_arrays,
                                                batch_size=batch_size)

  def batch(self,
            batch_size,
            shuffle=False,
            num_threads=1,
            queue_capacity=None,
            min_after_dequeue=None,
            seed=None):
    """Resize the batches in the `DataFrame` to the given `batch_size`.

    Args:
      batch_size: desired batch size.
      shuffle: whether records should be shuffled. Defaults to true.
      num_threads: the number of enqueueing threads.
      queue_capacity: capacity of the queue that will hold new batches.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` with `batch_size` rows.
    """
    column_names = list(self._columns.keys())
    if shuffle:
      batcher = batch.ShuffleBatch(batch_size,
                                   output_names=column_names,
                                   num_threads=num_threads,
                                   queue_capacity=queue_capacity,
                                   min_after_dequeue=min_after_dequeue,
                                   seed=seed)
    else:
      batcher = batch.Batch(batch_size,
                            output_names=column_names,
                            num_threads=num_threads,
                            queue_capacity=queue_capacity)

    batched_series = batcher(list(self._columns.values()))
    dataframe = type(self)()
    dataframe.assign(**(dict(zip(column_names, batched_series))))
    return dataframe

  @classmethod
  def _from_csv_base(cls, filepatterns, get_default_values, has_header,
                     column_names, num_threads, enqueue_size,
                     batch_size, queue_capacity, min_after_dequeue, shuffle,
                     seed):
    """Create a `DataFrame` from CSV files.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      get_default_values: a function that produces a list of default values for
        each column, given the column names.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """
    filenames = _expand_file_names(filepatterns)
    if not filenames:
      raise ValueError("No matching file names.")

    if column_names is None:
      if not has_header:
        raise ValueError("If column_names is None, has_header must be true.")
      with gfile.GFile(filenames[0]) as f:
        column_names = csv.DictReader(f).fieldnames

    if "index" in column_names:
      raise ValueError(
          "'index' is reserved and can not be used for a column name.")

    default_values = get_default_values(column_names)

    reader_kwargs = {"skip_header_lines": (1 if has_header else 0)}
    index, value = reader_source.TextFileSource(
        filenames,
        reader_kwargs=reader_kwargs,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        num_threads=num_threads,
        seed=seed)()
    parser = csv_parser.CSVParser(column_names, default_values)
    parsed = parser(value)

    column_dict = parsed._asdict()
    column_dict["index"] = index

    dataframe = cls()
    dataframe.assign(**column_dict)
    return dataframe

  @classmethod
  def from_csv(cls,
               filepatterns,
               default_values,
               has_header=True,
               column_names=None,
               num_threads=1,
               enqueue_size=None,
               batch_size=32,
               queue_capacity=None,
               min_after_dequeue=None,
               shuffle=True,
               seed=None):
    """Create a `DataFrame` from CSV files.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      default_values: a list of default values for each column.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """

    def get_default_values(column_names):
      # pylint: disable=unused-argument
      return default_values

    return cls._from_csv_base(filepatterns, get_default_values, has_header,
                              column_names, num_threads,
                              enqueue_size, batch_size, queue_capacity,
                              min_after_dequeue, shuffle, seed)

  @classmethod
  def from_csv_with_feature_spec(cls,
                                 filepatterns,
                                 feature_spec,
                                 has_header=True,
                                 column_names=None,
                                 num_threads=1,
                                 enqueue_size=None,
                                 batch_size=32,
                                 queue_capacity=None,
                                 min_after_dequeue=None,
                                 shuffle=True,
                                 seed=None):
    """Create a `DataFrame` from CSV files, given a feature_spec.

    If `has_header` is false, then `column_names` must be specified. If
    `has_header` is true and `column_names` are specified, then `column_names`
    overrides the names in the header.

    Args:
      filepatterns: a list of file patterns that resolve to CSV files.
      feature_spec: a dict mapping column names to `FixedLenFeature` or
          `VarLenFeature`.
      has_header: whether or not the CSV files have headers.
      column_names: a list of names for the columns in the CSV files.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed lines.
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with examples from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """

    def get_default_values(column_names):
      return [_get_default_value(feature_spec[name]) for name in column_names]

    dataframe = cls._from_csv_base(filepatterns, get_default_values, has_header,
                                   column_names, num_threads,
                                   enqueue_size, batch_size, queue_capacity,
                                   min_after_dequeue, shuffle, seed)

    # replace the dense columns with sparse ones in place in the dataframe
    for name in dataframe.columns():
      if name != "index" and isinstance(feature_spec[name],
                                        parsing_ops.VarLenFeature):
        strip_value = _get_default_value(feature_spec[name])
        (dataframe[name],) = sparsify.Sparsify(strip_value)(dataframe[name])

    return dataframe

  @classmethod
  def from_examples(cls,
                    filepatterns,
                    features,
                    reader_cls=io_ops.TFRecordReader,
                    num_threads=1,
                    enqueue_size=None,
                    batch_size=32,
                    queue_capacity=None,
                    min_after_dequeue=None,
                    shuffle=True,
                    seed=None):
    """Create a `DataFrame` from `tensorflow.Example`s.

    Args:
      filepatterns: a list of file patterns containing `tensorflow.Example`s.
      features: a dict mapping feature names to `VarLenFeature` or
        `FixedLenFeature`.
      reader_cls: a subclass of `tensorflow.ReaderBase` that will be used to
        read the `Example`s.
      num_threads: the number of readers that will work in parallel.
      enqueue_size: block size for each read operation.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.

    Returns:
      A `DataFrame` that has columns corresponding to `features` and is filled
      with `Example`s from `filepatterns`.

    Raises:
      ValueError: no files match `filepatterns`.
      ValueError: `features` contains the reserved name 'index'.
    """
    filenames = _expand_file_names(filepatterns)
    if not filenames:
      raise ValueError("No matching file names.")

    if "index" in features:
      raise ValueError(
          "'index' is reserved and can not be used for a feature name.")

    index, record = reader_source.ReaderSource(
        reader_cls,
        filenames,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        num_threads=num_threads,
        seed=seed)()
    parser = example_parser.ExampleParser(features)
    parsed = parser(record)

    column_dict = parsed._asdict()
    column_dict["index"] = index

    dataframe = cls()
    dataframe.assign(**column_dict)
    return dataframe

  @classmethod
  def from_pandas(cls,
                  pandas_dataframe,
                  num_threads=None,
                  enqueue_size=None,
                  batch_size=None,
                  queue_capacity=None,
                  min_after_dequeue=None,
                  shuffle=True,
                  seed=None,
                  data_name="pandas_data"):
    """Create a `tf.learn.DataFrame` from a `pandas.DataFrame`.

    Args:
      pandas_dataframe: `pandas.DataFrame` that serves as a data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given
      `pandas_dataframe`.
    """
    pandas_source = in_memory_source.PandasSource(
        pandas_dataframe,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(pandas_source()._asdict()))
    return dataframe

  @classmethod
  def from_numpy(cls,
                 numpy_array,
                 num_threads=None,
                 enqueue_size=None,
                 batch_size=None,
                 queue_capacity=None,
                 min_after_dequeue=None,
                 shuffle=True,
                 seed=None,
                 data_name="numpy_data"):
    """Creates a `tf.learn.DataFrame` from a `numpy.ndarray`.

    The returned `DataFrame` contains two columns: 'index' and 'value'. The
    'value' column contains a row from the array. The 'index' column contains
    the corresponding row number.

    Args:
      numpy_array: `numpy.ndarray` that serves as a data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given
      array.
    """
    numpy_source = in_memory_source.NumpySource(
        numpy_array,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(numpy_source()._asdict()))
    return dataframe

  @classmethod
  def from_ordereddict(cls,
                       ordered_dict_of_arrays,
                       num_threads=None,
                       enqueue_size=None,
                       batch_size=None,
                       queue_capacity=None,
                       min_after_dequeue=None,
                       shuffle=True,
                       seed=None,
                       data_name="numpy_data"):
    """Creates a `tf.learn.DataFrame` from an `OrderedDict` of `numpy.ndarray`.

    The returned `DataFrame` contains a column for each key of the dict plus an
    extra 'index' column. The 'index' column contains the row number. Each of
    the other columns contains a row from the corresponding array.

    Args:
      ordered_dict_of_arrays: `OrderedDict` of `numpy.ndarray` that serves as a
          data source.
      num_threads: the number of threads to use for enqueueing.
      enqueue_size: the number of rows to enqueue per step.
      batch_size: desired batch size.
      queue_capacity: capacity of the queue that will store parsed `Example`s
      min_after_dequeue: minimum number of elements that can be left by a
        dequeue operation. Only used if `shuffle` is true.
      shuffle: whether records should be shuffled. Defaults to true.
      seed: passed to random shuffle operations. Only used if `shuffle` is true.
      data_name: a scope name identifying the data.

    Returns:
      A `tf.learn.DataFrame` that contains batches drawn from the given arrays.

    Raises:
      ValueError: `ordered_dict_of_arrays` contains the reserved name 'index'.
    """
    numpy_source = in_memory_source.OrderedDictNumpySource(
        ordered_dict_of_arrays,
        num_threads=num_threads,
        enqueue_size=enqueue_size,
        batch_size=batch_size,
        queue_capacity=queue_capacity,
        shuffle=shuffle,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        data_name=data_name)
    dataframe = cls()
    dataframe.assign(**(numpy_source()._asdict()))
    return dataframe


import matplotlib.pyplot as plt
import numpy as np
from scipy.fftpack import fft, fftshift
import sys

sys.path.append('../../../software/models/')
import utilFunctions as UF

(fs, x) = UF.wavread('../../../sounds/oboe-A4.wav')
N = 512
M = 401
hN = N/2
hM = (M+1)/2
start = .8*fs
xw = x[start-hM:start+hM-1] * np.hamming(M)

plt.figure(1, figsize=(9.5, 6.5))
plt.subplot(411)
plt.plot(np.arange(-hM, hM-1), xw, lw=1.5)
plt.axis([-hN, hN-1, min(xw), max(xw)])
plt.title('x (oboe-A4.wav), M = 401')

fftbuffer = np.zeros(N)
fftbuffer[:hM] = xw[hM-1:]
fftbuffer[N-hM+1:] = xw[:hM-1]
plt.subplot(412)
plt.plot(np.arange(0, N), fftbuffer, lw=1.5)
plt.axis([0, N, min(xw), max(xw)])
plt.title('fftbuffer: N = 512')

X = fftshift(fft(fftbuffer))
mX = 20 * np.log10(abs(X)/N)
pX = np.unwrap(np.angle(X))
plt.subplot(413)
plt.plot(np.arange(-hN, hN), mX, 'r', lw=1.5)
plt.axis([-hN,hN-1,-100,max(mX)])
plt.title('mX')

plt.subplot(414)
plt.plot(np.arange(-hN, hN), pX, 'c', lw=1.5)
plt.axis([-hN,hN-1,min(pX),max(pX)])
plt.title('pX')

plt.tight_layout()
plt.savefig('fft-zero-phase.png')
plt.show()


"""
testscript.py
-------------

This is a test script for decu.

"""

from decu import Script, experiment, figure, run_parallel
import numpy as np
import matplotlib.pyplot as plt

class TestScript(Script):

    @experiment(data_param='data')
    def exp(self, data, param, param2):
        """Compute x**param for each data point."""
        self.log.info('Working hard for {}..'.format(TestScript.exp.run))
        return np.power(data, param) + param2

    @figure()
    def plot_result(self, data, result):
        """Plot results of experiment."""
        plt.plot(data, result)

    @figure()
    def plot_many_results(self, data, results):
        """Plot results of experiment."""
        plt.figure()
        for res in results:
            plt.plot(data, res)

    def main(self):
        """Run some experiments and make some figures."""
        data = np.arange(5)
        result1 = self.exp(data, param=4, param2=10)
        self.plot_result(data, result1)

        param_list = [(data, x, y) for x, y in
                      zip(np.arange(5), np.arange(5, 10))]
        result2 = run_parallel(self.exp, param_list)
        self.plot_many_results(data, result2, suffix='parallel')