sampling 0.3.1

use {
    crate::{rewl::log_density_to_log10_density, *},
    rand::Rng,
    std::{marker::PhantomData, mem::*, num::*, sync::*},
};

#[cfg(feature = "sweep_time_optimization")]
use std::time::*;

#[cfg(feature = "serde_support")]
use serde::{Deserialize, Serialize};

/// # Walker for Replica exchange entropic sampling
/// * performes the random walk in its respective domain
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
pub struct ReesWalker<R, Hist, Energy, S, Res> {
    id: usize,
    sweep_size: NonZeroUsize,
    pub(crate) rng: R,
    hist: Hist,
    log_density: Vec<f64>,
    step_count: u64,
    re: u64,
    proposed_re: u64,
    rejected_markov_steps: u64,
    old_energy: Energy,
    bin: usize,
    markov_steps: Vec<S>,
    marker_res: PhantomData<Res>,
    step_size: usize,
    step_threshold: u64,
    #[cfg(feature = "sweep_time_optimization")]
    duration: Duration,
}

impl<R, Hist, Energy, S, Res> From<RewlWalker<R, Hist, Energy, S, Res>>
    for ReesWalker<R, Hist, Energy, S, Res>
where
    Hist: Histogram,
{
    fn from(mut rewl_walker: RewlWalker<R, Hist, Energy, S, Res>) -> Self {
        rewl_walker.hist.reset();
        Self {
            id: rewl_walker.id,
            sweep_size: rewl_walker.sweep_size,
            markov_steps: rewl_walker.markov_steps,
            step_size: rewl_walker.step_size,
            marker_res: PhantomData::<Res>,
            rng: rewl_walker.rng,
            log_density: rewl_walker.log_density,
            bin: rewl_walker.bin,
            hist: rewl_walker.hist,
            old_energy: rewl_walker.old_energy,
            step_count: 0,
            step_threshold: rewl_walker.step_count,
            re: 0,
            proposed_re: 0,
            rejected_markov_steps: 0,
            #[cfg(feature = "sweep_time_optimization")]
            duration: rewl_walker.duration,
        }
    }
}

impl<R, Hist, Energy, S, Res> ReesWalker<R, Hist, Energy, S, Res> {
    /// # Returns id of walker
    /// * important for mapping the ensemble to the walker
    #[inline(always)]
    pub fn id(&self) -> usize {
        self.id
    }

    /// # Returns duration of last sweep that was performed
    #[cfg(feature = "sweep_time_optimization")]
    pub fn duration(&self) -> Duration {
        self.duration
    }

    /// Returns reference of current energy
    #[inline(always)]
    pub fn energy(&self) -> &Energy {
        &self.old_energy
    }

    /// Returns current energy
    #[inline(always)]
    pub fn energy_copy(&self) -> Energy
    where
        Energy: Copy,
    {
        self.old_energy
    }

    /// # Reference to internal histogram
    #[inline(always)]
    pub fn hist(&self) -> &Hist {
        &self.hist
    }

    /// # how many steps per sweep
    #[inline(always)]
    pub fn sweep_size(&self) -> NonZeroUsize {
        self.sweep_size
    }

    /// # change how many steps per sweep are performed
    pub fn sweep_size_change(&mut self, sweep_size: NonZeroUsize) {
        self.sweep_size = sweep_size;
    }

    /// # step size for markov steps
    #[inline(always)]
    pub fn step_size(&self) -> usize {
        self.step_size
    }

    /// # Change step sitze for markov steps
    pub fn step_size_change(&mut self, step_size: usize) {
        self.step_size = step_size;
    }

    /// # How many entropic steps were performed until now?
    #[inline(always)]
    pub fn step_count(&self) -> u64 {
        self.step_count
    }

    /// # How many successful replica exchanges were performed until now?
    #[inline(always)]
    pub fn replica_exchanges(&self) -> u64 {
        self.re
    }

    /// # How many replica exchanges were proposed until now?
    #[inline(always)]
    pub fn proposed_replica_exchanges(&self) -> u64 {
        self.proposed_re
    }

    /// fraction of how many replica exchanges were accepted and how many were proposed
    #[inline(always)]
    pub fn replica_exchange_frac(&self) -> f64 {
        self.re as f64 / self.proposed_re as f64
    }

    /// # How many markov steps were rejected until now
    #[inline(always)]
    pub fn rejected_markov_steps(&self) -> u64 {
        self.rejected_markov_steps
    }

