spin-sim 0.2.0

use super::equilibration::EquilCheckpoint;
use super::overlap::OverlapStats;

pub struct ClusterSnapshot {
    pub sweep_id: usize,
    pub mode_idx: usize,
    pub cluster_ids: Vec<Vec<u32>>,
    pub blue_ids: Option<Vec<Vec<u32>>>,
    pub spins: Vec<[Vec<i8>; 2]>,
    pub system_ids: Vec<[usize; 2]>,
}

pub struct ClusterStats {
    /// FK cluster size histogram per temperature: `hist[s]` = count of size-`s` clusters.
    pub fk_csd: Vec<Vec<u64>>,
    /// Per-mode overlap cluster size histogram: `[n_modes][n_temps][n_spins+1]`.
    pub overlap_csd: Vec<Vec<Vec<u64>>>,
    /// Per-mode average relative size of k-th largest overlap cluster: `[n_modes][n_temps][4]`.
    pub top_cluster_sizes: Vec<Vec<[f64; 4]>>,
}

pub struct Diagnostics {
    /// Integrated autocorrelation time τ_int(m²) per temperature.
    /// Empty if autocorrelation_max_lag is None.
    pub mags2_tau: Vec<f64>,
    /// Integrated autocorrelation time τ_int(q²) per temperature.
    /// Empty if autocorrelation_max_lag is None or n_replicas < 2.
    pub overlap2_tau: Vec<f64>,
    /// Equilibration diagnostic checkpoints (energy + link overlap running averages).
    /// Empty if equilibration_diagnostic is false.
    pub equil_checkpoints: Vec<EquilCheckpoint>,
}

/// Per-temperature observables averaged over measurement sweeps and replicas.
///
/// All vectors are indexed by temperature index and have length `n_temps`.
/// Overlap vectors are empty when `n_replicas < 2`.
pub struct SweepResult {
    /// ⟨m⟩ — mean magnetization per spin.
    pub mags: Vec<f64>,
    /// ⟨m²⟩.
    pub mags2: Vec<f64>,
    /// ⟨m⁴⟩.
    pub mags4: Vec<f64>,
    /// ⟨E⟩ — mean energy per spin.
    pub energies: Vec<f64>,
    /// ⟨E²⟩.
    pub energies2: Vec<f64>,
    pub overlap_stats: OverlapStats,
    pub cluster_stats: ClusterStats,
    pub diagnostics: Diagnostics,
    pub cluster_snapshots: Vec<ClusterSnapshot>,
}

