"""
Implements several extended-ensemble Monte Carlo sampling algorithms.

Here is a short example which shows how to sample from a PDF using the replica
exchange with non-equilibrium switches (RENS) method. It draws 5000 samples from
a 1D normal distribution using the RENS algorithm working on three Markov chains
being generated by the HMC algorithm:


    >>> import numpy
    >>> from numpy import sqrt
    >>> from csb.io.plots import Chart
    >>> from csb.statistics.pdf import Normal
    >>> from csb.statistics.samplers import State
    >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENSSwapParameterInfo
    >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENS, AlternatingAdjacentSwapScheme
    >>> from csb.statistics.samplers.mc.singlechain import HMCSampler

    >>> # Pick some initial state for the different Markov chains:
    >>> initial_state = State(numpy.array([1.]))

    >>> # Set standard deviations:
    >>> std_devs = [1./sqrt(5), 1. / sqrt(3), 1.]

    >>> # Set HMC timesteps and trajectory length:
    >>> hmc_timesteps = [0.6, 0.7, 0.6]
    >>> hmc_trajectory_length = 20
    >>> hmc_gradients = [lambda q, t: 1 / (std_dev ** 2) * q for std_dev in std_devs]

    >>> # Set parameters for the thermostatted RENS algorithm:
    >>> rens_trajectory_length = 30
    >>> rens_timesteps = [0.3, 0.5]

    >>> # Set interpolation gradients as a function of the work parameter l:
    >>> rens_gradients = [lambda q, l, i=i: (l / (std_devs[i + 1] ** 2) + (1 - l) / (std_devs[i] ** 2)) * q 
                          for i in range(len(std_devs)-1)]

    >>> # Initialize HMC samplers:
    >>> samplers = [HMCSampler(Normal(sigma=std_devs[i]), initial_state, hmc_gradients[i], hmc_timesteps[i],
                    hmc_trajectory_length) for i in range(len(std_devs))]

    >>> # Create swap parameter objects:
    params = [ThermostattedMDRENSSwapParameterInfo(samplers[0], samplers[1], rens_timesteps[0],
              rens_trajectory_length, rens_gradients[0]),
              ThermostattedMDRENSSwapParameterInfo(samplers[1], samplers[2], rens_timesteps[1],
              rens_trajectory_length, rens_gradients[1])]

    >>> # Initialize thermostatted RENS algorithm:
    >>> algorithm = ThermostattedMDRENS(samplers, params)

    >>> # Initialize swapping scheme:
    >>> swapper = AlternatingAdjacentSwapScheme(algorithm)

    >>> # Initialize empty list which will store the samples:
    >>> states = []
    >>> for i in range(5000):
            if i % 5 == 0:
                swapper.swap_all()
            states.append(algorithm.sample())

    >>> # Print acceptance rates:
    >>> print('HMC acceptance rates:', [s.acceptance_rate for s in samplers])
    >>> print('swap acceptance rates:', algorithm.acceptance_rates)

    >>> # Create and plot histogram for first sampler and numpy.random.normal reference:
    >>> chart = Chart()
    >>> rawstates = [state[0].position[0] for state in states]
    >>> chart.plot.hist([numpy.random.normal(size=5000, scale=std_devs[0]), rawstates], bins=30, normed=True)
    >>> chart.plot.legend(['numpy.random.normal', 'RENS + HMC'])
    >>> chart.show()


For L{ReplicaExchangeMC} (RE), the procedure is easier because apart from the
two sampler instances the corresponding L{RESwapParameterInfo} objects take
no arguments.

Every replica exchange algorithm in this module (L{ReplicaExchangeMC}, L{MDRENS},
L{ThermostattedMDRENS}) is used in a similar way. A simulation is always
initialized with a list of samplers (instances of classes derived from
L{AbstractSingleChainMC}) and a list of L{AbstractSwapParameterInfo} objects
suited for the algorithm under consideration. Every L{AbstractSwapParameterInfo}
object holds all the information needed to perform a swap between two samplers.
The usual scheme is to swap only adjacent replicae in a scheme::

    1 <--> 2, 3 <--> 4, ...
    2 <--> 3, 4 <--> 5, ...
    1 <--> 2, 3 <--> 4, ...
    
This swapping scheme is implemented in the L{AlternatingAdjacentSwapScheme} class,
but different schemes can be easily implemented by deriving from L{AbstractSwapScheme}.
Then the simulation is run by looping over the number of samples to be drawn
and calling the L{AbstractExchangeMC.sample} method of the algorithm. By calling
the L{AbstractSwapScheme.swap_all} method of the specific L{AbstractSwapScheme}
implementation, all swaps defined in the list of L{AbstractSwapParameterInfo}
objects are performed according to the swapping scheme. The
L{AbstractSwapScheme.swap_all} method may be called for example after sampling
intervals of a fixed length or randomly.
"""

import numpy

import csb.numeric

from abc import ABCMeta, abstractmethod

from csb.statistics.samplers import EnsembleState
from csb.statistics.samplers.mc import AbstractMC, Trajectory, MCCollection, augment_state
from csb.statistics.samplers.mc.propagators import MDPropagator, ThermostattedMDPropagator
from csb.statistics.samplers.mc.neqsteppropagator import NonequilibriumStepPropagator
from csb.statistics.samplers.mc.neqsteppropagator import Protocol, Step, ReducedHamiltonian
from csb.statistics.samplers.mc.neqsteppropagator import ReducedHamiltonianPerturbation
from csb.statistics.samplers.mc.neqsteppropagator import HMCPropagation, HMCPropagationParam
from csb.statistics.samplers.mc.neqsteppropagator import HamiltonianSysInfo, NonequilibriumTrajectory
from csb.numeric.integrators import AbstractGradient, FastLeapFrog