    /// # rate/fraction of acceptance
    pub fn acceptance_rate_markov(&self) -> f64 {
        let rej = self.rejected_markov_steps() as f64 / self.step_count() as f64;
        1.0 - rej
    }

    /// * Old non normalized estimate of the natural logarithm of the probability density function
    /// * for refined density use `self.log_density_refined()`
    #[inline(always)]
    pub fn log_density(&self) -> &[f64] {
        &self.log_density
    }

    /// * Current non normalized estimate of the natural logarithm of the probability density function
    /// * calculated by refining old density with current histogram
    ///
    /// # How does the refining work?
    /// * Let P(i) be the current probability density function (non normalized) at some index i
    /// * Let H(i) be the histogram at some index i
    ///
    /// We will now calculate the refined density P', which is calculated as follows:
    ///
    /// P'(i) = P(i) * H(i) (if H(i) != 0)
    ///
    /// P'(i) = P(i) (if H(i) == 0)
    ///
    /// Or in log space, which is what is actually calculated here:
    ///
    /// ln(P'(i)) = ln(P(i)) + ln(H(i)) (if H(i) != 0)
    ///
    /// ln(P'(i)) = ln(P(i)) (if H(i)=0)
    ///
    /// # for more information see
    /// > J. Lee,
    /// > “New Monte Carlo algorithm: Entropic sampling,”
    /// > Phys. Rev. Lett. 71, 211–214 (1993),
    /// > DOI: [10.1103/PhysRevLett.71.211](https://doi.org/10.1103/PhysRevLett.71.211)
    pub fn log_density_refined(&self) -> Vec<f64>
    where
        Hist: Histogram,
    {
        let mut refined_log_density = Vec::with_capacity(self.log_density.len());

        refined_log_density.extend(self.log_density.iter().zip(self.hist.hist().iter()).map(
            |(&log_d, &h)| {
                if h == 0 {
                    log_d
                } else {
                    log_d + (h as f64).ln()
                }
            },
        ));
        refined_log_density
    }

    /// # Old estimate of log10 of probability density
    /// * normalized (sum over non log values is 1 (within numerical precision))
    pub fn log10_density(&self) -> Vec<f64> {
        log_density_to_log10_density(self.log_density())
    }

    /// # Current estimate of log10 of probability density
    /// * normalized (sum over non log values is 1 (within numerical precision))
    pub fn log10_density_refined(&self) -> Vec<f64>
    where
        Hist: Histogram,
    {
        let density = self.log_density_refined();
        log_density_to_log10_density(&density)
    }

    /// # is the simulation finished?
    /// * true, if more (or equal) steps than the step threshold are performed
    #[inline(always)]
    pub fn is_finished(&self) -> bool {
        self.step_count >= self.step_threshold
    }

    /// # Return step threshold
    #[inline(always)]
    pub fn step_threshold(&self) -> u64 {
        self.step_threshold
    }

    /// # Refine current probability density estimate
    /// * refines the current probability estimate by setting it to [self.log_density_refined](`Self::log_density_refined`)
    pub fn refine(&mut self)
    where
        Hist: Histogram,
    {
        let refined = self.log_density_refined();
        self.log_density = refined;
        self.hist.reset();
        self.step_count = 0;
    }

    #[inline(always)]
    fn count_rejected(&mut self) {
        self.rejected_markov_steps += 1;
    }
}

