import numpy as np
from pysbmy.snapshot import read_tile
from pysbmy.field import Field
from pysbmy.density import get_density_pm, density_to_delta
from pysbmy import c_float_t
import tqdm
import argparse

def get_indices(tile, N_TILES0, N_TILES1, N_TILES2):
    """Get the indices of the tile in a 3D grid."""
    i = (tile // (N_TILES1 * N_TILES2))%N_TILES0
    j = ((tile - N_TILES1 * N_TILES2 * i) // N_TILES2)%N_TILES1
    k = (tile - N_TILES2 * j - N_TILES1 * N_TILES2 * i)%N_TILES2
    return i, j, k


def tile_to_density_with_buffer(T, N, d, buffer):
    """
    Convert a tile to density with a buffer.
    """
    # Create a buffer for the density
    A = np.zeros((N+2*buffer, N+2*buffer, N+2*buffer), dtype=np.float32)
    
    # Get the density of the tile
    A = get_density_pm(T.pos-T.corner_position.astype(c_float_t)+d*buffer, A, d)
    return A


def add_density_tile(A,A_tile,corner_indices):
    N = A.shape[0]
    Np = A_tile.shape[0]
    i0, j0, k0 = corner_indices

    # Create index arrays for the tile
    i_idx = (np.arange(Np) + i0) % N
    j_idx = (np.arange(Np) + j0) % N
    k_idx = (np.arange(Np) + k0) % N

    # Use np.ix_ to get all combinations of indices
    A[np.ix_(i_idx, j_idx, k_idx)] += A_tile


def gather_density(A, folder, tile_base, Np_tile, dpm, buffer, N_TILES):
    """
    Gather the density from the tiles.
    """
    for tile in tqdm.tqdm(range(N_TILES**3), desc="Density of tiles", unit="tiles"):
        T=read_tile(folder+tile_base+str(tile+1)+".h5", int(Np_tile**3))
        # Get the corner position of the tile
        corner_position = T.corner_position
        # Get the corner indices of the tile
        i,j,k = get_indices(tile, N_TILES, N_TILES, N_TILES)
        corner_grid_indices = (i*Np_tile-buffer, j*Np_tile-buffer, k*Np_tile-buffer)
        # Get the density of the tile with a buffer
        A_tile = tile_to_density_with_buffer(T, Np_tile, dpm, buffer)
        # Add the density of the tile to the grid
        add_density_tile(A, A_tile, corner_grid_indices)



def gather_tiles(folder, tile_base, L, Np, N_TILES, buffer):
    """
    Gather sCOLA tiles into a single density field.

    Parameters
    ----------
    folder : str
        Folder containing the tiles.
    tile_base : str
        Base name of the tiles.
    L : float
        Size of the box in Mpc/h.
    Np : int
        Number of cells per dimension for the full box.
    N_TILES : int
        Number of tiles per dimension.
    buffer : int
        Buffer size for the density field of tiles.
    """
    Np_tile = Np//N_TILES
    dpm = L/Np_tile

    print("Memory allocation for the grid...")
    A=np.zeros((Np,Np,Np), dtype=np.float32)

    print("Starting to read the tiles...")
    gather_density(A, folder, tile_base, Np_tile, dpm, buffer, N_TILES)

    print("Finished reading the tiles.")
    A=density_to_delta(A,-1)

    print("Converting to field...")
    F=Field(L,L,L, 0.,0.,0., 1, Np,Np,Np, 1., A)

    print("Saving field...")
    F.write(folder+"../results/final_density_sCOLA.h5")
    print("Density field saved to", folder+"../results/final_density_sCOLA.h5")
    print("Done.")




if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Gather density from tiles.")
    parser.add_argument("-d","--folder", type=str, default="./", help="Folder containing the tiles")
    parser.add_argument("--tile_base", type=str, default="sCOLA_tile", help="Base name of the tiles")
    parser.add_argument("-L","--L", type=int, default=1920, help="Size of the box in Mpc/h")
    parser.add_argument("-Np","--Np", type=int, default=80, help="Number of cells per dimension for the full box")
    parser.add_argument("-Nt","--N_tiles", type=int, default=4, help="Number of tiles per dimension.")
    parser.add_argument("--buffer", type=int, default=40, help="Buffer size for the density field of tiles")
    args = parser.parse_args()

    L = args.L
    Np = args.Np
    N_TILES = args.N_tiles
    buffer = args.buffer

    folder = args.folder
    tile_base = args.tile_base

    Np_tile = Np//N_TILES
    dpm = L/Np_tile

    gather_tiles(folder, tile_base, L, Np, N_TILES, buffer)