class AbstractEnsembleMC(AbstractMC):
    """
    Abstract class for Monte Carlo sampling algorithms simulating several ensembles.

    @param samplers: samplers which sample from their respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}    
    """

    __metaclass__ = ABCMeta

    def __init__(self, samplers):
        
        self._samplers = MCCollection(samplers)
        state = EnsembleState([x.state for x in self._samplers])
        
        super(AbstractEnsembleMC, self).__init__(state)        

    def sample(self):
        """
        Draw an ensemble sample.
        
        @rtype: L{EnsembleState}
        """
        
        sample = EnsembleState([sampler.sample() for sampler in self._samplers])
        self.state = sample

        return sample

    @property
    def energy(self):
        """
        Total ensemble energy.
        """ 
        return sum([x.energy for x in self._samplers])


class AbstractExchangeMC(AbstractEnsembleMC):
    """
    Abstract class for Monte Carlo sampling algorithms employing some replica exchange method.

    @param samplers: samplers which sample from their respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}

    @param param_infos: list of ParameterInfo instances providing information needed
        for performing swaps
    @type param_infos: list of L{AbstractSwapParameterInfo}
    """

    __metaclass__ = ABCMeta

    def __init__(self, samplers, param_infos):
        super(AbstractExchangeMC, self).__init__(samplers)
        
        self._swaplist1 = []
        self._swaplist2 = []        
        self._currentswaplist = self._swaplist1
        
        self._param_infos = param_infos
        self._statistics = SwapStatistics(self._param_infos)
        
    def _checkstate(self, state):
        if not isinstance(state, EnsembleState):
            raise TypeError(state)        

    def swap(self, index):
        """
        Perform swap between sampler pair described by param_infos[index]
        and return outcome (true = accepted, false = rejected).

        @param index: index of swap pair in param_infos
        @type index: int

        @rtype: boolean
        """
        param_info = self._param_infos[index]
        swapcom = self._propose_swap(param_info)
        swapcom = self._calc_pacc_swap(swapcom)
        result = self._accept_swap(swapcom)
        
        self.state = EnsembleState([x.state for x in self._samplers])

        self.statistics.stats[index].update(result)
        
        return result

    @abstractmethod
    def _propose_swap(self, param_info):
        """
        Calculate proposal states for a swap between two samplers.
        
        @param param_info: ParameterInfo instance holding swap parameters
        @type param_info: L{AbstractSwapParameterInfo}
        
        @rtype: L{AbstractSwapCommunicator}
        """ 
        pass

    @abstractmethod
    def _calc_pacc_swap(self, swapcom):
        """
        Calculate probability to accept a swap given initial and proposal states.

        @param swapcom: SwapCommunicator instance holding information to be communicated
                        between distinct swap substeps
        @type swapcom: L{AbstractSwapCommunicator}

        @rtype: L{AbstractSwapCommunicator}
        """
        pass

    def _accept_swap(self, swapcom):
        """
        Accept / reject an exchange between two samplers given proposal states and
        the acceptance probability and returns the outcome (true = accepted, false = rejected).

        @param swapcom: SwapCommunicator instance holding information to be communicated
            between distinct swap substeps
        @type swapcom: L{AbstractSwapCommunicator}

        @rtype: boolean
        """

        if numpy.random.random() < swapcom.acceptance_probability:
            if swapcom.sampler1.state.momentum is None and swapcom.sampler2.state.momentum is None:
                swapcom.traj12.final.momentum = None
                swapcom.traj21.final.momentum = None
            swapcom.sampler1.state = swapcom.traj21.final
            swapcom.sampler2.state = swapcom.traj12.final
            return True
        else:
            return False

    @property
    def acceptance_rates(self):
        """
        Return swap acceptance rates.

        @rtype: list of floats
        """        
        return self.statistics.acceptance_rates

    @property
    def param_infos(self):
        """
        List of SwapParameterInfo instances holding all necessary parameters.

        @rtype: list of L{AbstractSwapParameterInfo}
        """
        return self._param_infos
    
    @property
    def statistics(self):
        return self._statistics

    def _update_statistics(self, index, accepted):
        """
        Update statistics of a given swap process.
        
        @param index: position of swap statistics to be updated
        @type index: int
        
        @param accepted: outcome of the swap
        @type accepted: boolean
        """
        
        self._stats[index][0] += 1
        self._stats[index][1] += int(accepted)


class AbstractSwapParameterInfo(object):
    """
    Subclass instances hold all parameters necessary for performing a swap
    between two given samplers.
    """

    __metaclass__ = ABCMeta

    def __init__(self, sampler1, sampler2):
        """
        @param sampler1: First sampler
        @type sampler1: L{AbstractSingleChainMC}

        @param sampler2: Second sampler
        @type sampler2: L{AbstractSingleChainMC}
        """

        self._sampler1 = sampler1
        self._sampler2 = sampler2

    @property
    def sampler1(self):
        return self._sampler1

    @property
    def sampler2(self):
        return self._sampler2


class AbstractSwapCommunicator(object):
    """
    Holds all the information which needs to be communicated between
    distinct swap substeps.

    @param param_info: ParameterInfo instance holding swap parameters
    @type param_info: L{AbstractSwapParameterInfo}

    @param traj12: Forward trajectory
    @type traj12: L{Trajectory}

    @param traj21: Reverse trajectory
    @type traj21: L{Trajectory}
    """

    __metaclass__ = ABCMeta
    
    def __init__(self, param_info, traj12, traj21):
        
        self._sampler1 = param_info.sampler1
        self._sampler2 = param_info.sampler2

        self._traj12 = traj12
        self._traj21 = traj21
        
        self._param_info = param_info
        
        self._acceptance_probability = None
        self._accepted = False
        
    @property
    def sampler1(self):
        return self._sampler1

    @property
    def sampler2(self):
        return self._sampler2
    
    @property
    def traj12(self):
        return self._traj12    

    @property
    def traj21(self):
        return self._traj21

    @property
    def acceptance_probability(self):
        return self._acceptance_probability
    @acceptance_probability.setter
    def acceptance_probability(self, value):
        self._acceptance_probability = value

    @property
    def accepted(self):
        return self._accepted
    @accepted.setter
    def accepted(self, value):
        self._accepted = value

    @property
    def param_info(self):
        return self._param_info


