helena 0.1.0

use std::num::NonZeroUsize;

use crate::latent::Continuous;
use crate::{Error, Latent, Result};

use super::super::common::{l2_distance, shape};

/// Mean L2 distance between matched cells of two condition-by-seed grids.
///
/// The grids must align: same condition count, same per-condition seed budget, and
/// a single per-cell latent shape shared across *every* cell of *both* grids (so the
/// averaged L2 scores are commensurable). Cell `(c, s)` of `candidate` is compared to
/// cell `(c, s)` of `reference`, so the seeds must be drawn in the same order in both.
/// The latents are continuous, so their payloads are finite and rank-2 by
/// construction; only non-empty payloads and the grid-wide shape remain to check.
///
/// This is the pooled mean of [`grid_cell_distances`]; the per-condition split that
/// feeds the uncertainty bootstrap lives there.
pub fn grid_distance(
    reference: &[Vec<Latent<Continuous>>],
    candidate: &[Vec<Latent<Continuous>>],
) -> Result<f32> {
    let clusters = grid_cell_distances(reference, candidate)?;
    // Flatten in the same row-major (condition, then seed) order the per-cell loop
    // produced, so the f64 accumulation is bit-for-bit the former `total / cells`.
    let total: f64 = clusters.iter().flatten().sum();
    let cells: usize = clusters.iter().map(Vec::len).sum();
    Ok((total / cells as f64) as f32)
}

/// Per-condition L2 distances between matched cells of two condition-by-seed grids.
///
/// The clustered-bootstrap input behind [`grid_distance`]: each inner vector holds one
/// condition's per-seed L2 distances to the reference (the coherent unit the
/// uncertainty bootstrap resamples within), so the outer length is the condition count
/// and each inner length is that condition's seed budget. Alignment is validated
/// exactly as [`grid_distance`] requires (this is its shared core), so the two never
/// disagree on what a well-formed grid pair is.
pub fn grid_cell_distances(
    reference: &[Vec<Latent<Continuous>>],
    candidate: &[Vec<Latent<Continuous>>],
) -> Result<Vec<Vec<f64>>> {
    if reference.is_empty() {
        return Err(Error::validation(
            "step-gap grids need at least one condition",
        ));
    }
    if reference.len() != candidate.len() {
        return Err(Error::validation(format!(
            "grids must share a condition count; reference {} vs candidate {}",
            reference.len(),
            candidate.len()
        )));
    }

    let mut clusters: Vec<Vec<f64>> = Vec::with_capacity(reference.len());
    // The first cell pins the shape every other cell of both grids must match, so a
    // ragged-but-pairwise-consistent pair cannot average incomparable L2 scales.
    let mut expected: Option<[usize; 2]> = None;
    for (c, (ref_row, cand_row)) in reference.iter().zip(candidate).enumerate() {
        if ref_row.is_empty() {
            return Err(Error::validation(format!("condition {c} has no samples")));
        }
        if ref_row.len() != cand_row.len() {
            return Err(Error::validation(format!(
                "condition {c} must share a seed budget; reference {} vs candidate {}",
                ref_row.len(),
                cand_row.len()
            )));
        }
        let mut row = Vec::with_capacity(ref_row.len());
        for (s, (r, k)) in ref_row.iter().zip(cand_row).enumerate() {
            let rt = r.repr();
            let kt = k.repr();
            if rt.values().is_empty() {
                return Err(Error::validation("samples must have at least one element"));
            }
            let want = *expected.get_or_insert_with(|| shape(rt));
            if shape(rt) != want || shape(kt) != want {
                return Err(Error::validation(format!(
                    "all cells must share a shape {want:?}; cell ({c}, {s}) has \
                     reference {:?} / candidate {:?}",
                    shape(rt),
                    shape(kt)
                )));
            }
            row.push(l2_distance(rt, kt));
        }
        clusters.push(row);
    }
    Ok(clusters)
}

/// One point of a [`StepGapCurve`]: the gap to the reference at a step budget.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct StepGap {
    steps: NonZeroUsize,
    gap: f32,
}

impl StepGap {
    /// The sampler step budget this gap was measured at.
    pub fn steps(&self) -> NonZeroUsize {
        self.steps
    }

    /// The mean matched-cell distance to the many-step reference.
    pub fn gap(&self) -> f32 {
        self.gap
    }
}

/// The few-step sampling gap: how far each step budget's output sits from a
/// many-step reference at matched seeds.
///
/// helena samples in a small fixed number of steps (target 4-16), not a many-step
/// regime, so the operative quality question is how much a few-step output deviates
/// from the converged many-step result at the *same* seed. At a fixed seed the
/// chains share an initial noise canvas (its size is set by the clip, not the step
/// count), so the gap isolates integration error. A gap that shrinks as the budget
/// grows is a sampler converging toward its schedule's limit; whether an objective
/// lever (`GEN-9`/`GEN-10`) shrinks the 4-step gap is the few-step win the
/// scarce-data brief predicts.
///
/// Unlike [`sample_grid`], there is no generic driver here: the step budget is a
/// generator-construction property, not a per-call [`GenerationParams`] field, so
/// sweeping it is the (diffusion-specific) caller's job. This type scores the grids
/// the caller produces.
///
/// [`sample_grid`]: super::sample_grid
/// [`GenerationParams`]: crate::GenerationParams
#[derive(Clone, Debug, PartialEq)]
pub struct StepGapCurve {
    reference_steps: NonZeroUsize,
    gaps: Vec<StepGap>,
}

impl StepGapCurve {
    /// Score each `(steps, grid)` candidate against the `reference` grid.
    ///
    /// `reference_steps` is the budget the reference grid was sampled at (by
    /// convention the largest, the converged anchor). Returns
    /// [`Error::Validation`] if `candidates` is empty, repeats a step budget, or
    /// any candidate grid fails to align with `reference` (see [`grid_distance`]).
    /// The resulting gaps are ordered by ascending step budget.
    pub fn measure(
        reference_steps: NonZeroUsize,
        reference: &[Vec<Latent<Continuous>>],
        candidates: &[(NonZeroUsize, &[Vec<Latent<Continuous>>])],
    ) -> Result<Self> {
        if candidates.is_empty() {
            return Err(Error::validation(
                "step-gap curve needs at least one candidate budget",
            ));
        }
        let mut gaps: Vec<StepGap> = Vec::with_capacity(candidates.len());
        for &(steps, grid) in candidates {
            if gaps.iter().any(|g| g.steps == steps) {
                return Err(Error::validation(format!(
                    "duplicate candidate step budget {steps}"
                )));
            }
            let gap = grid_distance(reference, grid)?;
            gaps.push(StepGap { steps, gap });
        }
        gaps.sort_by_key(|g| g.steps);
        Ok(Self {
            reference_steps,
            gaps,
        })
    }

    /// The step budget the reference grid was sampled at.
    pub fn reference_steps(&self) -> NonZeroUsize {
        self.reference_steps
    }

    /// The per-budget gaps, ordered by ascending step budget.
    pub fn gaps(&self) -> &[StepGap] {
        &self.gaps
    }

    /// Whether the gap to the reference shrinks (non-strictly) as the step budget
    /// grows: a sampler converging toward its schedule's many-step limit. The
    /// decisive read of few-step behavior. Vacuously true for a single candidate.
    pub fn is_converging(&self) -> bool {
        self.gaps.windows(2).all(|w| w[1].gap <= w[0].gap)
    }
}