ALLEN INSTITUTE

# Source code for bmtk.builder.auxi.node_params

    # Copyright 2017. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANwhile (itercount < max_iter) and (n_pass < np.ceil(ndraws_tot / 5000)):Y EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import numpy as np
    import math
    # import nrrd
    import matplotlib.pyplot as plt
    from collections import defaultdict
    from types import SimpleNamespace

    try:
        from sklearn.neighbors import KDTree
    except ImportError:
        pass

    try:
        import nrrd
    except ImportError:
        pass

    class CellLocations(object):
        # Added class to facilitate checking of new locations against already placed locations in other populations for
        # the purpose of ensuring a minimum distance.
        # Only a unique minimum distance is allowed per network (populations sharing the same physical space) and must be
        # set before adding placements
        # Original functions outside of class are retained for compatibility with continued usage of non-class approach
        # for simpler situations.
        def __init__(self, name, dmin=0.00):
            # Currently only allowing a unique dmin
            self._name = name
            self._dmin = dmin
            self._CCF_orientation = True
            #self._all_positions = [np.array([]).reshape(0, 3)]
            self._all_positions =[]
            self._all_pop_names = []

        @property
        def dmin(self):
            return self._dmin

        @property
        def CCF_orientation(self):
            return self._CCF_orientation

        '''
        @property
        def pop_positions(self,pop_name):

            return
        '''

        @dmin.setter
        def dmin(self, value):
            self._dmin = value

        @CCF_orientation.setter
        def CCF_orientation(self, value):
            self._CCF_orientation = value

        def add_positions_nrrd(self, nrrd_filename, max_dens_per_mm3, pop_names, partitions=[1], split_bilateral=None,
                              method='prog', verbose = False):

            if self._all_positions:
                existing_positions = np.vstack(self._all_positions)
            else:
                existing_positions = np.array([]).reshape(0, 3)
            positions = positions_nrrd(nrrd_filename, max_dens_per_mm3, split_bilateral, self._dmin,
                                       CCF_orientation=self._CCF_orientation, plot=False,
                                       method=method, existing_positions=existing_positions, verbose=verbose)
            # list of position arrays
            positions_pop = partition_locations(positions, partitions)
            self.add_positions(pop_names, positions_pop)

        def add_positions_columnar(self, pop_names, partitions=[1], N=1, center=[0.0, 50.0, 0.0], height = 100.0,
                                   min_radius = 0.0, max_radius = 1.0, method='prog', verbose=False):
            if sum(partitions)!=1:
                raise ValueError('The elements of `partitions` must add up to 1.')

            if self._all_positions:
                existing_positions = np.vstack(self._all_positions)
            else:
                existing_positions = np.array([]).reshape(0, 3)
            vol = np.pi * (max_radius ** 2 - min_radius ** 2) * height / 1000**3

            if method == 'prog':
                def sampling_func(ndraws):
                    positions = positions_columinar(N=ndraws, center=center, height=height, min_radius=min_radius,
                                                    max_radius=max_radius)
                    return positions
                positions = positions_dmin_prog(N=N, vol_tot=vol, sampling_func=sampling_func,
                                dmin=self._dmin, existing_positions=existing_positions, verbose=verbose)
            elif method == 'lattice':
                box_dims = np.array([2*max_radius, height, 2*max_radius])
                position_scale = np.eye(3)*box_dims
                mat = np.array([[[N/vol]]])
                def filter_func(x, y, z):
                    inside = (np.sqrt((x-box_dims[0]/2)**2 + (z-box_dims[2]/2)**2) <= max_radius) & \
                             (np.sqrt((x-box_dims[0]/2)**2 + (z-box_dims[2]/2)**2) >= min_radius) & \
                             (np.abs(y-box_dims[1]/2) < height/2)
                    return inside

                positions = positions_dmin_lattice(mat, position_scale=position_scale, N=N, vol_tot=vol, dmin=self._dmin,
                                                   existing_positions=existing_positions, filter_func=filter_func,
                                                   verbose=verbose)
                positions = positions - box_dims/2 + center
            positions_pop = partition_locations(positions, partitions)
            self.add_positions(pop_names, positions_pop)

        def add_positions_rect_prism(self, pop_names, partitions=[1], N=1, center=[0.0, 50.0, 0.0], height=20.0,
                                   x_length=100.0, z_length=100.0, method='prog', verbose=False):
            if sum(partitions) != 1:
                raise ValueError('The elements of `partitions` must add up to 1.')

            if self._all_positions:
                existing_positions = np.vstack(self._all_positions)
            else:
                existing_positions = np.array([]).reshape(0, 3)
            vol = x_length * z_length * height / 1000**3

            if method == 'prog':
                def sampling_func(ndraws):
                    positions = positions_rect_prism(N=ndraws, center=center, height=height, x_length=x_length,
                                                     z_length=z_length)
                    return positions
                positions = positions_dmin_prog(N=N, vol_tot=vol, sampling_func=sampling_func,
                                                dmin=self._dmin, existing_positions=existing_positions,
                                                verbose=verbose)
            elif method == 'lattice':
                box_dims = np.array([x_length, height, z_length])
                position_scale = np.eye(3) * box_dims
                mat = np.array([[[N / vol]]])

                positions = positions_dmin_lattice(mat, position_scale=position_scale, N=N, vol_tot=vol, dmin=self._dmin,
                                                   existing_positions=existing_positions)
                positions = positions - box_dims / 2 + center
            positions_pop = partition_locations(positions, partitions)
            self.add_positions(pop_names, positions_pop)

        def add_positions_ellipsoid(self, pop_names, partitions=[1], N=1, center=[0.0, 50.0, 0.0], height=50.0,
                                   x_length=100.0, z_length=200.0, method='prog', verbose=False):
            if sum(partitions) != 1:
                raise ValueError('The elements of `partitions` must add up to 1.')

            if self._all_positions:
                existing_positions = np.vstack(self._all_positions)
            else:
                existing_positions = np.array([]).reshape(0, 3)
            vol = 4 / 3 * np.pi * (x_length/2) * (z_length/2) * (height/2) / 1000**3

            if method == 'prog':
                def sampling_func(ndraws):
                    positions = positions_ellipsoid(N=ndraws, center=center, height=height, x_length=x_length,
                                                    z_length=z_length)
                    return positions
                positions = positions_dmin_prog(N=N, vol_tot=vol, sampling_func=sampling_func,
                                                dmin=self._dmin, existing_positions=existing_positions, verbose=verbose)
            elif method == 'lattice':
                box_dims = np.array([x_length, height, z_length])
                position_scale = np.eye(3) * box_dims
                mat = np.array([[[N / vol]]])

                def filter_func(x, y, z):
                    inside = ((x - box_dims[0] / 2) ** 2 / (x_length/2) ** 2
                              + (y - box_dims[1] / 2) ** 2 / (height/2) ** 2
                              + (z - box_dims[2] / 2) ** 2 / (z_length/2) ** 2) <= 1
                    return inside
                positions = positions_dmin_lattice(mat, position_scale=position_scale, N=N, vol_tot=vol, dmin=self._dmin,
                                                   existing_positions=existing_positions, filter_func=filter_func)
                positions = positions - box_dims / 2 + center
            '''
            elif method == 'lattice':
                box_dims = np.array([2*max_radius, height, 2*max_radius])
                position_scale = np.eye(3)*box_dims
                mat = np.array([[[N/(np.prod(box_dims)/1000**3)]]])
                def filter_func(x, y, z):
                    inside = (np.sqrt((x-box_dims[0]/2)**2 + (z-box_dims[2]/2)**2) <= max_radius) & \
                             (np.sqrt((x-box_dims[0]/2)**2 + (z-box_dims[2]/2)**2) >= min_radius) & \
                             (np.abs(y-box_dims[1]/2) < height/2)
                    return inside

                positions = positions_dmin_lattice(mat, position_scale=position_scale, dmin=self._dmin,
                                                   existing_positions=existing_positions, filter_func=filter_func)
                positions = positions - box_dims/2 + center
            '''
            positions_pop = partition_locations(positions, partitions)
            self.add_positions(pop_names, positions_pop)

        def plot_locs (self):
            fig, ax = plt.subplots(1,2,figsize=(20,15), subplot_kw={'projection':'3d'})
            plt.style.use('seaborn-bright')
            if self._CCF_orientation:
                ax[0].invert_zaxis()
            for p, pop_name in enumerate(self._all_pop_names):
                # x,y,z axis orientations depending on function and orientation parameters
                x = self._all_positions[p][:,0]
                y = self._all_positions[p][:,1]
                z = self._all_positions[p][:,2]

                if self._CCF_orientation:
                    plot_positions(x, z, y, ax[0], labels=['X', 'Z', 'Y'], pop_name=pop_name)
                    ax[0].set_title('Posterior view')
                    ax[0].view_init(elev=10., azim=0)
                    plot_positions(x, z, y, ax[1], labels=['X', 'Z', 'Y'], pop_name=pop_name)
                    ax[1].set_title('Top view')
                    ax[1].view_init(elev=90., azim=0)
                    ax[1].legend(loc="upper right", markerscale=2, prop={'size': 15})
                else:
                    plot_positions(x, y, z, ax[0], labels=['X','Y','Z'], pop_name=pop_name)
                    ax.view_init(elev=10., azim=0)
                    ax[0].set_title('Side view')
                    ax[0].view_init(elev=10., azim=-90)
                    plot_positions(x, y, z, ax[1], labels=['X', 'Y', 'Z'], pop_name=pop_name)
                    ax[1].set_title('Top view')
                    ax[1].view_init(elev=90., azim=0)
                    ax[1].legend(loc="upper right", markerscale=2, prop={'size': 15})

            plt.tight_layout()
            plt.show()

        def add_positions(self, pop_names, positions_pop):
            if isinstance(pop_names, str):
                pop_names = [pop_names]
            for p, pop_name in enumerate(pop_names):
                self._all_positions.append(positions_pop[p])
                self._all_pop_names.append(pop_name)
                myvars = vars(self)
                myvars[pop_name] = SimpleNamespace()
                myvars[pop_name].positions = self._all_positions[p]
                myvars[pop_name].N = self._all_positions[p].shape[0]

    def positions_columinar(N=1, center=[0.0, 50.0, 0.0], height=100.0, min_radius=0.0, max_radius=1.0, plot=False):
        """Returns a set of random x,y,z coordinates within a given cylinder or cylindrical ring.
        Height is given as the y (index 1) coordinates.

        :param N: Number of points to return
        :param center: center of sphere
        :param height: maximum length of sphere (y coord)
        :param min_radius: minimum horizontal radius on x-z plane
        :param max_radius: maximum horizontal radius on x-z plane
        :return: A (N, 3) matrix
        """
        phi = 2.0 * math.pi * np.random.random([N])
        r = np.sqrt((min_radius**2 - max_radius**2) * np.random.random([N]) + max_radius**2)
        x = center[0] + r * np.cos(phi)
        z = center[2] + r * np.sin(phi)
        y = center[1] + height * (np.random.random([N]) - 0.5)

        if plot:
            fig, ax = plt.subplots(1,2,figsize=(9,12), subplot_kw={'projection':'3d'})
            plot_positions(x, z, y, ax[0], labels=['X','Z','Y'])
            ax[0].set_title('Side view')
            ax[0].view_init(elev=5., azim=-90)

            plot_positions(x, z, y, ax[1], labels=['X','Z','Y'])
            ax[1].set_title('Top view')
            ax[1].view_init(elev=90., azim=-90)

        return np.column_stack((x, y, z))

    positions_columnar = positions_columinar

    def positions_rect_prism(N=1, center=[0.0, 50.0, 0.0], height=20.0, x_length=100.0, z_length=100.0, plot=False):
        """Returns a set of random x,y,z coordinates within a rectangular prism.
        Height is given as the y (index 1) coordinates.

        :param N: Number of points to return
        :param center: center of rectangular prism
        :param height: height of prism (y-coordinate)
        :param x_length: length of prism in x
        :param z_length: length of prism in z
        :return: A (N, 3) matrix
        """

        x = center[0] + x_length * (np.random.random([N]) - 0.5)
        z = center[2] + z_length * (np.random.random([N]) - 0.5)
        y = center[1] + height * (np.random.random([N]) - 0.5)
        if plot:
            fig, ax = plt.subplots(1,2,figsize=(9,12), subplot_kw={'projection':'3d'})
            plot_positions(x, z, y, ax[0], labels=['X','Z','Y'])
            ax[0].set_title('Side view')
            ax[0].view_init(elev=5., azim=-90)

            plot_positions(x, z, y, ax[1], labels=['X','Z','Y'])
            ax[1].set_title('Top view')
            ax[1].view_init(elev=90., azim=-90)

        return np.column_stack((x, y, z))

    def positions_ellipsoid (N=1, center=[0.0, 50.0, 0.0], height=50, x_length=100.0, z_length=200.0, plot=False):
        """Returns a set of random x,y,z coordinates within an ellipsoid. Height is given as the y (index 1) coordinates.

        :param N: Number of points to return
        :param center: center of ellipsoid
        :param height: height of ellipsoid (y-coordinate)
        :param x_length: length of ellipsoid in x
        :param z_length: length of ellipsoid in z
        :return: A (N, 3) matrix
        """
        # Generate prism bounding the ellipsoid
        positions = positions_rect_prism (3*N, center, height, x_length, z_length, plot=False)
        val = ((positions[:,0]-center[0])**2/(x_length/2)**2)\
              +((positions[:,1]-center[1])**2/(height/2)**2)\
              +((positions[:,2]-center[2])**2/(z_length/2)**2)
        positions = np.squeeze(positions [np.where(val<1),:])
        positions = positions [:N,:]

        if plot:
            fig, ax = plt.subplots(1,2,figsize=(9,12), subplot_kw={'projection':'3d'})
            plot_positions(positions[:,0], positions[:,2], positions[:,1], ax[0], labels=['X','Z','Y'])
            ax[0].set_title('Side view')
            ax[0].view_init(elev=5., azim=-90)

            plot_positions(positions[:,0], positions[:,2], positions[:,1], ax[1], labels=['X','Z','Y'])
            ax[1].set_title('Top view')
            ax[1].view_init(elev=90., azim=-90)

        return positions

    def positions_cuboid(N=1, center=[0.0, 0.0, 0.0], height=100.0, xside_length=100.0, yside_length=100.0, min_dist=20.0,
                         plot = False):
        """This function distributes the cells in a 3D cuboid (x,y,z sides may have different lengths). The method used
        assures cells cannot be placed too close to one another (must be > min_dist apart)
        WARNING: If cell density is high and there is more than 1 population of cells, there is a high chance cells will be
        placed on top of one another. You can use positions_list() to avoid this...

        Written by Ben Latimer at University of Missouri (latimerb@missouri.edu)

        :return: A (N, 3) matrix
        """

        # Create the possible x,y,z coordinates
        x_grid = np.arange(center[0], xside_length+min_dist, min_dist)
        y_grid = np.arange(center[1], yside_length+min_dist, min_dist)
        z_grid = np.arange(center[2], height+min_dist, min_dist)
        xx, yy, zz = np.meshgrid(x_grid, y_grid, z_grid)
        positions = np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T

        # Pick N indices for the positions matrix randomly, without replacement (so coordinates are unique)
        inds = np.random.choice(np.arange(0, np.size(positions, 0)), N, replace=False)

        # Assign positions
        x = positions[inds][:, 0]
        y = positions[inds][:, 1]
        z = positions[inds][:, 2]

        if plot:
            fig, ax = plt.subplots(1,2,figsize=(9,12), subplot_kw={'projection':'3d'})
            plot_positions(x, y, z, ax[0], labels=['X','Y','Z'])
            ax[0].set_title('Side view')
            ax[0].view_init(elev=5., azim=-90)

            plot_positions(x, y, z, ax[1], labels=['X','Y','Z'])
            ax[1].set_title('Top view')
            ax[1].view_init(elev=90., azim=-90)

        return np.column_stack((x, y, z))

    def positions_list(positions=np.array([(0, 0, 0), (0, 0, 1)])):
        """This function is designed to be used with an externally supplied array of x,y,z coordinates. It is useful to
        avoid cell overlap in high density situations or to make some unique geometry. After you create each population,
        delete those positions from the "master" list of coordinates. This will assure that no cells are placed on top of
        one another.

        Written by Ben Latimer at University of Missouri (latimerb@missouri.edu)

        :param positions:
        :return:
        """

        x = positions[:, 0]
        y = positions[:, 1]
        z = positions[:, 2]

        return np.column_stack((x, y, z))

    def positions_density_matrix(mat, position_scale=np.array([[1,0,0],[0,1,0],[0,0,1]]), origin=np.array([0,0,0]),
                                 plot=False, CCF_orientation=False, dmin=0.0, method='prog',
                                 existing_positions=np.array([]).reshape(0, 3), verbose=False):
        """This function places random x,y,z coordinates according to a supplied 3D array of densities (cells/mm^3).
        The optional position_scale parameter defines a transformation matrix A to physical space such that::

            [x_phys, y_phys, z_phys] = A * [x_mat, y_mat, z_mat]

        Note: position_scale and output coordinates are in units of microns, while density is specified in mm^3.

        :param mat: A 3-dimensional matrix of densities (cells/mm^3)
        :param position_scale: A (3, 3) matrix (microns)
        :param plot: Generates a plot of the cell locations when set to True
        :param CCF_orientation: if True, plot will be oriented for Allen atlas common coordinate framework
               See https://help.brain-map.org/display/mouseconnectivity/API#API-DownloadAtlas3-DReferenceModels
        :param dmin: Minimum distance allowed between cell centers - using this function can severely slow down cell
               placement, especially as we approach the maximal possible density for this minimum distance
        :param verbose: If set to true, prints every 100 units placed for monitoring progress

        :return: A (N, 3) matrix (microns)
        """

        ind_nz = np.nonzero(mat)
        if method == 'prog':
            # Use batch progressive sampling
            def sampling_func(ndraws):
                rand_pos_new = np.random.uniform(size=(3, ndraws))  # Or for simple geometries, draw over max_dims
                rand_nz_vox_new = np.random.choice(range(ind_nz[0].shape[0]), ndraws)

                positions = position_scale.dot([ind_nz[0][rand_nz_vox_new].astype(float) + rand_pos_new[0, :],
                                                     ind_nz[1][rand_nz_vox_new].astype(float) + rand_pos_new[1, :],
                                                     ind_nz[2][rand_nz_vox_new].astype(float) + rand_pos_new[2, :]])
                positions = positions.T

                return positions

            positions = positions_dmin_prog(mat, position_scale=position_scale, sampling_func=sampling_func, dmin=dmin,
                                            existing_positions=existing_positions, verbose=verbose)
        elif method == 'lattice':
            # Use lattice jittering
            vol_per_voxel = np.abs(np.linalg.det(position_scale.astype(float))) / (1000) ** 3
            vol = vol_per_voxel * len(ind_nz[0])
            N = math.floor(np.max(mat)*vol)  # Final desired N
            positions = positions_dmin_lattice(mat, position_scale=position_scale, N=N, vol_tot=vol, dmin=dmin,
                                               existing_positions=existing_positions, verbose=verbose)
        else:
            raise ValueError("The 'method' argument can be either 'prog' for progressive sampling "
                             "or 'lattice' for lattice jittering")

        # Remove cells to achieve non-uniform density, if applicable

        temp = set(mat.flat)
        temp.discard(0)

        print('positions:', positions.shape)
        inds = np.matmul(np.linalg.inv(position_scale),(positions.T)).T
        inds = np.floor(inds).astype(int)
        indx = inds[:,0]
        indy = inds[:,1]
        indz = inds[:,2]
        max_dens = np.max(mat)

        if len(temp) > 1:
            rand_keep = np.random.uniform(size=positions.shape[0])

            for j in range(len(rand_keep)):
                vox_dens = mat[indx[j], indy[j], indz[j]]
                if rand_keep[j] >= vox_dens / max_dens:
                    positions[j,:] = [np.nan, np.nan, np.nan]
            positions = positions[~np.isnan(positions[:,0]),:]

        # Shift re: origin
        positions += origin

        x = positions[:,0]
        y = positions[:,1]
        z = positions[:,2]

        if plot:
            fig, ax = plt.subplots(figsize=(11, 11),  subplot_kw={'projection':'3d'})
            if CCF_orientation:
                plot_positions(x, z, y, ax, labels=['X','Z','Y'])
                ax.invert_zaxis()
                ax.view_init(elev=10., azim=0)
            else:
                plot_positions(x, y, z, ax, labels=['X','Y','Z'])
                ax.view_init(elev=10., azim=0)

        return positions

    def positions_dmin_prog(mat=None, position_scale=np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), N=None, vol_tot=None, sampling_func=None,
                                  dmin=0.0, existing_positions=np.array([]).reshape(0, 3), verbose=False):
        """This function places random x,y,z coordinates with a minimal distance according to a supplied 3D array of
        densities (cells/mm^3) using progressive sampling.
        The optional position_scale parameter defines a transformation matrix A to physical space such that:

            [x_phys, y_phys, z_phys] = A * [x_mat, y_mat, z_mat]

        Note: position_scale and output coordinates are in units of microns, while density is specified in mm^3.

        :param mat: A 3-dimensional matrix of densities (cells/mm^3)
        :param position_scale: A (3, 3) matrix (microns)
        :param plot: Generates a plot of the cell locations when set to True
        :param CCF_orientation: if True, plot will be oriented for Allen atlas common coordinate framework
               See https://help.brain-map.org/display/mouseconnectivity/API#API-DownloadAtlas3-DReferenceModels
        :param dmin: Minimum distance allowed between cell centers (microns) - using this function can severely slow down cell
               placement, especially as we approach the maximal possible density for this minimum distance

        :return: A (N, 3) matrix (microns)
        """

        # Can later put this inside check function?
        if (dmin < 0):
            raise ValueError('Minimum distance between cell centers (dmin) must not be negative')

        if isinstance(mat, np.ndarray):    # Density matrix
            max_dens=np.max(mat)
            vol_per_voxel = np.abs(np.linalg.det(position_scale.astype(float))) / (1000) ** 3
            mat = mat.astype(float) * vol_per_voxel
            max_dens_n = np.max(mat)
            n_vox_nz = np.nonzero(mat)[0].shape[0]
            ndraws_tot = int(max_dens_n * n_vox_nz)
            vol_tot = vol_per_voxel * n_vox_nz
        elif N is not None:     # Desired total N
            ndraws_tot = N
            max_dens = N/(vol_tot)
        else:
            raise Exception('Must specify either a cell density matrix (mat) or a total number of cells (N)')

        if dmin != 0:
            # Number of spheres at FCC/HCP packing: 0.74*V/V_sphere
            # Random close packing: 0.64*V/V_sphere
            V_sphere = 4 / 3 * np.pi * (dmin/1000/2) ** 3
            dens_lim = 0.63 / V_sphere   # in points per mm^3
            print('Density limit:{:.3f}'.format(dens_lim))
            print('Max density requested:{:.3f}'.format(max_dens))
            if max_dens > dens_lim:
                raise ValueError('Requested density is incompatible with minimum distance: either reduce the density or dmin.')
            thresh = np.ceil(vol_tot * 0.0001 / V_sphere)
        else:
            thresh = 1

        # First distribute uniformly in non zero voxels at max_dens
        ndraws = ndraws_tot
        positions = np.array([]).reshape(0, 3)
        len_existing = existing_positions.shape[0]
        max_iter = 1
        past_iters = np.full(3, np.nan)
        pi_ind = 0

        while positions.shape[0] < ndraws_tot:
            positions_new = []
            itercount = 0
            n_pass = 0
            while (itercount < max_iter) and (n_pass < thresh):
                positions_new.append(sampling_func(ndraws))
                if dmin>0:
                    # Check against existing tree
                    # Find nearest neighbor
                    positions_all = np.vstack((existing_positions, positions))
                    if positions_all.shape[0] > 0:
                        tree = KDTree(positions_all, leaf_size=2)
                        dists, d_inds = tree.query(positions_new[-1], k=1)
                        conf_inds = np.squeeze(dists < dmin)
                        positions_new[-1][conf_inds,:] = [np.nan, np.nan, np.nan]
                        notna = ~np.isnan(positions_new[-1][:, 0])
                        positions_new[-1] = positions_new[-1][notna, :]

                itercount += 1
                n_pass += positions_new[-1].shape[0]

            if pi_ind >= len(past_iters):
                # Get mean number of iterations for past 3 points and use to set new maxiter
                max_iter = 1.5 * np.max(past_iters)
                pi_ind = 0
            past_iters[pi_ind] = itercount
            pi_ind += 1

            if n_pass < int(ndraws_tot / 2000):
                # Replace some of existing points
                # Draw a small fraction of the total number of points

                positions_new.append(sampling_func(ndraws))

                # Remove fraction of points in conflict with new points (exclude pre-existing points from other populations)
                inds2 = tree.query_radius(positions_new[-1], r=dmin)
                lens = np.array([np.nan if np.any(a<len_existing) else len(a) for a in inds2])
                k = np.min([math.ceil(0.005*positions_all.shape[0]), np.count_nonzero(lens <= 1)])
                noconf = np.nonzero(lens == 0)[0]
                swap = np.argpartition(lens, k)[:k]   # k indices with smallest number of conflicts, nan's sorted to largest
                inds3 = np.unique(np.hstack(inds2[swap]))
                inds3 = inds3[inds3>=len_existing]-len_existing
                positions[inds3, :] = [np.nan, np.nan, np.nan]
                positions = positions[~np.isnan(positions[:, 0]),:]
                positions_new[-1] = positions_new[-1][np.concatenate((swap, noconf)),:]

            positions_new = np.vstack(positions_new)

            # Check new points that are retained against each other
            if (dmin > 0) and positions_new.shape[0] > 1:
                tree = KDTree(positions_new, leaf_size=2)
                dists, d_inds = tree.query(positions_new, k=2)
                # Overremoves by checking against removed cells, but tracking removed cells is not faster
                conf_inds = np.squeeze(dists[:,1] < dmin)
                positions_new[conf_inds, :] = [np.nan, np.nan, np.nan]
                notna = ~np.isnan(positions_new[:, 0])
                positions_new = positions_new[notna, :]

            positions = np.concatenate((positions, positions_new), axis=0)  # Add new positions
            if verbose:
                print(f'{positions.shape[0]}/{ndraws_tot} cells placed')

            ndraws = ndraws_tot
        positions = positions[:ndraws,:]

        if verbose:
            print(f'{positions.shape[0]} cells')

        return positions

    def positions_dmin_lattice(mat, position_scale=np.array([0,0,0]), N=1, vol_tot=None, dmin=0.0,
                               existing_positions=np.array([]).reshape(0, 3), filter_func=None, verbose=False):
        '''Packing with a minimum distance, starting from a hexagonal close packing lattice

        :param x_box, y_box, z_box: size of each dimension of lattice box to be generated (microns)
        :param N: number of points to be placed
        '''
        if (dmin < 0):
            raise ValueError('Minimum distance between cell centers (dmin) must not be negative')
        if existing_positions.shape[0] > 0:
            raise RuntimeError("The 'lattice' method cannot be used with already placed points. "
                               "Use the partition inputs to place multiple populations of cells into the same space, "
                               "or use the 'prog' method.")

        max_dens=N/(vol_tot)
        V_sphere = 4 / 3 * np.pi * (dmin / 1000 / 2) ** 3
        dens_lim = 0.74 / V_sphere  # in points per mm^3
        print('Density limit:{:.3f}'.format(dens_lim))
        print('Max density requested:{:.3f}'.format(max_dens))
        if max_dens > dens_lim:
            raise ValueError('Requested density is incompatible with minimum distance: either reduce the density or dmin.')

        ind_nz = np.nonzero(mat)
        minx = np.min(ind_nz[0])
        miny = np.min(ind_nz[1])
        minz = np.min(ind_nz[2])
        maxx = np.max(ind_nz[0])
        maxy = np.max(ind_nz[1])
        maxz = np.max(ind_nz[2])

        corners = [v.flatten() for v in (np.meshgrid([minx, maxx+1], [miny, maxy+1], [minz, maxz+1]))]
        corners = np.vstack(corners)
        corners_t = (position_scale.astype(float).dot(corners))

        minx2 = np.min(corners_t[0])
        miny2 = np.min(corners_t[1])
        minz2 = np.min(corners_t[2])
        maxx2 = np.max(corners_t[0])
        maxy2 = np.max(corners_t[1])
        maxz2 = np.max(corners_t[2])

        # max_dims
        # Calc additional origin
        origin2 = [minx2, miny2, minz2]
        x_box = maxx2 - minx2
        y_box = maxy2 - miny2
        z_box = maxz2 - minz2

        # Calc box N
        N_box = math.floor(x_box * y_box * z_box * np.max(mat) / (1000 ** 3))              # N for lattice box

        vol_box = x_box * y_box * z_box
        r = (vol_box * 0.74 / N_box * 3 / 4 / np.pi) ** (1. / 3)

        x_n = int((x_box-r) / (2 * r)) + 1
        y_n = int(y_box / (r) / np.sqrt(3) - 1/3) + 1
        z_n = int(z_box / (2 * r) / np.sqrt(6) * 3) + 1

        x, y, z = hcp(x_n, y_n, z_n, r)

        # Randomly remove units in excess of N
        # Set to NaN rather than drop to keep lattice in register

        x_vec = x.flatten()
        y_vec = y.flatten()
        z_vec = z.flatten()

        todrop = np.random.choice(range(len(x_vec)), np.max([0,len(x_vec) - N_box]), replace=False)
        x_vec[todrop] = np.nan
        y_vec[todrop] = np.nan
        z_vec[todrop] = np.nan

        positions = np.vstack((x_vec, y_vec, z_vec))
        # Remove unneeded units by mat or geometry function
        position_scale_inv = np.linalg.inv(position_scale)
        inds = np.matmul(position_scale_inv, (positions+np.reshape(origin2, (-1, 1)))).T
        inds = np.floor(inds)
        mask = ~np.isnan(inds)
        mat_vals = np.full(inds.shape[0], fill_value=np.nan)
        mat_vals[mask[:,0]] = mat[inds[mask[:, 0],0].astype(int),inds[mask[:, 0],1].astype(int),inds[mask[:, 0],2].astype(int)]
        mat_vals = mat_vals==0
        mask = ~np.isnan(mat_vals)
        positions = positions.T
        positions [mat_vals[mask],:] = [np.nan, np.nan, np.nan]

        x_vec = positions[:,0]
        y_vec = positions[:,1]
        z_vec = positions[:,2]

        x = x_vec.reshape((x_n, y_n, z_n))
        y = y_vec.reshape((x_n, y_n, z_n))
        z = z_vec.reshape((x_n, y_n, z_n))

        # Jitter by up to 2*r-dmin and check if okay,if not, redraw jitter
        jitr_max = 2 * r - dmin

        for i in range(x.shape[0]):
            for j in range(x.shape[1]):
                for k in range(x.shape[2]):
                    # Don't need it to be uniformly distributed - the outer points are less likely to qualify
                    if np.isnan(x[i, j, k]):
                        # Eliminated unit, skip
                        continue
                    pt_old = [x[i, j, k], y[i, j, k], z[i, j, k]]
                    x[i, j, k] = float('inf')
                    y[i, j, k] = float('inf')
                    z[i, j, k] = float('inf')

                    keep = False

                    q = 2
                    indsi = np.arange(max(i - q, 0), min(i + q, x.shape[0]))
                    indsj = np.arange(max(j - q, 0), min(j + q, x.shape[1]))
                    indsk = np.arange(max(k - q, 0), min(k + q, x.shape[2]))

                    near_inds = np.meshgrid(indsi, indsj, indsk)

                    while not keep:
                        jitx = np.random.uniform(-jitr_max, jitr_max)
                        jity = np.random.uniform(-jitr_max, jitr_max)
                        jitz = np.random.uniform(-jitr_max, jitr_max)

                        x_new = pt_old[0] + jitx
                        y_new = pt_old[1] + jity
                        z_new = pt_old[2] + jitz

                        dists = np.sqrt((x_new - x[near_inds[0], near_inds[1], near_inds[2]]) ** 2 +
                                        (y_new - y[near_inds[0], near_inds[1], near_inds[2]]) ** 2 +
                                        (z_new - z[near_inds[0], near_inds[1], near_inds[2]]) ** 2)
                        d = np.nanmin(np.append(dists.flatten(), float('inf')))
                        keep = d > dmin

                    # Replace existing
                    x[i, j, k] = x_new
                    y[i, j, k] = y_new
                    z[i, j, k] = z_new

        x, y, z = x.flatten(), y.flatten(), z.flatten()
        x = x[~np.isnan(x)]
        y = y[~np.isnan(y)]
        z = z[~np.isnan(z)]

        before_trim_len = x.shape[0]
        if filter_func is None:
            # Trim any points that have been jittered out of permitted region
            inside = (x <= x_box) & (x >= 0) & (y <= y_box) & (y >= 0) & (z <= z_box) & (z >= 0)
        else:
            inside = filter_func(x, y, z)
        x = x[inside]
        y = y[inside]
        z = z[inside]

        positions = np.vstack((x, y, z)).T

        if positions.shape[0] > N:
            # Trim excess N, points are ordered, so pick randomly
            todrop = np.random.choice(range(len(x_vec)), positions.shape[0]-N, replace=False)
            positions[todrop,:] = np.nan

        while positions.shape[0] < N:
            if verbose==True:
                print(f'{positions.shape[0]}/{N} cells placed')
            tree = KDTree(positions, leaf_size=2)

            rand_pos_new = np.random.uniform(size=(3, before_trim_len))
            rand_nz_vox_new = np.random.choice(range(ind_nz[0].shape[0]), before_trim_len)

            positions_new = position_scale.dot([(ind_nz[0]-minx)[rand_nz_vox_new].astype(float) + rand_pos_new[0, :],
                                                (ind_nz[1]-miny)[rand_nz_vox_new].astype(float) + rand_pos_new[1, :],
                                                (ind_nz[2]-minz)[rand_nz_vox_new].astype(float) + rand_pos_new[2, :]])
            positions_new = positions_new.T
            dists, d_inds = tree.query(positions_new, k=1)
            conf_inds = np.squeeze(dists < dmin)
            positions_new[conf_inds, :] = [np.nan, np.nan, np.nan]
            positions_new = positions_new[~np.isnan(positions_new[:, 0]), :]
            if positions_new.shape[0] > 1:
                tree = KDTree(positions_new, leaf_size=2)
                dists, d_inds = tree.query(positions_new, k=2)
                conf_inds = np.squeeze(dists[:,1] < dmin)
                positions_new[conf_inds, :] = [np.nan, np.nan, np.nan]
                notna = ~np.isnan(positions_new[:, 0])
                positions_new = positions_new[notna, :]
                if filter_func is not None:
                    inside = filter_func(positions_new[:,0], positions_new[:,1], positions_new[:,2])
                    positions_new = positions_new[inside]
            positions = np.concatenate((positions, positions_new), axis=0)  # Add new positions

        positions = positions[:N, :]
        # If N rather than density, replace any missing

        x = x[~np.isnan(x)]
        y = y[~np.isnan(y)]
        z = z[~np.isnan(z)]

        # Shift re: origin
        x += origin2[0]
        y += origin2[1]
        z += origin2[2]

        return positions

    def positions_nrrd (nrrd_filename, max_dens_per_mm3, split_bilateral=None, dmin = 0.0, CCF_orientation=True, plot=False, method='prog',
                        existing_positions=np.array([]).reshape(0, 3), verbose=False):
        '''Generates random cell positions based on a .nrrd file. Matrix values are interpreted as cell densities
        for each voxel. The maximum density is scaled to max_dens_per_mm3 (cells/mm^3).
        If the .nrrd file is a structural mask, cells will be placed uniformly at max_dens_per_mm3 within the
        structure geometry.

        By default, only one hemisphere is shown, but this can be disabled by setting bilateral=True.

        :param nrrd_filename: path to .nrrd file
        :param max_dens_per_mm3: desired density at maximum value of nrrd array (cells/mm^3)
        :param split_bilateral: return only unilateral structure by removing half of the array along the given axis.
                If no splitting is desired, pass in None
        :return: A (N, 3) matrix (microns)
        '''

        readdata, header = nrrd.read(nrrd_filename)
        space_dir = header['space directions']
        origin = header['space origin']

        readdata_scaled = readdata / np.max(readdata) * max_dens_per_mm3

        if split_bilateral is not None:
            if split_bilateral=='x':
                readdata_scaled = readdata_scaled[:int(readdata_scaled.shape[0]/2), :, :]
            elif split_bilateral=='y':
                readdata_scaled = readdata_scaled[:, :int(readdata_scaled.shape[1]/2), :]
            elif split_bilateral=='z':
                readdata_scaled = readdata_scaled[:, :, :int(readdata_scaled.shape[2]/2)]
            else:
                error("split_bilateral must take values of 'x', 'y', or 'z', or 'None' if no processing is desired")

        # For testing non-uniform density
        #t = np.arange(0, readdata_scaled.shape[2], 1)
        #s = 1+np.sin(2*np.pi*0.1*t)
        #readdata_scaled = readdata_scaled * s

        positions = positions_density_matrix(readdata_scaled, position_scale=space_dir, origin = origin, plot=plot,
                                             CCF_orientation=CCF_orientation, dmin=dmin, method=method,
                                             existing_positions=existing_positions, verbose=verbose)

        return positions

    def xiter_random(N=1, min_x=0.0, max_x=1.0):
        return np.random.uniform(low=min_x, high=max_x, size=(N,))

    def plot_positions(x,y,z,ax,labels,pop_name=None):
        # Do plot a verification for each of these helper functions

        ax.scatter(x,y,z, marker='.', s=5, label=pop_name)

        ax.set_xlabel(labels[0])
        ax.set_ylabel(labels[1])
        ax.set_zlabel(labels[2])

    def hcp(x_n, y_n, z_n, r):
        # Hexagonal close packing lattice generation
        i = np.arange(x_n)
        j = np.arange(y_n)
        k = np.arange(z_n)

        get_x = lambda i, j, k, r: 2 * r * i + r * ((j + k) % 2)
        x = get_x(i[:, None, None], j[None, :, None], k[None, None, :], r)
        get_y = lambda i, j, k, r: r * np.sqrt(3) * (j + 1 / 3 * (k % 2)) + 0 * i
        y = get_y(i[:, None, None], j[None, :, None], k[None, None, :], r)
        get_z = lambda i, j, k, r: 2 * r * np.sqrt(6) / 3 * k + 0 * i + 0 * j
        z = get_z(i[:, None, None], j[None, :, None], k[None, None, :], r)

        return x, y, z

    def partition_locations(positions, partitions):
        # Assign positions to subpopulations based on probabilities
        pop_positions = []
        pop_assigned = np.random.choice(range(len(partitions)), positions.shape[0], p=partitions)

        for i in range(len(partitions)):
            pop_positions.append(positions[pop_assigned==i,:])

        return pop_positions