class ReplicaExchangeMC(AbstractExchangeMC):
    """
    Replica Exchange (RE, Swendsen & Yang 1986) implementation.
    """
        
    def _propose_swap(self, param_info):
        
        return RESwapCommunicator(param_info, Trajectory([param_info.sampler1.state,
                                                          param_info.sampler1.state]),
                                              Trajectory([param_info.sampler2.state,
                                                          param_info.sampler2.state]))
    
    def _calc_pacc_swap(self, swapcom):
        
        E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x)
        E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x)
            
        T1 = swapcom.sampler1.temperature
        T2 = swapcom.sampler2.temperature
        
        state1 = swapcom.traj12.initial
        state2 = swapcom.traj21.initial
        
        proposal1 = swapcom.traj21.final
        proposal2 = swapcom.traj12.final

        swapcom.acceptance_probability = csb.numeric.exp(-E1(proposal1.position) / T1 
                                                         + E1(state1.position) / T1 
                                                         - E2(proposal2.position) / T2 
                                                         + E2(state2.position) / T2)
                                                         
        return swapcom


class RESwapParameterInfo(AbstractSwapParameterInfo):
    """
    Holds parameters for a standard Replica Exchange swap.
    """
    pass


class RESwapCommunicator(AbstractSwapCommunicator):
    """
    Holds all the information which needs to be communicated between distinct
    RE swap substeps.

    See L{AbstractSwapCommunicator} for constructor signature.
    """
    pass


class AbstractRENS(AbstractExchangeMC):
    """
    Abstract Replica Exchange with Nonequilibrium Switches
    (RENS, Ballard & Jarzynski 2009) class.
    Subclasses implement various ways of generating trajectories
    (deterministic or stochastic).
    """

    __metaclass__ = ABCMeta

    def _propose_swap(self, param_info):

        init_state1 = param_info.sampler1.state
        init_state2 = param_info.sampler2.state
        
        trajinfo12 = RENSTrajInfo(param_info, init_state1, direction="fw")
        trajinfo21 = RENSTrajInfo(param_info, init_state2, direction="bw")
        
        traj12 = self._run_traj_generator(trajinfo12)
        traj21 = self._run_traj_generator(trajinfo21)

        return RENSSwapCommunicator(param_info, traj12, traj21)

    def _setup_protocol(self, traj_info):
        """
        Sets the protocol lambda(t) to either the forward or the reverse protocol.

        @param traj_info: TrajectoryInfo object holding information neccessary to
                          calculate the rens trajectories.
        @type traj_info: L{RENSTrajInfo}
        """

        if traj_info.direction == "fw":
            return traj_info.param_info.protocol
        else:
            return lambda t, tau: traj_info.param_info.protocol(tau - t, tau)
        
        return protocol

    def _get_init_temperature(self, traj_info):
        """
        Determine the initial temperature of a RENS trajectory.

        @param traj_info: TrajectoryInfo object holding information neccessary to
                          calculate the RENS trajectory.
        @type traj_info: L{RENSTrajInfo}
        """
        
        if traj_info.direction == "fw":
            return traj_info.param_info.sampler1.temperature
        else:
            return traj_info.param_info.sampler2.temperature

    @abstractmethod
    def _calc_works(self, swapcom):
        """
        Calculates the works expended during the nonequilibrium
        trajectories.

        @param swapcom: Swap communicator object holding all the
                        neccessary information.
        @type swapcom: L{RENSSwapCommunicator}

        @return: The expended during the forward and the backward
                 trajectory.
        @rtype: 2-tuple of floats
        """
        
        pass

    def _calc_pacc_swap(self, swapcom):

        work12, work21 = self._calc_works(swapcom)
        swapcom.acceptance_probability = csb.numeric.exp(-work12 - work21)

        return swapcom

    @abstractmethod
    def _propagator_factory(self, traj_info):
        """
        Factory method which produces the propagator object used to calculate
        the RENS trajectories.

        @param traj_info: TrajectoryInfo object holding information neccessary to
                          calculate the rens trajectories.
        @type traj_info: L{RENSTrajInfo}
        @rtype: L{AbstractPropagator}
        """
        pass

    def _run_traj_generator(self, traj_info):
        """
        Run the trajectory generator which generates a trajectory
        of a given length between the states of two samplers.

        @param traj_info: TrajectoryInfo instance holding information
                          needed to generate a nonequilibrium trajectory   
        @type traj_info: L{RENSTrajInfo}
        
        @rtype: L{Trajectory}
        """

        init_temperature = self._get_init_temperature(traj_info)
        
        init_state = traj_info.init_state.clone()

        if init_state.momentum is None:
            init_state = augment_state(init_state,
                                       init_temperature,
                                       traj_info.param_info.mass_matrix)
            
        gen = self._propagator_factory(traj_info)

        traj = gen.generate(init_state, int(traj_info.param_info.traj_length))
        
        return traj


class AbstractRENSSwapParameterInfo(RESwapParameterInfo):
    """
    Holds parameters for a RENS swap.
    """

    __metaclass__ = ABCMeta

    def __init__(self, sampler1, sampler2, protocol):

        super(AbstractRENSSwapParameterInfo, self).__init__(sampler1, sampler2)

        ## Can't pass the linear protocol as a default argument because of a reported bug
        ## in epydoc parsing which makes it fail building the docs.
        self._protocol = None
        if protocol is None:
            self._protocol = lambda t, tau: t / tau
        else:
            self._protocol = protocol

    @property
    def protocol(self):
        """
        Switching protocol determining the time dependence
        of the switching parameter.
        """
        return self._protocol
    @protocol.setter
    def protocol(self, value):
        self._protocol = value


class RENSSwapCommunicator(AbstractSwapCommunicator):
    """
    Holds all the information which needs to be communicated between distinct
    RENS swap substeps.

    See L{AbstractSwapCommunicator} for constructor signature.
    """
    
    pass


