from __future__ import print_function
from numpy import asarray, unravel_index, prod, mod, ndarray, ceil, where, \
    r_, sort, argsort, array, random, arange, ones, expand_dims, sum
from itertools import groupby

from bolt.base import BoltArray
from bolt.spark.stack import StackedArray
from bolt.spark.utils import zip_with_index
from bolt.spark.statcounter import StatCounter
from bolt.utils import slicify, listify, tupleize, argpack, inshape, istransposeable, isreshapeable

class BoltArraySpark(BoltArray):

    _metadata = {
        '_shape': None,
        '_split': None,
        '_dtype': None,
        '_ordered': True
    }

    def __init__(self, rdd, shape=None, split=None, dtype=None, ordered=True):
        self._rdd = rdd
        self._shape = shape
        self._split = split
        self._dtype = dtype
        self._mode = 'spark'
        self._ordered = ordered

    @property
    def _constructor(self):
        return BoltArraySpark

    def __array__(self):
        return self.toarray()

    def cache(self):
        """
        Cache the underlying RDD in memory.
        """
        self._rdd.cache()

    def unpersist(self):
        """
        Remove the underlying RDD from memory.
        """
        self._rdd.unpersist()

    def repartition(self, npartitions):
        """
        Repartitions the underlying RDD

        Parameters
        ----------
        npartitions : int
            Number of partitions to repartion the underlying RDD to
        """

        rdd = self._rdd.repartition(npartitions)
        return self._constructor(rdd, ordered=False).__finalize__(self)

    def stack(self, size=None):
        """
        Aggregates records of a distributed array.

        Stacking should improve the performance of vectorized operations,
        but the resulting StackedArray object only exposes a restricted set
        of operations (e.g. map, reduce). The unstack method can be used
        to restore the full bolt array.

        Parameters
        ----------
        size : int, optional, default=None
            The maximum size for each stack (number of original records),
            will aggregate groups of records per partition up to this size,
            if None will aggregate all records on each partition.

        Returns
        -------
        StackedArray
        """
        stk = StackedArray(self._rdd, shape=self.shape, split=self.split)
        return stk.stack(size)

    def _align(self, axis):
        """
        Align spark bolt array so that axes for iteration are in the keys.

        This operation is applied before most functional operators.
        It ensures that the specified axes are valid, and swaps
        key/value axes so that functional operators can be applied
        over the correct records.

        Parameters
        ----------
        axis: tuple[int]
            One or more axes that wil be iterated over by a functional operator

        Returns
        -------
        BoltArraySpark
        """
        # ensure that the specified axes are valid
        inshape(self.shape, axis)

        # find the value axes that should be moved into the keys (axis >= split)
        tokeys = [(a - self.split) for a in axis if a >= self.split]

        # find the key axes that should be moved into the values (axis < split)
        tovalues = [a for a in range(self.split) if a not in axis]

        if tokeys or tovalues:
            return self.swap(tovalues, tokeys)
        else:
            return self

    def first(self):
        """
        Return the first element of an array
        """
        from bolt.local.array import BoltArrayLocal
        rdd = self._rdd if self._ordered else self._rdd.sortByKey()
        return BoltArrayLocal(rdd.values().first())

    def map(self, func, axis=(0,), value_shape=None, dtype=None, with_keys=False):
        """
        Apply a function across an axis.

        Array will be aligned so that the desired set of axes
        are in the keys, which may incur a swap.

        Parameters
        ----------
        func : function
            Function of a single array to apply. If with_keys=True,
            function should be of a (tuple, array) pair.

        axis : tuple or int, optional, default=(0,)
            Axis or multiple axes to apply function along.

        value_shape : tuple, optional, default=None
            Known shape of values resulting from operation

        dtype: numpy.dtype, optional, default=None
            Known dtype of values resulting from operation

        with_keys : bool, optional, default=False
            Include keys as an argument to the function

        Returns
        -------
        BoltArraySpark
        """
        axis = tupleize(axis)
        swapped = self._align(axis)

        if with_keys:
            test_func = lambda x: func(((0,), x))
        else:
            test_func = func

        if value_shape is None or dtype is None:
            # try to compute the size of each mapped element by applying func to a random array
            try:
                mapped = test_func(random.randn(*swapped.values.shape).astype(self.dtype))
            except Exception:
                first = swapped._rdd.first()
                if first:
                    # eval func on the first element
                    mapped = test_func(first[1])
            if value_shape is None:
                value_shape = mapped.shape
            if dtype is None:
                dtype = mapped.dtype

        shape = tuple([swapped._shape[ax] for ax in range(len(axis))]) + tupleize(value_shape)

        if with_keys:
            rdd = swapped._rdd.map(lambda kv: (kv[0], func(kv)))
        else:
            rdd = swapped._rdd.mapValues(func)

        # reshaping will fail if the elements aren't uniformly shaped
        def check(v):
            if len(v.shape) > 0 and v.shape != tupleize(value_shape):
                raise Exception("Map operation did not produce values of uniform shape.")
            return v

        rdd = rdd.mapValues(lambda v: check(v))

        return self._constructor(rdd, shape=shape, dtype=dtype, split=swapped.split).__finalize__(swapped)

    def filter(self, func, axis=(0,), sort=False):
        """
        Filter array along an axis.

        Applies a function which should evaluate to boolean,
        along a single axis or multiple axes. Array will be
        aligned so that the desired set of axes are in the keys,
        which may incur a swap.

        Parameters
        ----------
        func : function
            Function to apply, should return boolean

        axis : tuple or int, optional, default=(0,)
            Axis or multiple axes to filter along.

        sort: bool, optional, default=False
            Whether or not to sort by key before reindexing

        Returns
        -------
        BoltArraySpark
        """
        axis = tupleize(axis)

        swapped = self._align(axis)
        def f(record):
            return func(record[1])
        rdd = swapped._rdd.filter(f)
        if sort:
            rdd = rdd.sortByKey().values()
        else:
            rdd = rdd.values()

        # count the resulting array in order to reindex (linearize) the keys
        count, zipped = zip_with_index(rdd)
        if not count:
            count = zipped.count()
        reindexed = zipped.map(lambda kv: (tupleize(kv[1]), kv[0]))

        # since we can only filter over one axis, the remaining shape is always the following
        remaining = list(swapped.shape[len(axis):])
        if count != 0:
            shape = tuple([count] + remaining)
        else:
            shape = (0,)

        return self._constructor(reindexed, shape=shape, split=1).__finalize__(swapped)

    def reduce(self, func, axis=(0,), keepdims=False):
        """
        Reduce an array along an axis.

        Applies a commutative/associative function of two
        arguments cumulatively to all arrays along an axis.
        Array will be aligned so that the desired set of axes
        are in the keys, which may incur a swap.

        Parameters
        ----------
        func : function
            Function of two arrays that returns a single array

        axis : tuple or int, optional, default=(0,)
            Axis or multiple axes to reduce along.

        Returns
        -------
        BoltArraySpark
        """
        from bolt.local.array import BoltArrayLocal
        from numpy import ndarray

        axis = tupleize(axis)
        swapped = self._align(axis)
        arr = swapped._rdd.values().treeReduce(func, depth=3)

        if keepdims:
            for i in axis:
                arr = expand_dims(arr, axis=i)

        if not isinstance(arr, ndarray):
            # the result of a reduce can also be a scalar
            return arr
        elif arr.shape == (1,):
            # ndarrays with single values in them should be converted into scalars
            return arr[0]

        return BoltArrayLocal(arr)

    def _stat(self, axis=None, func=None, name=None, keepdims=False):
        """
        Compute a statistic over an axis.

        Can provide either a function (for use in a reduce)
        or a name (for use by a stat counter).

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        func : function, optional, default=None
            Function for reduce, see BoltArraySpark.reduce

        name : str
            A named statistic, see StatCounter

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        if axis is None:
            axis = list(range(len(self.shape)))
        axis = tupleize(axis)

        if func and not name:
            return self.reduce(func, axis, keepdims)

        if name and not func:
            from bolt.local.array import BoltArrayLocal

            swapped = self._align(axis)

            def reducer(left, right):
                return left.combine(right)

            counter = swapped._rdd.values()\
                             .mapPartitions(lambda i: [StatCounter(values=i, stats=name)])\
                             .treeReduce(reducer, depth=3)

            arr = getattr(counter, name)

            if keepdims:
                for i in axis:
                    arr = expand_dims(arr, axis=i)

            return BoltArrayLocal(arr).toscalar()

        else:
            raise ValueError('Must specify either a function or a statistic name.')

    def mean(self, axis=None, keepdims=False):
        """
        Return the mean of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        return self._stat(axis, name='mean', keepdims=keepdims)

    def var(self, axis=None, keepdims=False):
        """
        Return the variance of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        return self._stat(axis, name='variance', keepdims=keepdims)

    def std(self, axis=None, keepdims=False):
        """
        Return the standard deviation of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        return self._stat(axis, name='stdev', keepdims=keepdims)

    def sum(self, axis=None, keepdims=False):
        """
        Return the sum of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        from operator import add
        return self._stat(axis, func=add, keepdims=keepdims)

    def max(self, axis=None, keepdims=False):
        """
        Return the maximum of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        from numpy import maximum
        return self._stat(axis, func=maximum, keepdims=keepdims)

    def min(self, axis=None, keepdims=False):
        """
        Return the minimum of the array over the given axis.

        Parameters
        ----------
        axis : tuple or int, optional, default=None
            Axis to compute statistic over, if None
            will compute over all axes

        keepdims : boolean, optional, default=False
            Keep axis remaining after operation with size 1.
        """
        from numpy import minimum
        return self._stat(axis, func=minimum, keepdims=keepdims)

    def concatenate(self, arry, axis=0):
        """
        Join this array with another array.

        Paramters
        ---------
        arry : ndarray, BoltArrayLocal, or BoltArraySpark
            Another array to concatenate with

        axis : int, optional, default=0
            The axis along which arrays will be joined.

        Returns
        -------
        BoltArraySpark
        """
        if isinstance(arry, ndarray):
            from bolt.spark.construct import ConstructSpark
            arry = ConstructSpark.array(arry, self._rdd.context, axis=range(0, self.split))
        else:
            if not isinstance(arry, BoltArraySpark):
                raise ValueError("other must be local array or spark array, got %s" % type(arry))

        if not all([x == y if not i == axis else True
                    for i, (x, y) in enumerate(zip(self.shape, arry.shape))]):
            raise ValueError("all the input array dimensions except for "
                             "the concatenation axis must match exactly")

        if not self.split == arry.split:
            raise NotImplementedError("two arrays must have the same split ")

        if axis < self.split:
            shape = self.keys.shape

            def key_func(key):
                key = list(key)
                key[axis] += shape[axis]
                return tuple(key)

            rdd = self._rdd.union(arry._rdd.map(lambda kv: (key_func(kv[0]), kv[1])))

        else:
            from numpy import concatenate as npconcatenate
            shift = axis - self.split
            rdd = self._rdd.join(arry._rdd).map(lambda kv: (kv[0], npconcatenate(kv[1], axis=shift)))

        shape = tuple([x + y if i == axis else x
                      for i, (x, y) in enumerate(zip(self.shape, arry.shape))])

        return self._constructor(rdd, shape=shape, ordered=False).__finalize__(self)

    def _getbasic(self, index):
        """
        Basic indexing (for slices or ints).
        """
        key_slices = index[0:self.split]
        value_slices = index[self.split:]

        def key_check(key):
            def inrange(k, s):
                if s.step > 0:
                    return s.start <= k < s.stop
                else:
                    return s.stop < k <= s.start
            def check(k, s):
                return inrange(k, s) and mod(k - s.start, s.step) == 0
            out = [check(k, s) for k, s in zip(key, key_slices)]
            return all(out)

        def key_func(key):
            return tuple([(k - s.start)/s.step for k, s in zip(key, key_slices)])

        filtered = self._rdd.filter(lambda kv: key_check(kv[0]))

        if self._split == self.ndim:
            rdd = filtered.map(lambda kv: (key_func(kv[0]), kv[1]))
        else:
            # handle use of use slice.stop = -1 for a special case (see utils.slicify)
            value_slices = [s if s.stop != -1 else slice(s.start, None, s.step) for s in value_slices]
            rdd = filtered.map(lambda kv: (key_func(kv[0]), kv[1][value_slices]))

        shape = tuple([int(ceil((s.stop - s.start) / float(s.step))) for s in index])
        split = self.split
        return rdd, shape, split

    def _getadvanced(self, index):
        """
        Advanced indexing (for sets, lists, or ndarrays).
        """
        index = [asarray(i) for i in index]
        shape = index[0].shape
        if not all([i.shape == shape for i in index]):
            raise ValueError("shape mismatch: indexing arrays could not be broadcast "
                             "together with shapes " + ("%s " * self.ndim)
                             % tuple([i.shape for i in index]))

        index = tuple([listify(i, d) for (i, d) in zip(index, self.shape)])

        # build tuples with target indices
        key_tuples = list(zip(*index[0:self.split]))
        value_tuples = list(zip(*index[self.split:]))

        # build dictionary to look up targets in values
        d = {}
        for k, g in groupby(zip(value_tuples, key_tuples), lambda x: x[1]):
            d[k] = map(lambda x: x[0], list(g))

        def key_check(key):
            return key in key_tuples

        def key_func(key):
            return unravel_index(key, shape)

        # filter records based on key targets
        filtered = self._rdd.filter(lambda kv: key_check(kv[0]))

        # subselect and flatten records based on value targets (if they exist)
        if len(value_tuples) > 0:
            flattened = filtered.flatMap(lambda kv: [(kv[0], kv[1][i]) for i in d[kv[0]]])
        else:
            flattened = filtered

        # reindex
        indexed = flattened.zipWithIndex()
        rdd = indexed.map(lambda kkv: (key_func(kkv[1]), kkv[0][1]))
        split = len(shape)

        return rdd, shape, split

    def _getmixed(self, index):
        """
        Mixed indexing (combines basic and advanced indexes)

        Assumes that only a single advanced index is used, due to the complicated
        behavior needed to be compatible with NumPy otherwise.
        """
        # find the single advanced index
        loc = where([isinstance(i, (tuple, list, ndarray)) for i in index])[0][0]
        idx = list(index[loc])

        if isinstance(idx[0], (tuple, list, ndarray)):
            raise ValueError("When mixing basic and advanced indexing, "
                             "advanced index must be one-dimensional")

        # single advanced index is on a key -- filter and update key
        if loc < self.split:
            def newkey(key):
                newkey = list(key)
                newkey[loc] = idx.index(key[loc])
                return tuple(newkey)
            rdd = self._rdd.filter(lambda kv: kv[0][loc] in idx).map(lambda kv: (newkey(kv[0]), kv[1]))
        # single advanced index is on a value -- use NumPy indexing
        else:
            slices = [slice(0, None, None) for _ in self.values.shape]
            slices[loc - self.split] = idx
            rdd = self._rdd.map(lambda kv: (kv[0], kv[1][slices]))
        newshape = list(self.shape)
        newshape[loc] = len(idx)
        barray = self._constructor(rdd, shape=tuple(newshape)).__finalize__(self)

        # apply the rest of the simple indices
        new_index = index[:]
        new_index[loc] = slice(0, None, None)
        barray = barray[tuple(new_index)]
        return barray._rdd, barray.shape, barray.split

    def __getitem__(self, index):
        """
        Get an item from the array through indexing.

        Supports basic indexing with slices and ints, or advanced
        indexing with lists or ndarrays of integers.
        Mixing basic and advanced indexing across axes is currently supported
        only for a single advanced index amidst multiple basic indices.

        Parameters
        ----------
        index : tuple of slices, ints, list, tuple, or ndarrays
            One or more index specifications

        Returns
        -------
        BoltSparkArray
        """
        if isinstance(index, tuple):
            index = list(index)
        else:
            index = [index]
        int_locs = where([isinstance(i, int) for i in index])[0]

        if len(index) > self.ndim:
            raise ValueError("Too many indices for array")

        if not all([isinstance(i, (slice, int, list, tuple, ndarray)) for i in index]):
            raise ValueError("Each index must either be a slice, int, list, set, or ndarray")

        # fill unspecified axes with full slices
        if len(index) < self.ndim:
            index += tuple([slice(0, None, None) for _ in range(self.ndim - len(index))])

        # standardize slices and bounds checking
        for n, idx in enumerate(index):
            size = self.shape[n]
            if isinstance(idx, (slice, int)):
                slc = slicify(idx, size)
                # throw an error if this would lead to an empty dimension in numpy
                if slc.step > 0:
                    minval, maxval = slc.start, slc.stop
                else:
                    minval, maxval = slc.stop, slc.start
                if minval > size-1 or maxval < 1 or minval >= maxval:
                    raise ValueError("Index {} in dimension {} with shape {} would "
                                     "produce an empty dimension".format(idx, n, size))
                index[n] = slc
            else:
                adjusted = array(idx)
                inds = where(adjusted<0)
                adjusted[inds] += size
                if adjusted.min() < 0 or adjusted.max() > size-1:
                    raise ValueError("Index {} out of bounds in dimension {} with "
                                     "shape {}".format(idx, n, size))
                index[n] = adjusted

        # select basic or advanced indexing
        if all([isinstance(i, slice) for i in index]):
            rdd, shape, split = self._getbasic(index)
        elif all([isinstance(i, (tuple, list, ndarray)) for i in index]):
            rdd, shape, split = self._getadvanced(index)
        elif sum([isinstance(i, (tuple, list, ndarray)) for i in index]) == 1:
            rdd, shape, split = self._getmixed(index)
        else:
            raise NotImplementedError("When mixing basic indexing (slices and int) with "
                                      "with advanced indexing (lists, tuples, and ndarrays), "
                                      "can only have a single advanced index")

        # if any key indices used negative steps, records are no longer ordered
        if self._ordered is False or any([isinstance(s, slice) and s.step<0 for s in index[:self.split]]):
            ordered = False
        else:
            ordered = True

        result = self._constructor(rdd, shape=shape, split=split, ordered=ordered).__finalize__(self)

        # squeeze out int dimensions (and squeeze to singletons if all ints)
        if len(int_locs) == self.ndim:
            return result.squeeze().toarray()[()]
        else:
            return result.squeeze(tuple(int_locs))

    def chunk(self, size="150", axis=None, padding=None):
        """
        Chunks records of a distributed array.

        Chunking breaks arrays into subarrays, using an specified
        size of chunks along each value dimension. Can alternatively
        specify an average chunk byte size (in kilobytes) and the size of
        chunks (as ints) will be computed automatically.

        Parameters
        ----------
        size : tuple, int, or str, optional, default = "150"
            A string giving the size in kilobytes, or a tuple with the size
            of chunks along each dimension.

        axis : int or tuple, optional, default = None
            One or more axis to chunk array along, if None
            will use all axes,

        padding: tuple or int, default = None
            Number of elements per dimension that will overlap with the adjacent chunk.
            If a tuple, specifies padding along each chunked dimension; if a int, same
            padding will be applied to all chunked dimensions.

        Returns
        -------
        ChunkedArray
        """
        if type(size) is not str:
            size = tupleize((size))
        axis = tupleize((axis))
        padding = tupleize((padding))

        from bolt.spark.chunk import ChunkedArray

        chnk = ChunkedArray(rdd=self._rdd, shape=self._shape, split=self._split, dtype=self._dtype)
        return chnk._chunk(size, axis, padding)

    def swap(self, kaxes, vaxes, size="150"):
        """
        Swap axes from keys to values.

        This is the core operation underlying shape manipulation
        on the Spark bolt array. It exchanges an arbitrary set of axes
        between the keys and the valeus. If either is None, will only
        move axes in one direction (from keys to values, or values to keys).
        Keys moved to values will be placed immediately after the split;
        values moved to keys will be placed immediately before the split.

        Parameters
        ----------
        kaxes : tuple
            Axes from keys to move to values

        vaxes : tuple
            Axes from values to move to keys

        size : tuple or int, optional, default = "150"
            Can either provide a string giving the size in kilobytes,
            or a tuple with the number of chunks along each
            value dimension being moved

        Returns
        -------
        BoltArraySpark
        """
        kaxes = asarray(tupleize(kaxes), 'int')
        vaxes = asarray(tupleize(vaxes), 'int')
        if type(size) is not str:
            size = tupleize(size)

        if len(kaxes) == self.keys.ndim and len(vaxes) == 0:
            raise ValueError('Cannot perform a swap that would '
                             'end up with all data on a single key')

        if len(kaxes) == 0 and len(vaxes) == 0:
            return self

        from bolt.spark.chunk import ChunkedArray

        chunks = self.chunk(size)

        swapped = chunks.keys_to_values(kaxes).values_to_keys([v+len(kaxes) for v in vaxes])
        barray = swapped.unchunk()

        return barray

    def transpose(self, *axes):
        """
        Return an array with the axes transposed.

        This operation will incur a swap unless the
        desiured permutation can be obtained
        only by transpoing the keys or the values.

        Parameters
        ----------
        axes : None, tuple of ints, or n ints
            If None, will reverse axis order.
        """
        if len(axes) == 0:
            p = arange(self.ndim-1, -1, -1)
        else:
            p = asarray(argpack(axes))

        istransposeable(p, range(self.ndim))

        split = self.split

        # compute the keys/value axes that need to be swapped
        new_keys, new_values = p[:split], p[split:]
        swapping_keys = sort(new_values[new_values < split])
        swapping_values = sort(new_keys[new_keys >= split])
        stationary_keys = sort(new_keys[new_keys < split])
        stationary_values = sort(new_values[new_values >= split])

        # compute the permutation that the swap causes
        p_swap = r_[stationary_keys, swapping_values, swapping_keys, stationary_values]

        # compute the extra permutation (p_x)  on top of this that
        # needs to happen to get the full permutation desired
        p_swap_inv = argsort(p_swap)
        p_x = p_swap_inv[p]
        p_keys, p_values = p_x[:split], p_x[split:]-split

        # perform the swap and the the within key/value permutations
        arr = self.swap(swapping_keys, swapping_values-split)
        arr = arr.keys.transpose(tuple(p_keys.tolist()))
        arr = arr.values.transpose(tuple(p_values.tolist()))

        return arr

    @property
    def T(self):
        """
        Transpose by reversing the order of the axes.
        """
        return self.transpose()

    def swapaxes(self, axis1, axis2):
        """
        Return the array with two axes interchanged.

        Parameters
        ----------
        axis1 : int
            The first axis to swap

        axis2 : int
            The second axis to swap
        """
        p = list(range(self.ndim))
        p[axis1] = axis2
        p[axis2] = axis1

        return self.transpose(p)

    def reshape(self, *shape):
        """
        Return an array with the same data but a new shape.

        Currently only supports reshaping that independently
        reshapes the keys, or the values, or both.

        Parameters
        ----------
        shape :  tuple of ints, or n ints
            New shape
        """
        new = argpack(shape)
        isreshapeable(new, self.shape)

        if new == self.shape:
            return self

        i = self._reshapebasic(new)
        if i == -1:
            raise NotImplementedError("Currently no support for reshaping between "
                                      "keys and values for BoltArraySpark")
        else:
            new_key_shape, new_value_shape = new[:i], new[i:]
            return self.keys.reshape(new_key_shape).values.reshape(new_value_shape)

    def _reshapebasic(self, shape):
        """
        Check if the requested reshape can be broken into independant reshapes
        on the keys and values. If it can, returns the index in the new shape
        separating keys from values, otherwise returns -1
        """
        new = tupleize(shape)
        old_key_size = prod(self.keys.shape)
        old_value_size = prod(self.values.shape)

        for i in range(len(new)):
            new_key_size = prod(new[:i])
            new_value_size = prod(new[i:])
            if new_key_size == old_key_size and new_value_size == old_value_size:
                return i

        return -1

    def squeeze(self, axis=None):
        """
        Remove one or more single-dimensional axes from the array.

        Parameters
        ----------
        axis : tuple or int
            One or more singleton axes to remove.
        """
        if not any([d == 1 for d in self.shape]):
            return self

        if axis is None:
            drop = where(asarray(self.shape) == 1)[0]
        elif isinstance(axis, int):
            drop = asarray((axis,))
        elif isinstance(axis, tuple):
            drop = asarray(axis)
        else:
            raise ValueError("an integer or tuple is required for the axis")

        if any([self.shape[i] > 1 for i in drop]):
            raise ValueError("cannot select an axis to squeeze out which has size greater than one")

        if any(asarray(drop) < self.split):
            kmask = set([d for d in drop if d < self.split])
            kfunc = lambda k: tuple([kk for ii, kk in enumerate(k) if ii not in kmask])
        else:
            kfunc = lambda k: k

        if any(asarray(drop) >= self.split):
            vmask = tuple([d - self.split for d in drop if d >= self.split])
            vfunc = lambda v: v.squeeze(vmask)
        else:
            vfunc = lambda v: v

        rdd = self._rdd.map(lambda kv: (kfunc(kv[0]), vfunc(kv[1])))
        shape = tuple([ss for ii, ss in enumerate(self.shape) if ii not in drop])
        split = len([d for d in range(self.keys.ndim) if d not in drop])
        return self._constructor(rdd, shape=shape, split=split).__finalize__(self)

    def astype(self, dtype, casting='unsafe'):
        """
        Cast the array to a specified type.

        Parameters
        ----------
        dtype : str or dtype
            Typecode or data-type to cast the array to (see numpy)
        """
        rdd = self._rdd.mapValues(lambda v: v.astype(dtype, 'K', casting))
        return self._constructor(rdd, dtype=dtype).__finalize__(self)

    def clip(self, min=None, max=None):
        """
        Clip values above and below.

        Parameters
        ----------
        min : scalar or array-like
            Minimum value. If array, will be broadcasted

        max : scalar or array-like
            Maximum value. If array, will be broadcasted.
        """
        rdd = self._rdd.mapValues(lambda v: v.clip(min=min, max=max))
        return self._constructor(rdd).__finalize__(self)

    @property
    def shape(self):
        """
        Size of each dimension.
        """
        return self._shape

    @property
    def size(self):
        """
        Total number of elements.
        """
        return prod(self._shape)

    @property
    def ndim(self):
        """
        Number of dimensions.
        """
        return len(self._shape)

    @property
    def split(self):
        """
        Axis at which the array is split into keys/values.
        """
        return self._split

    @property
    def dtype(self):
        """
        Data-type of array.
        """
        return self._dtype

    @property
    def mask(self):
        return tuple([1] * len(self.keys.shape) + [0] * len(self.values.shape))

    @property
    def keys(self):
        """
        Returns a restricted keys.
        """
        from bolt.spark.shapes import Keys
        return Keys(self)

    @property
    def values(self):
        from bolt.spark.shapes import Values
        return Values(self)

    def tolocal(self):
        """
        Returns a local bolt array by first collecting as an array.
        """
        from bolt.local.array import BoltArrayLocal
        return BoltArrayLocal(self.toarray())

    def toarray(self):
        """
        Returns the contents as a local array.

        Will likely cause memory problems for large objects.
        """
        rdd = self._rdd if self._ordered else self._rdd.sortByKey()
        x = rdd.values().collect()
        return asarray(x).reshape(self.shape)

    def tordd(self):
        """
        Return the underlying RDD of the bolt array.
        """
        return self._rdd

    def display(self):
        """
        Show a pretty-printed representation of this BoltArrayLocal.
        """
        for x in self._rdd.take(10):
            print(x)


