spin-sim 0.3.0

use crate::config::PtSchedule;
use crate::geometry::Lattice;
use crate::spins::model::{Real, Spin};
use rand::SeedableRng;
use rand_xoshiro::Xoshiro256StarStar;

const SYSTEM_SEED_DOMAIN: u64 = 0x53A9_17E1_4C2D_8B6F;
const PAIR_SEED_DOMAIN: u64 = 0xA147_5EED_91C3_02D7;

fn splitmix64(mut value: u64) -> u64 {
    value = value.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut mixed = value;
    mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    mixed ^ (mixed >> 31)
}

fn child_seed(root: u64, domain: u64, index: usize) -> u64 {
    splitmix64(root ^ domain ^ splitmix64(index as u64))
}

/// Position in the interval schedule, carried across sampling calls so that an
/// event due every `k` sweeps stays every `k` sweeps however the run is split.
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct ScheduleState {
    pub sweeps: usize,
    /// Overlap events so far, which selects the next of several build modes.
    pub overlap_calls: usize,
}

pub(super) struct PtState {
    edge_attempts: Vec<u64>,
    edge_acceptances: Vec<u64>,
    round_trips: Vec<u64>,
    trip_state: Vec<u8>,
    /// The end of the ladder each system visited last: `NO_END` until it reaches one.
    last_end: Vec<u8>,
    /// Per temperature slot, visits by systems that last visited the cold end.
    flow_up: Vec<u64>,
    /// Per temperature slot, visits by systems that have visited either end.
    flow_labelled: Vec<u64>,
    cold_slot: usize,
    hot_slot: usize,
}

const NO_END: u8 = 0;
const COLD_END: u8 = 1;
const HOT_END: u8 = 2;

impl PtState {
    fn new<T: PartialOrd>(
        n_replicas: usize,
        n_temps: usize,
        system_ids: &[usize],
        temperatures: &[T],
    ) -> Self {
        let mut state = Self {
            edge_attempts: Vec::new(),
            edge_acceptances: Vec::new(),
            round_trips: Vec::new(),
            trip_state: Vec::new(),
            last_end: Vec::new(),
            flow_up: Vec::new(),
            flow_labelled: Vec::new(),
            cold_slot: 0,
            hot_slot: 0,
        };
        state.reset(n_replicas, n_temps, system_ids, temperatures);
        state
    }

    fn reset<T: PartialOrd>(
        &mut self,
        n_replicas: usize,
        n_temps: usize,
        system_ids: &[usize],
        temperatures: &[T],
    ) {
        let n_edges = n_temps.saturating_sub(1);
        self.edge_attempts.resize(n_edges, 0);
        self.edge_attempts.fill(0);
        self.edge_acceptances.resize(n_edges, 0);
        self.edge_acceptances.fill(0);
        self.round_trips.resize(n_replicas * n_temps, 0);
        self.round_trips.fill(0);
        self.trip_state.resize(n_replicas * n_temps, 0);
        self.trip_state.fill(0);
        self.last_end.resize(n_replicas * n_temps, NO_END);
        self.last_end.fill(NO_END);
        self.flow_up.resize(n_temps, 0);
        self.flow_up.fill(0);
        self.flow_labelled.resize(n_temps, 0);
        self.flow_labelled.fill(0);
        (self.cold_slot, self.hot_slot) = Self::extreme_temperature_slots(temperatures, n_temps);

        if n_temps == 0 {
            return;
        }
        for replica in 0..n_replicas {
            let row = replica * n_temps;
            self.last_end[system_ids[row + self.cold_slot]] = COLD_END;
            self.last_end[system_ids[row + self.hot_slot]] = HOT_END;
            self.trip_state[system_ids[row + self.hot_slot]] = 1;
        }
    }

    pub(super) fn record_attempt(&mut self, attempt: crate::mcmc::tempering::TemperingAttempt) {
        self.edge_attempts[attempt.edge] += 1;
        if !attempt.accepted {
            return;
        }
        self.edge_acceptances[attempt.edge] += 1;

        self.record_arrival(attempt.left_system, attempt.edge + 1);
        self.record_arrival(attempt.right_system, attempt.edge);
    }

    fn extreme_temperature_slots<T: PartialOrd>(
        temperatures: &[T],
        n_temps: usize,
    ) -> (usize, usize) {
        if n_temps == 0 {
            return (0, 0);
        }
        let mut cold_slot = 0;
        let mut hot_slot = 0;
        for slot in 1..n_temps {
            if temperatures[slot] < temperatures[cold_slot] {
                cold_slot = slot;
            }
            if temperatures[slot] > temperatures[hot_slot] {
                hot_slot = slot;
            }
        }
        (cold_slot, hot_slot)
    }

