import cosmotool as ct
import numpy as np
import cosmolopy as cpy
from cosmogrowth import *
import borgadaptor as ba

def gen_posgrid(N, L):
    """ Generate an ordered lagrangian grid"""

    ix = np.arange(N)*L/N

    x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
    y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
    z = ix[None,None,:].repeat(N, axis=0).repeat(N, axis=1)

    return x.flatten(), y.flatten(), z.flatten()

def run_generation(input_borg, a_borg, a_ic, **cosmo):
    """ Generate particles and velocities from a BORG snapshot. Returns a tuple of 
    (positions,velocities,N,BoxSize,scale_factor)."""

    borg_vol = ct.read_borg_vol(input_borg)
    N = borg_vol.density.shape[0]

    cgrowth = CosmoGrowth(**cosmo)

    density_hat, L = ba.half_pixel_shift(borg_vol)


    lpt = LagrangianPerturbation(density_hat, L, fourier=True)

    # Generate grid
    posq = gen_posgrid(N, L)
    vel= []
    posx = []
    
    # Compute LPT scaling coefficient
    D1 = cgrowth.D(a_ic)
    D1_0 = D1/cgrowth.D(a_borg)
    velmul = cgrowth.compute_velmul(a_ic)*D1_0
    
    D2 = 3./7 * D1**2

    for j in xrange(3):
        # Generate psi_j (displacement along j)
        psi = D1_0*lpt.lpt1(j).flatten()*(N/L)**3
        # Generate posx
        posx.append(((posq[j] + psi)%L).astype(np.float32))
        # Generate vel
        vel.append((psi*velmul).astype(np.float32)) 

    density = np.fft.irfftn(density_hat*D1_0)*(N/L)**3

    return posx,vel,density,N,L,a_ic
    
def write_icfiles(*generated_ic, **cosmo):
    """Write the initial conditions from the tuple returned by run_generation"""
    posx,vel,density,N,L,a_ic = generated_ic

    ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) 
    for i,c in enumerate(['x','y','z']):
        ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)

    ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
    ct.writeWhitePhase("white.dat", density)