from numpy import ndarray, asarray, prod, concatenate, expand_dims
from numpy import any as npany
from collections import Iterable

def tupleize(arg):
    """
    Coerce singletons and lists and ndarrays to tuples.

    Parameters
    ----------
    arg : tuple, list, ndarray, or singleton
        Item to coerce
    """
    if arg is None:
        return None
    if not isinstance(arg, (tuple, list, ndarray, Iterable)):
        return tuple((arg,))
    elif isinstance(arg, (list, ndarray)):
        return tuple(arg)
    elif isinstance(arg, Iterable) and not isinstance(arg, str):
        return tuple(arg)
    else:
        return arg

def argpack(args):
    """
    Coerce a list of arguments to a tuple.

    Parameters
    ----------
    args : tuple or nested tuple
        Pack arguments into a tuple, converting ((,...),) or (,) -> (,)
    """
    if isinstance(args[0], (tuple, list, ndarray)):
        return tupleize(args[0])
    elif isinstance(args[0], Iterable) and not isinstance(args[0], str):
        # coerce any iterable into a list before calling tupleize (Python 3 compatibility)
        return tupleize(list(args[0]))
    else:
        return tuple(args)

def inshape(shape, axes):
    """
    Checks to see if a list of axes are contained within an array shape.

    Parameters
    ----------
    shape : tuple[int]
        the shape of a BoltArray

    axes : tuple[int]
        the axes to check against shape
    """
    valid = all([(axis < len(shape)) and (axis >= 0) for axis in axes])
    if not valid:
        raise ValueError("axes not valid for an ndarray of shape: %s" % str(shape))

