# @auto-fold regex /^\s*if/ /^\s*else/ /^\s*def/
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__version__ = '0.5'
__all__ = ['PTSampler']

import numpy as np
import emcee
from emcee.utils import deprecation_warning

from tqdm import tqdm

from typing import Dict, List, Optional, Union

from .reddwrappers import PTWrapper

temp_table = np.array([25.2741, 7., 4.47502, 3.5236, 3.0232,
                      2.71225, 2.49879, 2.34226, 2.22198, 2.12628,
                      2.04807, 1.98276, 1.92728, 1.87946, 1.83774,
                      1.80096, 1.76826, 1.73895, 1.7125, 1.68849,
                      1.66657, 1.64647, 1.62795, 1.61083, 1.59494,
                      1.58014, 1.56632, 1.55338, 1.54123, 1.5298,
                      1.51901, 1.50881, 1.49916, 1.49, 1.4813,
                      1.47302, 1.46512, 1.45759, 1.45039, 1.4435,
                      1.4369, 1.43056, 1.42448, 1.41864, 1.41302,
                      1.40761, 1.40239, 1.39736, 1.3925, 1.38781,
                      1.38327, 1.37888, 1.37463, 1.37051, 1.36652,
                      1.36265, 1.35889, 1.35524, 1.3517, 1.34825,
                      1.3449, 1.34164, 1.33847, 1.33538, 1.33236,
                      1.32943, 1.32656, 1.32377, 1.32104, 1.31838,
                      1.31578, 1.31325, 1.31076, 1.30834, 1.30596,
                      1.30364, 1.30137, 1.29915, 1.29697, 1.29484,
                      1.29275, 1.29071, 1.2887, 1.28673, 1.2848,
                      1.28291, 1.28106, 1.27923, 1.27745, 1.27569,
                      1.27397, 1.27227, 1.27061, 1.26898, 1.26737,
                      1.26579, 1.26424, 1.26271, 1.26121,
                      1.25973])

dmax = temp_table.shape[0]


def set_temp_ladder(ntemps, ndims, temp_table=temp_table):
    # returns betas
    dmax = temp_table.shape[0]

    tstep = temp_table[ndims-1]
    if ndims > dmax:
        # An approximation to the temperature step at large dimension
        tstep = 1.0 + 2.0*np.sqrt(np.log(4.0))/np.sqrt(ndims)

    return np.exp(np.linspace(0, -(ntemps-1)*np.log(tstep), ntemps))


