import numpy as np
import aquila_borg as borg
import configparser
import matplotlib.pyplot as plt

import borg_velocity.likelihood as likelihood
import borg_velocity.utils as utils
import borg_velocity.forwards as forwards

ini_file = '../conf/basic_ini.ini'

# Setup config
config = configparser.ConfigParser()
config.read(ini_file)

# Cosmology
cosmo = utils.get_cosmopar(ini_file)

# Input box
box_in = borg.forward.BoxModel()
box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))

# Load forward model
fwd_model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file)
fwd_model.setCosmoParams(cosmo)

# Make some initial conditions
s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)

# Get the real version of this
s_real = np.fft.irfftn(s_hat, norm="ortho")

# Run BORG density field
output_density = np.zeros(box_in.N)
fwd_model.forwardModel_v2(s_hat)
fwd_model.getDensityFinal(output_density)

# Get growth rate
cosmology = borg.cosmo.Cosmology(cosmo)
af = float(config['model']['af'])
f = cosmology.gplus(af)  # dD / da
f *= af / cosmology.d_plus(af)  # f = dlnD / dlna

# Get velocity
smooth_R = float(config['model']['smooth_R'])
output_vel = forwards.dens2vel_linear(output_density, f, box_in.L[0], smooth_R)
print(output_vel.shape)


# Plot the IC and dens fields
fig, axs = plt.subplots(2, 3, figsize=(15,10))
for i, field in enumerate([s_real, np.log10(2 + output_density)]):
    vmin, vmax = field.min(), field.max()
    pc = axs[i,0].pcolor(field[box_in.N[0]//2], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,0])
    pc = axs[i,1].pcolor(field[:,box_in.N[1]//2,:], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,1])
    pc = axs[i,2].pcolor(field[:,:,box_in.N[2]//2], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,2])
axs[0,1].set_title('Initial Conditions')
axs[1,1].set_title(r'$\log_{10} (2 + \delta)$')
for ax in axs.flatten():
    ax.set_aspect('equal')
fig.savefig('../figs/test_dens_ic.png')


# Plot the velocity fields
fig, axs = plt.subplots(3, 3, figsize=(15,15))
vmin, vmax = output_vel.min(), output_vel.max()
for i in range(3):
    pc = axs[i,0].pcolor(output_vel[i,box_in.N[0]//2], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,0])
    pc = axs[i,1].pcolor(output_vel[i,:,box_in.N[1]//2,:], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,1])
    pc = axs[i,2].pcolor(output_vel[i,:,:,box_in.N[2]//2], vmin=vmin, vmax=vmax)
    fig.colorbar(pc, ax=axs[i,2])
axs[0,1].set_title(r'$v_x$')
axs[1,1].set_title(r'$v_y$')
axs[2,1].set_title(r'$v_z$')
for ax in axs.flatten():
    ax.set_aspect('equal')
fig.savefig('../figs/test_vel.png')