def allclose(a, b):
    """
    Test that a and b are close and match in shape.

    Parameters
    ----------
    a : ndarray
        First array to check

    b : ndarray
        First array to check
    """
    from numpy import allclose
    return (a.shape == b.shape) and allclose(a, b)

def tuplesort(seq):
    """
    Sort a list by a sequence.

    Parameters
    ----------
    seq : tuple
        Sequence to sort by
    """

    return sorted(range(len(seq)), key=seq.__getitem__)

def listify(lst, dim):
    """
    Flatten lists of indices and ensure bounded by a known dim.

    Parameters
    ----------
    lst : list
        List of integer indices

    dim : tuple
        Bounds for indices
    """
    if not all([l.dtype == int for l in lst]):
        raise ValueError("indices must be integers")

    if npany(asarray(lst) >= dim):
        raise ValueError("indices out of bounds for axis with size %s" % dim)

    return lst.flatten()

def slicify(slc, dim):
    """
    Force a slice to have defined start, stop, and step from a known dim.
    Start and stop will always be positive. Step may be negative.

    There is an exception where a negative step overflows the stop needs to have
    the default value set to -1. This is the only case of a negative start/stop
    value.

    Parameters
    ----------
    slc : slice or int
        The slice to modify, or int to convert to a slice

    dim : tuple
        Bound for slice
    """
    if isinstance(slc, slice):

        # default limits
        start = 0 if slc.start is None else slc.start
        stop = dim if slc.stop is None else slc.stop
        step = 1 if slc.step is None else slc.step
        # account for negative indices
        if start < 0: start += dim
        if stop < 0: stop += dim
        # account for over-flowing the bounds
        if step > 0:
            if start < 0: start = 0
            if stop > dim: stop = dim
        else:
            if stop < 0: stop = -1
            if start > dim: start = dim-1

        return slice(start, stop, step)

    elif isinstance(slc, int):
        if slc < 0:
            slc += dim
        return slice(slc, slc+1, 1)

    else:
        raise ValueError("Type for slice %s not recongized" % type(slc))