class RENSTrajInfo(object):
    """
    Holds information necessary for calculating a RENS trajectory.

    @param param_info: ParameterInfo instance holding swap parameters
    @type param_info: L{AbstractSwapParameterInfo}

    @param init_state: state from which the trajectory is supposed to start
    @type init_state: L{State}

    @param direction: Either "fw" or "bw", indicating a forward or backward
                      trajectory. This is neccessary to pick the protocol or
                      the reversed protocol, respectively.
    @type direction: string, either "fw" or "bw"
    """
    
    def __init__(self, param_info, init_state, direction):
        
        self._param_info = param_info
        self._init_state = init_state
        self._direction = direction
        
    @property
    def param_info(self):
        return self._param_info

    @property
    def init_state(self):
        return self._init_state

    @property
    def direction(self):
        return self._direction


class MDRENS(AbstractRENS):
    """
    Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
    with Molecular Dynamics (MD) trajectories.

    @param samplers: Samplers which sample their
                         respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}

    @param param_infos: ParameterInfo instance holding
                        information required to perform a MDRENS swap
    @type param_infos: list of L{MDRENSSwapParameterInfo}

    @param integrator: Subclass of L{AbstractIntegrator} to be used to
                       calculate the non-equilibrium trajectories
    @type integrator: type
    """

    def __init__(self, samplers, param_infos,
                 integrator=csb.numeric.integrators.FastLeapFrog):
        
        super(MDRENS, self).__init__(samplers, param_infos)
        
        self._integrator = integrator

    def _propagator_factory(self, traj_info):
        
        protocol = self._setup_protocol(traj_info)
        tau = traj_info.param_info.traj_length * traj_info.param_info.timestep
        factory = InterpolationFactory(protocol, tau)
        gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient),
                           traj_info.param_info.timestep,
                           mass_matrix=traj_info.param_info.mass_matrix,
                           integrator=self._integrator)

        return gen

    def _calc_works(self, swapcom):

        T1 = swapcom.param_info.sampler1.temperature
        T2 = swapcom.param_info.sampler2.temperature
        
        heat12 = swapcom.traj12.heat
        heat21 = swapcom.traj21.heat
        
        proposal1 = swapcom.traj21.final
        proposal2 = swapcom.traj12.final
        
        state1 = swapcom.traj12.initial
        state2 = swapcom.traj21.initial
        
        if swapcom.param_info.mass_matrix.is_unity_multiple:
            inverse_mass_matrix = 1.0 / swapcom.param_info.mass_matrix[0][0]
        else:
            inverse_mass_matrix = swapcom.param_info.mass_matrix.inverse
        
        E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x)
        E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x)
        K = lambda x: 0.5 * numpy.dot(x.T, numpy.dot(inverse_mass_matrix, x))

        w12 = (K(proposal2.momentum) + E2(proposal2.position)) / T2 - \
              (K(state1.momentum) + E1(state1.position)) / T1 - heat12 
        w21 = (K(proposal1.momentum) + E1(proposal1.position)) / T1 - \
              (K(state2.momentum) + E2(state2.position)) / T2 - heat21

        return w12, w21


class MDRENSSwapParameterInfo(RESwapParameterInfo):
    """
    Holds parameters for a MDRENS swap.

    @param sampler1: First sampler
    @type sampler1: L{AbstractSingleChainMC}

    @param sampler2: Second sampler
    @type sampler2: L{AbstractSingleChainMC}

    @param timestep: Integration timestep
    @type timestep: float

    @param traj_length: Trajectory length in number of timesteps
    @type traj_length: int

    @param gradient: Gradient which determines the dynamics during a trajectory
    @type gradient: L{AbstractGradient}
    
    @param protocol: Switching protocol determining the time dependence of the
                     switching parameter. It is a function M{f} taking the running
                     time t and the switching time tau to yield a value in M{[0, 1]}
                     with M{f(0, tau) = 0} and M{f(tau, tau) = 1}. Default is a linear
                     protocol, which is being set manually due to an epydoc bug
    @type protocol: callable
    
    @param mass_matrix: Mass matrix
    @type mass_matrix: n-dimensional matrix of type L{InvertibleMatrix} with n being the dimension
                               of the configuration space, that is, the dimension of
                               the position / momentum vectors
    """

    def __init__(self, sampler1, sampler2, timestep, traj_length, gradient,
                 protocol=None, mass_matrix=None):
        
        super(MDRENSSwapParameterInfo, self).__init__(sampler1, sampler2)

        self._mass_matrix = mass_matrix
        if self.mass_matrix is None:
            d = len(sampler1.state.position)
            self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d))

        self._traj_length = traj_length
        self._gradient = gradient
        self._timestep = timestep

        ## Can't pass the linear protocol as a default argument because of a reported bug
        ## in epydoc parsing which makes it fail building the docs.
        self._protocol = None
        if protocol is None:
            self._protocol = lambda t, tau: t / tau
        else:
            self._protocol = protocol
    
    @property
    def timestep(self):
        """
        Integration timestep.
        """
        return self._timestep
    @timestep.setter
    def timestep(self, value):
        self._timestep = float(value)

    @property
    def traj_length(self):
        """
        Trajectory length in number of integration steps.
        """
        return self._traj_length
    @traj_length.setter
    def traj_length(self, value):
        self._traj_length = int(value)

    @property
    def gradient(self):
        """
        Gradient which governs the equations of motion.
        """
        return self._gradient

    @property
    def mass_matrix(self):
        return self._mass_matrix
    @mass_matrix.setter
    def mass_matrix(self, value):
        self._mass_matrix = value

    @property
    def protocol(self):
        """
        Switching protocol determining the time dependence
        of the switching parameter.
        """
        return self._protocol
    @protocol.setter
    def protocol(self, value):
        self._protocol = value