impl SweepResult {
    /// Average [`SweepResult`]s across disorder realizations.
    pub fn aggregate(results: &[Self]) -> Self {
        let n = results.len() as f64;
        let n_temps = results[0].mags.len();
        let n_fk_csd = results[0].cluster_stats.fk_csd.len();

        let fk_len = results[0]
            .cluster_stats
            .fk_csd
            .first()
            .map_or(0, |v| v.len());

        let n_modes = results[0].cluster_stats.overlap_csd.len();
        let ov_inner_len = results[0]
            .cluster_stats
            .overlap_csd
            .first()
            .and_then(|v| v.first())
            .map_or(0, |v| v.len());
        let ov_inner_temps = results[0]
            .cluster_stats
            .overlap_csd
            .first()
            .map_or(0, |v| v.len());

        let n_top_modes = results[0].cluster_stats.top_cluster_sizes.len();
        let n_top_temps = results[0]
            .cluster_stats
            .top_cluster_sizes
            .first()
            .map_or(0, |v| v.len());

        let m2_tau_len = results[0].diagnostics.mags2_tau.len();
        let q2_tau_len = results[0].diagnostics.overlap2_tau.len();
        let n_ckpts = results[0].diagnostics.equil_checkpoints.len();

        let overlap_stats =
            OverlapStats::aggregate(&results.iter().map(|r| &r.overlap_stats).collect::<Vec<_>>());

        let mut agg = SweepResult {
            mags: vec![0.0; n_temps],
            mags2: vec![0.0; n_temps],
            mags4: vec![0.0; n_temps],
            energies: vec![0.0; n_temps],
            energies2: vec![0.0; n_temps],
            overlap_stats,
            cluster_stats: ClusterStats {
                fk_csd: (0..n_fk_csd).map(|_| vec![0u64; fk_len]).collect(),
                overlap_csd: (0..n_modes)
                    .map(|_| {
                        (0..ov_inner_temps)
                            .map(|_| vec![0u64; ov_inner_len])
                            .collect()
                    })
                    .collect(),
                top_cluster_sizes: (0..n_top_modes)
                    .map(|_| vec![[0.0; 4]; n_top_temps])
                    .collect(),
            },
            diagnostics: Diagnostics {
                mags2_tau: vec![0.0; m2_tau_len],
                overlap2_tau: vec![0.0; q2_tau_len],
                equil_checkpoints: (0..n_ckpts)
                    .map(|i| EquilCheckpoint {
                        sweep: results[0].diagnostics.equil_checkpoints[i].sweep,
                        energy_avg: vec![0.0; n_temps],
                        link_overlap_avg: vec![0.0; n_temps],
                    })
                    .collect(),
            },
            cluster_snapshots: Vec::new(),
        };

        for r in results {
            for (a, &v) in agg.mags.iter_mut().zip(r.mags.iter()) {
                *a += v;
            }
            for (a, &v) in agg.mags2.iter_mut().zip(r.mags2.iter()) {
                *a += v;
            }
            for (a, &v) in agg.mags4.iter_mut().zip(r.mags4.iter()) {
                *a += v;
            }
            for (a, &v) in agg.energies.iter_mut().zip(r.energies.iter()) {
                *a += v;
            }
            for (a, &v) in agg.energies2.iter_mut().zip(r.energies2.iter()) {
                *a += v;
            }
            for (a, s) in agg
                .cluster_stats
                .fk_csd
                .iter_mut()
                .zip(r.cluster_stats.fk_csd.iter())
            {
                for (ah, &sh) in a.iter_mut().zip(s.iter()) {
                    *ah += sh;
                }
            }
            for (am, sm) in agg
                .cluster_stats
                .overlap_csd
                .iter_mut()
                .zip(r.cluster_stats.overlap_csd.iter())
            {
                for (a, s) in am.iter_mut().zip(sm.iter()) {
                    for (ah, &sh) in a.iter_mut().zip(s.iter()) {
                        *ah += sh;
                    }
                }
            }
            for (am, sm) in agg
                .cluster_stats
                .top_cluster_sizes
                .iter_mut()
                .zip(r.cluster_stats.top_cluster_sizes.iter())
            {
                for (a, &s) in am.iter_mut().zip(sm.iter()) {
                    for k in 0..4 {
                        a[k] += s[k];
                    }
                }
            }
            for (a, &v) in agg
                .diagnostics
                .mags2_tau
                .iter_mut()
                .zip(r.diagnostics.mags2_tau.iter())
            {
                *a += v;
            }
            for (a, &v) in agg
                .diagnostics
                .overlap2_tau
                .iter_mut()
                .zip(r.diagnostics.overlap2_tau.iter())
            {
                *a += v;
            }
            for (ac, rc) in agg
                .diagnostics
                .equil_checkpoints
                .iter_mut()
                .zip(r.diagnostics.equil_checkpoints.iter())
            {
                for (a, &v) in ac.energy_avg.iter_mut().zip(rc.energy_avg.iter()) {
                    *a += v;
                }
                for (a, &v) in ac
                    .link_overlap_avg
                    .iter_mut()
                    .zip(rc.link_overlap_avg.iter())
                {
                    *a += v;
                }
            }
        }

        for v in agg
            .mags
            .iter_mut()
            .chain(agg.mags2.iter_mut())
            .chain(agg.mags4.iter_mut())
            .chain(agg.energies.iter_mut())
            .chain(agg.energies2.iter_mut())
        {
            *v /= n;
        }

        for mode_tops in agg.cluster_stats.top_cluster_sizes.iter_mut() {
            for arr in mode_tops.iter_mut() {
                for v in arr.iter_mut() {
                    *v /= n;
                }
            }
        }

        for v in agg.diagnostics.mags2_tau.iter_mut() {
            *v /= n;
        }
        for v in agg.diagnostics.overlap2_tau.iter_mut() {
            *v /= n;
        }
        for ckpt in agg.diagnostics.equil_checkpoints.iter_mut() {
            for v in ckpt.energy_avg.iter_mut() {
                *v /= n;
            }
            for v in ckpt.link_overlap_avg.iter_mut() {
                *v /= n;
            }
        }

        agg
    }
}