class PTSampler(object):
    def __init__(self, nwalkers, ndim, log_likelihood_fn, log_prior_fn,
                 pool=None, moves=None, backend=None, vectorize=False,
                 blobs_dtype=None, parameter_names=None, logl_args=None,
                 logl_kwargs=None, logp_args=None, logp_kwargs=None,
                 ntemps=None, betas=None, adaptative=True, a=None,
                 postargs=None, threads=None, live_dangerously=None,
                 runtime_sortingfn=None):
        # Default arguments
        logl_args = logl_args or []
        logp_args = logp_args or []
        logl_kwargs = logl_kwargs or {}
        logp_kwargs = logp_kwargs or {}

        ######################################################
        # Warn about deprecated arguments
        deprecated_args = [postargs, threads, runtime_sortingfn, live_dangerously]
        deprecated_args_str = ['postargs', 'threads', 'runtime_sortingfn', 'live_dangerously']

        if a is not None:
            deprecation_warning("The 'a' argument is deprecated, use 'moves' instead")
        for i in range(len(deprecated_args)):
            if deprecated_args[i] is not None:
                deprecation_warning(f"The '{deprecated_args_str[i]}' argument is deprecated")
        ######################################################
        # Parse the move schedule

        # Beta ladder initialization
        self.ntemps = ntemps or 5
        self.betas = betas or set_temp_ladder(self.ntemps, ndim)
        #####################################################
        # Initialize instance variables
        self.nwalkers = nwalkers
        self.ndim = ndim
        self.lp_ = log_prior_fn
        self.ll_ = log_likelihood_fn

        self.lp_args_ = logp_args
        self.ll_args_ = logl_args

        self.n_swap_accept = np.zeros(ntemps-1)

        self.pool = pool
        self.adaptative = adaptative
        self.config_adaptation_lag = 10000
        self.config_adaptation_time = 100
        self.nglobal = 0

        # Default values for lists
        default_list = lambda lst, size, default: [default for _ in range(size)] if lst is None else lst
        self.vectorize = [False for _ in range(self.ntemps)] if vectorize is False else vectorize
        self.moves = default_list(moves, self.ntemps, None)
        self.blobs_dtype = default_list(blobs_dtype, self.ntemps, None)
        self.backend = default_list(backend, self.ntemps, None)

        self.ratios = None

        ## BACKEND
        #self.backend = Backend() if backend is None else backend
        # Deal with re-used backends
        # Check the backend shape
        #########
        # Probability function wrappers
        self.my_probs_fn = [PTWrapper(self.ll_, self.lp_, b, loglargs=self.ll_args_,
                                  logpargs=self.lp_args_, loglkwargs={},
                                  logpkwargs={}) for b in self.betas]

        self.sampler = [emcee.EnsembleSampler(self.nwalkers, self.ndim,
                        self.my_probs_fn[t], pool=self.pool,
                        moves=self.moves[t], backend=self.backend[t],
                        vectorize=self.vectorize[t], blobs_dtype=self.blobs_dtype[t],
                        ) for t in range(self.ntemps)]


    def __str__(self):
        return 'My sampler, ntemps = %i' % self.ntemps


    def __getitem__(self, n):
        return self.sampler[n]


    def __setitem__(self, n, thing):
        self.sampler[n] = thing


    def sample(
        self,
        initial_state,
        iterations=1,
        tune=False,
        skip_initial_state_check=False,
        thin_by=1,
        thin=None,
        store=True,
        progress=False,
        progress_kwargs=None):

        self.samp = initial_state
        for t in range(self.ntemps):
            for self.samp[t] in self[t].sample(self.samp[t],
                                               iterations=iterations,
                                               tune=tune,
                                               skip_initial_state_check=skip_initial_state_check,
                                               thin_by=thin_by,
                                               thin=thin,
                                               store=store,
                                               progress=progress,
                                               progress_kwargs=progress_kwargs):
                pass


        self.temp_swaps_()
        if self.adaptative:
            dbetas = self.ladder_adjustment()
            self.betas += dbetas

            # Mutate my_probs_fn
            for t in range(self.ntemps-1, 0, -1):
                self.my_probs_fn[t].beta = self.betas[t]
                self.samp[t].log_prob += self.samp[t].blobs * dbetas[t]

        yield self.samp


    def temp_swaps_(self):
        for t in range(self.ntemps-1, 0, -1):
            dbeta = self.betas[t-1] - self.betas[t]

            #iperm = np.ones(self.nwalkers, int)
            #i1perm = np.ones(self.nwalkers, int)

            ll1 = self.samp[t].blobs  # hot
            ll2 = self.samp[t-1].blobs  # cold

            raccept = np.log(np.random.uniform(0, 1, self.nwalkers))
            paccept = dbeta * (ll1 - ll2)

            asel = paccept > raccept

            self.n_swap_accept[t-1] = np.sum(asel)

            self.samp[t].coords[asel], self.samp[t-1].coords[asel] = self.samp[t-1].coords[asel], self.samp[t].coords[asel]
            self.samp[t].log_prob[asel], self.samp[t-1].log_prob[asel] = self.samp[t-1].log_prob[asel] - dbeta*ll2[asel], self.samp[t].log_prob[asel] + dbeta*ll1[asel]
            self.samp[t].blobs[asel], self.samp[t-1].blobs[asel] = ll2[asel], ll1[asel]

        self.ratios = self.n_swap_accept / self.nwalkers


    def ladder_adjustment(self):
        """
        Execute temperature adjustment according to dynamics outlined in
        `arXiv:1501.05823 <http://arxiv.org/abs/1501.05823>`_.
        """

        betas = self.betas.copy()
        time = self[0].iteration

        # Modulate temperature adjustments with a hyperbolic decay.
        decay = self.config_adaptation_lag / (time + self.config_adaptation_lag)
        kappa = decay / self.config_adaptation_time

        # Construct temperature adjustments.
        dSs = kappa * (self.ratios[:-1] - self.ratios[1:])

        # Compute new ladder (hottest and coldest chains don't move).
        deltaTs = np.diff(1 / betas[:-1])
        deltaTs *= np.exp(dSs)
        betas[1:-1] = 1 / (np.cumsum(deltaTs) + 1 / betas[0])

        return betas - self.betas
        # Mutate Ladder
        #self.betas = betas
        # Mutate my_probs_fn
        #for t in range(self.ntemps):
        #    self.my_probs_fn[t].beta = self.betas[t]


    def thermodynamic_integration(self):
        sampler_dict = {'flat':False, 'discard':discard}
        logls0 = self.get_func('get_blobs', kwargs=sampler_dict)
        logls = self.get_mean_logls(logls0, coef=coef)

        if self.betas[-1] != 0:
            betas1 = np.concatenate((self.betas, [0]))
            betas2 = np.concatenate((self.betas[::2], [0]))

            logls1 = np.concatenate((logls, [logls[-1]]))
            logls2 = np.concatenate((logls[::2], [logls[-1]]))

        else:
            betas1 = self.betas
            betas2 = np.concatenate((self.betas[:-1:2], [0]))

            logls1 = logls
            logls2 = np.concatenate((logls1[:-1:2], [logls1[-1]]))

        logZ1 = -np.trapz(logls1, betas1)
        logZ2 = -np.trapz(logls2, betas2)

        return logZ1, np.abs(logZ1 - logZ2)


    def hand_calc_z(self, logls):
        betas = self.betas.copy()
        order = np.argsort(betas)[::-1]
        betas = betas[order]
        logls0 = logls[order]
        
        if betas[-1] != 0:
            betas1 = np.concatenate((betas, [0]))
            logls1 = np.concatenate((logls0, [logls0[-1]]))            
        else:
            betas1 = betas
            logls1 = logls0

        # MANUAL INTEGRATION 3, INTERPOLATED
        # INTERPOLATION
        x0 = betas1
        y0 = logls1
        
        index = 0
        for i in range(len(y0)-1):
            if y0[i]*y0[i+1] < 0:
                index = i+1
                m = (y0[i+1] - y0[i]) / (x0[i+1] - x0[i])
                insertx = x0[i] - y0[i]/m
                inserty = 0

        if index:
            y0 = np.insert(y0, index, inserty)
            x0 = np.insert(x0, index, insertx)
        
        # INTEGRAL CALC
        hand_Z = np.array([])
        hand_Zerr = np.array([])
        
        for i in range(len(y0)-1):
            xa, xb = x0[i], x0[i+1]
            A, B = y0[i], y0[i+1]
            
            ctex = (xb-xa)
            quad = A * ctex
            tria = (B-A) * ctex / 2
            suma = quad + tria
            
            hand_Z = np.append(hand_Z, suma)
            hand_Zerr = np.append(hand_Zerr, tria)
            
        suma = np.sum(hand_Z)
        suma_err = np.sqrt(np.sum(np.array(hand_Zerr)**2))

        return -suma, suma_err


    def extract_chunks(self, array, chunk_size=1000):
        chunks = []
        end_index = len(array[0])
        while end_index > 0:
            start_index = max(end_index - chunk_size, 0)  # Ensuring the start index is not negative
            chunk = array[:, start_index:end_index]
            if len(chunk) > 1:
                chunks.append(chunk)
                end_index = start_index  # Update the end_index for the next iteration
        return chunks


    def get_Z(self, discard=1, coef=3, largo=100):
        if True:
            sampler_dict={'flat':True, 'discard':discard}
            logl_r = np.array(self.get_func('get_blobs', kwargs=sampler_dict))

            chunks = self.extract_chunks(logl_r, chunk_size=largo)

            zz_ = []
            zze_ = []
            for subset in chunks:
                logls = self.get_mean_logls(subset, coef=coef)
                
                z, zerr = self.hand_calc_z(logls)
                zz_.append(z)
                zze_.append(zerr)
        if False:
            print('Try larger sample')
            return [0], [0]
        return zz_, zze_


    def get_mean_logls(self, logls, coef=3):
        ll0 = np.array([])
        for i in range(self.ntemps):
            loglrow = logls[i]
            mask = (np.mean(loglrow) + coef*np.std(loglrow) > loglrow) & (loglrow > np.mean(loglrow) - coef*np.std(loglrow))
            loglrow0 = np.mean(loglrow[mask])
            ll0 = np.append(ll0, loglrow0)

        return ll0


    def run_mcmc(self, initial_state, nsteps, iterations=1, progress=True):
        if initial_state is None:
            print('Initial state is none')
            if self.__previous_state[0] is None:
                raise ValueError(
                    "Cannot have `initial_state=None` if run_mcmc has never "
                    "been called.")
            initial_state = self.__previous_state

        self.sampler = [emcee.EnsembleSampler(self.nwalkers, self.ndim,
                        self.my_probs_fn[t], pool=self.pool,
                        moves=self.moves[t], backend=self.backend[t],
                        vectorize=self.vectorize[t], blobs_dtype=self.blobs_dtype[t],
                        ) for t in range(self.ntemps)]

        results = None
        pbar = tqdm(total=nsteps*iterations, disable=not progress)
        for _ in range(nsteps):
            for results in self.sample(initial_state, iterations=iterations):
                pbar.update(iterations)

        pbar.close()

        return results


    def get_attr(self, x):
        return [getattr(sampler_instance, x) for sampler_instance in self]


    def get_func(self, x, kwargs=None):
        if kwargs is None:
            kwargs = {}
        return [getattr(sampler_instance, x)(**kwargs) for sampler_instance in self]


    pass
