import numpy as np
import aquila_borg as borg
import jax.numpy as jnp
import configparser
import warnings
import aquila_borg as borg
#import symbolic_pofk.linear
import jax
from jax import lax
import jaxlib
from functools import partial
import ast
import numbers
import h5py
import re
import os

import borg_velocity.utils as utils
from borg_velocity.utils import myprint, compute_As, get_sigma_bulk
import borg_velocity.forwards as forwards
import borg_velocity.mock_maker as mock_maker
import borg_velocity.projection as projection
from borg_velocity.samplers import HMCBiasSampler, derive_prior, MVSliceBiasSampler, BlackJaxBiasSampler, TransformedBlackJaxBiasSampler
import borg_velocity.samplers as samplers

class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
    """
    HADES likelihood for distance-tracers
    """

    def __init__(self, fwd: borg.forward.BaseForwardModel, param_model: forwards.NullForward, fwd_vel: borg.forward.BaseForwardModel, ini_file: str) -> None:
        """
        Initialises the VelocityBORGLikelihood class

        Args:
            - fwd (borg.forward.BaseForwardModel): The forward model to be used in the likelihood.
            - param_model (forwards.NullForward): An empty forward model for storing model parameters.
            - ini_file (str): The name of the ini file containing the model and borg parameters. 
            ADD FW_VEL HERE
        """
        
        self.ini_file = ini_file
        
        myprint("Reading from configuration file: " + ini_file)
        config = configparser.ConfigParser()
        config.read(ini_file)
        
        # Grid parameters
        self.N = [int(config['system'][f'N{i}']) for i in range(3)]
        self.L = [float(config['system'][f'L{i}']) for i in range(3)]
        
        # Catalogues
        self.nsamp = int(config['run']['nsamp'])
        assert self.nsamp > 0, "Need at least one sample to run"
        self.Nt = [int(config[f'sample_{i}']['Nt']) for i in range(self.nsamp)]
        self.alpha = [float(config[f'sample_{i}']['alpha']) for i in range(self.nsamp)]
        self.lam = [float(config[f'sample_{i}']['lam']) for i in range(self.nsamp)]
        self.muA = [float(config[f'sample_{i}']['muA']) for i in range(self.nsamp)]
        self.frac_sig_rhMpc = [float(config[f'sample_{i}']['frac_sig_rhMpc']) for i in range(self.nsamp)]
        
        # What type of run we're doing
        self.run_type = config['run']['run_type']
        
        # Seed if creating mocks
        self.mock_seed = int(config['mock']['seed'])
        self.R_max = float(config['mock']['r_max'])
        
        # Model parameters
        self.ai = float(config['model']['ai'])
        self.af = float(config['model']['af'])
        self.sig_v = float(config['model']['sig_v'])
        if config['model']['R_lim'] == 'none':
            self.R_lim = self.L[0]/2
        else:
            self.R_lim = float(config['model']['R_lim'])
        self.Nint_points = int(config['model']['nint_points'])
        self.Nsig = float(config['model']['Nsig'])
        self.sig_v = float(config['model']['sig_v'])
        self.smooth_R = float(config['model']['smooth_R'])
        self.bias_epsilon = float(config['model']['bias_epsilon'])
        self.interp_order = int(config['model']['interp_order'])
        
        # Load bulk flow
        self.bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
        
        # For log-likelihood values
        self.bignum = float(config['mcmc']['bignum'])
        
        myprint(f" Init {self.N}, {self.L}")
        super().__init__(fwd, self.N, self.L)
        
        # Define the forward models
        self.fwd = fwd
        self.fwd_param = param_model
        self.fwd_vel = fwd_vel
        
        # Initialise model parameters
        self.model_params = {
            **{f'mua{i}':self.muA[i] for i in range(self.nsamp)},
            **{f'alpha{i}':self.alpha[i] for i in range(self.nsamp)},
            **{f'lam{i}':self.lam[i] for i in range(self.nsamp)},
            **{f'bulk_flow_{d}':self.bulk_flow[i] for i, d in enumerate(['x', 'y', 'z'])},
            'sig_v':self.sig_v
        }
        self.fwd_param.setModelParams(self.model_params)
        
        # Initialise cosmological parameters
        cpar = utils.get_cosmopar(self.ini_file)
        self.fwd.setCosmoParams(cpar)
        self.fwd_param.setCosmoParams(cpar)
        self.updateCosmology(cpar)
        myprint(f"Original cosmological parameters: {self.fwd.getCosmoParams()}")
        
        # Initialise derivative
        self.grad_like = jax.grad(self.dens2like, argnums=(0,1))
        
        # Find out what kind of run this is
        if borg.EMBEDDED:
            self.action = utils.get_action() 
        else:
            self.action = None
        
        
    def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
        """
        Initialise the likelihood internal variables and MarkovState variables.

        Args:
            - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
        """
        
        myprint("Init likelihood")

        # for i in range(self.nsamp):
        #     state.newArray1d(f"tracer_vr_{i}", self.Nt[i], True)
        #     if self.run_type != 'data':
        #         state.newArray1d(f"data_x_{i}", self.Nt[i], True)
        #         state.newArray1d(f"data_y_{i}", self.Nt[i], True)
        #         state.newArray1d(f"data_z_{i}", self.Nt[i], True)
        #         state.newArray1d(f"true_x_{i}", self.Nt[i], True)
        #         state.newArray1d(f"true_y_{i}", self.Nt[i], True)
        #         state.newArray1d(f"true_z_{i}", self.Nt[i], True)
        # self.data = [state[f"tracer_vr_{i}"] for i in range(self.nsamp)]
        state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)
        
        if self.run_type == 'data':
            self.loadObservedData(make_plot=False)
        elif borg.EMBEDDED:
            if self.action == 'INIT':
                pass  # Data will be loaded later
            elif self.action == 'RESUME':
                self.loadMockData(state)  # load data from mock_data.h5
            else:
                myprint(f'Unknown action: {self.action}') 
                raise NotImplementedError 
        
        
    def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None:
        """
        Update the meta parameters of the sampler (not sampled) from the MarkovState.

        Args:
            - state (borg.likelihood.MarkovState): The state object to be used in the likelihood.

        """
        cpar = state['cosmology']
        cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
        self.fwd.setCosmoParams(cpar)
        self.fwd_param.setCosmoParams(cpar)

        
    def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None:
        """
        Updates the forward model's cosmological parameters with the given values.

        Args:
            - cosmo (borg.cosmo.CosmologicalParameters): The cosmological parameters.

        """
        cpar = cosmo
        
        # Convert sigma8 to As
        cpar = compute_As(cpar)
        # cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
            # cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
        myprint(f"Updating cosmology Om = {cosmo.omega_m}, sig8 = {cosmo.sigma8}, As = {cosmo.A_s}")
        
        cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
        self.fwd.setCosmoParams(cpar)
        self.fwd_param.setCosmoParams(cpar)
        
        self.sigma_bulk = get_sigma_bulk(self.L[0], self.fwd.getCosmoParams())
        
        
    def generateMBData(self) -> None:
        """
        Generate points along the line of sight of each tracer to enable marginalisation
        over distance uncertainty. The central distance is given such that the observed
        redshift equals the cosmological redshift at this distance. The range is then
        +/- Nsig * sig, where 
        sig^2 = (sig_v/100)^2 + sig_r^2
        and sig_v is the velocity uncertainty in km/s
        """
        
        self.MB_pos = [None] * self.nsamp
        self.r_integration = [None] * self.nsamp
        self.RA = [None] * self.nsamp
        self.DEC = [None] * self.nsamp
        
        for i in range(self.nsamp):
            myprint(f"Making MB data for sample {i}")
            
            # Get angular coordinates of all points
            r_hat = projection.get_radial_vectors(self.coord_meas[i])
            
            # Get min and max distance to integrate over
            robs = self.cz_obs[i] / 100
            sigr = np.sqrt((self.sig_v / 100) ** 2 + (self.frac_sig_rhMpc[i] * robs)**2)
            rmin = robs - self.Nsig * sigr
            rmin = rmin.at[rmin <= 0].set(self.L[0] / self.N[0] / 100.)
            rmax = robs + self.Nsig * sigr
            rmax = rmax.at[rmax > self.R_lim].set(self.R_lim)
            
            # Compute coordinates of integration points
            self.r_integration[i] = np.linspace(rmin, rmax, self.Nint_points)
            cartesian_pos_MB = np.expand_dims(self.r_integration[i], axis=2) * r_hat
            self.r_integration[i] = self.r_integration[i].T
            self.MB_pos[i] = cartesian_pos_MB
            self.MB_pos[i] = jnp.transpose(self.MB_pos[i], (2, 1, 0))
        
        
    def generateMockData(self, s_hat: np.ndarray, state: borg.likelihood.MarkovState,) -> None:
        """
        Generates mock data by simulating the forward model with the given white noise,
        drawing distance tracers from the density field, computing their distance
        moduli and radial velocities, and adding Gaussian noise to the appropriate
        variables. Also calculates the initial negative log-likelihood of the data.

        Args:
             - s_hat (np.ndarray): The input (initial) density field.
             - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
             - make_plot (bool, default=True): Whether to make diagnostic plots for the mock data generation
        """
        
        if self.run_type == 'data':
            raise NotImplementedError
        elif self.run_type == 'velmass':
            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
                mock_maker.velmass_mock(self.ini_file, seed=self.mock_seed)
        elif self.run_type == 'mock':
            self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
                mock_maker.borg_mock(s_hat, state, self.fwd, self.fwd_vel, self.ini_file, seed=self.mock_seed)
        else:
            raise NotImplementedError
            
        # Save mock to file
        with h5py.File(f'tracer_data_{self.run_type}.h5', 'w') as h5file:
            for i in range(self.nsamp):
                sample_group = h5file.create_group(f'sample_{i}')
                sample_group.create_dataset('coord_true', data=self.coord_true[i])
                sample_group.create_dataset('coord_meas', data=self.coord_meas[i])
                sample_group.create_dataset('sig_mu', data=self.sig_mu[i])
                sample_group.create_dataset('vr_true', data=self.vr_true[i])
                sample_group.create_dataset('cz_obs', data=self.cz_obs[i])
            
        self.r_hMpc = [np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0)) for i in range(self.nsamp)]
        
        # Check that the measured coordinates all lie within the box
        for i in range(self.nsamp):
            if np.amax(self.r_hMpc[i]) > self.R_lim:
                raise ValueError('All tracers must have measured coordinates within R_lim')
        
        self.generateMBData()
        myprint('From mock')
        self.saved_s_hat = s_hat.copy()
        self.logLikelihoodComplex(s_hat, False)
        
    def loadMockData(self, state: borg.likelihood.MarkovState) -> None:
        
        myprint('Loading previously made mock data')
        
        self.coord_true = [None] * self.nsamp
        self.coord_meas = [None] * self.nsamp
        self.sig_mu = [None] * self.nsamp
        self.vr_true = [None] * self.nsamp
        self.cz_obs = [None] * self.nsamp
        self.r_hMpc = [None] * self.nsamp

        with h5py.File(f'tracer_data_{self.run_type}.h5', 'r') as f:
            for i in range(self.nsamp):
                self.coord_true[i] = jnp.array(f[f'sample_{i}/coord_true'][:])
                self.coord_meas[i] = jnp.array(f[f'sample_{i}/coord_meas'][:])
                self.sig_mu[i] = f[f'sample_{i}/sig_mu'][()]
                self.vr_true[i] = jnp.array(f[f'sample_{i}/vr_true'][:])
                self.cz_obs[i] = jnp.array(f[f'sample_{i}/cz_obs'][:])
                self.r_hMpc[i] = np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0))
                
        # Check that the measured coordinates all lie within the box
        for i in range(self.nsamp):
            if np.amax(self.r_hMpc[i]) > self.R_lim:
                raise ValueError('All tracers must have measured coordinates within R_lim')

        self.generateMBData()


    def dens2like(self, output_density: np.ndarray, output_velocity: np.ndarray):
        """
        Given stored distance tracer data, computes the negative log-likelihood of the data
        for this final density field.

        Args:
            output_density (np.ndarray): The z=0 density field.
        Return:
            lkl (float): The negative log-likelihood of the data.
        """
        
        lkl = 0
        
        sig_v = self.model_params['sig_v']
        
        # Compute velocity field
        self.bulk_flow = jnp.array([self.model_params['bulk_flow_x'],
                          self.model_params['bulk_flow_y'],
                          self.model_params['bulk_flow_z']])
        v = output_velocity + self.bulk_flow.reshape((3, 1, 1, 1))
        
        omega_m = self.fwd.getCosmoParams().omega_m
        