class ThermostattedMDRENS(MDRENS):
    """
    Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009)
    with Andersen-thermostatted Molecular Dynamics (MD) trajectories.

    @param samplers: Samplers which sample their
                         respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}

    @param param_infos: ParameterInfo instance holding
                        information required to perform a MDRENS swap
    @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo}

    @param integrator: Subclass of L{AbstractIntegrator} to be used to
                       calculate the non-equilibrium trajectories
    @type integrator: type
    """

    def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
        
        super(ThermostattedMDRENS, self).__init__(samplers, param_infos, integrator)

    def _propagator_factory(self, traj_info):

        protocol = self._setup_protocol(traj_info)
        tau = traj_info.param_info.traj_length * traj_info.param_info.timestep
        factory = InterpolationFactory(protocol, tau)
        
        grad = factory.build_gradient(traj_info.param_info.gradient)
        temp = factory.build_temperature(traj_info.param_info.temperature)

        gen = ThermostattedMDPropagator(grad,
                                        traj_info.param_info.timestep, temperature=temp, 
                                        collision_probability=traj_info.param_info.collision_probability,
                                        update_interval=traj_info.param_info.collision_interval,
                                        mass_matrix=traj_info.param_info.mass_matrix,
                                        integrator=self._integrator)

        return gen

class ThermostattedMDRENSSwapParameterInfo(MDRENSSwapParameterInfo):
    """
    @param sampler1: First sampler
    @type sampler1: subclass instance of L{AbstractSingleChainMC}

    @param sampler2: Second sampler
    @type sampler2: subclass instance of L{AbstractSingleChainMC}

    @param timestep: Integration timestep
    @type timestep: float

    @param traj_length: Trajectory length in number of timesteps
    @type traj_length: int

    @param gradient: Gradient which determines the dynamics during a trajectory
    @type gradient: subclass instance of L{AbstractGradient}

    @param mass_matrix: Mass matrix
    @type mass_matrix: n-dimensional L{InvertibleMatrix} with n being the dimension
                       of the configuration space, that is, the dimension of
                       the position / momentum vectors

    @param protocol: Switching protocol determining the time dependence of the
                     switching parameter. It is a function f taking the running
                     time t and the switching time tau to yield a value in [0, 1]
                     with f(0, tau) = 0 and f(tau, tau) = 1
    @type protocol: callable

    @param temperature: Temperature interpolation function.
    @type temperature: Real-valued function mapping from [0,1] to R.
        T(0) = temperature of the ensemble sampler1 samples from, T(1) = temperature
        of the ensemble sampler2 samples from

    @param collision_probability: Probability for a collision with the heatbath during one timestep
    @type collision_probability: float

    @param collision_interval: Interval during which collision may occur with probability
        collision_probability
    @type collision_interval: int
    """
        
    def __init__(self, sampler1, sampler2, timestep, traj_length, gradient, mass_matrix=None,
                 protocol=None, temperature=lambda l: 1.0,
                 collision_probability=0.1, collision_interval=1):
        
        super(ThermostattedMDRENSSwapParameterInfo, self).__init__(sampler1, sampler2, timestep,
                                                                   traj_length, gradient,
                                                                   mass_matrix=mass_matrix,
                                                                   protocol=protocol)
        
        self._collision_probability = None
        self._collision_interval = None
        self._temperature = temperature
        self.collision_probability = collision_probability
        self.collision_interval = collision_interval

    @property
    def collision_probability(self):
        """
        Probability for a collision with the heatbath during one timestep.
        """
        return self._collision_probability
    @collision_probability.setter
    def collision_probability(self, value):
        self._collision_probability = float(value)

    @property
    def collision_interval(self):
        """
        Interval during which collision may occur with probability
        C{collision_probability}.
        """
        return self._collision_interval
    @collision_interval.setter
    def collision_interval(self, value):
        self._collision_interval = int(value)

    @property
    def temperature(self):
        return self._temperature