    /// Counts the systems of one replica row by the end they last visited. The
    /// fraction `flow_up / flow_labelled` falls from 1 at the cold end to 0 at the
    /// hot end, linearly along a well-spaced ladder and steeply across a bottleneck
    /// (Katzgraber, Trebst, Huse & Troyer, J. Stat. Mech. P03018 (2006)).
    fn record_visits(&mut self, row: &[usize]) {
        for (slot, &system) in row.iter().enumerate() {
            let end = self.last_end[system];
            self.flow_labelled[slot] += u64::from(end != NO_END);
            self.flow_up[slot] += u64::from(end == COLD_END);
        }
    }

    fn record_arrival(&mut self, system: usize, slot: usize) {
        if slot == self.cold_slot {
            self.last_end[system] = COLD_END;
        }
        if slot == self.hot_slot {
            self.last_end[system] = HOT_END;
        }
        if slot == self.hot_slot {
            if self.trip_state[system] == 2 {
                self.round_trips[system] += 1;
            }
            self.trip_state[system] = 1;
            return;
        }
        if slot == self.cold_slot && self.trip_state[system] == 1 {
            self.trip_state[system] = 2;
        }
    }
}

/// Mutable state for one disorder realization.
///
/// Holds the coupling array (fixed after construction), spin configurations for
/// every replica at every temperature, and bookkeeping for parallel tempering.
///
/// With `n_replicas` replicas and `n_temps` temperatures there are
/// `n_systems = n_replicas * n_temps` independent spin configurations.
/// Spins are stored in a single flat `Vec` of length `n_systems * n_spins`,
/// where system `i` occupies `spins[i*n_spins .. (i+1)*n_spins]`.
pub struct ModelRealization<S: Spin> {
    /// Forward couplings, length `n_spins * n_neighbors`.
    pub couplings: Vec<S::Value>,
    /// All spin configurations, length `n_systems * n_spins` (±1 for Ising, unit vectors for XY).
    pub spins: Vec<S>,
    /// Temperature assigned to each system slot, length `n_systems`.
    pub temperatures: Vec<S::Value>,
    /// Parallel-tempering permutation: `system_ids[slot]` is the system index
    /// currently occupying temperature slot `slot`.
    pub system_ids: Vec<usize>,
    /// One PRNG per system.
    pub rngs: Vec<Xoshiro256StarStar>,
    /// One PRNG per overlap-update pair slot, length `n_temps * (n_replicas / 2)`.
    pub pair_rngs: Vec<Xoshiro256StarStar>,
    pub(super) pt: PtState,
    pub(crate) schedule: ScheduleState,
    /// Scratch cache of the interaction energy per system (-H/N), length
    /// `n_systems`. It is current only right after a measurement or tempering step.
    pub energies: Vec<S::Value>,
}

pub type Realization = ModelRealization<i8>;
pub type XyRealization = ModelRealization<[f64; 2]>;

impl<S: Spin> ModelRealization<S> {
    /// Initialize a realization with random spins (±1, or uniform unit vectors for XY).
    ///
    /// Seeds independent system and pair streams from domain-separated child seeds.
    pub fn new(
        lattice: &Lattice,
        couplings: Vec<S::Value>,
        temps: &[S::Value],
        n_replicas: usize,
        base_seed: u64,
    ) -> Self {
        let n_spins = lattice.n_spins;
        let n_temps = temps.len();
        let n_systems = n_replicas * n_temps;

        let temperatures = temps.repeat(n_replicas);

        let mut rngs = Vec::with_capacity(n_systems);
        for i in 0..n_systems {
            rngs.push(Xoshiro256StarStar::seed_from_u64(child_seed(
                base_seed,
                SYSTEM_SEED_DOMAIN,
                i,
            )));
        }

        let mut spins = vec![S::default(); n_systems * n_spins];
        for (i, rng) in rngs.iter_mut().enumerate() {
            for j in 0..n_spins {
                spins[i * n_spins + j] = S::random(rng);
            }
        }

        let system_ids: Vec<usize> = (0..n_systems).collect();

        let n_pairs = n_replicas / 2;
        let mut pair_rngs = Vec::with_capacity(n_temps * n_pairs);
        for i in 0..n_temps * n_pairs {
            pair_rngs.push(Xoshiro256StarStar::seed_from_u64(child_seed(
                base_seed,
                PAIR_SEED_DOMAIN,
                i,
            )));
        }

        let mut energies = vec![S::Value::default(); n_systems];
        S::interactions(lattice, &spins, &couplings, &mut energies);

        let pt = PtState::new(n_replicas, n_temps, &system_ids, &temperatures);
        Self {
            couplings,
            spins,
            temperatures,
            system_ids,
            rngs,
            pair_rngs,
            pt,
            schedule: ScheduleState::default(),
            energies,
        }
    }

