import numpy as np
import scipy.stats

def sample_uniform(N: int, Nt: int, L: float, origin: np.ndarray):
    """
    Generate Nt points uniformly sampled from a box of side length L.

    Args:
        - N (int): The number of grid points per side.
        - Nt (int): The number of tracers to generate.
        - L (float): The side-length of the box (Mpc/h).
        - origin (np.ndarray): The coordinates of the origin of the box (Mpc/h).

    Returns:
        - xtrue (np.ndarray): The true coordinates (Mpc/h) of the tracers.
    """

    h = 1

    xtrue = np.random.uniform(low=0.0, high=N+1, size=Nt)
    ytrue = np.random.uniform(low=0.0, high=N+1, size=Nt)
    ztrue = np.random.uniform(low=0.0, high=N+1, size=Nt)

    # Convert to coordinates, and move relative to origin
    xtrue *= L / N  # Mpc/h
    ytrue *= L / N  # Mpc/h
    ztrue *= L / N  # Mpc/h

    xtrue += origin[0]
    ytrue += origin[1]
    ztrue += origin[2]

    xtrue = np.array([xtrue, ytrue, ztrue])

    return xtrue


def draw_linear(nsamp: int, alpha: float, beta: float, u0: float, u1: float) -> np.ndarray:
    """
    Draw a sample from the probability distribution:
    p(u) \propto alpha (u1 - u) + beta (u - u0)
    for u0 <= u <= u1 and p(u) = 0 otherwise.

    Args:
        - nsamp (int): Number of samples to draw.
        - alpha (float): The coefficient of (u1 - u) in p(u).
        - beta (float): The coefficient of (u - u0) in p(u).
        - u0 (float): The minimum allowed value of u.
        - u1 (float): The maximum allowed value of u.

    Return:
        - np.ndarray: The samples from p(u).
    """
    n = scipy.stats.uniform(0, 1).rvs(nsamp)
    if isinstance(alpha, np.ndarray):
        res = np.zeros(alpha.shape)
        m = alpha != beta
        res[m] = ((u1 - u0) * np.sqrt(n * (beta ** 2 - alpha ** 2) + alpha ** 2) - u1 * alpha + u0 * beta)[m] / (beta - alpha)[m]
        res[~m] = (u0 + (u1 - u0) * n)[~m]
        return res
    else:
        if alpha != beta:
            return ((u1 - u0) * np.sqrt(n * (beta ** 2 - alpha ** 2) + alpha ** 2) - u1 * alpha + u0 * beta) / (beta - alpha)
        else:
            return u0 + (u1 - u0) * n

        
def periodic_index(index: np.ndarray, shape: tuple) -> np.ndarray:
    """
    Apply periodic boundary conditions to an array of indices.

    Args:
        - index (np.ndarray): The indices to transform. Shape = (ndim, nvals). 
        - shape (tuple): The shape of the box used for periodic boundary conditions (N0, N1, ...)

    Returns:
        - new_index (np.ndarray): The values in index after applying periodic boundary conditions, such that for dimension i, the values are in the range [0, Ni)
    """
    assert index.shape[0] == len(shape)
    new_index = index.copy()
    for i in range(len(shape)):
        new_index[i,:] = np.mod(new_index[i,:], shape[i])
    return new_index


def get_new_index(index: np.ndarray, shape: tuple, subscript: tuple) -> np.ndarray:
    """
    If each entry of index corresponds to (0,0,0), find the index corresponding to the point given by subscript.

    Args:
        - index (np.ndarray): The indices to transform. Shape = (ndim, nvals). 
        - shape (tuple): The shape of the box used (N0, N1, ...).
        - subscript (tuple): The coordinate to find, relative to the values given in index.

    Returns:
        - new_index (np.ndarray): The new index values.

    """
    new_index = index.copy()
    for i in range(len(subscript)):
        new_index[i,:] += subscript[i]
    new_index = periodic_index(new_index, shape)
    return new_index