class AbstractStepRENS(AbstractRENS):
    """
    Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
    with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate
    Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011.
    The switching parameter dependence of the Hamiltonian is a linear interpolation
    between the PDFs of the sampler objects, 
    M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}.
    The perturbation kernel is a thermodynamic perturbation and the propagation is subclass
    responsibility.
    Note that due to the linear interpolations between the two Hamiltonians, the
    log-probability has to be evaluated four times per perturbation step which can be
    costly. In this case it is advisable to define the intermediate log probabilities
    in _run_traj_generator differently.

    @param samplers: Samplers which sample their respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}

    @param param_infos: ParameterInfo instances holding
                        information required to perform a HMCStepRENS swaps
    @type param_infos: list of L{AbstractSwapParameterInfo}
    """

    __metaclass__ = ABCMeta

    def __init__(self, samplers, param_infos):

        super(AbstractStepRENS, self).__init__(samplers, param_infos)

        self._evaluate_im_works = True

    @abstractmethod
    def _setup_propagations(self, im_sys_infos, param_info):
        """
        Set up the propagation steps using the information about the current system
        setup and parameters from the SwapParameterInfo object.

        @param im_sys_infos: Information about the intermediate system setups
        @type im_sys_infos: List of L{AbstractSystemInfo}

        @param param_info: SwapParameterInfo object containing parameters for the
                           propagations like timesteps, trajectory lengths etc.
        @type param_info: L{AbstractSwapParameterInfo}
        """
        
        pass

    @abstractmethod
    def _add_gradients(self, im_sys_infos, param_info, t_prot):
        """
        If needed, set im_sys_infos.hamiltonian.gradient.

        @param im_sys_infos: Information about the intermediate system setups
        @type im_sys_infos: List of L{AbstractSystemInfo}

        @param param_info: SwapParameterInfo object containing parameters for the
                           propagations like timesteps, trajectory lengths etc.
        @type param_info: L{AbstractSwapParameterInfo}

        @param t_prot: Switching protocol defining the time dependence of the switching
                       parameter.
        @type t_prot: callable
        """
        
        pass

    def _setup_stepwise_protocol(self, traj_info):
        """
        Sets up the stepwise protocol consisting of perturbation and relaxation steps.

        @param traj_info: TrajectoryInfo instance holding information
                          needed to generate a nonequilibrium trajectory   
        @type traj_info: L{RENSTrajInfo}
        
        @rtype: L{Protocol}
        """

        pdf1 = traj_info.param_info.sampler1._pdf
        pdf2 = traj_info.param_info.sampler2._pdf
        T1 = traj_info.param_info.sampler1.temperature
        T2 = traj_info.param_info.sampler2.temperature
        traj_length = traj_info.param_info.intermediate_steps
        prot = self._setup_protocol(traj_info)
        t_prot = lambda i: prot(float(i), float(traj_length))

        im_log_probs = [lambda x, i=i: pdf2.log_prob(x) * t_prot(i) + \
                                       (1 - t_prot(i)) * pdf1.log_prob(x)
                        for i in range(traj_length + 1)]
        
        im_temperatures = [T2 * t_prot(i) + (1 - t_prot(i)) * T1 
                           for i in range(traj_length + 1)]
        im_reduced_hamiltonians = [ReducedHamiltonian(im_log_probs[i],
                                                      temperature=im_temperatures[i]) 
                                   for i in range(traj_length + 1)]
        im_sys_infos = [HamiltonianSysInfo(im_reduced_hamiltonians[i])
                        for i in range(traj_length + 1)]
        perturbations = [ReducedHamiltonianPerturbation(im_sys_infos[i], im_sys_infos[i+1])
                        for i in range(traj_length)]
        if self._evaluate_im_works == False:
            for p in perturbations:
                p.evaluate_work = False
        im_sys_infos = self._add_gradients(im_sys_infos, traj_info.param_info, t_prot)
        propagations = self._setup_propagations(im_sys_infos, traj_info.param_info)
        
        steps = [Step(perturbations[i], propagations[i]) for i in range(traj_length)]

        return Protocol(steps)

    def _propagator_factory(self, traj_info):

        protocol = self._setup_stepwise_protocol(traj_info)
        gen = NonequilibriumStepPropagator(protocol)
        
        return gen

    def _run_traj_generator(self, traj_info):

        init_temperature = self._get_init_temperature(traj_info)

        gen = self._propagator_factory(traj_info)

        traj = gen.generate(traj_info.init_state)
        
        return NonequilibriumTrajectory([traj_info.init_state, traj.final], jacobian=1.0,
                                         heat=traj.heat, work=traj.work, deltaH=traj.deltaH)       


class HMCStepRENS(AbstractStepRENS):
    """
    Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
    with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate
    Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011.
    The switching parameter dependence of the Hamiltonian is a linear interpolation
    between the PDFs of the sampler objects, 
    M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}.
    The perturbation kernel is a thermodynamic perturbation and the propagation is done using HMC.

    Note that due to the linear interpolations between the two Hamiltonians, the
    log-probability and its gradient has to be evaluated four times per perturbation step which
    can be costly. In this case it is advisable to define the intermediate log probabilities
    in _run_traj_generator differently.

    @param samplers: Samplers which sample their respective equilibrium distributions
    @type samplers: list of L{AbstractSingleChainMC}

    @param param_infos: ParameterInfo instances holding
                        information required to perform a HMCStepRENS swaps
    @type param_infos: list of L{HMCStepRENSSwapParameterInfo}
    """

    def __init__(self, samplers, param_infos):

        super(HMCStepRENS, self).__init__(samplers, param_infos)

    def _add_gradients(self, im_sys_infos, param_info, t_prot):

        im_gradients = [lambda x, t, i=i: param_info.gradient(x, t_prot(i))
                        for i in range(param_info.intermediate_steps + 1)]

        for i, s in enumerate(im_sys_infos):
            s.hamiltonian.gradient = im_gradients[i]

        return im_sys_infos
            
    def _setup_propagations(self, im_sys_infos, param_info):
                        
        propagation_params = [HMCPropagationParam(param_info.timestep,
                                                  param_info.hmc_traj_length,
                                                  im_sys_infos[i+1].hamiltonian.gradient,
                                                  param_info.hmc_iterations,
                                                  mass_matrix=param_info.mass_matrix,
                                                  integrator=param_info.integrator)
                              for i in range(param_info.intermediate_steps)]

        propagations = [HMCPropagation(im_sys_infos[i+1], propagation_params[i], evaluate_heat=False)
                        for i in range(param_info.intermediate_steps)]

        return propagations        

    def _calc_works(self, swapcom):

        return swapcom.traj12.work, swapcom.traj21.work