def istransposeable(new, old):
    """
    Check to see if a proposed tuple of axes is a valid permutation
    of an old set of axes. Checks length, axis repetion, and bounds.

    Parameters
    ----------
    new : tuple
        tuple of proposed axes

    old : tuple
        tuple of old axes
    """

    new, old = tupleize(new), tupleize(old)

    if not len(new) == len(old):
        raise ValueError("Axes do not match axes of keys")

    if not len(set(new)) == len(set(old)):
        raise ValueError("Repeated axes")

    if any(n < 0 for n in new) or max(new) > len(old) - 1:
        raise ValueError("Invalid axes")

def isreshapeable(new, old):
    """
    Check to see if a proposed tuple of axes is a valid reshaping of
    the old axes by ensuring that they can be factored.

    Parameters
    ----------
    new : tuple
        tuple of proposed axes

    old : tuple
        tuple of old axes
    """

    new, old = tupleize(new), tupleize(old)

    if not prod(new) == prod(old):
        raise ValueError("Total size of new keys must remain unchanged")

def allstack(vals, depth=0):
    """
    If an ndarray has been split into multiple chunks by splitting it along
    each axis at a number of locations, this function rebuilds the
    original array from chunks.

    Parameters
    ----------
    vals : nested lists of ndarrays
        each level of nesting of the lists representing a dimension of
        the original array.
    """
    if type(vals[0]) is ndarray:
        return concatenate(vals, axis=depth)
    else:
        return concatenate([allstack(x, depth+1) for x in vals], axis=depth)

