helena 0.1.0

//! Per-dimension standardization for continuous latents (GEN-8).
//!
//! [`LatentNorm`] is a frozen provenance artifact for one codec/dataset pair:
//! it stores a per-dimension mean and a strictly positive standard deviation,
//! standardizes training targets, and inverts generated latents before decode.
//! It is *fitted by construction* (refoundation A5): the only ways to obtain
//! one are the validating [`new`](LatentNorm::new), the streaming
//! [`NormFitter`], or the iterator convenience [`fit`](LatentNorm::fit) — an
//! accumulating builder and a finished statistic are distinct types, so
//! "standardize with an unfinished fit" does not type-check.
//!
//! Everything here operates on [`FrameSeq`], the continuous latent payload.
//! Two whole bug classes the old `AudioLatent`-enum version guarded against at
//! runtime are now unrepresentable: a [`FrameSeq`] is *continuous by type* (a
//! token latent can never reach these functions), and it is *finite and rank-2
//! `[frames, dim]` by construction*, so the NaN scans and the rank/`dim ≥ 1`
//! re-derivation are the payload type's invariants, not this module's checks.
//! The only structural check that remains is dimension agreement between a
//! sequence and the fitted statistics.

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};
use crate::latent::FrameSeq;

/// Per-dimension standardization statistics for continuous latents.
///
/// Invariants, enforced by every constructor and re-checked on
/// deserialization: `mean` and `std` have equal, non-zero length, every value
/// is finite, and every standard deviation is strictly positive (floored at an
/// `eps` when fitted, so a constant dimension yields a usable — not
/// zero-division — scale). The fields are private and there is no unchecked
/// constructor, so [`dim`](Self::dim), [`mean`](Self::mean), and
/// [`std`](Self::std) can never disagree, and a degenerate statistic cannot
/// reach [`standardize`](Self::standardize).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "LatentNormRepr")]
pub struct LatentNorm {
    mean: Vec<f32>,
    std: Vec<f32>,
}

/// Wire form for [`LatentNorm`], routed back through invariant validation on
/// deserialization. The stored `std` is already floored (a persisted statistic
/// is a *finished* one), so the decode gate re-proves the invariants — equal
/// length, finite, strictly positive — rather than re-flooring against an
/// `eps` the wire form does not carry.
#[derive(Deserialize)]
struct LatentNormRepr {
    mean: Vec<f32>,
    std: Vec<f32>,
}

impl TryFrom<LatentNormRepr> for LatentNorm {
    type Error = Error;

    fn try_from(repr: LatentNormRepr) -> Result<Self> {
        Self::from_stats(repr.mean, repr.std)
    }
}

impl LatentNorm {
    /// Default standard-deviation floor for constant dimensions.
    pub const DEFAULT_EPS: f32 = 1e-5;

    /// Reject a standard-deviation floor that cannot anchor a valid statistic.
    ///
    /// Exposed so a caller that drives an expensive or side-effectful fit (e.g.
    /// streaming encoded sidecars from disk) can fail fast on a known-bad `eps`
    /// before doing the work — the same fail-fast guard [`fit`](Self::fit) and
    /// [`new`](Self::new) apply internally.
    pub fn validate_eps(eps: f32) -> Result<()> {
        if !eps.is_finite() || eps <= 0.0 {
            return Err(Error::validation(format!(
                "eps must be finite and positive, got {eps}"
            )));
        }
        Ok(())
    }

    /// Build from explicit per-dimension statistics, flooring each standard
    /// deviation at `eps`.
    ///
    /// Returns [`Error::Validation`] if `eps` is non-finite or non-positive, if
    /// `mean` and `std` differ in length, if either is empty, or if any `mean`
    /// value is non-finite. A non-finite standard deviation is likewise
    /// rejected; a finite but too-small (or non-positive) one is raised to
    /// `eps` rather than rejected, so a constant dimension produces a usable
    /// scale.
    pub fn new(mean: impl Into<Vec<f32>>, std: impl Into<Vec<f32>>, eps: f32) -> Result<Self> {
        Self::validate_eps(eps)?;
        let mean = mean.into();
        let mut std = std.into();
        if mean.len() != std.len() {
            return Err(Error::validation(format!(
                "latent normalization mean/std length mismatch: {} vs {}",
                mean.len(),
                std.len()
            )));
        }
        if mean.is_empty() {
            return Err(Error::validation(
                "latent normalization needs at least one dimension",
            ));
        }
        if let Some(d) = mean.iter().position(|m| !m.is_finite()) {
            return Err(Error::validation(format!(
                "non-finite mean at dimension {d}"
            )));
        }
        if let Some(d) = std.iter().position(|s| !s.is_finite()) {
            return Err(Error::validation(format!(
                "non-finite std at dimension {d}"
            )));
        }
        for s in &mut std {
            *s = s.max(eps);
        }
        Ok(Self { mean, std })
    }