class HMCStepRENSSwapParameterInfo(AbstractRENSSwapParameterInfo):
    """
    Holds all required information for performing HMCStepRENS swaps.

    @param sampler1: First sampler
    @type sampler1: subclass instance of L{AbstractSingleChainMC}

    @param sampler2: Second sampler
    @type sampler2: subclass instance of L{AbstractSingleChainMC}

    @param timestep: integration timestep
    @type timestep: float

    @param hmc_traj_length: HMC trajectory length
    @type hmc_traj_length: int

    @param hmc_iterations: number of HMC iterations in the propagation step
    @type hmc_iterations: int

    @param gradient: gradient governing the equations of motion, function of
                     position array and switching protocol
    @type gradient: callable
    
    @param intermediate_steps: number of steps in the protocol; this is a discrete version
                               of the switching time in "continuous" RENS implementations
    @type intermediate_steps: int

    @param protocol: Switching protocol determining the time dependence of the
                     switching parameter. It is a function f taking the running
                     time t and the switching time tau to yield a value in [0, 1]
                     with f(0, tau) = 0 and f(tau, tau) = 1
    @type protocol: callable

    @param mass_matrix: mass matrix for kinetic energy definition
    @type mass_matrix: L{InvertibleMatrix}

    @param integrator: Integration scheme to be utilized
    @type integrator: l{AbstractIntegrator}
    """

    def __init__(self, sampler1, sampler2, timestep, hmc_traj_length, hmc_iterations, 
                 gradient, intermediate_steps, parametrization=None, protocol=None,
                 mass_matrix=None, integrator=FastLeapFrog):

        super(HMCStepRENSSwapParameterInfo, self).__init__(sampler1, sampler2, protocol)

        self._mass_matrix = None
        self.mass_matrix = mass_matrix
        if self.mass_matrix is None:
            d = len(sampler1.state.position)
            self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d))

        self._hmc_traj_length = None
        self.hmc_traj_length = hmc_traj_length
        self._gradient = None
        self.gradient = gradient
        self._timestep = None
        self.timestep = timestep
        self._hmc_iterations = None
        self.hmc_iterations = hmc_iterations
        self._intermediate_steps = None
        self.intermediate_steps = intermediate_steps
        self._integrator = None
        self.integrator = integrator
    
    @property
    def timestep(self):
        """
        Integration timestep.
        """
        return self._timestep
    @timestep.setter
    def timestep(self, value):
        self._timestep = float(value)

    @property
    def hmc_traj_length(self):
        """
        HMC trajectory length in number of integration steps.
        """
        return self._hmc_traj_length
    @hmc_traj_length.setter
    def hmc_traj_length(self, value):
        self._hmc_traj_length = int(value)

    @property
    def gradient(self):
        """
        Gradient which governs the equations of motion.
        """
        return self._gradient
    @gradient.setter
    def gradient(self, value):
        self._gradient = value

    @property
    def mass_matrix(self):
        return self._mass_matrix
    @mass_matrix.setter
    def mass_matrix(self, value):
        self._mass_matrix = value

    @property
    def hmc_iterations(self):
        return self._hmc_iterations
    @hmc_iterations.setter
    def hmc_iterations(self, value):
        self._hmc_iterations = value

    @property
    def intermediate_steps(self):
        return self._intermediate_steps
    @intermediate_steps.setter
    def intermediate_steps(self, value):
        self._intermediate_steps = value

    @property
    def integrator(self):
        return self._integrator
    @integrator.setter
    def integrator(self, value):
        self._integrator = value


class AbstractSwapScheme(object):
    """
    Provides the interface for classes defining schemes according to which swaps in
    Replica Exchange-like simulations are performed.

    @param algorithm: Exchange algorithm that performs the swaps
    @type algorithm: L{AbstractExchangeMC}
    """

    __metaclass__ = ABCMeta

    def __init__(self, algorithm):

        self._algorithm = algorithm

    @abstractmethod
    def swap_all(self):
        """
        Advises the Replica Exchange-like algorithm to perform swaps according to
        the schedule defined here.
        """
        
        pass


class AlternatingAdjacentSwapScheme(AbstractSwapScheme):
    """
    Provides a swapping scheme in which tries exchanges between neighbours only
    following the scheme 1 <-> 2, 3 <-> 4, ... and after a sampling period 2 <-> 3, 4 <-> 5, ...

    @param algorithm: Exchange algorithm that performs the swaps
    @type algorithm: L{AbstractExchangeMC}
    """

    def __init__(self, algorithm):

        super(AlternatingAdjacentSwapScheme, self).__init__(algorithm)
        
        self._current_swap_list = None
        self._swap_list1 = []
        self._swap_list2 = []
        self._create_swap_lists()
    
    def _create_swap_lists(self):

        if len(self._algorithm.param_infos) == 1:
            self._swap_list1.append(0)
            self._swap_list2.append(0)
        else:
            i = 0
            while i < len(self._algorithm.param_infos):
                self._swap_list1.append(i)
                i += 2
                
            i = 1
            while i < len(self._algorithm.param_infos):
                self._swap_list2.append(i)
                i += 2

        self._current_swap_list = self._swap_list1

    def swap_all(self):
        
        for x in self._current_swap_list:
            self._algorithm.swap(x)

        if self._current_swap_list == self._swap_list1:
            self._current_swap_list = self._swap_list2
        else:
            self._current_swap_list = self._swap_list1


class SingleSwapStatistics(object):
    """
    Tracks swap statistics of a single sampler pair.

    @param param_info: ParameterInfo instance holding swap parameters
    @type param_info: L{AbstractSwapParameterInfo}
    """
    
    def __init__(self, param_info):
        self._total_swaps = 0
        self._accepted_swaps = 0

    @property
    def total_swaps(self):
        return self._total_swaps
    
    @property
    def accepted_swaps(self):
        return self._accepted_swaps
    
    @property
    def acceptance_rate(self):
        """
        Acceptance rate of the sampler pair.
        """
        if self.total_swaps > 0:
            return float(self.accepted_swaps) / float(self.total_swaps)
        else:
            return 0.

    def update(self, accepted):
        """
        Updates swap statistics.
        """        
        self._total_swaps += 1
        self._accepted_swaps += int(accepted)


class SwapStatistics(object):
    """
    Tracks swap statistics for an AbstractExchangeMC subclass instance.

    @param param_infos: list of ParameterInfo instances providing information
                        needed for performing swaps
    @type param_infos: list of L{AbstractSwapParameterInfo}
    """
    
    def __init__(self, param_infos):        
        self._stats = [SingleSwapStatistics(x) for x in param_infos]
        
    @property
    def stats(self):
        return tuple(self._stats)

    @property
    def acceptance_rates(self):
        """
        Returns acceptance rates for all swaps.
        """        
        return [x.acceptance_rate for x in self._stats]
        