#         self.lkl_ind = [None] * self.nsamp
        
#         global temp_lkl_ind
        
        for i in range(self.nsamp):
            muA = self.model_params[f'mua{i}']
            alpha = self.model_params[f'alpha{i}']
            lam = self.model_params[f'lam{i}']
            lkl += vel2like(
                self.cz_obs[i], 
                v, 
                output_density, 
                self.MB_pos[i], 
                self.r_integration[i], 
                self.r_hMpc[i], 
                self.sig_mu[i], 
                sig_v, 
                omega_m, 
                muA, 
                alpha, 
                lam,
                self.L[0], 
                jnp.array(self.fwd.getOutputBoxModel().xmin), 
                self.interp_order, 
                self.bias_epsilon,
                self.R_max
            )
            
            # self.lkl_ind[i] = temp_lkl_ind.copy()
            
        # Add in bulk flow prior
        lkl += jnp.sum(0.5 * jnp.log(2 * np.pi) + jnp.log(self.sigma_bulk / jnp.sqrt(3)) + self.bulk_flow ** 2 / 2. / (self.sigma_bulk / jnp.sqrt(3)) ** 2)
            
        # lkl = jnp.clip(lkl, -self.bignum, self.bignum)
        lkl = lax.cond(
            jnp.isfinite(lkl),
            lambda _: lkl,         # If True (finite), return lkl
            lambda _: self.bignum, # If False (not finite), return self.bignum
            operand=None           # No additional operand needed here
        )
            
        return lkl
        
    def logLikelihoodComplex(self, s_hat: np.ndarray, gradientIsNext: bool, skip_density: bool=False, update_from_model: bool=True):
        """
        Calculates the negative log-likelihood of the data.

        Args:
            - s_hat (np.ndarray): The input white noise.
            - gradientIsNext (bool): If True, prepares the forward model for gradient calculations.
            - skip_density (bool, default=False): If True, do not repeat the s_hat -> density, velocity computation
                and use the stored result
            - update_from_model (bool, default=True): If True, update self.model_params with self.fwd_param.getModelParam

        Returns:
            The negative log-likelihood value.

        """

        N = self.fwd.getBoxModel().N[0]
        L = self.fwd.getOutputBoxModel().L[0]
        
        if update_from_model:
            for k in self.model_params.keys():
                self.model_params[k] = self.fwd_param.getModelParam('nullforward', k)
            self.updateCosmology(self.fwd.getCosmoParams())  # for sigma8 -> As
            
        # fname = f'{os.getcwd()}/s_hat.npy'
        # myprint(f'Saving s_hat field to {fname}')
        # np.save(fname, s_hat)
        # myprint('Done')

        if not skip_density:
            
            myprint('Getting density field now')
            
            # Run BORG density field
            output_density = np.zeros((N,N,N))
            self.fwd.forwardModel_v2(s_hat)
            self.fwd.getDensityFinal(output_density)
            
            myprint('Getting velocity field now')
        
            # Get velocity field
            output_velocity = self.fwd_vel.getVelocityField()
            
        else:
            output_density = self.delta.copy()
            output_velocity = self.vel.copy()

        L = self.dens2like(output_density, output_velocity)
        if isinstance(L, numbers.Number) or isinstance(L, jaxlib.xla_extension.ArrayImpl):
            myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")
            myprint(self.model_params)
            myprint(self.fwd.getCosmoParams())
                    
        self.delta = output_density.copy()
        self.vel = output_velocity.copy()

        return L
    
    
    def gradientLikelihoodComplex(self, s_hat: np.ndarray):
        """
        Calculates the adjoint negative log-likelihood of the data.

        Args:
            - s_hat (np.ndarray): The input density field.

        Returns:
            The adjoint negative log-likelihood gradient.

        """

        N = self.fwd.getBoxModel().N[0]
        L = self.fwd.getOutputBoxModel().L[0]

        # Run BORG density field
        output_density = np.zeros((N,N,N))
        self.fwd.forwardModel_v2(s_hat)
        self.fwd.getDensityFinal(output_density)
        
        # Get velocity field
        output_velocity = self.fwd_vel.getVelocityField()

        # getlike(dens, vel)
        mygradient, velgradient = self.grad_like(output_density, output_velocity)
        
        mygradient = np.array(mygradient, dtype=np.float64)
        vgradient = np.array(velgradient, dtype=np.float64)
        
        self.fwd_vel.computeAdjointModel(vgradient)
        self.fwd.adjointModel_v2(mygradient)
        mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128)
        self.fwd.getAdjointModel(mygrad_hat)
        self.fwd.clearAdjointGradient()
        
        return mygrad_hat


    def commitAuxiliaryFields(self, state: borg.likelihood.MarkovState) -> None:
        """
        Commits the final density field to the Markov state.
        Args:
            - state (borg.state.State): The state object containing the final density field.
        """
        self.updateCosmology(self.fwd.getCosmoParams())
        self.dens2like(self.delta, self.vel)
        state["BORG_final_density"][:] = self.delta
        
        
