#!/usr/local/bin/env python

# ==============================================================================
# MODULE DOCSTRING
# ==============================================================================

"""
ParallelTemperingSampler
========================

Derived multi-thermodynamic state multistate class with exchanging configurations between replicas of different
temperatures. This is a special case which accepts a single thermodynamic_state and different temperatures to sample.
If you want different temperatures and Hamiltonians, use ReplicaExchangeSampler with temperatures pre-set.

COPYRIGHT

Current version by Andrea Rizzi <andrea.rizzi@choderalab.org>, Levi N. Naden <levi.naden@choderalab.org> and
John D. Chodera <john.chodera@choderalab.org> while at Memorial Sloan Kettering Cancer Center.

Original version by John D. Chodera <jchodera@gmail.com> while at the University of
California Berkeley.

LICENSE

This code is licensed under the latest available version of the MIT License.

"""

import copy
import math
import logging
import numpy as np
import openmmtools as mmtools

from .replicaexchange import ReplicaExchangeSampler, ReplicaExchangeAnalyzer
from .multistatereporter import MultiStateReporter

logger = logging.getLogger(__name__)


# ==============================================================================
# PARALLEL TEMPERING
# ==============================================================================

class ParallelTemperingSampler(ReplicaExchangeSampler):
    """Parallel tempering simulation facility.

    This class provides a facility for parallel tempering simulations. It
    is a subclass of :class:`ReplicaExchange`, but provides efficiency improvements
    for parallel tempering simulations, so should be preferred for this type
    of simulation. In particular, this makes use of the fact that the reduced
    potentials are linear in inverse temperature.

    Examples
    --------

    Create the system.

    >>> from simtk import unit
    >>> from openmmtools import testsystems, states, mcmc
    >>> import tempfile
    >>> testsystem = testsystems.AlanineDipeptideImplicit()

    Create thermodynamic states for parallel tempering with exponentially-spaced schedule.

    >>> n_replicas = 3  # Number of temperature replicas.
    >>> T_min = 298.0 * unit.kelvin  # Minimum temperature.
    >>> T_max = 600.0 * unit.kelvin  # Maximum temperature.
    >>> reference_state = states.ThermodynamicState(system=testsystem.system, temperature=T_min)

    Initialize simulation object with options. Run with a GHMC integrator.

    >>> move = mcmc.GHMCMove(timestep=2.0*unit.femtoseconds, n_steps=50)
    >>> simulation = ParallelTemperingSampler(mcmc_moves=move, number_of_iterations=2)

    Create simulation with its storage file (in a temporary directory) and run.

    >>> storage_path = tempfile.NamedTemporaryFile(delete=False).name + '.nc'
    >>> reporter = MultiStateReporter(storage_path, checkpoint_interval=10)
    >>> simulation.create(reference_state,
    ...                   states.SamplerState(testsystem.positions),
    ...                   reporter, min_temperature=T_min,
    ...                   max_temperature=T_max, n_temperatures=n_replicas)
    >>> simulation.run(n_iterations=1)

    Clean up.

    >>> os.remove(storage_path)


    See Also
    --------
    MultiStateSampler
    ReplicaExchangeSampler

    """

    _TITLE_TEMPLATE = ('Parallel tempering simulation created using ParallelTempering '
                       'class of yank.multistate on {}')

    def create(self, thermodynamic_state, sampler_states: list, storage,
               min_temperature=None, max_temperature=None, n_temperatures=None,
               temperatures=None, **kwargs):
        """Initialize a parallel tempering simulation object.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            Reference thermodynamic state that will be simulated at the given
            temperatures.

            WARNING: This is a SINGLE state, not a list of states!

        sampler_states : openmmtools.states.SamplerState or list
            One or more sets of initial sampler states. If a list of SamplerStates,
            they will be assigned to replicas in a round-robin fashion.
        storage : str or Reporter
            If str: path to the storage file, checkpoint options are default
            If Reporter: Instanced :class:`Reporter` class, checkpoint information is read from
            In the future this will be able to take a Storage class as well.
        min_temperature : simtk.unit.Quantity, optional
           Minimum temperature (units of temperature, default is None).
        max_temperature : simtk.unit.Quantity, optional
           Maximum temperature (units of temperature, default is None).
        n_temperatures : int, optional
           Number of exponentially-spaced temperatures between ``min_temperature``
           and ``max_temperature`` (default is None).
        temperatures : list of simtk.unit.Quantity, optional
           If specified, this list of temperatures will be used instead of
           ``min_temperature``, ``max_temperature``, and ``n_temperatures`` (units of temperature,
           default is None).
        metadata : dict, optional
           Simulation metadata to be stored in the file.

        Notes
        -----
        Either (``min_temperature``, ``max_temperature``, ``n_temperatures``) must all be
        specified or the list of '`temperatures`' must be specified.

        """
        # Create thermodynamic states from temperatures.
        if not isinstance(thermodynamic_state, mmtools.states.ThermodynamicState):
            raise ValueError("ParallelTempering only accepts a single ThermodynamicState!\n"
                             "If you have already set temperatures in your list of states, please use the "
                             "standard ReplicaExchange class with your list of states.")
        if temperatures is not None:
            logger.debug("Using provided temperatures")
        elif min_temperature is not None and max_temperature is not None and n_temperatures is not None:
            temperatures = [min_temperature + (max_temperature - min_temperature) *
                            (math.exp(i / n_temperatures-1) - 1.0) / (math.e - 1.0)
                            for i in range(n_temperatures)]  # Python 3 uses true division for /
            logger.debug('using temperatures {}'.format(temperatures))
        else:
            raise ValueError("Either 'temperatures' or ('min_temperature', 'max_temperature', "
                             "and 'n_temperatures') must be provided.")

        thermodynamic_states = [copy.deepcopy(thermodynamic_state) for _ in range(n_temperatures)]
        for state, temperature in zip(thermodynamic_states, temperatures):
            state.temperature = temperature

        # Initialize replica-exchange simulation.
        super(ParallelTemperingSampler, self).create(thermodynamic_states, sampler_states, storage=storage, **kwargs)

    def _compute_replica_energies(self, replica_id):
        """Compute the energy for the replica at every temperature.

        Because only the temperatures differ among replicas, we replace the generic O(N^2)
        replica-exchange implementation with an O(N) implementation.

        """
        # Initialize replica energies for each thermodynamic state.
        energy_thermodynamic_states = np.zeros(self.n_states)
        energy_unsampled_states = np.zeros(len(self._unsampled_states))

        # Determine neighborhood
        state_index = self._replica_thermodynamic_states[replica_id]
        neighborhood = self._neighborhood(state_index)
        # Only compute energies over neighborhoods
        energy_neighborhood_states = energy_thermodynamic_states[neighborhood]  # Array, can be indexed like this
        neighborhood_thermodynamic_states = [self._thermodynamic_states[n] for n in neighborhood]  # List

        # Retrieve sampler states associated to this replica.
        sampler_state = self._sampler_states[replica_id]

        # Thermodynamic state differ only by temperatures.
        reference_thermodynamic_state = self._thermodynamic_states[0]

        # Get the context, any Integrator works.
        context, integrator = mmtools.cache.global_context_cache.get_context(reference_thermodynamic_state)

        # Update positions and box vectors.
        sampler_state.apply_to_context(context)

        # Compute energy.
        reference_reduced_potential = reference_thermodynamic_state.reduced_potential(context)

        # Strip reference potential of reference state's beta.
        reference_beta = 1.0 / (mmtools.constants.kB * reference_thermodynamic_state.temperature)
        reference_reduced_potential /= reference_beta

        # Update potential energy by temperature.
        for thermodynamic_state_id, thermodynamic_state in enumerate(neighborhood_thermodynamic_states):
            beta = 1.0 / (mmtools.constants.kB * thermodynamic_state.temperature)
            energy_neighborhood_states[thermodynamic_state_id] = beta * reference_reduced_potential

        # Since no assumptions can be made about the unsampled thermodynamic states, do it the hard way
        for unsampled_id, state in enumerate(self._unsampled_states):
            if unsampled_id == 0 or not state.is_state_compatible(context_state):
                context_state = state

                # Get the context, any Integrator works.
                context, integrator = mmtools.cache.global_context_cache.get_context(state)

                # Update positions and box vectors. We don't need
                # to set Context velocities for the potential.
                sampler_state.apply_to_context(context, ignore_velocities=True)
            else:
                # If this state is compatible with the context, just fix the
                # thermodynamic state as positions/box vectors are the same.
                state.apply_to_context(context)

            # Compute energy.
            energy_unsampled_states[unsampled_id] = state.reduced_potential(context)

        # Return the new energies.
        return energy_neighborhood_states, energy_unsampled_states


class ParallelTemperingAnalyzer(ReplicaExchangeAnalyzer):
    """
    The ParallelTemperingAnalyzer is the analyzer for a simulation generated from a Parallel Tempering sampler
    simulation, implemented as an instance of the :class:`ReplicaExchangeAnalyzer` as the sampler is a subclass of
    the :class:`yank.multistate.ReplicaExchangeSampler`

    See Also
    --------
    PhaseAnalyzer
    ReplicaExchangeAnalyzer

    """
    pass

# ==============================================================================
# MAIN AND TESTS
# ==============================================================================

if __name__ == "__main__":
    import doctest
    doctest.testmod()