def iterexpand(arry, extra):
    """
    Expand dimensions by iteratively append empty axes.

    Parameters
    ----------
    arry : ndarray
        The original array

    extra : int
        The number of empty axes to append
    """
    for d in range(arry.ndim, arry.ndim+extra):
        arry = expand_dims(arry, axis=d)
    return arry


def get_kv_shape(shape, key_axes):
    func = lambda axis: shape[axis]
    return _get_kv_func(func, shape, key_axes)

def get_kv_axes(shape, key_axes):
    func = lambda axis: axis
    return _get_kv_func(func, shape, key_axes)

def _get_kv_func(func, shape, key_axes):
    key_res = [func(axis) for axis in key_axes]
    value_res = [func(axis) for axis in range(len(shape)) if axis not in key_axes]
    return key_res, value_res

def zip_with_index(rdd):
    """
    Alternate version of Spark's zipWithIndex that eagerly returns count.
    """
    starts = [0]
    if rdd.getNumPartitions() > 1:
        nums = rdd.mapPartitions(lambda it: [sum(1 for _ in it)]).collect()
        count = sum(nums)
        for i in range(len(nums) - 1):
            starts.append(starts[-1] + nums[i])
    else:
        count = rdd.count()

    def func(k, it):
        for i, v in enumerate(it, starts[k]):
            yield v, i

    return count, rdd.mapPartitionsWithIndex(func)


from bolt.local.construct import ConstructLocal
from bolt.spark.construct import ConstructSpark

constructors = [
    ('local', ConstructLocal),
    ('spark', ConstructSpark)
]

def wrapped(f):
    """
    Decorator to append routed docstrings
    """
    import inspect

    def extract(func):
        append = ""
        args = inspect.getargspec(func)
        for i, a in enumerate(args.args):
            if i < (len(args) - len(args.defaults)):
                append += str(a) + ", "
            else:
                default = args.defaults[i-len(args.defaults)]
                if hasattr(default, "__name__"):
                    default = default.__name__
                else:
                    default = str(default)
                append += str(a) + "=" + default + ", "
        append = append[:-2] + ")"
        return append

    doc = f.__doc__ + "\n"
    doc += "    local -> array(" + extract(getattr(ConstructLocal, f.__name__)) + "\n"
    doc += "    spark -> array(" + extract(getattr(ConstructSpark, f.__name__)) + "\n"
    f.__doc__ = doc
    return f

def lookup(*args, **kwargs):
    """
    Use arguments to route constructor.

    Applies a series of checks on arguments to identify constructor,
    starting with known keyword arguments, and then applying
    constructor-specific checks
    """
    if 'mode' in kwargs:
        mode = kwargs['mode']
        if mode not in constructors:
            raise ValueError('Mode %s not supported' % mode)
        del kwargs['mode']
        return constructors[mode]
    else:
        for mode, constructor in constructors:
            if constructor._argcheck(*args, **kwargs):
                return constructor
    return ConstructLocal

@wrapped
def array(*args, **kwargs):
    """
    Create a bolt array.
    """
    return lookup(*args, **kwargs).dispatch('array', *args, **kwargs)

@wrapped
def ones(*args, **kwargs):
    """
    Create a bolt array of ones.
    """
    return lookup(*args, **kwargs).dispatch('ones', *args, **kwargs)

@wrapped
def zeros(*args, **kwargs):
    """
    Create a bolt array of zeros.
    """
    return lookup(*args, **kwargs).dispatch('zeros', *args, **kwargs)

@wrapped
def concatenate(*args, **kwargs):
    """
    Create a bolt array of ones.
    """
    return lookup(*args, **kwargs).dispatch('concatenate', *args, **kwargs)


from numpy import unravel_index, ravel_multi_index

from bolt.utils import argpack, istransposeable, isreshapeable
from bolt.spark.array import BoltArraySpark

class Shapes(object):
    """
    Base Shape class. These classes wrap a BoltArraySpark in their
    entirity, but implement the following attributes and methods as if
    they were only working on the keys or the values, depending which
    subclass is used.
    """
    @property
    def shape(self):
        raise NotImplementedError

    @property
    def ndim(self):
        return len(self.shape)

    def reshape(self):
        raise NotImplementedError

    def transpose(self):
        raise NotImplementedError