    /// Validate already-finished statistics into a [`LatentNorm`] without
    /// re-flooring: the deserialization gate. Requires equal, non-zero length,
    /// finite means, and strictly positive finite standard deviations.
    fn from_stats(mean: Vec<f32>, std: Vec<f32>) -> Result<Self> {
        if mean.len() != std.len() {
            return Err(Error::validation(format!(
                "latent normalization mean/std length mismatch: {} vs {}",
                mean.len(),
                std.len()
            )));
        }
        if mean.is_empty() {
            return Err(Error::validation(
                "latent normalization needs at least one dimension",
            ));
        }
        if let Some(d) = mean.iter().position(|m| !m.is_finite()) {
            return Err(Error::validation(format!(
                "non-finite mean at dimension {d}"
            )));
        }
        if let Some(d) = std.iter().position(|s| !s.is_finite() || *s <= 0.0) {
            return Err(Error::validation(format!(
                "std at dimension {d} must be finite and positive, got {}",
                std[d]
            )));
        }
        Ok(Self { mean, std })
    }

    /// Fit per-dimension statistics over continuous target latents.
    ///
    /// The iterator form of [`NormFitter`]; prefer the builder when streaming
    /// encoded targets through a frozen codec so the whole dataset need not be
    /// materialized at once. `eps` floors each fitted standard deviation. A
    /// known-bad `eps` is rejected before the iterator is consumed, so a
    /// side-effectful source is never touched for a fit that cannot succeed.
    pub fn fit<'a>(seqs: impl IntoIterator<Item = &'a FrameSeq>, eps: f32) -> Result<Self> {
        // Fail fast on a known-bad eps before consuming a (possibly expensive or
        // side-effectful) iterator.
        Self::validate_eps(eps)?;
        let mut fit = NormFitter::new();
        for seq in seqs {
            fit.push(seq)?;
        }
        fit.finish_with_eps(eps)
    }

    /// Number of latent dimensions these statistics cover.
    pub fn dim(&self) -> usize {
        self.mean.len()
    }

    /// Per-dimension means.
    pub fn mean(&self) -> &[f32] {
        &self.mean
    }

    /// Per-dimension standard deviations, each strictly positive.
    pub fn std(&self) -> &[f32] {
        &self.std
    }

    /// Map a continuous latent into standardized space.
    ///
    /// Returns [`Error::Shape`] if the sequence's width does not match
    /// [`dim`](Self::dim). Finiteness and rank are the [`FrameSeq`] invariant,
    /// so no other rejection is possible.
    pub fn standardize(&self, seq: &FrameSeq) -> Result<FrameSeq> {
        self.apply(seq, |x, mean, std| (x - mean) / std)
    }

    /// Map a standardized latent back to codec space — the exact inverse of
    /// [`standardize`](Self::standardize) (see the roundtrip test).
    ///
    /// Returns [`Error::Shape`] on a width mismatch, as
    /// [`standardize`](Self::standardize) does.
    pub fn unstandardize(&self, seq: &FrameSeq) -> Result<FrameSeq> {
        self.apply(seq, |x, mean, std| x * std + mean)
    }

    fn apply(&self, seq: &FrameSeq, f: impl Fn(f32, f32, f32) -> f32) -> Result<FrameSeq> {
        let dim = seq.dim();
        if dim != self.dim() {
            return Err(Error::shape(
                vec![seq.frames(), self.dim()],
                vec![seq.frames(), dim],
            ));
        }
        let mapped: Vec<f32> = seq
            .values()
            .iter()
            .enumerate()
            .map(|(i, &x)| {
                let d = i % dim;
                f(x, self.mean[d], self.std[d])
            })
            .collect();
        FrameSeq::new(seq.frames(), dim, mapped)
    }
}

/// Per-dimension Welford accumulator (the numeric core of [`NormFitter`]).
#[derive(Clone, Copy, Debug, Default)]
struct DimAccumulator {
    count: u64,
    mean: f64,
    m2: f64,
}

impl DimAccumulator {
    fn push(&mut self, v: f32) {
        self.count += 1;
        let x = v as f64;
        let delta = x - self.mean;
        self.mean += delta / self.count as f64;
        self.m2 += delta * (x - self.mean);
    }

    fn std(&self) -> f64 {
        if self.count == 0 {
            return 0.0;
        }
        (self.m2.max(0.0) / self.count as f64).sqrt()
    }
}

