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

cosmo={'omega_M_0':0.3175, 'h':0.6711}
cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
cosmo['omega_k_0'] = 0

a_start=0.001
z_start=1/a_start-1


def fourier_analysis(borg_vol):
  L = (borg_vol.ranges[1]-borg_vol.ranges[0])
  N = borg_vol.density.shape[0]
    
  return np.fft.rfftn(borg_vol.density)*(L/N)**3, L, N

def half_pixel_shift(borg):

    dhat,L,N = fourier_analysis(borg)

    ik = np.fft.fftfreq(N,d=L/N)*2*np.pi
    phi = 0.5*L/N*(ik[:,None,None]+ik[None,:,None]+ik[None,None,:(N/2+1)])
    phase = np.cos(phi)+1j*np.sin(phi)
    
    return dhat*phase, L

def new_shape(N, direction, q=3):
    return ((1,)*direction) + (N,) + ((1,)*(q-1-direction))

def build_dir(ik, direction, q=3):
    if direction != q-1:
        return ik.reshape(new_shape(ik.size, direction, q=q))
    else:
        N = ik.size/2+1
        return ik[:N].reshape(new_shape(N, direction, q=q))

def get_k2(ik, q=3):

    N = ik.size
    k2 = (ik.reshape(new_shape(N, 0, q=q))**2)
    for d in xrange(1,q):
        k2 = k2 + build_dir(ik, d, q=q)**2
    return k2

def get_displacement(dhat, L, direction=0):

    N = dhat.shape[0]
    ik = np.fft.fftfreq(N,d=1.0/N)*2*np.pi/L
    
    k2 = get_k2(ik)
    k2[0,0,0] = 1
    
    return -build_dir(ik, direction)*1j*dhat / k2

def gen_posgrid(N, L):

    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_ic):

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

    cgrowth = CosmoGrowth(**cosmo)

    density_hat, L = half_pixel_shift(borg_vol)

    posq = gen_posgrid(N, L)
    vel= []
    posx = []
    velmul = cgrowth.compute_velmul(a_start)

    D1 = cgrowth.D(a_ic)
    D1_0 = D1/cgrowth.D(a_start)
    D2 = 3./7 * D1**2

    for j in xrange(3):
        psi = D1_0*np.fft.irfftn(get_displacement(density_hat, L, direction=j)).flatten()*(N/L)**3
        posx.append(((posq[j] + psi)%L).astype(np.float32))
        vel.append((psi*velmul).astype(np.float32)) 

    return posx,vel,N,L,a_ic
    
def write_icfiles(*generated_ic):
    posx,vel,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("borg.ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
        
        
if __name__=="__main__":
    write_icfiles(*run_generation("initial_condition_borg.dat", 1.0))