helena 0.1.0

use std::num::NonZeroUsize;

use crate::diffusion::NoiseLevel;
use crate::eval::{ClusterBootstrap, ConfidenceInterval};
use crate::{Error, Result, Tensor};

use super::super::common::mean_squared_error;

/// A clean-target (`x0`) reconstruction loss measured at one noise level.
///
/// The per-element mean squared error between a head's clean-signal estimate and
/// the true target, tagged with the [`NoiseLevel`] it was measured at. Collect a
/// set of these across the schedule and fold them with [`PerSnrLossCurve`] to read
/// *where* in the SNR spectrum a head spends its accuracy. This is the
/// discriminator the scarce-data objective brief calls for: a v-head that only
/// drives the loss down in the easy high-SNR bins is overspending on cleanup the
/// few-step sampler skips, whereas a flat, low mid-band is the healthy shape.
///
/// The loss is held in `f64`: clean-target error spans orders of magnitude across
/// the schedule (it vanishes as `gamma -> 1`, where the noisy sample already *is*
/// the target), and the curve's value is in comparing those bins.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CleanTargetLoss {
    level: NoiseLevel,
    loss: f64,
}

impl CleanTargetLoss {
    /// Measure the loss of a clean-target estimate `x0_hat` against `target`.
    ///
    /// The x0-head parameterization (the bakeoff's Arm C), where the head emits
    /// the clean signal directly. `x0_hat` and `target` must share a non-empty
    /// shape and be finite.
    pub fn from_estimate(level: NoiseLevel, x0_hat: &Tensor, target: &Tensor) -> Result<Self> {
        let loss = mean_squared_error(x0_hat, target)?;
        Ok(Self { level, loss })
    }

    /// Recover the clean signal from a velocity prediction, then measure its loss.
    ///
    /// The v-head parameterization (the bakeoff's Arm A/B): `x0_hat` is
    /// [`NoiseLevel::signal_from_velocity`] of `(x_t, v_hat)`, the same recovery
    /// the sampler uses. `x_t`, `v_hat`, and `target` must share a shape.
    pub fn from_velocity(
        level: NoiseLevel,
        x_t: &Tensor,
        v_hat: &Tensor,
        target: &Tensor,
    ) -> Result<Self> {
        let x0_hat = level.signal_from_velocity(x_t, v_hat)?;
        Self::from_estimate(level, &x0_hat, target)
    }

    /// The noise level this loss was measured at.
    pub fn level(&self) -> NoiseLevel {
        self.level
    }

    /// The mean squared clean-target reconstruction error.
    pub fn loss(&self) -> f64 {
        self.loss
    }
}

/// One signal-variance bin of a [`PerSnrLossCurve`].
///
/// Covers `gamma in [lo, hi)` (the top bin closes at `1.0`). An empty bin reports
/// a `count` of zero and no mean.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SnrBin {
    lo: f32,
    hi: f32,
    sum: f64,
    count: usize,
}

impl SnrBin {
    /// The half-open signal-variance range `[lo, hi)` this bin covers.
    pub fn range(&self) -> (f32, f32) {
        (self.lo, self.hi)
    }

    /// The number of observations that fell in this bin.
    pub fn count(&self) -> usize {
        self.count
    }

    /// The mean clean-target loss in this bin, or `None` if it is empty.
    pub fn mean_loss(&self) -> Option<f64> {
        (self.count > 0).then(|| self.sum / self.count as f64)
    }
}

/// The clean-target reconstruction loss binned by signal variance (SNR).
///
/// `bins` equal-width bins partition `gamma in [0, 1]`; observation `o` lands in
/// bin `floor(gamma_o * bins)` (a `gamma` of `1.0` folds into the top bin). Read
/// low to high gamma, the per-bin mean loss is the curve the scarce-data
/// objective brief uses to settle its hypothesis H1: is a bakeoff arm's win from
/// per-SNR *weighting* (the mid-band falls) or x0 *parameterization*? The metric
/// is head-agnostic by construction (see [`CleanTargetLoss`]), so Arm A/B (v-head)
/// and Arm C (x0-head) are scored on one ruler.
///
/// A faithful curve wants each bin populated, so the caller should draw `gamma`
/// across the whole schedule (uniform-in-`gamma`, not the training density) when
/// producing the [`CleanTargetLoss`] set.
#[derive(Clone, Debug, PartialEq)]
pub struct PerSnrLossCurve {
    bins: Vec<SnrBin>,
}