def sample_3d(phi: np.ndarray, Nt: int, L: float, origin: np.ndarray) -> np.ndarray:
    """
    Sample Nt points, assuming that the points are drawn from a Poisson process given by the field phi.
    phi gives the value of the field at the grid points, and we assume linear interpolation between points.

    Args:
        - phi (np.ndarray): The field which defines the mean of the Poisson process.
        - Nt (int): The number of tracers to generate.
        - L (float): The side-length of the box (Mpc/h).
        - origin (np.ndarray): The coordinates of the origin of the box (Mpc/h).

    Returns:
        - xtrue (np.ndarray): The true coordinates (Mpc/h) of the tracers.
    """

    N = phi.shape[0]
    h = 1

    # (1) Find which cell each point lives in
    mean = phi + \
            np.roll(phi, -1, axis=0) + \
            np.roll(phi, -1, axis=1) + \
            np.roll(phi, -1, axis=2) + \
            np.roll(phi, -1, axis=(0,1)) + \
            np.roll(phi, -1, axis=(0,2)) + \
            np.roll(phi, -1, axis=(1,2)) + \
            np.roll(phi, -1, axis=(0,1,2))
    prob = mean.flatten() / mean.sum()
    i = np.arange(prob.shape[0])
    a1d = np.random.choice(i, Nt, p=prob)
    a3d = np.array(np.unravel_index(a1d, (N,N,N)))

    # (2) Find the x values
    shape = (N, N, N)
    alpha = np.zeros(Nt)
    for subscript in [(0,0,0), (0,0,1), (0,1,0), (0,1,1)]:
        idx = get_new_index(a3d, shape, subscript)
        alpha += phi[idx[0,:], idx[1,:], idx[2,:]]
    beta = np.zeros(Nt)
    for subscript in [(1,0,0), (1,0,1), (1,1,0), (1,1,1)]:
        idx = get_new_index(a3d, shape, subscript)
        beta += phi[idx[0,:], idx[1,:], idx[2,:]]
    u0 = a3d[0,:]
    u1 = a3d[0,:] + 1
    xtrue = draw_linear(Nt, alpha, beta, u0, u1)

    # (3) Find the y values
    shape = (N, N, N)
    alpha = np.zeros(Nt)
    for subscript in [(0,0,0), (0,0,1)]:
        idx = get_new_index(a3d, shape, subscript)
        alpha += phi[idx[0,:], idx[1,:], idx[2,:]] * (a3d[0,:] + 1 - xtrue)
    for subscript in [(1,0,0), (1,0,1)]:
        idx = get_new_index(a3d, shape, subscript)
        alpha += phi[idx[0,:], idx[1,:], idx[2,:]] * (xtrue - a3d[0,:])
    beta = np.zeros(Nt)
    for subscript in [(0,1,0), (0,1,1)]:
        idx = get_new_index(a3d, shape, subscript)
        beta += phi[idx[0,:], idx[1,:], idx[2,:]] * (a3d[0,:] + 1 - xtrue)
    for subscript in [(1,1,0), (1,1,1)]:
        idx = get_new_index(a3d, shape, subscript)
        beta += phi[idx[0,:], idx[1,:], idx[2,:]] * (xtrue - a3d[0,:])
    u0 = a3d[1,:]
    u1 = a3d[1,:] + 1
    ytrue = draw_linear(Nt, alpha, beta, u0, u1)

    # (4) Find the z values
    xd = (xtrue - a3d[0,:])  # x1-x0=1 so xd = x - x0
    yd = (ytrue - a3d[1,:])  # y1-y0=1 so yd = y - y0
    ia = get_new_index(a3d, shape, (0,0,0))
    ib = get_new_index(a3d, shape, (1,0,0))
    phi00 = phi[ia[0,:], ia[1,:], ia[2,:]] * (1 - xd) + \
            phi[ib[0,:], ib[1,:], ib[2,:]] * xd
    ia = get_new_index(a3d, shape, (0,0,1))
    ib = get_new_index(a3d, shape, (1,0,1))
    phi01 = phi[ia[0,:], ia[1,:], ia[2,:]] * (1 - xd) + \
            phi[ib[0,:], ib[1,:], ib[2,:]] * xd
    ia = get_new_index(a3d, shape, (0,1,0))
    ib = get_new_index(a3d, shape, (1,1,0))
    phi10 = phi[ia[0,:], ia[1,:], ia[2,:]] * (1 - xd) + \
            phi[ib[0,:], ib[1,:], ib[2,:]] * xd
    ia = get_new_index(a3d, shape, (0,1,1))
    ib = get_new_index(a3d, shape, (1,1,1))
    phi11 = phi[ia[0,:], ia[1,:], ia[2,:]] * (1 - xd) + \
            phi[ib[0,:], ib[1,:], ib[2,:]] * xd
    alpha = phi00 * (1 - yd) + phi10 * yd  # alpha = phi0
    beta = phi01 * (1 - yd) + phi11 * yd   # beta = phi1
    u0 = a3d[2,:]
    u1 = a3d[2,:] + 1
    ztrue = draw_linear(Nt, alpha, beta, u0, u1)

    # Convert to coordinates, and move relative to origin
    xtrue *= L / N  # Mpc/h
    ytrue *= L / N  # Mpc/h
    ztrue *= L / N  # Mpc/h

    xtrue += origin[0]
    ytrue += origin[1]
    ztrue += origin[2]
    
    xtrue = np.array([xtrue, ytrue, ztrue])

    return xtrue