import borg_velocity.likelihood as likelihood
import borg_velocity.forwards as forwards
import borg_velocity.utils as utils
from borg_velocity.utils import myprint

import aquila_borg as borg
import configparser
import h5py as h5
import numpy as np
from tqdm import tqdm
import glob
import os
import matplotlib.pyplot as plt

dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics'
ini_file = f'{dirname}/ini_file.ini'

wdir = os.getcwd()
os.chdir(dirname)

def get_mcmc_steps(nframe, iter_max, iter_min=0):
    """
    Obtain evenly-spaced sample of MCMC steps to make movie from
    """
    all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
    x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc]
    all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc])
    if iter_max >= 0:
        all_mcmc = all_mcmc[all_mcmc <= iter_max]
    all_mcmc = all_mcmc[all_mcmc >= iter_min]
    if nframe > 0:
        max_out = max(all_mcmc)
        min_out = min(all_mcmc)
        step = max(int((max_out - min_out+1) / nframe), 1)
        all_mcmc = all_mcmc[::step]
        if max_out not in all_mcmc:
            all_mcmc = np.concatenate([all_mcmc, [max_out]])
    return all_mcmc


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

    Args:
        - gravity (str): Which gravity model to use (in ['pm', 'cola', 'lpt', '2lpt']
        - 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.

    """

    global chain, fwd_param, fwd_vel
    myprint("Building gravity model")

    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'])
    
    if gravity in ['pm', 'cola']:
        if 'forcesampling' in config['model']:
            forcesampling = int(config['model']['forcesampling'])
        else:
            myprint('Could not find forcesampling, so will set to 4')
            forcesampling = 4
            
    # 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


def build_likelihood(gravity, box_in, ini_file):
    
    global fwd_vel
    
    model = build_gravity_model(gravity, None, box_in, ini_file=ini_file)
    cosmo = utils.get_cosmopar(ini_file)
    model.setCosmoParams(cosmo)
    fwd_param = forwards.NullForward(box_in)
    mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
    
    # Load the mock data
    mylike.loadMockData(None)
    
    return mylike

ref_gravity = 'lpt'
ref_gravity = '2lpt'


# Input box
box_in = borg.forward.BoxModel()
config = configparser.ConfigParser()
config.read(ini_file)
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']))

# Setup BORG forward model and likelihood
cola_like = build_likelihood('cola', box_in, ini_file)
lpt_like = build_likelihood(ref_gravity, box_in, ini_file)

# Load some MCMC samples and see how they compare
nframe = 100
iter_min = 1000
iter_max = -1
all_mcmc = get_mcmc_steps(nframe, iter_max, iter_min=iter_min)

all_lpt_like = np.empty(len(all_mcmc))
all_cola_like = np.empty(len(all_mcmc))

for i, mcmc in enumerate(all_mcmc):
    myprint(f'\nStarting MCMC {i+1} of {len(all_mcmc)}: {mcmc}')

    # Load an s_hat from file
    with h5.File(f'{dirname}/mcmc_{mcmc}.h5', 'r') as f:
        s_hat = f['scalars/s_hat_field'][:]
        all_lpt_like[i] = lpt_like.logLikelihoodComplex(s_hat, None)
        all_cola_like[i] = cola_like.logLikelihoodComplex(s_hat, None)

fig, axs = plt.subplots(1, 2, figsize=(15,4))
cmap = plt.cm.spring
markersize = 50

# # Remove first point
# all_lpt_like = all_lpt_like[1:]
# all_cola_like = all_cola_like[1:]
# all_mcmc = all_mcmc[1:]

# Normalise by the minimum value
lmin = all_lpt_like.min()
all_lpt_like -= lmin
all_cola_like -= lmin

sc0 = axs[0].scatter(all_lpt_like, all_cola_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
xlim = axs[0].get_xlim()
axs[0].plot(xlim, xlim, color='k')
axs[0].set_xlim(xlim)
cbar = fig.colorbar(sc0, ax=axs[0], location='right', aspect=40)
cbar.set_label("MCMC Step")

sc1 = axs[1].scatter(all_lpt_like, all_cola_like - all_lpt_like, c=all_mcmc, cmap=cmap, marker='.', s=markersize)
axs[1].axhline(0, color='k')
cbar = fig.colorbar(sc1, ax=axs[1], location='right', aspect=40)
cbar.set_label("MCMC Step")

for ax in axs:
    ax.set_xlabel(f'{ref_gravity} Log-Likelihood')
axs[0].set_ylabel('COLA Log-Likelihood')
axs[1].set_ylabel(f'COLA Log-Likelihood - {ref_gravity} Log-Likelihood')

# fig.align_xlabels()
# fig.align_ylabels()
# fig.tight_layout()
fig.savefig(f'{wdir}/likelihood_forwards.png', facecolor='white', bbox_inches='tight')