impl PerSnrLossCurve {
    /// Bin `losses` by signal variance into `bins` equal-width bins.
    ///
    /// Returns [`Error::Validation`] if `losses` is empty.
    pub fn measure(losses: &[CleanTargetLoss], bins: NonZeroUsize) -> Result<Self> {
        if losses.is_empty() {
            return Err(Error::validation("no clean-target losses to bin"));
        }
        let n = bins.get();
        let mut acc: Vec<SnrBin> = (0..n)
            .map(|b| SnrBin {
                lo: b as f32 / n as f32,
                hi: (b + 1) as f32 / n as f32,
                sum: 0.0,
                count: 0,
            })
            .collect();
        for loss in losses {
            let idx = bin_index(loss.level().signal_variance(), n);
            acc[idx].sum += loss.loss();
            acc[idx].count += 1;
        }
        Ok(Self { bins: acc })
    }

    /// Per-bin clustered-bootstrap confidence intervals over per-conditioning-item
    /// losses, one entry per bin of [`PerSnrLossCurve::measure`] called with the same
    /// `bins` (parallel to its [`bins()`](Self::bins) in order).
    ///
    /// The point estimate ([`SnrBin::mean_loss`]) pools every observation in a bin, but
    /// its *uncertainty* must respect that same-item losses are not independent, so the
    /// bootstrap clusters by conditioning item: `per_item[i]` is item `i`'s losses
    /// across the gamma sweep (and across its windows, when one item segments into
    /// several), and within a bin all of that item's in-bin losses form one cluster (the
    /// [`ClusterBootstrap`] unit). So a clip that segments into more windows is still one
    /// cluster, and the interval reads the bin mean against item-to-item, not
    /// window-to-window, sampling noise.
    ///
    /// The interval is the *pooled-mean* bootstrap ([`ClusterBootstrap::measure_pooled`]),
    /// so its point is exactly the bin's pooled `mean_loss` no matter how many windows
    /// each item carries; only the resampling sees the cluster structure. Because every
    /// item is scored at the same gammas, a bin is populated for every item or empty for
    /// all; an empty bin yields `None` (nothing to resample) and a populated bin's
    /// clusters are all non-empty.
    ///
    /// Returns [`Error::Validation`] if `per_item` holds no losses at all, or if a
    /// non-finite loss reaches the bootstrap (the [`ClusterBootstrap`] boundary).
    pub fn bin_intervals(
        per_item: &[Vec<CleanTargetLoss>],
        bins: NonZeroUsize,
        bootstrap: &ClusterBootstrap,
    ) -> Result<Vec<Option<ConfidenceInterval>>> {
        if per_item.iter().all(Vec::is_empty) {
            return Err(Error::validation("no clean-target losses to bootstrap"));
        }
        let n = bins.get();
        let mut intervals = Vec::with_capacity(n);
        for b in 0..n {
            // One cluster per item that has any loss in this bin; an item scored at the
            // same gammas as the rest either populates the bin or is absent from it.
            let clusters: Vec<Vec<f64>> = per_item
                .iter()
                .filter_map(|item| {
                    let in_bin: Vec<f64> = item
                        .iter()
                        .filter(|l| bin_index(l.level().signal_variance(), n) == b)
                        .map(CleanTargetLoss::loss)
                        .collect();
                    (!in_bin.is_empty()).then_some(in_bin)
                })
                .collect();
            intervals.push(if clusters.is_empty() {
                None
            } else {
                Some(bootstrap.measure_pooled(&clusters)?)
            });
        }
        Ok(intervals)
    }

    /// The bins, ordered low to high signal variance.
    pub fn bins(&self) -> &[SnrBin] {
        &self.bins
    }

    /// Total observations folded into the curve.
    pub fn observation_count(&self) -> usize {
        self.bins.iter().map(|b| b.count).sum()
    }
}

/// The bin a signal variance `gamma in [0, 1]` lands in, for `n` equal-width bins
/// partitioning the unit interval; the top edge (`gamma == 1`) folds into the last
/// bin. The single source of truth shared by [`PerSnrLossCurve::measure`] and
/// [`PerSnrLossCurve::bin_intervals`], so the point estimate and its interval never
/// disagree on bin membership.
fn bin_index(gamma: f32, n: usize) -> usize {
    ((gamma * n as f32) as usize).min(n - 1)
}