# @partial(jax.jit, static_argnames=['L_BOX', 'interp_order', 'bias_epsilon'])
def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA, alpha, lam, 
             L_BOX, X_MIN, interp_order, bias_epsilon, R_max):
    """
    Jitted part of dens2like
    """
    
    tracer_vel = projection.interp_field(v,
                                MB_pos,
                                L_BOX,
                                X_MIN,
                                interp_order,
                                use_jitted=True,
                                )
            
    tracer_vr = projection.project_radial(
                    tracer_vel,
                    MB_pos,
                    jnp.zeros(3,)
                    )
    
    # Remove the velocity term
    # tracer_vr = jnp.zeros(tracer_vr.shape)
    
    # Convert velocities to redshifts
    zco = utils.z_cos(r, omega_m)
    cz_pred = utils.speed_of_light * zco + (1 + zco) * tracer_vr
    delta_cz_sigv = (cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sig_v

    # Get p_r
    delta_mu = 5. * jnp.log10(r / jnp.expand_dims(r_hMpc,axis=1) * muA)

    # Get los biased density
    los_density = projection.interp_field(
                                MB_field,
                                MB_pos,
                                L_BOX,
                                X_MIN,
                                interp_order,
                                use_jitted=True,
                            )
    los_density = jax.nn.relu(1. + los_density)
    los_density = jnp.power(los_density + bias_epsilon, alpha)
    
    # Remove bias term
    # los_density = jnp.ones(los_density.shape)
    
    d2 = (delta_mu / sig_mu) ** 2
    best = jnp.amin(jnp.abs(d2), axis=1)

    # Multiply p_r by arbitrary number for numerical stability (cancels in p_r / p_r_norm)
    d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1)
    p_r = r ** 2 * jnp.exp(-0.5 * d2) * los_density * jnp.exp(- lam * r / R_max)
    p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1)
    
    # Integrate to get likelihood
    d2 = delta_cz_sigv**2
    scale = jnp.nanmin(d2, axis=1)
    d2 = d2 - jnp.expand_dims(scale, axis=1)
    exp_delta_cz = jnp.exp(-0.5 * d2)
    
    p_cz = jnp.trapezoid(exp_delta_cz * p_r / p_r_norm, r, axis=1)
    lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sig_v**2)
    lkl = - lkl_ind.sum()
    
    # global temp_lkl_ind
    # temp_lkl_ind = lkl_ind
        
    return lkl