class Keys(Shapes):
    """
    This class implements all the base shape attributes and methods
    for the keys of a BoltArraySpark.
    """
    def __init__(self, barray):
        self._barray = barray

    @property
    def shape(self):
        return self._barray.shape[:self._barray.split]

    def reshape(self, *shape):
        """
        Reshape just the keys of a BoltArraySpark, returning a
        new BoltArraySpark.

        Parameters
        ----------
        shape : tuple
              New proposed axes.
        """
        new = argpack(shape)
        old = self.shape
        isreshapeable(new, old)

        if new == old:
            return self._barray

        def f(k):
            return unravel_index(ravel_multi_index(k, old), new)

        newrdd = self._barray._rdd.map(lambda kv: (f(kv[0]), kv[1]))
        newsplit = len(new)
        newshape = new + self._barray.values.shape

        return BoltArraySpark(newrdd, shape=newshape, split=newsplit).__finalize__(self._barray)

    def transpose(self, *axes):
        """
        Transpose just the keys of a BoltArraySpark, returning a
        new BoltArraySpark.

        Parameters
        ----------
        axes : tuple
             New proposed axes.
        """
        new = argpack(axes)
        old = range(self.ndim)
        istransposeable(new, old)

        if new == old:
            return self._barray

        def f(k):
            return tuple(k[i] for i in new)

        newrdd = self._barray._rdd.map(lambda kv: (f(kv[0]), kv[1]))
        newshape = tuple(self.shape[i] for i in new) + self._barray.values.shape

        return BoltArraySpark(newrdd, shape=newshape, ordered=False).__finalize__(self._barray)

    def __str__(self):
        s = "BoltArray Keys\n"
        s += "shape: %s" % str(self.shape)
        return s

    def __repr__(self):
        return str(self)

class Values(Shapes):
    """
    This class implements all the base shape attributes and methods
    for the values of a BoltArraySpark.
    """
    def __init__(self, barray):
        self._barray = barray

    @property
    def shape(self):
        return self._barray.shape[self._barray.split:]

    def reshape(self, *shape):
        """
        Reshape just the values of a BoltArraySpark, returning a
        new BoltArraySpark.

        Parameters
        ----------
        shape : tuple
              New proposed axes.
        """
        new = argpack(shape)
        old = self.shape
        isreshapeable(new, old)

        if new == old:
            return self._barray

        def f(v):
            return v.reshape(new)

        newrdd = self._barray._rdd.mapValues(f)
        newshape = self._barray.keys.shape + new

        return BoltArraySpark(newrdd, shape=newshape).__finalize__(self._barray)

    def transpose(self, *axes):
        """
        Transpose just the values of a BoltArraySpark, returning a
        new BoltArraySpark.

        Parameters
        ----------
        axes : tuple
             New proposed axes.
        """
        new = argpack(axes)
        old = range(self.ndim)
        istransposeable(new, old)

        if new == old:
            return self._barray

        def f(v):
            return v.transpose(new)

        newrdd = self._barray._rdd.mapValues(f)
        newshape = self._barray.keys.shape + tuple(self.shape[i] for i in new)

        return BoltArraySpark(newrdd, shape=newshape).__finalize__(self._barray)

    def __str__(self):
        s = "BoltArray Values\n"
        s += "shape: %s" % str(self.shape)
        return s

    def __repr__(self):
        return str(self)


from numpy import float64, asarray

from bolt.construct import ConstructBase
from bolt.local.array import BoltArrayLocal

class ConstructLocal(ConstructBase):

    @staticmethod
    def array(a, dtype=None, order='C'):
        """
        Create a local bolt array.

        Parameters
        ----------
        a : array-like
            An array, any object exposing the array interface, an
            object whose __array__ method returns an array, or any
            (nested) sequence.

        dtype : data-type, optional, default=None
            The desired data-type for the array. If None, will
            be determined from the data. (see numpy)

        order : {'C', 'F', 'A'}, optional, default='C'
            The order of the array. (see numpy)

        Returns
        -------
        BoltArrayLocal
        """
        return BoltArrayLocal(asarray(a, dtype, order))

    @staticmethod
    def ones(shape, dtype=float64, order='C'):
        """
        Create a local bolt array of ones.

        Parameters
        ----------
        shape : tuple
            Dimensions of the desired array

        dtype : data-type, optional, default=float64
            The desired data-type for the array. (see numpy)

        order : {'C', 'F', 'A'}, optional, default='C'
            The order of the array. (see numpy)

        Returns
        -------
        BoltArrayLocal
        """
        from numpy import ones
        return ConstructLocal._wrap(ones, shape, dtype, order)

    @staticmethod
    def zeros(shape, dtype=float64, order='C'):
        """
        Create a local bolt array of zeros.

        Parameters
        ----------
        shape : tuple
            Dimensions of the desired array.

        dtype : data-type, optional, default=float64
            The desired data-type for the array. (see numpy)

        order : {'C', 'F', 'A'}, optional, default='C'
            The order of the array. (see numpy)

        Returns
        -------
        BoltArrayLocal
        """
        from numpy import zeros
        return ConstructLocal._wrap(zeros, shape, dtype, order)

    @staticmethod
    def _wrap(func, shape, dtype, order):
        return BoltArrayLocal(func(shape, dtype, order))

    @staticmethod
    def concatenate(arrays, axis=0):
        """
        Join a sequence of arrays together.

        Parameters
        ----------
        arrays : tuple
            A sequence of array-like e.g. (a1, a2, ...)

        axis : int, optional, default=0
            The axis along which the arrays will be joined.

        Returns
        -------
        BoltArrayLocal
        """
        if not isinstance(arrays, tuple):
            raise ValueError("data type not understood")
        arrays = tuple([asarray(a) for a in arrays])
        from numpy import concatenate
        return BoltArrayLocal(concatenate(arrays, axis))


# intended to implement a power-law fitting routine as specified in.....
# that paper Erik sent me in July 2009.
# http://www.santafe.edu/~aaronc/powerlaws/
#
# The MLE for the power-law alpha is very easy to derive given knowledge
# of the lowest value at which a power law holds, but that point is
# difficult to derive and must be acquired iteratively.
from pylab import *

# def plfit(x,xmin=None):
#   if xmin is None:
#     raise InputError("can't do that yet")
#     xguess = mean(x)
#     mle = 1+len(x) / (sum(log(x[x>xguess]/xguess)))
#
#   return 1+len(x) / (sum(log(x[x>xmin]/xmin)))

def plfit_lsq(x,y):
    """
    Returns A and B in y=Ax^B
    http://mathworld.wolfram.com/LeastSquaresFittingPowerLaw.html
    """
    n = len(x)
    btop = n * (log(x)*log(y)).sum() - (log(x)).sum()*(log(y)).sum()
    bbottom = n*(log(x)**2).sum() - (log(x).sum())**2
    b = btop / bbottom
    a = ( log(y).sum() - b * log(x).sum() ) / n

    A = exp(a)
    return A,b