    /// Re-randomize all spins and reset the tempering permutation.
    pub fn reset(&mut self, lattice: &Lattice, n_replicas: usize, n_temps: usize, base_seed: u64) {
        let n_spins = lattice.n_spins;
        let n_systems = n_replicas * n_temps;

        for i in 0..n_systems {
            self.rngs[i] =
                Xoshiro256StarStar::seed_from_u64(child_seed(base_seed, SYSTEM_SEED_DOMAIN, i));
            for j in 0..n_spins {
                self.spins[i * n_spins + j] = S::random(&mut self.rngs[i]);
            }
        }

        self.system_ids = (0..n_systems).collect();

        let n_pairs = n_replicas / 2;
        for i in 0..n_temps * n_pairs {
            self.pair_rngs[i] =
                Xoshiro256StarStar::seed_from_u64(child_seed(base_seed, PAIR_SEED_DOMAIN, i));
        }

        self.energies.resize(n_systems, S::Value::default());
        S::interactions(lattice, &self.spins, &self.couplings, &mut self.energies);
        self.pt
            .reset(n_replicas, n_temps, &self.system_ids, &self.temperatures);
        self.schedule = ScheduleState::default();
    }

    /// Validate externally mutable buffers before entering parallel kernels.
    pub fn validate(
        &self,
        lattice: &Lattice,
        n_replicas: usize,
        n_temps: usize,
    ) -> Result<(), String> {
        if lattice.n_spins == 0 {
            return Err("lattice has no sites".into());
        }
        if lattice.has_self_neighbors() {
            return Err(
                "a neighbor offset maps sites to themselves (lattice extent 1 along it)".into(),
            );
        }
        let n = n_replicas
            .checked_mul(n_temps)
            .ok_or("system count overflow")?;
        if n == 0
            || self.spins.len()
                != n.checked_mul(lattice.n_spins)
                    .ok_or("spin count overflow")?
            || self.rngs.len() != n
            || self.system_ids.len() != n
            || self.temperatures.len() != n
            || self.energies.len() != n
            || self.couplings.len() != lattice.n_spins * lattice.n_neighbors
            || self.pair_rngs.len() != n_temps * (n_replicas / 2)
        {
            return Err("invalid realization buffer dimensions".into());
        }
        if self
            .temperatures
            .chunks_exact(n_temps)
            .any(|ladder| ladder != &self.temperatures[..n_temps])
        {
            return Err("temperature ladders must agree across replicas".into());
        }
        let mut seen = vec![false; n];
        for (slot, &id) in self.system_ids.iter().enumerate() {
            if id >= n || seen[id] || id / n_temps != slot / n_temps {
                return Err("invalid temperature permutation".into());
            }
            seen[id] = true;
        }
        if self.couplings.iter().any(|x| !x.to_f64().is_finite()) {
            return Err("couplings must be finite".into());
        }
        // Kernels divide by temperature in the model's precision, so 1/T must be finite too.
        if self.temperatures.iter().any(|&x| {
            let inverse = S::Value::from_f64(1.0).to_f64() / x.to_f64();
            !x.to_f64().is_finite()
                || x.to_f64() <= 0.0
                || !S::Value::from_f64(inverse).to_f64().is_finite()
        }) {
            return Err("temperatures must be positive and finite with a finite inverse".into());
        }
        Ok(())
    }

    pub(crate) fn temper(
        &mut self,
        n_spins: usize,
        n_replicas: usize,
        n_temps: usize,
        schedule: PtSchedule,
    ) {
        for replica in 0..n_replicas {
            let offset = replica * n_temps;
            let ids = &mut self.system_ids[offset..offset + n_temps];
            let temps = &self.temperatures[offset..offset + n_temps];
            let mut record = |attempt| self.pt.record_attempt(attempt);
            match schedule {
                PtSchedule::SingleRandomEdge => crate::mcmc::tempering::parallel_tempering(
                    &self.energies,
                    temps,
                    ids,
                    n_spins,
                    &mut self.rngs[offset],
                    &mut record,
                ),
                PtSchedule::FullLadder => crate::mcmc::tempering::parallel_tempering_full_ladder(
                    &self.energies,
                    temps,
                    ids,
                    n_spins,
                    &mut self.rngs[offset],
                    &mut record,
                ),
            }
            self.pt
                .record_visits(&self.system_ids[offset..offset + n_temps]);
        }
    }