@derive_prior
def transform_mua(x):
    a, b = model_params_prior['mua']
    return samplers.transform_uniform(x, a, b)
    
def inv_transform_mua(alpha):
    a, b = model_params_prior['mua']
    return samplers.inv_transform_uniform(alpha, a, b)

@derive_prior
def transform_alpha(x):
    a, b = model_params_prior['alpha']
    return samplers.transform_uniform(x, a, b)
    
def inv_transform_alpha(alpha):
    a, b = model_params_prior['alpha']
    return samplers.inv_transform_uniform(alpha, a, b)

@derive_prior
def transform_lam(x):
    a, b = model_params_prior['lam']
    return samplers.transform_uniform(x, a, b)
    
def inv_transform_lam(lam):
    a, b = model_params_prior['lam']
    return samplers.inv_transform_uniform(lam, a, b)

@derive_prior
def transform_sig_v(x):
    a, b = model_params_prior['sig_v']
    return samplers.transform_uniform(x, a, b)
    
def inv_transform_sig_v(alpha):
    a, b = model_params_prior['sig_v']
    return samplers.inv_transform_uniform(alpha, a, b)

@derive_prior
def transform_bulk_flow(x):
    a, b = model_params_prior['bulk_flow']
    return samplers.transform_uniform(x, a, b)
    