def plfit(x,nosmall=False,finite=False):
    """
    A Python implementation of the Matlab code http://www.santafe.edu/~aaronc/powerlaws/plfit.m
    from http://www.santafe.edu/~aaronc/powerlaws/

    See A. Clauset, C.R. Shalizi, and M.E.J. Newman, "Power-law distributions
    in empirical data" SIAM Review, to appear (2009). (arXiv:0706.1062)
    http://arxiv.org/abs/0706.1062
    """
    xmins = unique(x)
    xmins = xmins[1:-1]
    dat = xmins * 0
    z = sort(x)
    for xm in arange(len(xmins)):
        xmin = xmins[xm]
        z    = z[z>=xmin]
        n    = float(len(z))
        # estimate alpha using direct MLE
        a    =  n / sum( log(z/xmin) )
        if nosmall:
            # 4. For continuous data, PLFIT can return erroneously large estimates of
            #    alpha when xmin is so large that the number of obs x >= xmin is very
            #    small. To prevent this, we can truncate the search over xmin values
            #    before the finite-size bias becomes significant by calling PLFIT as
            if (a-1)/sqrt(n) > 0.1:
                #dat(xm:end) = [];
                dat = dat[:xm]
                xm = len(xmins)+1
                break
        # compute KS statistic
        cx   = arange(n)/float(n)  #data
        cf   = 1-(xmin/z)**a  # fitted
        dat[xm] = max( abs(cf-cx) )
    D     = min(dat);
    #xmin  = xmins(find(dat<=D,1,'first'));
    xmin  = xmins[argmin(dat)]
    z     = x[x>=xmin]
    n     = len(z)
    alpha = 1 + n / sum( log(z/xmin) )
    if finite:
        alpha = alpha*(n-1)/n+1/n
    if n < 50 and ~finite:
        print '(PLFIT) Warning: finite-size bias may be present.'
    L = n*log((alpha-1)/xmin) - alpha*sum(log(z/xmin));
    return xmin,alpha,L,dat

def plotcdf(x,xmin,alpha):
    """
    Plots CDF and powerlaw
    """

    x=sort(x)
    n=len(x)
    xcdf = arange(n,0,-1,dtype='float')/float(n)

    q = x[x>=xmin]
    fcdf = (q/xmin)**(1-alpha)
    nc = xcdf[argmax(x>=xmin)]
    fcdf_norm = nc*fcdf

    loglog(x,xcdf)
    loglog(q,fcdf_norm)

def plotpdf(x,xmin,alpha,nbins=30,dolog=False):
    """
    Plots PDF and powerlaw....
    """

    x=sort(x)
    n=len(x)

    if dolog:
        hb = hist(x,bins=logspace(log10(min(x)),log10(max(x)),nbins),log=True)
        alpha += 1
    else:
        hb = hist(x,bins=linspace((min(x)),(max(x)),nbins))
    h,b=hb[0],hb[1]
    b = b[1:]

    q = x[x>=xmin]
    px = (alpha-1)/xmin * (q/xmin)**(-alpha)

    arg = argmin(abs(b-xmin))
    norm = mean( h[b>xmin] / ((alpha-1)/xmin * (b[b>xmin]/xmin)**(-alpha))  )
    px = px*norm

    loglog(q,px)

    gca().set_xlim(min(x),max(x))


# intended to implement a power-law fitting routine as specified in.....
# http://www.santafe.edu/~aaronc/powerlaws/
#
# The MLE for the power-law alpha is very easy to derive given knowledge
# of the lowest value at which a power law holds, but that point is
# difficult to derive and must be acquired iteratively.

"""
Pure-Python version of plfit.py
===============================

A *pure* python power-law distribution fitter based on code by Aaron Clauset.
This is the slowest implementation, but has no dependencies.

Example very simple use::

    from plfit_py import plfit

    MyPL = plfit(mydata)
    MyPL.plotpdf(log=True)

"""
from __future__ import print_function

import time
import random
import math

class plfit:
    """
    A Python implementation of the Matlab code http://www.santafe.edu/~aaronc/powerlaws/plfit.m
    from http://www.santafe.edu/~aaronc/powerlaws/

    See A. Clauset, C.R. Shalizi, and M.E.J. Newman, "Power-law distributions
    in empirical data" SIAM Review, 51, 661-703 (2009). (arXiv:0706.1062)
    http://arxiv.org/abs/0706.1062

    The output "alpha" is defined such that :math:`p(x) \sim (x/xmin)^{-alpha}`
    """

    def __init__(self,x,**kwargs):
        """
        Initializes and fits the power law.  Can pass "quiet" to turn off
        output (except for warnings; "silent" turns off warnings)
        """
        neg = [i<0 for i in x]
        if any(neg) > 0:
            print("Removed %i negative points" % (sum(neg)))
            x = [i for i in x if i > 0]
        self.data = x
        self.plfit(**kwargs)

    def alpha_(self,x):
        """ Create a mappable function alpha to apply to each xmin in a list of xmins.
        This is essentially the slow version of fplfit/cplfit, though I bet it could
        be speeded up with a clever use of parellel_map.  Not intended to be used by users."""
        def alpha(xmin,x=x):
            """
            given a sorted data set and a minimum, returns power law MLE fit
            data is passed as a keyword parameter so that it can be vectorized
            """
            x = [i for i in x if i>=xmin]
            n = sum(x)
            divsum = sum([math.log(i/xmin) for i in x])
            if divsum == 0:
                return float('inf')
            # the "1+" here is unimportant because alpha_ is only used for minimization
            a = 1 + float(n) / divsum
            return a
        return alpha

    def kstest_(self,x):
        def kstest(xmin,x=x):
            """
            given a sorted data set and a minimum, returns power law MLE ks-test w/data
            data is passed as a keyword parameter so that it can be vectorized

            The returned value is the "D" parameter in the ks test...
            """
            x = [i for i in x if i>=xmin]
            n = len(x)
            if n == 0: return float('inf')
            divsum = sum([math.log(i/xmin) for i in x])
            if divsum == 0: return float('inf')
            a = float(n) / divsum
            cx = [float(i)/float(n) for i in range(int(n))]
            cf = [1-(xmin/i)**a for i in x]
            ks = max([abs(a-b) for a,b in zip(cf,cx)])
            return ks
        return kstest

    def plfit(self,nosmall=True,finite=False,quiet=False,silent=False,
              xmin=None, verbose=False):
        """
        A pure-Python implementation of the Matlab code http://www.santafe.edu/~aaronc/powerlaws/plfit.m
        from http://www.santafe.edu/~aaronc/powerlaws/

        See A. Clauset, C.R. Shalizi, and M.E.J. Newman, "Power-law distributions
        in empirical data" SIAM Review, 51, 661-703 (2009). (arXiv:0706.1062)
        http://arxiv.org/abs/0706.1062

        nosmall is on by default; it rejects low s/n points
        can specify xmin to skip xmin estimation

        This is only for continuous distributions; I have not implemented a
        pure-python discrete distribution fitter
        """
        x = self.data
        z = sorted(x)
        t = time.time()
        possible_xmins = sorted(set(z))
        argxmins = [z.index(i) for i in possible_xmins]
        self._nunique = len(possible_xmins)
        if xmin is None:
            av  = map(self.alpha_(z),possible_xmins)
            dat = map(self.kstest_(z),possible_xmins)
            sigma = [(a-1)/math.sqrt(len(z)-i+1) for a,i in zip(av,argxmins)]
            if nosmall:
                # test to make sure the number of data points is high enough
                # to provide a reasonable s/n on the computed alpha
                goodvals = [s<0.1 for s in sigma]
                if False in goodvals:
                    nmax = goodvals.index(False)
                    dat = dat[:nmax]
                    possible_xmins = possible_xmins[:nmax]
                    av = av[:nmax]
                else:
                    print("Not enough data left after flagging - using all positive data.")
            if not quiet: print("PYTHON plfit executed in %f seconds" % (time.time()-t))
            self._av = av
            self._xmin_kstest = dat
            self._sigma = sigma
            # [:-1] to weed out the very last data point; it cannot be correct
            # (can't have a power law with 1 data point).
            # However, this should only be done if the ends have not previously
            # been excluded with nosmall
            if nosmall:
                xmin = possible_xmins[dat.index(min(dat))]
            else:
                xmin = possible_xmins[dat.index(min(dat[:-1]))]
        z     = [i for i in z if i >= xmin]
        n     = len(z)
        alpha = 1 + n / sum([math.log(a/xmin) for a in z])
        if finite:
            alpha = alpha*(n-1.)/n+1./n
        if n == 1 and not silent:
            print("Failure: only 1 point kept.  Probably not a power-law distribution.")
            self._alpha = 0
            self._alphaerr = 0
            self._likelihood = 0
            self._ks = 0
            self._ks_prob = 0
            self._xmin = xmin
            return xmin,0
        if n < 50 and not finite and not silent:
            print('(PLFIT) Warning: finite-size bias may be present. n=%i' % n)
        # ks = max(abs( numpy.arange(n)/float(n) - (1-(xmin/z)**(alpha-1)) ))
        ks = max( [abs( i/float(n) - (1-(xmin/b)**(alpha-1))) for i,b in zip(range(n),z)] )
        # Parallels Eqn 3.5 in Clauset et al 2009, but zeta(alpha, xmin) = (alpha-1)/xmin.  Really is Eqn B3 in paper.
        #L = n*log((alpha-1)/xmin) - alpha*sum(log(z/xmin))
        sl = sum([math.log(a/xmin) for a in z])
        L = (n*math.log((alpha-1)/xmin) - alpha*sl)
        #requires another map... Larr = arange(len(unique(x))) * log((av-1)/unique(x)) - av*sum
        self._likelihood = L
        self._xmin = xmin
        self._xmins = possible_xmins
        self._alpha= alpha
        self._alphaerr = (alpha-1)/math.sqrt(n)
        self._ks = ks  # this ks statistic may not have the same value as min(dat) because of unique()
        #if scipyOK: self._ks_prob = scipy.stats.kstwobign.sf(ks*numpy.sqrt(n))
        self._ngtx = n
        if math.isnan(L) or math.isnan(xmin) or math.isnan(alpha):
            raise ValueError("plfit failed; returned a nan")

        if not quiet:
            if verbose: print("The lowest value included in the power-law fit, ", end=' ')
            print("xmin: %g" % xmin, end=' ')
            if verbose: print("\nThe number of values above xmin, ", end=' ')
            print("n(>xmin): %i" % n, end=' ')
            if verbose: print("\nThe derived power-law alpha (p(x)~x^-alpha) with MLE-derived error, ", end=' ')
            print("alpha: %g +/- %g  " % (alpha,self._alphaerr), end=' ')
            if verbose: print("\nThe log of the Likelihood (the maximized parameter), ", end=' ')
            print("Log-Likelihood: %g  " % L, end=' ')
            if verbose: print("\nThe KS-test statistic between the best-fit power-law and the data, ", end=' ')
            print("ks: %g" % (ks))

        return xmin,alpha