class InterpolationFactory(object):
    """
    Produces interpolations for functions changed during non-equilibrium
    trajectories.
    
    @param protocol: protocol to be used to generate non-equilibrium trajectories
    @type protocol: function mapping t to [0...1] for fixed tau
    @param tau: switching time
    @type tau: float    
    """
    
    def __init__(self, protocol, tau):
        
        self._protocol = None
        self._tau = None
        
        self.protocol = protocol
        self.tau = tau
        
    @property
    def protocol(self):
        return self._protocol
    @protocol.setter
    def protocol(self, value):
        if not hasattr(value, '__call__'):
            raise TypeError(value)
        self._protocol = value
                    
    @property
    def tau(self):
        return self._tau
    @tau.setter
    def tau(self, value):
        self._tau = float(value)
        
    def build_gradient(self, gradient):
        """
        Create a gradient instance with according to given protocol
        and switching time.
        
        @param gradient: gradient with G(0) = G_1 and G(1) = G_2
        @type gradient: callable    
        """
        return Gradient(gradient, self._protocol, self._tau)
    
    def build_temperature(self, temperature):
        """
        Create a temperature function according to given protocol and
        switching time.

        @param temperature: temperature with T(0) = T_1 and T(1) = T_2
        @type temperature: callable        
        """
        return lambda t: temperature(self.protocol(t, self.tau))
        

class Gradient(AbstractGradient):
    
    def __init__(self, gradient, protocol, tau):
        
        self._protocol = protocol
        self._gradient = gradient
        self._tau = tau
    
    def evaluate(self, q, t):
        return self._gradient(q, self._protocol(t, self._tau))


class ReplicaHistory(object):
    '''
    Replica history object, works with both RE and RENS for
    the AlternatingAdjacentSwapScheme.

    @param samples: list holding ensemble states
    @type samples: list

    @param swap_interval: interval with which swaps were attempted, e.g.,
                          5 means that every 5th regular MC step is replaced
                          by a swap
    @type swap_interval: int

    @param first_swap: sample index of the first sample generated by a swap attempt.
                       If None, the first RE sampled is assumed to have sample index
                       swap_interval. If specified, it has to be greater than zero
    @type first_swap: int
    '''
    
    def __init__(self, samples, swap_interval, first_swap=None):
        self.samples = samples
        self.swap_interval = swap_interval
        if first_swap == None:
            self.first_swap = swap_interval - 1
        elif first_swap > 0:
            self.first_swap = first_swap - 1
        else:
            raise(ValueError("Sample index of first swap has to be greater than zero!"))
        self.n_replicas = len(samples[0])

    @staticmethod
    def _change_direction(x):
        if x == 1:
            return -1
        if x == -1:
            return 1

    def calculate_history(self, start_ensemble):
        '''
        Calculates the replica history of the first state of ensemble #start_ensemble.

        @param start_ensemble: index of the ensemble to start at, zero-indexed
        @type start_ensemble: int

        @return: replica history as a list of ensemble indices
        @rtype: list of ints
        '''
        
        sample_counter = 0

        # determine the direction (up = 1, down = -1) in the "temperature ladder" of
        # the first swap attempt. Remember: first swap series is always 0 <-> 1, 2 <-> 3, ...
        if start_ensemble % 2 == 0:
            direction = +1
        else:
            direction = -1

        # if number of replicas is not even and the start ensemble is the highest-temperature-
        # ensemble, the first swap will be attempted "downwards"
        if start_ensemble % 2 == 0 and start_ensemble == self.n_replicas - 1:
            direction = -1

        # will store the indices of the ensembles the state will visit in chronological order
        history = []

        # the ensemble the state is currently in
        ens = start_ensemble
        
        while sample_counter < len(self.samples):
            if self.n_replicas == 2:
                if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
                   sample_counter >= self.first_swap:
                    ## swap attempt: determine whether it was successfull or not
                    # state after swap attempt
                    current_sample = self.samples[sample_counter][ens]

                    # state before swap attempt
                    previous_sample = self.samples[sample_counter - 1][history[-1]]

                    # swap was accepted when position of the current state doesn't equal
                    # the position of the state before the swap attempt, that is, the last 
                    # state in the history
                    swap_accepted = not numpy.all(current_sample.position == 
                                                  previous_sample.position)

                    if swap_accepted:
                        if ens == 0:
                            ens = 1
                        else:
                            ens = 0
                    history.append(ens)                    
                else:
                    history.append(ens)

            else:
                
                if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
                   sample_counter >= self.first_swap:
                    # state after swap attempt
                    current_sample = self.samples[sample_counter][ens]

                    # state before swap attempt
                    previous_sample = self.samples[sample_counter - 1][ens]

                    # swap was accepted when position of the current state doesn't equal
                    # the position of the state before the swap attempt, that is, the last 
                    # state in the history
                    swap_accepted = not numpy.all(current_sample.position == previous_sample.position)

                    if swap_accepted:
                        ens += direction
                    else:
                        if ens == self.n_replicas - 1:
                            # if at the top of the ladder, go downwards again
                            direction = -1
                        elif ens == 0:
                            # if at the bottom of the ladder, go upwards
                            direction = +1
                        else:
                            # in between, reverse the direction of the trajectory
                            # in temperature space
                            direction = self._change_direction(direction)

                history.append(ens)

            sample_counter += 1

        return history

    def calculate_projected_trajectories(self, ensemble):
        '''
        Calculates sequentially correlated trajectories projected on a specific ensemble.

        @param ensemble: ensemble index of ensemble of interest, zero-indexed
        @type ensemble: int

        @return: list of Trajectory objects containg sequentially correlated trajectories
        @rtype: list of L{Trajectory} objects.
        '''
        
        trajectories = []

        for i in range(self.n_replicas):
            history = self.calculate_history(i)
            traj = [x[ensemble] for k, x in enumerate(self.samples) if history[k] == ensemble]
            trajectories.append(Trajectory(traj))

        return trajectories

    def calculate_trajectories(self):
        '''
        Calculates sequentially correlated trajectories.

        @return: list of Trajectory objects containg sequentially correlated trajectories
        @rtype: list of L{Trajectory} objects.
        '''
        
        trajectories = []

        for i in range(self.n_replicas):
            history = self.calculate_history(i)
            traj = [x[history[k]] for k, x in enumerate(self.samples)]
            trajectories.append(Trajectory(traj))

        return trajectories