import numpy as np
import jax.numpy as jnp
import configparser
import warnings
import aquila_borg as borg
import symbolic_pofk.linear
import jax
from functools import partial
import ast
import torch

import utils as utils
from utils import myprint, initial_pos
import forwards
import networks

class GaussianLikelihood(borg.likelihood.BaseLikelihood):
    """
    HADES Gaussian likelihood
    """

    def __init__(self, fwd: borg.forward.BaseForwardModel, param_model: forwards.NullForward, ini_file: str) -> None:
        """
        Initialises the GaussianLikelihood 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. 
        """
        
        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)]
        
        # 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'])
        
        # 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
        
        # Initialise model parameters
        model_params = {
            'logfR0':float(config['model']['logfr0']),
            'gauss_sigma':float(config['model']['gauss_sigma'])
        }
        self.fwd_param.setModelParams(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)
        
        
    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")
        state.newArray3d("mock", *self.fwd.getOutputBoxModel().N, False)
        state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)        
        
        
    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.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)
        
        
    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 == 'mock':
            output_density = np.zeros(self.fwd.getOutputBoxModel().N)
            self.fwd.forwardModel_v2(s_hat)
            self.fwd.getDensityFinal(output_density)
            state["BORG_final_density"][:] = output_density
            output_density += np.random.normal(size=self.fwd.getOutputBoxModel().N) * self.getModelParam('nullforward', 'gauss_sigma')
            state["mock"][:] = output_density
        else:
            raise NotImplementedError
        
    
    def dens2like(self, output_density: np.ndarray):
        """
        Given stored 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.
        """
        
        sigma = self.fwd_param.getModelParam('nullforward', 'gauss_sigma')
        lkl = (self.data - output_density)**2 / (2. * sigma ** 2) + 0.5 * jnp.log(2 * jnp.pi * sigma**2)
        lkl = jnp.sum(lkl)
        
        if not jnp.isfinite(lkl):
            lkl = self.bignum
            
        return lkl
        
    def logLikelihoodComplex(self, s_hat: np.ndarray, gradientIsNext: bool):
        """
        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.

        Returns:
            The negative log-likelihood value.

        """

        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)

        self.delta = output_density
        self.vel = output_velocity

        L = self.dens2like(output_density, output_velocity)
        myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")

        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)

        # getlike(dens, vel)
        mygradient = self.grad_like(output_density)
        mygradient = np.array(mygradient, dtype=np.float64)

        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)
        state["BORG_final_density"][:] = self.delta
    

@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
    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'])

    # Setup forward model
    chain = borg.forward.ChainForwardModel(box)
    chain.addModel(borg.forward.models.HermiticEnforcer(box))
    
    # Cosmological parameters
    if ini_file is None:
        cpar = utils.get_cosmopar(borg.getIniConfigurationFilename())
    else:
        cpar = utils.get_cosmopar(ini_file)
    chain.setCosmoParams(cpar)
  
    # 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'] == 'lpt':
        lpt = borg.forward.model_lib.M_LPT_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=1,
                      lightcone=False,
                      part_factor=1.01,))
    elif config['model']['gravity'] == '2lpt':
        lpt = borg.forward.model_lib.M_2LPT_CIC(
            box, 
            opts=dict(a_initial=af,
                      a_final=af,
                      do_rsd=False,
                      supersampling=1,
                      lightcone=False,
                      part_factor=1.01,))
    else:
        raise NotImplementedError(config['model']['gravity'])
        
    lpt.accumulateAdjoint(True)
    chain @= lpt
        
    if config['emulator']['use_emulator'].lower().strip() == 'true':
        myprint('Adding emulator to the chain')
        
        # Check device:
        cuda_avail = torch.cuda.is_available()
        device = torch.device('cuda' if cuda_avail else 'cpu')
        myprint(f'device = {device}')
        if cuda_avail:
            torch.cuda.empty_cache()

        # Load weights
        model_weights_path = config['emulator']['model_weights_path']
        myprint(f'Use model params from: {model_weights_path}')
        emu_weights = torch.load(model_weights_path, map_location=device)

        # Initialize model
        model = getattr(networks, config['emulator']['architecture'])
        model = model(int(config['emulator']['style_size']), int(config['emulator']['in_chan']), int(config['emulator']['out_chan']))
        model.load_state_dict(emu_weights['model'])
    
        use_float64 = config['emulator']['use_float64'].lower().strip() == 'true'
        if use_float64:
            model.double()
            dtype = torch.float64
        else:
            dtype = torch.float32
        model.to(device)

        # Extract omega as style param
        Om = cpar.omega_m
        Om = torch.tensor([Om],dtype=dtype) # style parameter

        # from emulator hacking:
        Om -= torch.tensor([0.3])
        Om *= torch.tensor([5.0])
        if cuda_avail:
            Om = Om.cuda()

        # Initial positions
        q = initial_pos(box.L[0],box.N[0])

        # Create module in BORG chain
        emu = forwards.Emulator(
            box, 
            lpt, 
            model, 
            Om, 
            config['emulator']['use_pad_and_NN'].strip().lower() == 'true', 
            use_float64, 
            q, 
            config['emulator']['requires_grad'].strip().lower() == 'true', 
            cuda_avail, 
            debug=False
        )
        chain.addModel(emu)
        
        # DIFFERENT TO LUDVIG - CHECK THIS
        chain @= emu
    
    # 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)

    return chain

_glob_model = None
_glob_cosmo = None
begin_model = None
begin_cosmo = None

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.

    Returns:
        List of samplers to use.

    """
    global _glob_model, _glob_cosmo, begin_model, begin_cosmo
    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 = {}
    for p in to_sample:
        if p in config['prior'].keys():
            myprint(f'Adding {p} sampler')
            params.append(p)
            initial_values[f'{p}'] = float(config[f'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]))
        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')
        model_sampler = borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior)
        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))
        
    print('Warmups:', begin_cosmo, begin_model)

    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, borg.getIniConfigurationFilename())
    return likelihood


"""
NOTES:
 - Currently if we update Om, then we don't seem to update the emulator
 - Few arguments missing from forwards.Emulator
"""
