import numpy as np
import aquila_borg as borg
from typing import List
import jax.numpy as jnp
import jax
import blackjax
from borg_velocity.utils import myprint
import matplotlib.pyplot as plt
import os


class HMCBiasSampler(borg.samplers.PyBaseSampler):
    def __init__(
        self,
        prefix: str,
        likelihood: borg.likelihood.BaseLikelihood,
        attributes: List[str],
        transform_attributes: List[object],
        inv_transform_attributes: List[object],
        prior: List[object],
        algo="DKD",
        refresh: float=0.1,
        Nsteps: int=50,
        epsilon: float=0.001,
    ):
        super().__init__()
        self.likelihood = likelihood
        self.attributes = attributes
        self.prefix = prefix
        self.transform_attributes = transform_attributes
        self.inv_transform_attributes = inv_transform_attributes
        self.prior_attributes = prior

        self.refresh = refresh
        self.p_old = None
        self.Nsteps = Nsteps
        self.epsilon = epsilon

        # Drift-Kick-Drift
        if algo == "DKD":
            self.coefs = [[0, 0.5], [1, 0.5]]
        # Kick-Drift-Kick
        elif algo == "KDK":
            self.coefs = [[0.5, 1], [0.5, 0]]
        else:
            raise ValueError(f"unknown algo  {algo=}")

    def initialize(self, state: borg.likelihood.MarkovState):
        self.restore(state)
        for i in range(len(self.attributes)):
            attr = self.likelihood.fwd_param.getModelParam(
                'mod_null',
                self.attributes[i]
            )
            self.y_real[i] = attr
            if isinstance(attr, float):
                attr = jnp.array(attr)
            if self.inv_transform_attributes[i]:
                self.y[i] = self.inv_transform_attributes[i](attr).item()
            else:
                self.y[i] = attr.item()

        myprint(f"initial {self.y=}")

    def _update_attrs(self, x):
        for i in range(len(self.attributes)):
            xnew = x[i]
            if self.transform_attributes[i]:
                xnew = self.transform_attributes[i](xnew)
            self.likelihood.model_params[self.attributes[i]] = xnew
        

    def _likelihood(self, x_hat, x, skip_density=False):
        self._update_attrs(x)
        return self.likelihood.logLikelihoodComplex(
            x_hat, False, skip_density=skip_density, update_from_model=False,
        )

    def _grad_likelihood(self, x_hat, x, skip_density=False):
        grad_fn = jax.grad(self._likelihood, argnums=1)
        grad_x = grad_fn(x_hat, x, skip_density=skip_density)
        return grad_x

    def _prior(self, x):
        prior_value = 0.0
        for i in range(len(self.attributes)):
            if self.prior_attributes[i]:
                prior_value += self.prior_attributes[i](x[i])
        return prior_value

    def _grad_prior(self, x):
        grad_x = jax.grad(self._prior, argnums=0)(x)
        return grad_x


    def sample(self, state: borg.likelihood.MarkovState):
        myprint(f"Sampling attributes {self.attributes}")
        x_hat = state["s_hat_field"]

        mass = 1.0
        inv_mass = 1 / mass

        params = jnp.array(self.y)

        # Initialise the PRNG key
        key = jax.random.PRNGKey(np.random.randint(0, 1e6))
        
        # Split the key for independent random number generation
        key, subkey = jax.random.split(key)
        p_new = jax.random.normal(key, shape=params.shape) * jnp.sqrt(mass)
        
        if self.p_old is not None:
            key, subkey = jax.random.split(key)  # Split the key again
            p = self.p_old * self.refresh + jnp.sqrt(1 - self.refresh**2) * p_new
        else:
            p = p_new

        Lstart = (
            self._likelihood(x_hat, params, skip_density=False)
            + self._prior(params)
            + 0.5 * jnp.sum(p**2) * inv_mass
        )

        new_gradient = True
        myprint('running model parameter hmc ...')
        for i in range(self.Nsteps):
            for c in self.coefs:
                if c[0] != 0:
                    if new_gradient:
                        gradient = self._grad_likelihood(
                            x_hat, params, skip_density=True
                        ) + self._grad_prior(params)
                        new_gradient = False
                    p -= self.epsilon * c[0] * gradient
                if c[1] != 0:
                    params += self.epsilon * c[1] * p * inv_mass
                    new_gradient = True

        Lend = self._likelihood(x_hat, params, skip_density=True)

        deltaH = (
            Lend - Lstart + self._prior(params) + 0.5 * jnp.sum(p**2) * inv_mass
        )
        myprint(f"# deltaH={deltaH}")
        key, subkey = jax.random.split(key)
        if jnp.log(jax.random.uniform(subkey)) <= -deltaH:
            # accept
            myprint(f"# accepting...")
            self.y[:] = params
            for i in range(len(self.attributes)):
                xnew = params[i]
                if self.transform_attributes[i]:
                    xnew = self.transform_attributes[i](xnew)
                self.y_real[i] = xnew
                
            to_set = {k:v for k, v in zip(self.attributes, self.y_real)}
            self.likelihood.fwd_param.setModelParams(to_set)
            
        else:
            # reject
            myprint(f"# rejecting...")
        
        self._update_attrs(params)
        self.p_old = p

    def restore(self, state: borg.likelihood.MarkovState):
        
        # Define attribute names
        attrname_real = self.prefix
        attrname_var = f"{self.prefix}_hmc_var"

        # Initialize attributes in the state
        state.newArray1d(attrname_real, len(self.attributes), True)
        self.y_real = state[attrname_real]

        state.newArray1d(attrname_var, len(self.attributes), True)
        self.y = state[attrname_var]

        def _loaded():
            myprint(f"Reinjecting parameters for bias sampler: {self.attributes}")
            # Ensure `self.y` is a JAX array
            self._update_attrs(jnp.array(self.y))

        # Subscribe to updates
        state.subscribeLoaded(attrname_var, _loaded)

        
