import aquila_borg as borg
import numpy as np
import configparser
import borg_velocity.utils as utils
from borg_velocity.utils import myprint
import borg_velocity.forwards as forwards
import Pk_library as PKL
import matplotlib.pyplot as plt
import symbolic_pofk.syrenhalofit as syrenhalofit

def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None, gravity='lpt') -> borg.forward.BaseForwardModel:
    """
    Builds the gravity model and returns the forward model chain.

    Args:
        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
        - box (borg.forward.BoxModel): The input box model.
        - ini_file (str, default=None): The location of the ini file. If None, use borg.getIniConfigurationFilename()

    Returns:
        borg.forward.BaseForwardModel: The forward model.

    """
    
    myprint(f"Building gravity model {gravity}")

    if ini_file is None:
        myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
        config = configparser.ConfigParser()
        config.read(borg.getIniConfigurationFilename())
    else:
        myprint("Reading from configuration file: " + ini_file)
        config = configparser.ConfigParser()
        config.read(ini_file)
    ai = float(config['model']['ai'])
    af = float(config['model']['af'])
    supersampling = int(config['model']['supersampling'])
    print('Supersampling:', supersampling)
    
    if gravity in ['pm', 'cola']:
        forcesampling = int(config['model']['forcesampling'])
        forcesampling = 2
        print('FORCESAMPLING:', forcesampling)

    # Setup forward model
    chain = borg.forward.ChainForwardModel(box)
    chain.addModel(borg.forward.models.HermiticEnforcer(box))
  
    # CLASS transfer function
    chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box)
    transfer_class = borg.forward.model_lib.M_TRANSFER_CLASS(box, opts=dict(a_transfer=1.0))
    transfer_class.setModelParams({"extra_class_arguments":{'YHe':'0.24'}})
    chain @= transfer_class
    
    if gravity == 'linear':
        raise NotImplementedError(gravity)
    elif gravity == 'lpt':
        mod = borg.forward.model_lib.M_LPT_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=supersampling,
                      lightcone=False,
                      part_factor=1.01,))
    elif gravity == '2lpt':
        mod = borg.forward.model_lib.M_2LPT_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=supersampling,
                      lightcone=False,
                      part_factor=1.01,))
    elif gravity == 'pm':
        mod =  borg.forward.model_lib.M_PM_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=supersampling,
                      part_factor=1.01,
                      forcesampling=forcesampling,
                      pm_start_z=1/ai - 1,
                      pm_nsteps=int(config['model']['nsteps']),
                      tcola=False))
    elif gravity == 'cola':
        mod = borg.forward.model_lib.M_PM_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=supersampling,
                      part_factor=1.01,
                      forcesampling=forcesampling,
                      pm_start_z=1/ai - 1,
                      pm_nsteps=int(config['model']['nsteps']),
                      tcola=True))
    else:
        raise NotImplementedError(gravity)
        
    mod.accumulateAdjoint(True)
    chain @= mod
        
    # Cosmological parameters
    if ini_file is None:
        cpar = utils.get_cosmopar(borg.getIniConfigurationFilename())
    else:
        cpar = utils.get_cosmopar(ini_file)
    chain.setCosmoParams(cpar)
    
    # This is the forward model for the model parameters
    fwd_param = borg.forward.ChainForwardModel(box)
    mod_null = forwards.NullForward(box)
    fwd_param.addModel(mod_null)
    fwd_param.setCosmoParams(cpar)
    
    # This is the forward model for velocity
    velmodel_name = config['model']['velocity']
    velmodel = getattr(borg.forward.velocity, velmodel_name)
    if velmodel_name == 'LinearModel':
        fwd_vel = velmodel(box, mod, af)
    elif velmodel_name == 'CICModel':
        rsmooth = float(config['model']['rsmooth'])  # Mpc/h
        fwd_vel = velmodel(box, mod, rsmooth)
    else:
        fwd_vel = velmodel(box, mod)
    
    return chain

ini_file = '../conf/velmass_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.N = (128, 128, 128)
# box_in.N = (256, 256, 256)
box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))

# Make some initial conditions
s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
s_real = np.fft.irfftn(s_hat, norm="ortho")

all_gravity = ['lpt', '2lpt', 'pm', 'cola']
all_dens = [None] * len(all_gravity)
all_pk = [None] * len(all_gravity)

fig, ax = plt.subplots(1, 1, figsize=(6,4))

for i, gravity in enumerate(all_gravity):
    # Run BORG density field
    output_density = np.zeros(box_in.N)
    fwd_model = build_gravity_model(None, box_in, ini_file=ini_file, gravity=gravity)
    fwd_model.forwardModel_v2(s_hat)
    fwd_model.getDensityFinal(output_density)
    all_dens[i] = output_density.copy()
    
    Pk = PKL.Pk(all_dens[i].astype(np.float32), box_in.L[0], axis=0, MAS='CIC', threads=1, verbose=True) 
    all_pk[i] = (Pk.k3D.copy(), Pk.Pk[:,0].copy())

    ax.loglog(all_pk[i][0], all_pk[i][1], label=gravity)
    
# Load VELMASS field
dirname = config['mock']['velmass_dirname']
field = np.load(dirname + '/../density.npz')['d']
field = field / np.mean(field) - 1
with open(dirname + '/auto-rockstar.cfg') as f:
    data = [r for r in f]
Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip()
Lbox = float(Lbox.split('=')[1])
print('VELMASS box', Lbox)
Pk = PKL.Pk(field.astype(np.float32), Lbox, axis=0, MAS='CIC', threads=1, verbose=True)
m = Pk.k3D > all_pk[-1][0].min()
ax.loglog(Pk.k3D[m], Pk.Pk[m,0], label='velmass')

# Syren
k = Pk.k3D[m]
pk_syrenhalofit = syrenhalofit.run_halofit(k, 
                       cosmo.sigma8, cosmo.omega_m, cosmo.omega_b, cosmo.h, cosmo.n_s, 1.0,
                       emulator='fiducial', extrapolate=True, which_params='Bartlett', add_correction=True)
ax.loglog(k, pk_syrenhalofit, label='syren')
ymin = max(ax.get_ylim()[0], 1e1)
ax.set_ylim(ymin, None)    
ax.legend()
fig.tight_layout()
fig.savefig('../figs/pk_gravity_models_test.png', facecolor='white', bbox_inches='tight')
plt.close(fig)