import numpy as np

# how to read the websky ICs
boxsize = 7700. # comoving Mpc

N=768
icsfname='Fvec_7700Mpc_n6144_nb30_nt16_no768'

# uncomment following two lines for the full resolution instead
# beware you need 864 GB of RAM to actually read this in
# N=6144
# icsfname='Fvec_7700Mpc_n6144_nb30_nt16'

icsfile = open(icsfname,"rb")

# read all N**3 values into an array
# delta will be a 1D array initially
delta = np.fromfile(icsfile,count=N**3,dtype=np.float32)

# reshape delta to be a 3D array to make bookkeeping easier
delta = delta.reshape((N,N,N))

grid_spacing = boxsize / N

# report value of delta at a specific lattice point
# (i,j,k) = (3,4,5)
i=3; j=4; k=5

def linear_growth_factor(z):
    # insert correct function here
    return 1./(1+z)

def redshift_of_chi(chi):
    # insert correct function here
    return 70. / 3e5 * chi - 1.

# loop over octants
for xd in (-1,0):
    x = (i+0.5)*grid_spacing + xd * boxsize
    for yd in (-1,0):
        y = (j+0.5)*grid_spacing + yd * boxsize
        for zd in (-1,0):
            z = (k+0.5)*grid_spacing + zd * boxsize

            chi = np.sqrt(x**2+y**2+z**2)
            redshift = redshift_of_chi(chi)
            delta_linear = delta[i,j,k] * linear_growth_factor(redshift)
            print('delta_linear at ({:+0.2e},{:+0.2e},{:+0.2e}) in octant ({:+d},{:+d},{:+d}) is {:+0.2e}'.format(x, y, z,
                                                                                 xd,yd,zd,delta_linear))