import numpy as np
import cosmotool as ct
import borgicgen as bic

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

TestCase=True
zstart=10
astart=1/(1.+zstart)

if TestCase:
  pos,_,density,N,L,_ = bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=True)

  dcic = ct.cicParticles(pos, L, N)
  dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0)
  dcic -= 1

  dcic_hat = np.fft.rfftn(dcic)*(L/N)**3

  Pcic, bcic = bic.bin_power(np.abs(dcic_hat)**2/L**3, L, bins=50)

  borg_evolved = ct.read_borg_vol("final_density_1380.dat")

if __name__=="__main__":
  if not TestCase:
    bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, do_lpt2=True), **cosmo)