def inv_transform_bulk_flow(alpha):
    a, b = model_params_prior['bulk_flow']
    return samplers.inv_transform_uniform(alpha, a, b)
    

@borg.registerGravityBuilder
def build_gravity_model(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:
        - 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 config['model']['gravity'] in ['pm', 'cola']:
        forcesampling = int(config['model']['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 config['model']['gravity'] == 'linear':
        raise NotImplementedError(config['model']['gravity'])
    elif config['model']['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 config['model']['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 config['model']['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 config['model']['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(config['model']['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

_glob_model = None
_glob_cosmo = None
begin_model = None
begin_cosmo = None
model_params_prior = {}

def check_model_sampling(loop):
    return loop.getStepID() > begin_model

def check_cosmo_sampling(loop):
    return loop.getStepID() > begin_cosmo
            

@borg.registerSamplerBuilder
def build_sampler(
    state: borg.likelihood.MarkovState,
    info: borg.likelihood.LikelihoodInfo,
    loop: borg.samplers.MainLoop
):
    """
    Builds the sampler and returns it.
    Which parameters to sample are given in the ini file.
    We assume all parameters are NOT meant to be sampled, unless we find "XX_sampler_blocked = false" in the ini file

    Args:
        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
        - info (borg.likelihood.LikelihoodInfo): The likelihood information.
        ADD LOOP - DEFINE CONDITIONS AND DEPENDENCIES

    Returns:
        List of samplers to use.

    """
    global _glob_model, _glob_cosmo, begin_model, begin_cosmo, model_params_prior
    borg.print_msg(borg.Level.std, "Hello sampler, loop is {l}, step_id={s}", l=loop, s=loop.getStepID())
    
    myprint("Building sampler")
    
    myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
    config = configparser.ConfigParser()
    config.read(borg.getIniConfigurationFilename())
    end = '_sampler_blocked'
    to_sample = [k[:-len(end)] for (k, v) in config['block_loop'].items() if k[-len(end):] == end and v.lower() == 'false']
    myprint(f'Parameters to sample: {to_sample}')
    nsamp = int(config['run']['nsamp'])

    all_sampler = []

    # Cosmology sampler arguments
    prefix = ""
    params = []
    initial_values = {}
    prior = {}
    for p in ["omega_m", "sigma8"]:
        if p not in to_sample:
            continue
        if p in config['prior'].keys() and p in config['cosmology'].keys():
            myprint(f'Adding {p} sampler')
            params.append(f"cosmology.{p}")
            initial_values[f"cosmology.{p}"] = float(config['cosmology'][p])
            prior[f"cosmology.{p}"] = np.array(ast.literal_eval(config['prior'][p]))
        else:
            s = f'Could not find {p} prior and/or default, so will not sample'
            warnings.warn(s, stacklevel=2)
        # Remove for later to prevent duplication 
        to_sample.remove(p)
       
    begin_cosmo = int(config['mcmc']['warmup_cosmo'])
    
    if len(params) > 0:
        myprint('Adding cosmological parameter sampler')
        cosmo_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, chain, initial_values, prior)
        cosmo_sampler.setName("cosmo_sampler")
        _glob_cosmo = cosmo_sampler
        all_sampler.append(cosmo_sampler)
        loop.push(cosmo_sampler)
        loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler")
        loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop))

    # Model parameter sampler
    prefix = ""
    params = []
    initial_values = {}
    prior = {}
    transform_attributes = []
    inv_transform_attributes = []
    
    for p in to_sample:
        if p in config['prior'].keys():
            if p == 'sig_v':
                myprint(f'Adding {p} sampler')
                params.append(p)
                initial_values[p] = float(config['model'][p])
                if 'inf' in config['prior'][p]:
                    x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
                    prior[p] = np.array([xx if xx != 'inf' else np.inf for xx in x])
                else:
                    prior[p] = np.array(ast.literal_eval(config['prior'][p]))
                model_params_prior[p] = prior[p]
                transform_attributes.append(globals().get(f'transform_{p}'))
                inv_transform_attributes.append(globals().get(f'inv_transform_{p}'))                            
            elif p == 'bulk_flow':
                for i, d in enumerate(['_x', '_y', '_z']):
                    myprint(f'Adding {p}{d} sampler')
                    params.append(f'{p}{d}')
                    initial_values[f'{p}{d}'] = np.array(ast.literal_eval(config['model']['bulk_flow']))[i]
                    if 'inf' in config['prior'][p]:
                        x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
                        prior[f'{p}{d}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
                    else:
                        prior[f'{p}{d}'] = np.array(ast.literal_eval(config['prior'][p]))
                    transform_attributes.append(globals().get(f'transform_{p}'))
                    inv_transform_attributes.append(globals().get(f'inv_transform_{p}')) 
                model_params_prior[p] = prior[f'{p}_x']
            else:
                for i in range(nsamp):
                    myprint(f'Adding {p}{i} sampler')
                    params.append(f'{p}{i}')
                    initial_values[f'{p}{i}'] = float(config[f'sample_{i}'][p])
                    if 'inf' in config['prior'][p]:
                        x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))    
                        prior[f'{p}{i}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
                    else:
                        prior[f'{p}{i}'] = np.array(ast.literal_eval(config['prior'][p]))
                    transform_attributes.append(globals().get(f'transform_{p}'))
                    inv_transform_attributes.append(globals().get(f'inv_transform_{p}')) 
                model_params_prior[p] = prior[f'{p}{0}']
        else:
            s = f'Could not find {p} prior, so will not sample'
            warnings.warn(s, stacklevel=2)
            
    begin_model = int(config['mcmc']['warmup_model'])
   
    if len(params) > 0:
        myprint('Adding model parameter sampler')
        if config['sampling']['algorithm'].lower() == 'slice':
            model_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior)
        elif config['sampling']['algorithm'].lower() == 'hmc':
            model_sampler = HMCBiasSampler(
                "model_params",
                likelihood,
                params,
                transform_attributes=transform_attributes,
                inv_transform_attributes=inv_transform_attributes,
                prior = [p.prior for p in transform_attributes],
                Nsteps = int(config['sampling']['nsteps']),
                epsilon = float(config['sampling']['epsilon']),
                refresh = float(config['sampling']['refresh']),
            )
            with open('model_params.txt', 'w') as file:
                for item in params:
                    file.write(f"{item}\n")
        elif config['sampling']['algorithm'].lower() == 'mvslice':
            lam = [None] * len(params)
            for i, p in enumerate(params):
                if p.startswith('bulk_flow'):
                    lam[i] = float(config['sampling'][f'mvlam_bulk_flow'])
                elif p[-1].isdigit():
                    # Search for a non-digit before end of reversed string
                    match = re.search(r'\D(?=\d*$)', p[::-1])
                    lam[i] = float(config['sampling'][f'mvlam_{p[:match.start()]}'])
                else:
                    lam[i] = float(config['sampling'][f'mvlam_{p}'])
            model_sampler =  MVSliceBiasSampler(
                "model_params",
                likelihood,
                params,
                lam
            )
            with open('model_params.txt', 'w') as file:
                for item in params:
                    file.write(f"{item}\n")
        elif config['sampling']['algorithm'].lower() == 'blackjax':
            # state.newArray1d("inverse_mass_matrix", len(params), True)
            # state.newScalar('step_size', 0., True)
            model_sampler = BlackJaxBiasSampler(
                "model_params",
                likelihood,
                params,
                int(config['sampling']['rng_seed']),
                int(config['sampling']['warmup_nsteps']),
                float(config['sampling']['warmup_target_acceptance_rate']),
            )
            with open('model_params.txt', 'w') as file:
                for item in params:
                    file.write(f"{item}\n")
        elif config['sampling']['algorithm'].lower() == 'transformedblackjax':
            model_sampler = TransformedBlackJaxBiasSampler(
                "model_params",
                likelihood,
                params,
                transform_attributes=transform_attributes,
                inv_transform_attributes=inv_transform_attributes,
                prior = [p.prior for p in transform_attributes],
                rng_seed = int(config['sampling']['rng_seed']),
                warmup_nsteps = int(config['sampling']['warmup_nsteps']),
                warmup_target_acceptance_rate = float(config['sampling']['warmup_target_acceptance_rate']),
            )
            with open('model_params.txt', 'w') as file:
                for item in params:
                    file.write(f"{item}\n")
        else:
            raise NotImplementedError
        model_sampler.setName("model_sampler")
        _glob_model = model_sampler
        loop.push(model_sampler)
        all_sampler.append(model_sampler)
        loop.addToConditionGroup("warmup_model", "model_sampler")
        loop.addConditionToConditionGroup("warmup_model", partial(check_model_sampling, loop))
        myprint('Done')
        
    return []


@borg.registerLikelihoodBuilder
def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood:
    """
    Builds the likelihood object and returns it.

    Args:
        - state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
        - info (borg.likelihood.LikelihoodInfo): The likelihood information.

    Returns:
        borg.likelihood.BaseLikelihood: The likelihood object.

    """
    global likelihood, fwd_param
    myprint("Building likelihood")
    myprint(chain.getCosmoParams())
    boxm = chain.getBoxModel()
    likelihood = VelocityBORGLikelihood(chain, fwd_param, fwd_vel, borg.getIniConfigurationFilename())
    return likelihood