def derive_prior(f):
    def _prior(x):
        alpha = f(x)
        grad_fn = jax.grad(f)
        grad = grad_fn(x)
        return -jnp.log(jnp.abs(grad))
    f.prior = _prior
    return f


def transform_uniform(x, a, b):
    """
    Take a x value, from ]-inf,inf[ and map onto alpha in ]a,b[
    """
    return jax.nn.sigmoid(x) * (b - a) + a


def inv_transform_uniform(alpha, a, b):
    """
    Take an alpha value, from ]a,b[ and map onto x in ]-inf,inf[
    """
    x = (alpha - a) / (b - a)
    return - jnp.log(1/x - 1)


class MVSliceBiasSampler(borg.samplers.PyBaseSampler):
    def __init__(
        self,
        prefix: str,
        likelihood: borg.likelihood.Likelihood3d,
        attributes: List[str],
        lam: List[float],
    ):
        super().__init__()
        self.likelihood = likelihood
        self.attributes = attributes
        self.prefix = prefix
        self.lam = np.array(lam)

    def initialize(self, state: borg.likelihood.MarkovState):
        self.restore(state)
        for i in range(len(self.attributes)):
            self.y[i] = self.likelihood.fwd_param.getModelParam(
                'mod_null',
                self.attributes[i]
            )
        self.mvs_state[:] = 0  # np.random.uniform(size=len(self.attributes))

    def _update_attrs(self, x):
        for i in range(len(self.attributes)):
            self.likelihood.model_params[self.attributes[i]] = x[i]
        self.likelihood.fwd_param.setModelParams(self.likelihood.model_params)
            
    def sample(self, state: borg.likelihood.MarkovState):
        myprint(f"Sampling attributes {self.attributes}")
        x_hat = state["s_hat_field"]

        def _callback(x):
            self._update_attrs(x)
            return -self.likelihood.logLikelihoodComplex(x_hat, False)

        self.y[:] = borg.samplers.mpi_mv_slice_sampler(
            state, _callback, self.y, self.mvs_state, self.lam, mpi=None
        )
        self._update_attrs(self.y)

    def restore(self, state: borg.likelihood.MarkovState):
        attrname = f"{self.prefix}attributes"
        state.newArray1d(attrname, len(self.attributes), True)
        self.y = state[attrname]

        def _loaded():
            myprint(f"Reinjecting parameters for bias sampler: {self.attributes}")
            self._update_attrs(self.y)

        state.subscribeLoaded(attrname, _loaded)

        mvsname = f"{self.prefix}mvs_state"
        state.newArray1d(mvsname, len(self.attributes), True)
        self.mvs_state = state[mvsname]
        
        