def plexp(x,xm=1,a=2.5):
    """
    CDF(x) for the piecewise distribution exponential x<xmin, powerlaw x>=xmin
    This is the CDF version of the distributions drawn in fig 3.4a of Clauset et al.
    """

    C = 1/(-xm/(1 - a) - xm/a + math.exp(a)*xm/a)
    Ppl = lambda X: 1+C*(xm/(1-a)*(X/xm)**(1-a))
    Pexp = lambda X: C*xm/a*math.exp(a)-C*(xm/a)*math.exp(-a*(X/xm-1))
    d=Ppl(x)
    d[x<xm]=Pexp(x)
    return d

def plexp_inv(P,xm,a):
    """
    Inverse CDF for a piecewise PDF as defined in eqn. 3.10
    of Clauset et al.
    """

    C = 1/(-xm/(1 - a) - xm/a + math.exp(a)*xm/a)
    Pxm = 1+C*(xm/(1-a))
    pp = P
    x = xm*(pp-1)*(1-a)/(C*xm)**(1/(1-a)) if pp >= Pxm else (math.log( ((C*xm/a)*math.exp(a)-pp)/(C*xm/a)) - a) * (-xm/a)
    #x[P>=Pxm] = xm*( (P[P>=Pxm]-1) * (1-a)/(C*xm) )**(1/(1-a)) # powerlaw
    #x[P<Pxm] = (math.log( (C*xm/a*math.exp(a)-P[P<Pxm])/(C*xm/a) ) - a) * (-xm/a) # exp

    return x

def pl_inv(P,xm,a):
    """
    Inverse CDF for a pure power-law
    """

    x = (1-P)**(1/(1-a)) * xm
    return x

def test_fitter(xmin=1.0, alpha=2.5, niter=500, npts=1000, invcdf=plexp_inv,
        quiet=True, silent=True):
    """
    Tests the power-law fitter

    Examples
    ========
    Example (fig 3.4b in Clauset et al.)::

        xminin=[0.25,0.5,0.75,1,1.5,2,5,10,50,100]
        xmarr,af,ksv,nxarr = plfit.test_fitter(xmin=xminin,niter=1,npts=50000)
        loglog(xminin,xmarr.squeeze(),'x')

    Example 2::

        xminin=[0.25,0.5,0.75,1,1.5,2,5,10,50,100]
        xmarr,af,ksv,nxarr = plfit.test_fitter(xmin=xminin,niter=10,npts=1000)
        loglog(xminin,xmarr.mean(axis=0),'x')

    Example 3::

        xmarr,af,ksv,nxarr = plfit.test_fitter(xmin=1.0,niter=1000,npts=1000)
        hist(xmarr.squeeze());
        # Test results:
        # mean(xmarr) = 0.70, median(xmarr)=0.65 std(xmarr)=0.20
        # mean(af) = 2.51 median(af) = 2.49  std(af)=0.14
        # biased distribution; far from correct value of xmin but close to correct alpha

    Example 4::

        xmarr,af,ksv,nxarr = plfit.test_fitter(xmin=1.0,niter=1000,npts=1000,invcdf=pl_inv)
        print("mean(xmarr): %0.2f median(xmarr): %0.2f std(xmarr): %0.2f" % (mean(xmarr),median(xmarr),std(xmarr)))
        print("mean(af): %0.2f median(af): %0.2f std(af): %0.2f" % (mean(af),median(af),std(af)))
        # mean(xmarr): 1.19 median(xmarr): 1.03 std(xmarr): 0.35
        # mean(af): 2.51 median(af): 2.50 std(af): 0.07

    """
    sz = niter
    xmarr,alphaf_v,ksv,nxarr = ([0]*sz,)*4
    for i in range(niter):
        randarr = [random.random() for k in range(npts)]
        fakedata = [invcdf(r,xmin,alpha) for r in randarr]
        TEST = plfit(fakedata,quiet=quiet,silent=silent,nosmall=True)
        alphaf_v[i] = TEST._alpha
        ksv[i] = TEST._ks
        nxarr[i] = TEST._ngtx
        xmarr[i] = TEST._xmin

    return xmarr,alphaf_v,ksv,nxarr