/// Single-pass, bounded-memory builder for a [`LatentNorm`].
///
/// Folds continuous latents one at a time with per-dimension Welford
/// accumulators, so a caller can stream encoded targets through a frozen codec
/// without retaining the whole encoded dataset: state is bounded by the latent
/// dimension, not the number of frames. This is the "fit stats" stage the
/// gradient-flow brief (`DOCS-6`) recommends running once per codec/dataset
/// pair before prior training. [`LatentNorm::fit`] is this builder driven over
/// an iterator.
///
/// The first [`push`](Self::push) fixes the dimension; later sequences must
/// match it. [`finish`](Self::finish) is total: it rejects a builder that saw
/// no sequences or no frames, so a degenerate statistic cannot reach an apply
/// path.
#[derive(Clone, Debug, Default)]
pub struct NormFitter {
    /// Empty until the first `push` fixes the dimension.
    acc: Vec<DimAccumulator>,
}

impl NormFitter {
    /// A builder that has seen no sequences.
    pub fn new() -> Self {
        Self::default()
    }

    /// Whether no sequence has been folded in yet.
    pub fn is_empty(&self) -> bool {
        self.acc.is_empty()
    }

    /// Number of latent dimensions fixed by the first [`push`](Self::push), or
    /// `None` while the builder is still empty.
    pub fn dim(&self) -> Option<usize> {
        (!self.acc.is_empty()).then_some(self.acc.len())
    }

    /// Fold one continuous latent into the running statistics.
    ///
    /// The first call fixes the dimension; a later sequence of a different
    /// width is an [`Error::Shape`]. Non-finite values cannot occur —
    /// finiteness is the [`FrameSeq`] invariant — so a width mismatch is the
    /// only failure. The mismatch is detected before any accumulator is
    /// touched, so an errored push leaves the builder unchanged.
    pub fn push(&mut self, seq: &FrameSeq) -> Result<()> {
        let dim = seq.dim();
        if self.acc.is_empty() {
            let mut acc = vec![DimAccumulator::default(); dim];
            accumulate(seq, dim, &mut acc);
            self.acc = acc;
        } else if self.acc.len() != dim {
            return Err(Error::shape(
                vec![seq.frames(), self.acc.len()],
                vec![seq.frames(), dim],
            ));
        } else {
            accumulate(seq, dim, &mut self.acc);
        }
        Ok(())
    }

    /// Finish into per-dimension statistics, flooring each standard deviation
    /// at [`LatentNorm::DEFAULT_EPS`].
    ///
    /// Rejects a builder that saw no sequences and one whose sequences
    /// contributed no frames.
    pub fn finish(self) -> Result<LatentNorm> {
        self.finish_with_eps(LatentNorm::DEFAULT_EPS)
    }

    /// The one numeric finish path, parameterized by the standard-deviation
    /// floor. Public [`finish`](Self::finish) fixes `eps` at
    /// [`LatentNorm::DEFAULT_EPS`]; [`LatentNorm::fit`] threads the caller's
    /// `eps` through the same code.
    fn finish_with_eps(self, eps: f32) -> Result<LatentNorm> {
        LatentNorm::validate_eps(eps)?;
        if self.acc.is_empty() {
            return Err(Error::validation(
                "cannot fit normalization over an empty set",
            ));
        }
        let count = self.acc[0].count;
        if count == 0 {
            return Err(Error::validation(
                "cannot fit normalization: latents contributed no frames",
            ));
        }
        // Every dimension is folded once per frame, so unequal counts can only
        // mean corrupt state; reject rather than emit a lopsided statistic.
        if self.acc.iter().any(|a| a.count != count) {
            return Err(Error::validation(
                "cannot fit normalization: inconsistent per-dimension frame counts",
            ));
        }

        let dim = self.acc.len();
        let mut mean = Vec::with_capacity(dim);
        let mut std = Vec::with_capacity(dim);
        for a in &self.acc {
            mean.push(a.mean as f32);
            std.push((a.std().max(eps as f64)) as f32);
        }
        LatentNorm::from_stats(mean, std)
    }
}

/// Fold a continuous latent into per-dimension accumulators.
///
/// Finiteness is the [`FrameSeq`] invariant, so — unlike the old
/// `AudioLatent`-enum version — there is no NaN scan and this is infallible.
/// Each value is folded into its dimension's accumulator by column position.
fn accumulate(seq: &FrameSeq, dim: usize, acc: &mut [DimAccumulator]) {
    for (i, &x) in seq.values().iter().enumerate() {
        acc[i % dim].push(x);
    }
}

#[cfg(test)]
mod tests;