class BlackJaxBiasSampler(borg.samplers.PyBaseSampler):
    def __init__(
        self,
        prefix: str,
        likelihood: borg.likelihood.Likelihood3d,
        attributes: List[str],
        rng_seed: int,
        warmup_nsteps: int,
        warmup_target_acceptance_rate: float
    ):
        super().__init__()
        self.likelihood = likelihood
        self.attributes = attributes
        self.prefix = prefix
        self.rng_key = jax.random.key(rng_seed)
        self.warmup_nsteps = warmup_nsteps
        self.warmup_target_acceptance_rate = warmup_target_acceptance_rate
        self.parameters = None

    def initialize(self, state: borg.likelihood.MarkovState):
        print(dir(state))
        self.restore(state)
        for i in range(len(self.attributes)):
            self.y[i] = self.likelihood.fwd_param.getModelParam(
                'mod_null',
                self.attributes[i]
            )
        self.mvs_state[:] = 0  # np.random.uniform(size=len(self.attributes))

    def _update_attrs(self, x):
        for i in range(len(self.attributes)):
            self.likelihood.model_params[self.attributes[i]] = x[i]
        # self.likelihood.fwd_param.setModelParams(self.likelihood.model_params)
        
    def _likelihood(self, x_hat, x, skip_density=False):
        self._update_attrs(x)
        return self.likelihood.logLikelihoodComplex(
            x_hat, False, skip_density=skip_density, update_from_model=False,
        )
 
    def sample(self, state: borg.likelihood.MarkovState):
        myprint(f"Sampling attributes {self.attributes}")
        
        # Update the sampling hyperparameters if possible
        if (self.parameters is None) and (state["step_size"] > 0) and all(state["inverse_mass_matrix"] > 0):
            myprint('Using loaded hyperparameters for model parameter sampling')
            myprint(f'Step size: {state["step_size"]}')
            myprint(f'Inverse mass matrix: {state["inverse_mass_matrix"]}')
            self.parameters = {"step_size": state["step_size"], "inverse_mass_matrix":state["inverse_mass_matrix"]}
        
        x_hat = state["s_hat_field"]
        
        # Run forward model to compute correct density field
        _ = self.likelihood.logLikelihoodComplex(x_hat, False, skip_density=False)
        
        @jax.custom_vjp
        def logdensity_fn(x):
            return - self._likelihood(x_hat, x, skip_density=True)

        def logdensity_fn_fwd(x):
            y = logdensity_fn(x)
            return y, x  # Save `x` for the backward pass

        def logdensity_fn_bwd(x, cotangent_y):
            grad_fn = jax.grad(self._likelihood, argnums=1)
            grad_x = - grad_fn(x_hat, x, skip_density=True)
            cotangent_x = grad_x * cotangent_y
            return (cotangent_x,)

        logdensity_fn.defvjp(logdensity_fn_fwd, logdensity_fn_bwd)
        
        if self.parameters is None:
            get_info = blackjax.adaptation.base.get_filter_adapt_info_fn(
                state_keys={'position', 'logdensity'},
                info_keys={'acceptance_rate'},
                adapt_state_keys={'step_size', 'inverse_mass_matrix'}
            )
            logdensity_fn_jitted = jax.jit(logdensity_fn)
            warmup = blackjax.window_adaptation(
                blackjax.nuts, 
                logdensity_fn_jitted, 
                target_acceptance_rate=self.warmup_target_acceptance_rate,
                progress_bar=True, 
                adaptation_info_fn=get_info
            )
            rng_key, warmup_key, sample_key = jax.random.split(self.rng_key, 3)
            (bj_state, self.parameters), info = warmup.run(warmup_key, self.y, num_steps=self.warmup_nsteps)
            x = info.state.position
            
            # Make and save summary figures
            dirname = os.getcwd()
            trace_plot(self.attributes, info.state.position, f'{dirname}/tuning_params.png', ncol=4)
            trace_plot(self.attributes, info.adaptation_state.inverse_mass_matrix, f'{dirname}/tuning_inverse_mass_matrix.png', ncol=4)
            trace_plot(['stepsize'], info.adaptation_state.step_size, f'{dirname}/tuning_step_size.png', ncol=4)

            self.y[:] = bj_state.position
            
            # Save results to the borg state object
            state["inverse_mass_matrix"][:] = self.parameters['inverse_mass_matrix']
            state["step_size"] = float(self.parameters['step_size'])

        nuts = blackjax.nuts(logdensity_fn, **self.parameters)
        bias_state = nuts.init(self.y)
        nuts_step = nuts.step

        self.rng_key, sample_key = jax.random.split(self.rng_key)
        bias_state, _ = nuts_step(sample_key, bias_state)
        
        self.y[:] = bias_state.position
        self._update_attrs(self.y)
        
        to_set = {k:v for k, v in zip(self.attributes, self.y)}
        self.likelihood.fwd_param.setModelParams(to_set)
        

    def restore(self, state: borg.likelihood.MarkovState):
        
        state.newArray1d("inverse_mass_matrix", len(self.attributes), True)
        state.newScalar('step_size', 0., True)
        
        attrname = f"{self.prefix}attributes"
        state.newArray1d(attrname, len(self.attributes), True)
        self.y = state[attrname]

        def _loaded():
            myprint(f"Reinjecting parameters for bias sampler: {self.attributes}")
            self._update_attrs(self.y)

        state.subscribeLoaded(attrname, _loaded)

        mvsname = f"{self.prefix}mvs_state"
        state.newArray1d(mvsname, len(self.attributes), True)
        self.mvs_state = state[mvsname]
        
        
