import numpy as np
from scipy.special import gammainc
from skimage.transform import resize


# Uniform sampling in a hyperspere
# Based on Matlab implementation by Roger Stafford
# Can be optimized for Bridson algorithm by excluding all points within the r/2 sphere
def hypersphere_volume_sample(center, radius, k=1, rng_new=np.random.default_rng):
    ndim = center.size
    x = rng_new.normal(size=(k, ndim))
    ssq = np.sum(x ** 2, axis=1)
    fr = radius * gammainc(ndim / 2, ssq / 2) ** (1 / ndim) / np.sqrt(ssq)
    frtiled = np.tile(fr.reshape(k, 1), (1, ndim))
    p = center + np.multiply(x, frtiled)
    return p


# Uniform sampling on the sphere's surface
def hypersphere_surface_sample(center, radius, k=1):
    ndim = center.size
    vec = np.random.standard_normal(size=(k, ndim))
    vec /= np.linalg.norm(vec, axis=1)[:, None]
    p = center + np.multiply(vec, radius)
    return p


def squared_distance(p0, p1):
    return np.sum(np.square(p0 - p1))


def Bridson_sampling(dims=np.array([1.0, 1.0]), radius=0.05, k=30, hypersphere_sample=hypersphere_volume_sample, rng_new=None):
    # References: Fast Poisson Disk Sampling in Arbitrary Dimensions
    #             Robert Bridson, SIGGRAPH, 2007

    ndim = dims.size

    # size of the sphere from which the samples are drawn relative to the size of a disc (radius)
    sample_factor = 2
    if hypersphere_sample == hypersphere_volume_sample:
        sample_factor = 2

    # for the surface sampler, all new points are almost exactly 1 radius away from at least one existing sample
    # eps to avoid rejection
    if hypersphere_sample == hypersphere_surface_sample:
        eps = 0.001
        sample_factor = 1 + eps

    def in_limits(p):
        return np.all(np.zeros(ndim) <= p) and np.all(p < dims)

    # Check if there are samples closer than "squared_radius" to the candidate "p"
    def in_neighborhood(p, n=2):
        indices = (p / cellsize).astype(int)
        indmin = np.maximum(indices - n, np.zeros(ndim, dtype=int))
        indmax = np.minimum(indices + n + 1, gridsize)

        # Check if the center cell is empty
        if not np.isnan(P[tuple(indices)][0]):
            return True
        a = []
        for i in range(ndim):
            a.append(slice(indmin[i], indmax[i]))
        with np.errstate(invalid='ignore'):
            if np.any(np.sum(np.square(p - P[tuple(a)]), axis=ndim) < squared_radius):
                return True

    def add_point(p):
        points.append(p)
        indices = (p / cellsize).astype(int)
        P[tuple(indices)] = p

    cellsize = radius / np.sqrt(ndim)
    gridsize = (np.ceil(dims / cellsize)).astype(int)

    # Squared radius because we'll compare squared distance
    squared_radius = radius * radius

    # Positions of cells
    P = np.empty(np.append(gridsize, ndim), dtype=np.float32)  # n-dim value for each grid cell
    # Initialise empty cells with NaNs
    P.fill(np.nan)

    points = []
    add_point(rng_new.uniform(np.zeros(ndim), dims))
    while len(points):
        i = rng_new.integers(len(points))
        p = points[i]
        del points[i]
        Q = hypersphere_sample(np.array(p), radius * sample_factor, k, rng_new)
        for q in Q:
            if in_limits(q) and not in_neighborhood(q):
                add_point(q)
    return P[~np.isnan(P).any(axis=ndim)]


def create_circular_mask(h, w, center=None, radius=None):
    if center is None:  # use the middle of the image
        center = (int(w / 2), int(h / 2))
    if radius is None:  # use the smallest distance between the center and image walls
        radius = min(center[0], center[1], w - center[0], h - center[1])

    Y, X = np.ogrid[:h, :w]
    dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)

    mask = (dist_from_center <= radius).astype(float)
    return mask


class CircularPoissonMaskGenerator:
    def __init__(self, shape: np.array, central_sampling: bool, accel: float, radius: float = 0.0125, rng_new=None):
        if rng_new is None:
            rng_new = np.random.default_rng(42)
        self.shape = shape
        self.central_sampling = central_sampling
        self.accel = accel
        self.radius = radius
        self.mask = np.zeros(tuple(self.shape))
        self.rng_new = rng_new

    def __call__(self):
        while True:
            shape_bs = self.shape / max(self.shape)
            bs = Bridson_sampling(shape_bs, radius=self.radius, rng_new=self.rng_new)
            bs = (bs * max(self.shape)).astype(int)
            self.mask *= 0
            if self.central_sampling:
                cent_x = round(self.shape[0] // 2)
                cent_y = round(self.shape[1] // 2)
                perc_x = round(self.shape[0] * 0.15)
                perc_y = round(self.shape[1] * 0.05)
                self.mask[cent_x - perc_x:cent_x + perc_x, cent_y - perc_y:cent_y + perc_y] = 1
            self.mask[bs[:, 0], bs[:, 1]] = 1
            cm = create_circular_mask(100, 100)
            cm = resize(cm, tuple(self.shape))
            self.mask *= cm
            R = 1 / (np.sum(self.mask) / np.prod(self.mask.shape) + 1e-8)
            if abs(R - self.accel) < 0.1:
                return self.mask
            elif R > self.accel:
                self.radius -= 0.0001
            elif R < self.accel:
                self.radius += 0.0001


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    gen = CircularPoissonMaskGenerator(np.array([32, 148]), True, 6, rng_new=np.random.default_rng())
    mask1 = gen()
    mask2 = gen()
    plt.subplot(1, 3, 1)
    plt.imshow(mask1)
    plt.subplot(1, 3, 2)
    plt.imshow(mask2)
    plt.subplot(1, 3, 3)
    plt.imshow(abs(mask1 - mask2))
    plt.savefig("test.png")