impl<R, Hist, Energy, S, Res> ReesWalker<R, Hist, Energy, S, Res>
where
    Hist: HistogramVal<Energy>,
    R: Rng,
{
    pub(crate) fn check_energy_fn<F, Ensemble>(
        &self,
        ensemble_vec: &[RwLock<Ensemble>],
        energy_fn: F,
    ) -> bool
    where
        Energy: PartialEq,
        F: Fn(&mut Ensemble) -> Option<Energy>,
    {
        let mut e = ensemble_vec[self.id].write().expect(
            "Fatal Error encountered; ERRORCODE 0x1 - this should be \
                impossible to reach. If you are using the latest version of the \
                'sampling' library, please contact the library author via github by opening an \
                issue! https://github.com/Pardoxa/sampling/issues",
        );

        let energy = match energy_fn(&mut e) {
            Some(energy) => energy,
            None => {
                return false;
            }
        };
        energy == self.old_energy
    }

    pub(crate) fn sweep<Ensemble, F, Extra, P>(
        &mut self,
        ensemble_vec: &[RwLock<Ensemble>],
        extra: &mut Extra,
        extra_fn: P,
        energy_fn: F,
    ) where
        F: Fn(&mut Ensemble) -> Option<Energy>,
        P: Fn(&Self, &mut Ensemble, &mut Extra),
        Ensemble: MarkovChain<S, Res>,
        Hist: Histogram,
    {
        #[cfg(feature = "sweep_time_optimization")]
        let start = Instant::now();

        let mut e = ensemble_vec[self.id].write().expect(
            "Fatal Error encountered; ERRORCODE 0x3 - this should be \
                impossible to reach. If you are using the latest version of the \
                'sampling' library, please contact the library author via github by opening an \
                issue! https://github.com/Pardoxa/sampling/issues",
        );

        for _ in 0..self.sweep_size.get() {
            self.step_count += 1;
            e.m_steps(self.step_size, &mut self.markov_steps);

            let energy = match energy_fn(&mut e) {
                Some(energy) => energy,
                None => {
                    self.count_rejected();
                    e.undo_steps_quiet(&self.markov_steps);
                    continue;
                }
            };

            match self.hist.get_bin_index(&energy) {
                Ok(current_bin) => {
                    // metropolis hastings
                    let acception_prob =
                        (self.log_density[self.bin] - self.log_density[current_bin]).exp();
                    if self.rng.random::<f64>() > acception_prob {
                        self.count_rejected();
                        e.steps_rejected(&self.markov_steps);
                        e.undo_steps_quiet(&self.markov_steps);
                    } else {
                        self.old_energy = energy;
                        self.bin = current_bin;
                        e.steps_accepted(&self.markov_steps);
                    }
                }
                _ => {
                    self.count_rejected();
                    e.steps_rejected(&self.markov_steps);
                    e.undo_steps_quiet(&self.markov_steps);
                }
            }

            self.hist
                .increment_index(self.bin)
                .expect("Histogram index Error, ERRORCODE 0x4");

            extra_fn(self, &mut e, extra);
        }

        #[cfg(feature = "sweep_time_optimization")]
        {
            self.duration = start.elapsed();
        }
    }
}

pub(crate) fn replica_exchange<R, Hist, Energy, S, Res>(
    walker_a: &mut ReesWalker<R, Hist, Energy, S, Res>,
    walker_b: &mut ReesWalker<R, Hist, Energy, S, Res>,
) where
    Hist: HistogramVal<Energy> + Histogram,
    R: Rng,
{
    walker_a.proposed_re += 1;
    walker_b.proposed_re += 1;
    // check if exchange is even possible
    let new_bin_a = match walker_a.hist.get_bin_index(&walker_b.old_energy) {
        Ok(bin) => bin,
        _ => return,
    };

    let new_bin_b = match walker_b.hist.get_bin_index(&walker_a.old_energy) {
        Ok(bin) => bin,
        _ => return,
    };

    // see paper equation 1
    let log_gi_x = walker_a.log_density[walker_a.bin];
    let log_gi_y = walker_a.log_density[new_bin_a];

    let log_gj_y = walker_b.log_density[walker_b.bin];
    let log_gj_x = walker_b.log_density[new_bin_b];

    let log_prob = log_gi_x + log_gj_y - log_gi_y - log_gj_x;

    let prob = log_prob.exp();

    // if exchange is accepted
    if walker_b.rng.random::<f64>() < prob {
        swap(&mut walker_b.id, &mut walker_a.id);
        swap(&mut walker_b.old_energy, &mut walker_a.old_energy);
        walker_b.bin = new_bin_b;
        walker_a.bin = new_bin_a;
        walker_a.re += 1;
        walker_b.re += 1;
    }

    walker_a.hist.increment_index(walker_a.bin).unwrap();
    walker_b.hist.increment_index(walker_b.bin).unwrap();
}

pub(crate) fn get_merged_refined_walker_prob<R, Hist, Energy, S, Res>(
    walker: &[ReesWalker<R, Hist, Energy, S, Res>],
) -> Vec<f64>
where
    Hist: Histogram,
{
    let log_len = walker[0].log_density.len();
    debug_assert!(walker.iter().all(|w| w.log_density.len() == log_len));

    let mut averaged_log_density = walker[0].log_density_refined();

    if walker.len() > 1 {
        walker[1..]
            .iter()
            .map(|w| w.log_density_refined())
            .for_each(|log_density| {
                averaged_log_density
                    .iter_mut()
                    .zip(log_density)
                    .for_each(|(average, other)| {
                        *average += other;
                    })
            });

        let number_of_walkers = walker.len() as f64;
        averaged_log_density
            .iter_mut()
            .for_each(|average| *average /= number_of_walkers);
    }

    averaged_log_density
}