class TransformedBlackJaxBiasSampler(borg.samplers.PyBaseSampler):
    def __init__(
        self,
        prefix: str,
        likelihood: borg.likelihood.BaseLikelihood,
        attributes: List[str],
        transform_attributes: List[object],
        inv_transform_attributes: List[object],
        prior: List[object],
        rng_seed: int,
        warmup_nsteps: int,
        warmup_target_acceptance_rate: float
    ):
        super().__init__()
        self.likelihood = likelihood
        self.attributes = attributes
        self.prefix = prefix
        self.transform_attributes = transform_attributes
        self.inv_transform_attributes = inv_transform_attributes
        self.prior_attributes = prior
        self.rng_key = jax.random.key(rng_seed)
        self.warmup_nsteps = warmup_nsteps
        self.warmup_target_acceptance_rate = warmup_target_acceptance_rate
        self.parameters = None
       

    def initialize(self, state: borg.likelihood.MarkovState):
        self.restore(state)
        for i in range(len(self.attributes)):
            attr = self.likelihood.fwd_param.getModelParam(
                'mod_null',
                self.attributes[i]
            )
            self.y_real[i] = attr
            if isinstance(attr, float):
                attr = jnp.array(attr)
            if self.inv_transform_attributes[i]:
                self.y[i] = self.inv_transform_attributes[i](attr).item()
            else:
                self.y[i] = attr.item()

        myprint(f"initial {self.y=}")

    def _update_attrs(self, x):
        for i in range(len(self.attributes)):
            xnew = x[i]
            if self.transform_attributes[i]:
                xnew = self.transform_attributes[i](xnew)
            self.likelihood.model_params[self.attributes[i]] = xnew
        

    def _likelihood(self, x_hat, x, skip_density=False):
        self._update_attrs(x)
        return self.likelihood.logLikelihoodComplex(
            x_hat, False, skip_density=skip_density, update_from_model=False,
        )

    def _grad_likelihood(self, x_hat, x, skip_density=False):
        grad_fn = jax.grad(self._likelihood, argnums=1)
        grad_x = grad_fn(x_hat, x, skip_density=skip_density)
        return grad_x

    def _prior(self, x):
        prior_value = 0.0
        for i in range(len(self.attributes)):
            if self.prior_attributes[i]:
                prior_value += self.prior_attributes[i](x[i])
        return prior_value

    def _grad_prior(self, x):
        grad_x = jax.grad(self._prior, argnums=0)(x)
        return grad_x
    
    
    def sample(self, state: borg.likelihood.MarkovState):
        myprint(f"Sampling attributes {self.attributes}")
        
        # Update the sampling hyperparameters if possible
        if (self.parameters is None) and (state["step_size"] > 0) and all(state["inverse_mass_matrix"] > 0):
            myprint('Using loaded hyperparameters for model parameter sampling')
            myprint(f'Step size: {state["step_size"]}')
            myprint(f'Inverse mass matrix: {state["inverse_mass_matrix"]}')
            self.parameters = {"step_size": state["step_size"], "inverse_mass_matrix":state["inverse_mass_matrix"]}
        
        x_hat = state["s_hat_field"]
        
        # Run forward model to compute correct density field
        _ = self.likelihood.logLikelihoodComplex(x_hat, False, skip_density=False)
        
        @jax.custom_vjp
        def logdensity_fn(x):
            self._update_attrs(x)
            p = self._likelihood(x_hat, x, skip_density=True) + self._prior(x)
            return - p 

        def logdensity_fn_fwd(x):
            y = logdensity_fn(x)
            return y, x  # Save `x` for the backward pass

        def logdensity_fn_bwd(x, cotangent_y):
            grad_x = - (self._grad_likelihood(x_hat, x, skip_density=True) + self._grad_prior(x))
            cotangent_x = grad_x * cotangent_y
            return (cotangent_x,)

        logdensity_fn.defvjp(logdensity_fn_fwd, logdensity_fn_bwd)
        
        if self.parameters is None:
            get_info = blackjax.adaptation.base.get_filter_adapt_info_fn(
                state_keys={'position', 'logdensity'},
                info_keys={'acceptance_rate'},
                adapt_state_keys={'step_size', 'inverse_mass_matrix'}
            )
            logdensity_fn_jitted = jax.jit(logdensity_fn)
            warmup = blackjax.window_adaptation(
                blackjax.nuts, 
                logdensity_fn_jitted, 
                target_acceptance_rate=self.warmup_target_acceptance_rate,
                progress_bar=True, 
                adaptation_info_fn=get_info
            )
            rng_key, warmup_key, sample_key = jax.random.split(self.rng_key, 3)
            (bj_state, self.parameters), info = warmup.run(warmup_key, self.y, num_steps=self.warmup_nsteps)
            # x = info.state.position
            self.y[:] = bj_state.position
            
            # Save results to the borg state object
            state["inverse_mass_matrix"][:] = self.parameters['inverse_mass_matrix']
            state["step_size"] = float(self.parameters['step_size'])
            
            # Make and save summary figures
            dirname = os.getcwd()
            trace_plot(self.attributes, info.state.position, f'{dirname}/tuning_params.png', ncol=4)
            trace_plot(self.attributes, info.adaptation_state.inverse_mass_matrix, f'{dirname}/tuning_inverse_mass_matrix.png', ncol=4)
            trace_plot(['stepsize'], info.adaptation_state.step_size, f'{dirname}/tuning_step_size.png', ncol=4)


        nuts = blackjax.nuts(logdensity_fn, **self.parameters)
        bias_state = nuts.init(self.y)
        nuts_step = nuts.step

        self.rng_key, sample_key = jax.random.split(self.rng_key)
        bias_state, _ = nuts_step(sample_key, bias_state)
        
        self.y[:] = bias_state.position
        self._update_attrs(self.y)

        for i in range(len(self.attributes)):
            self.y_real[i] = self.likelihood.model_params[self.attributes[i]]
        
        to_set = {k:v for k, v in zip(self.attributes, self.y_real)}
        self.likelihood.fwd_param.setModelParams(to_set)
        

    def restore(self, state: borg.likelihood.MarkovState):
        
        state.newArray1d("inverse_mass_matrix", len(self.attributes), True)
        state.newScalar('step_size', 0., True)
        
        # Define attribute names
        attrname_real = self.prefix
        attrname_var = f"{self.prefix}_hmc_var"

        # Initialize attributes in the state
        state.newArray1d(attrname_real, len(self.attributes), True)
        self.y_real = state[attrname_real]

        state.newArray1d(attrname_var, len(self.attributes), True)
        self.y = state[attrname_var]

        def _loaded():
            myprint(f"Reinjecting parameters for bias sampler: {self.attributes}")
            # Ensure `self.y` is a JAX array
            self._update_attrs(jnp.array(self.y))

        # Subscribe to updates
        state.subscribeLoaded(attrname_var, _loaded)



def trace_plot(names, samples, savename, ncol=4):
    
    if len(names) == 1:
        
        fig, ax = plt.subplots(1, 1, figsize=(4,3))
        ax.plot(samples[:])
        ax.set_title(names[0])
        ax.set_xlabel('Step Number')
        
    else:
   
        nrow = int(np.ceil(len(names) / ncol))

        fig, axs = plt.subplots(nrow, ncol, figsize=(ncol*4, nrow*3))
        axs = np.atleast_2d(axs)
        axs = axs.flatten()

        for i in range(len(names)):
            axs[i].plot(samples[:,i])
            axs[i].set_title(names[i])
            axs[i].set_xlabel('Step Number')
        for i in range(len(names), len(axs)):
            axs[i].remove()
            
    fig.tight_layout()
    myprint(f'Saving tuning diagnosis figure: {savename}')
    fig.savefig(savename, facecolor='white', bbox_inches='tight')
    fig.clf()
    plt.close(fig)
    
    return
