# -*- coding: utf-8 -*-
from abc import ABC
from copy import deepcopy
import numpy as np
import warnings

from ..state import BranchSupplemental, State
from .move import Move


__all__ = ["RedBlueMove"]


class RedBlueMove(Move, ABC):
    """
    An abstract red-blue ensemble move with parallelization as described in
    `Foreman-Mackey et al. (2013) <https://arxiv.org/abs/1202.3665>`_.

    This class is heavily based on the original from ``emcee``.

    Args:
        nsplits (int, optional): The number of sub-ensembles to use. Each
            sub-ensemble is updated in parallel using the other sets as the
            complementary ensemble. The default value is ``2`` and you
            probably won't need to change that.
        randomize_split (bool, optional): Randomly shuffle walkers between
            sub-ensembles. The same number of walkers will be assigned to
            each sub-ensemble on each iteration. (default: ``True``)
        live_dangerously (bool, optional): By default, an update will fail with
            a ``RuntimeError`` if the number of walkers is smaller than twice
            the dimension of the problem because the walkers would then be
            stuck on a low dimensional subspace. This can be avoided by
            switching between the stretch move and, for example, a
            Metropolis-Hastings step. If you want to do this and suppress the
            error, set ``live_dangerously = True``. Thanks goes (once again)
            to @dstndstn for this wonderful terminology. (default: ``False``)
        **kwargs (dict, optional): Kwargs for parent classes.

    """

    def __init__(
        self, nsplits=2, randomize_split=True, live_dangerously=False, **kwargs
    ):
        super(RedBlueMove, self).__init__(**kwargs)
        self.nsplits = int(nsplits)
        self.live_dangerously = live_dangerously
        self.randomize_split = randomize_split

    def setup(self, branches_coords):
        """Any setup for the proposal.

        Args:
            branches_coords (dict): Keys are ``branch_names``. Values are
                np.ndarray[ntemps, nwalkers, nleaves_max, ndim]. These are the curent
                coordinates for all the walkers.

        """

    @classmethod
    def get_proposal(self, sample, complement, random, gibbs_ndim=None):
        """Make a proposal

        Args:
            sample (dict): Keys are ``branch_names``. Values are
                np.ndarray[ntemps, nwalkers, nleaves_max, ndim].
            complement (dict): Keys are ``branch_names``. Values are lists of other
                other np.ndarray[ntemps, nwalkers - subset size, nleaves_max, ndim] from
                all other subsets. This is the compliment whose ``coords`` are
                used to form the proposal for the ``sample`` subset.
            random (object): Current random state of the sampler.
            gibbs_ndim (int or np.ndarray, optional): If Gibbs sampling, this indicates
                the true dimension. If given as an array, must have shape ``(ntemps, nwalkers)``.
                See the tutorial for more information.
                (default: ``None``)

        Returns:
            tuple: Tuple contained proposal information.
                First entry is the new coordinates as a dictionary with keys
                as ``branch_names`` and values as np.ndarray[ntemps, nwalkers, nleaves_max, ndim]
                of new coordinates. Second entry is the factors associated with the
                proposal necessary for detailed balance. This is effectively
                any term in the detailed balance fraction. +log of factors if
                in the numerator. -log of factors if in the denominator.

        """

        raise NotImplementedError("The proposal must be implemented by " "subclasses")

    def propose(self, model, state):
        """Use the move to generate a proposal and compute the acceptance

        Args:
            model (:class:`eryn.model.Model`): Carrier of sampler information.
            state (:class:`State`): Current state of the sampler.

        Returns:
            tuple: (state, accepted)
                The first return is the state of the sampler after the move.
                The second return value is the accepted count array.

        """
        # Check that the dimensions are compatible.
        ndim_total = 0
        for branch in state.branches.values():
            ntemps, nwalkers, nleaves_, ndim_ = branch.shape
            ndim_total += ndim_ * nleaves_

        if nwalkers < 2 * ndim_total and not self.live_dangerously:
            raise RuntimeError(
                "It is unadvisable to use a red-blue move "
                "with fewer walkers than twice the number of "
                "dimensions. If you would like to do this, please set live_dangerously"
                "to True."
            )

        # Run any move-specific setup.
        self.setup(state.branches)

        # Split the ensemble in half and iterate over these two halves.
        accepted = np.zeros((ntemps, nwalkers), dtype=bool)
        all_inds = np.tile(np.arange(nwalkers), (ntemps, 1))
        inds = all_inds % self.nsplits
        if self.randomize_split:
            [np.random.shuffle(x) for x in inds]

        all_branch_names = list(state.branches.keys())

        ntemps, nwalkers, _, _ = state.branches[all_branch_names[0]].shape

        # get gibbs sampling information
        for branch_names_run, inds_run in self.gibbs_sampling_setup_iterator(
            all_branch_names
        ):
            # setup proposals based on Gibbs sampling
            (
                coords_going_for_proposal,
                inds_going_for_proposal,
                at_least_one_proposal,
            ) = self.setup_proposals(
                branch_names_run, inds_run, state.branches_coords, state.branches_inds
            )

            if not at_least_one_proposal:
                continue

            # prepare accepted fraction
            accepted_here = np.zeros((ntemps, nwalkers), dtype=bool)
            for split in range(self.nsplits):
                # get split information
                S1 = inds == split
                num_total_here = np.sum(inds == split)
                nwalkers_here = np.sum(S1[0])

                all_inds_shaped = all_inds[S1].reshape(ntemps, nwalkers_here)

                # inds including gibbs information
                new_inds = {
                    name: np.take_along_axis(
                        state.branches[name].inds,
                        all_inds_shaped[:, :, None],
                        axis=1,
                    )
                    for name in state.branches
                }

                # the actual inds for the subset
                real_inds_subset = {
                    name: new_inds[name] for name in inds_going_for_proposal
                }

                # actual coordinates of subset
                temp_coords = {
                    name: np.take_along_axis(
                        state.branches_coords[name],
                        all_inds_shaped[:, :, None, None],
                        axis=1,
                    )
                    for name in state.branches_coords
                }

                # prepare the sets for each model
                # goes into the proposal as (ntemps * (nwalkers / subset size), nleaves_max, ndim)
                sets = {
                    key: [
                        np.take_along_axis(
                            state.branches[key].coords,
                            all_inds[inds == j].reshape(ntemps, -1)[:, :, None, None],
                            axis=1,
                        )
                        for j in range(self.nsplits)
                    ]
                    for key in branch_names_run
                }

                # setup s and c based on splits
                s = {key: sets[key][split] for key in sets}
                c = {key: sets[key][:split] + sets[key][split + 1 :] for key in sets}

                # need to trick stretch proposal into using the dimenionality associated
                # with Gibbs sampling if it is being used
                gibbs_ndim = 0
                for brn, ir in zip(branch_names_run, inds_run):
                    if ir is not None:
                        gibbs_ndim += ir.sum()
                    else:
                        # nleaves * ndim
                        gibbs_ndim += np.prod(state.branches[brn].shape[-2:])

                # Get the move-specific proposal.
                q, factors = self.get_proposal(
                    s, c, model.random, gibbs_ndim=gibbs_ndim
                )

                # account for gibbs sampling
                self.cleanup_proposals_gibbs(branch_names_run, inds_run, q, temp_coords)

                # setup supplemental information
                if state.supplemental is not None:
                    # TODO: should there be a copy?
                    new_supps = BranchSupplemental(
                        state.supplemental.take_along_axis(all_inds_shaped, axis=1),
                        base_shape=(ntemps, nwalkers),
                        copy=False,
                    )

                else:
                    new_supps = None

                # default for removing inds info from supp
                if not np.all(
                    np.asarray(list(state.branches_supplemental.values())) == None
                ):
                    new_branch_supps_tmp = {
                        name: state.branches[name].branch_supplemental.take_along_axis(
                            all_inds_shaped[:, :, None], axis=1
                        )
                        for name in state.branches
                        if state.branches[name].branch_supplemental is not None
                    }

                    new_branch_supps = {
                        name: BranchSupplemental(
                            new_branch_supps_tmp[name],
                            base_shape=new_inds[name].shape,
                            copy=False,
                        )
                        for name in new_branch_supps_tmp
                    }

                else:
                    new_branch_supps = None

                # order everything properly
                q, new_inds, new_branch_supps = self.ensure_ordering(
                    list(state.branches.keys()), q, new_inds, new_branch_supps
                )

                # Compute prior of the proposed position
                # new_inds_prior is adjusted if product-space is used
                logp = model.compute_log_prior_fn(
                    q,
                    inds=new_inds,
                    supps=new_supps,
                    branch_supps=new_branch_supps,
                )

                self.fix_logp_gibbs(branch_names_run, inds_run, logp, real_inds_subset)

                # Compute the lnprobs of the proposed position.
                logl, new_blobs = model.compute_log_like_fn(
                    q,
                    inds=new_inds,
                    logp=logp,
                    supps=new_supps,
                    branch_supps=new_branch_supps,
                )

                # catch and warn about nans
                if np.any(np.isnan(logl)):
                    logl[np.isnan(logl)] = -1e300
                    warnings.warn("Getting Nan in likelihood computation.")

                logP = self.compute_log_posterior(logl, logp)

                prev_logl = np.take_along_axis(state.log_like, all_inds_shaped, axis=1)

                prev_logp = np.take_along_axis(state.log_prior, all_inds_shaped, axis=1)

                # takes care of tempering
                prev_logP = self.compute_log_posterior(prev_logl, prev_logp)

                lnpdiff = factors + logP - prev_logP

                keep = lnpdiff > np.log(model.random.rand(ntemps, nwalkers_here))

                # if gibbs sampling, this will say it is accepted if
                # any of the gibbs proposals were accepted
                np.put_along_axis(
                    accepted_here,
                    all_inds_shaped,
                    keep,
                    axis=1,
                )

                # readout for total
                accepted = (accepted.astype(int) + accepted_here.astype(int)).astype(
                    bool
                )
                # new state
                new_state = State(
                    q,
                    log_like=logl,
                    log_prior=logp,
                    blobs=new_blobs,
                    inds=new_inds,
                    supplemental=new_supps,
                    branch_supplemental=new_branch_supps,
                )

                # update state
                state = self.update(
                    state, new_state, accepted_here, subset=all_inds_shaped
                )

            # add to move-specific accepted information
            self.accepted += accepted
            self.num_proposals += 1

        # temp swaps
        if self.temperature_control is not None:
            state = self.temperature_control.temper_comps(state)

        return state, accepted