    /// Swap attempts per temperature edge since creation or the last reset. Like the
    /// acceptances and round trips, they count warmup sweeps too.
    pub fn pt_edge_attempts(&self) -> &[u64] {
        &self.pt.edge_attempts
    }

    pub fn pt_edge_acceptances(&self) -> &[u64] {
        &self.pt.edge_acceptances
    }

    pub fn pt_round_trips(&self) -> &[u64] {
        &self.pt.round_trips
    }

    /// Per temperature slot, visits after each tempering event by systems that last
    /// visited the cold end of the ladder.
    pub fn pt_flow_up(&self) -> &[u64] {
        &self.pt.flow_up
    }

    /// Per temperature slot, visits by systems that have visited either end; the
    /// denominator of the flow fraction `pt_flow_up / pt_flow_labelled`.
    pub fn pt_flow_labelled(&self) -> &[u64] {
        &self.pt.flow_labelled
    }
}

/// Stable disorder stream mapping shared by the Rust and Python interfaces.
pub fn realization_seed(root: u64, realization: usize) -> u64 {
    splitmix64(root ^ splitmix64(realization as u64))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcmc::tempering::TemperingAttempt;

    #[test]
    fn rejects_lattices_with_self_neighbors() {
        let lattice = Lattice::new(vec![4, 1]);
        let couplings = vec![1.0; lattice.n_spins * lattice.n_neighbors];
        let realization = Realization::new(&lattice, couplings, &[1.0], 1, 0);
        let error = realization.validate(&lattice, 1, 1).unwrap_err();
        assert!(error.contains("themselves"), "{error}");
    }

    #[test]
    fn reset_replays_spins_and_pt_state() {
        let lattice = Lattice::new(vec![3, 3]);
        let couplings = vec![1.0; lattice.n_spins * lattice.n_neighbors];
        let mut realization = Realization::new(&lattice, couplings, &[1.0, 2.0], 2, 17);
        let initial_spins = realization.spins.clone();
        realization.system_ids.swap(0, 1);
        realization.pt.edge_attempts.fill(9);

        realization.reset(&lattice, 2, 2, 17);

        assert_eq!(realization.spins, initial_spins);
        assert_eq!(realization.system_ids, vec![0, 1, 2, 3]);
        assert_eq!(realization.pt_edge_attempts(), &[0]);
    }

    /// With every swap accepted, the full ladder must move walkers across it; a
    /// schedule whose consecutive passes undo each other completes no round trips.
    #[test]
    fn full_ladder_moves_walkers_across_the_ladder() {
        let lattice = Lattice::new(vec![2, 2]);
        let couplings = vec![1.0; lattice.n_spins * lattice.n_neighbors];
        let temps = [1.0, 2.0, 3.0, 4.0, 5.0];
        let mut realization = Realization::new(&lattice, couplings, &temps, 1, 3);
        realization.energies.fill(0.0);
        for _ in 0..40 {
            realization.temper(lattice.n_spins, 1, temps.len(), PtSchedule::FullLadder);
        }
        assert!(realization.pt_edge_acceptances().iter().all(|&a| a == 40));
        assert!(realization.pt_round_trips().iter().sum::<u64>() > 0);
        // Whoever sits at an end has just visited it.
        let (up, labelled) = (realization.pt_flow_up(), realization.pt_flow_labelled());
        assert_eq!((up[0], up[4]), (labelled[0], 0));
        assert_eq!((labelled[0], labelled[4]), (40, 40));
        assert!(up.iter().zip(labelled).all(|(up, labelled)| up <= labelled));
    }

    #[test]
    fn tracks_hot_cold_hot_across_attempts() {
        let lattice = Lattice::new(vec![2, 2]);
        let couplings = vec![1.0; lattice.n_spins * lattice.n_neighbors];
        let mut realization = Realization::new(&lattice, couplings, &[0.5, 1.0, 2.0], 1, 7);

        for (edge, left_system, right_system) in [(1, 1, 2), (0, 0, 2), (0, 2, 0), (1, 2, 1)] {
            realization.pt.record_attempt(TemperingAttempt {
                edge,
                accepted: true,
                left_system,
                right_system,
            });
        }

        assert_eq!(realization.pt_edge_attempts(), &[2, 2]);
        assert_eq!(realization.pt_edge_acceptances(), &[2, 2]);
        assert_eq!(realization.pt_round_trips()[2], 1);
    }
}