# Path: mjhmc/misc/utils.py
# def overrides(interface_class):
#     """
#     override decorator. need to specify the super class, unfortunately
#     """
#     def overrider(method):
#         assert(method.__name__ in dir(interface_class))
#         return method
#     return overrider
# 
# def min_idx(draws):
#     """
#     selects the argmin of the row wise minimum
#     returns a list of rows where each input draw was min
#     draws: a list of the various draws
# 
#     eg: draws = [fl_draws, f_draws] where fl_draws.shape = f_draws.shape = (1, N)
#     returns fl_idx, f_idx where fl_idx (and the others) are
#     arrays of shape (nbatch,) containing the indices of the rows where each variable was
#     the minimum
#     """
#     cdraws = np.concatenate(draws, axis=0)
#     transition_idx = np.argmin(cdraws, axis=0)
#     return [np.where(transition_idx == i)[0] for i in xrange(len(draws))]
# 
# def normalize_by_row(matrix):
#     row_sums = matrix.sum(axis=1)
#     # replaces 0 with 1 to avoid divide by 0
#     np.place(row_sums, row_sums == 0, 1)
#     norm_matrix = matrix / row_sums[:, np.newaxis]
#     return norm_matrix


import numpy as np
import copy
import itertools
from mjhmc.misc.utils import overrides, min_idx, normalize_by_row
        p_f = 1 - p_fl

        p_transitions = [p_fl, p_f]
        transition_idxs = [fl_idx, f_idx]
        for t_idx, p_t in zip(transition_idxs, p_transitions):
            T_full[idx, t_idx] = p_t

class AlgebraicReducedFlip(AlgebraicContinuous):
    """The continuous algebraic sampler with reduced flip rate
    """

    @overrides(AlgebraicContinuous)
    def sampling_iteration(self):
        # states
        pre_state = self.ladder_states.copy()
        f_state = self.ladder_states.copy().F()
        l_state = self.ladder_states.copy().F().FL()
        # aka l^-1 state
        flf_state = self.ladder_states.copy().FL().F()

        # rates
        l_rates = self.acceptance_rate(self.ladder_states, l_state)
        flf_rates = self.acceptance_rate(self.ladder_states, flf_state)
        f_rates = flf_rates - np.min((flf_rates, l_rates), axis=0)

        # draws from exponential distribution
        l_draws = self.draw_from(l_rates)
        f_draws = self.draw_from(f_rates)

        # select minimum waiting times

# Path: mjhmc/misc/utils.py
# def overrides(interface_class):
#     """
#     override decorator. need to specify the super class, unfortunately
#     """
#     def overrider(method):
#         assert(method.__name__ in dir(interface_class))
#         return method
#     return overrider
# 
# def min_idx(draws):
#     """
#     selects the argmin of the row wise minimum
#     returns a list of rows where each input draw was min
#     draws: a list of the various draws
# 
#     eg: draws = [fl_draws, f_draws] where fl_draws.shape = f_draws.shape = (1, N)
#     returns fl_idx, f_idx where fl_idx (and the others) are
#     arrays of shape (nbatch,) containing the indices of the rows where each variable was
#     the minimum
#     """
#     cdraws = np.concatenate(draws, axis=0)
#     transition_idx = np.argmin(cdraws, axis=0)
#     return [np.where(transition_idx == i)[0] for i in xrange(len(draws))]
# 
# def normalize_by_row(matrix):
#     row_sums = matrix.sum(axis=1)
#     # replaces 0 with 1 to avoid divide by 0
#     np.place(row_sums, row_sums == 0, 1)
#     norm_matrix = matrix / row_sums[:, np.newaxis]
#     return norm_matrix


# Path: mjhmc/samplers/algebraic_hmc.py
import numpy as np
import copy
import itertools
from mjhmc.misc.utils import overrides, min_idx, normalize_by_row

        p_f = 1 - p_fl

        p_transitions = [p_fl, p_f]
        transition_idxs = [fl_idx, f_idx]
        for t_idx, p_t in zip(transition_idxs, p_transitions):
            T_full[idx, t_idx] = p_t

class AlgebraicReducedFlip(AlgebraicContinuous):
    """The continuous algebraic sampler with reduced flip rate
    """

    @overrides(AlgebraicContinuous)
    def sampling_iteration(self):
        # states
        pre_state = self.ladder_states.copy()
        f_state = self.ladder_states.copy().F()
        l_state = self.ladder_states.copy().F().FL()
        # aka l^-1 state
        flf_state = self.ladder_states.copy().FL().F()

        # rates
        l_rates = self.acceptance_rate(self.ladder_states, l_state)
        flf_rates = self.acceptance_rate(self.ladder_states, flf_state)
        f_rates = flf_rates - np.min((flf_rates, l_rates), axis=0)

        # draws from exponential distribution
        l_draws = self.draw_from(l_rates)
        f_draws = self.draw_from(f_rates)

        # select minimum waiting times

        l_idx, f_idx = min_idx([l_draws, f_draws])