helena 0.1.0

//! The conditioning bundle: slot-typed data latents (PRD §10.1, §13.1).
//!
//! Each PRD latent slot is its own type, because the slots are *not
//! interchangeable*: a consumer that treats `z_global` and `z_temporal` as
//! one vector destroys the encoder contract. The mask lives *inside*
//! [`TemporalLatent`] with its length tied to the sequence at construction —
//! a mask without a temporal sequence, or of the wrong length, is
//! unrepresentable (the old `is_empty`-forgets-the-mask bug class is gone
//! with the field that hosted it).

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};
use crate::tensor::Tensor;

/// Free-form provenance and labels. `BTreeMap` so serialization order — and
/// therefore fingerprinting — is deterministic.
pub type Metadata = BTreeMap<String, serde_json::Value>;

/// Clip-level conditioning: identity, style, class (PRD §10.1 `z_global`).
///
/// Invariants: rank-2 `[items, dim]`, `dim ≥ 1`, finite.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "Tensor")]
pub struct GlobalLatent(Tensor);

impl GlobalLatent {
    /// Adopt a `[items, dim]` tensor, proving the slot's invariants.
    pub fn new(tensor: Tensor) -> Result<Self> {
        let &[_, dim] = tensor.shape() else {
            return Err(Error::shape(vec![0, 0], tensor.shape().to_vec()));
        };
        if dim == 0 {
            return Err(Error::validation("z_global needs dim ≥ 1"));
        }
        if !tensor.is_finite() {
            return Err(Error::NonFinite {
                context: "z_global",
            });
        }
        Ok(Self(tensor))
    }

    /// Number of items (rows).
    pub fn items(&self) -> usize {
        self.0.shape()[0]
    }

    /// Latent width (columns), ≥ 1.
    pub fn dim(&self) -> usize {
        self.0.shape()[1]
    }

    /// The backing tensor.
    pub fn tensor(&self) -> &Tensor {
        &self.0
    }
}

impl TryFrom<Tensor> for GlobalLatent {
    type Error = Error;

    fn try_from(tensor: Tensor) -> Result<Self> {
        Self::new(tensor)
    }
}

/// A per-step validity mask for a temporal sequence. Constructed only by
/// [`TemporalLatent::with_mask`], which ties its length to the sequence.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Mask(Vec<bool>);

impl Mask {
    /// Per-step validity flags, `true` = attend, `false` = padding.
    pub fn steps(&self) -> &[bool] {
        &self.0
    }

    /// Number of steps covered.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Whether the mask covers zero steps.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

/// Time-varying conditioning (PRD §10.1 `z_temporal`): a `[steps, channels]`
/// sequence with an optional per-step mask whose length is *proven equal* to
/// `steps` at construction.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "TemporalRepr")]
pub struct TemporalLatent {
    features: Tensor,
    mask: Option<Mask>,
}

impl TemporalLatent {
    /// Adopt a `[steps, channels]` tensor, unmasked.
    pub fn new(features: Tensor) -> Result<Self> {
        let &[_, channels] = features.shape() else {
            return Err(Error::shape(vec![0, 0], features.shape().to_vec()));
        };
        if channels == 0 {
            return Err(Error::validation("z_temporal needs channels ≥ 1"));
        }
        if !features.is_finite() {
            return Err(Error::NonFinite {
                context: "z_temporal",
            });
        }
        Ok(Self {
            features,
            mask: None,
        })
    }

    /// Attach a per-step mask; its length must equal the step count.
    pub fn with_mask(mut self, steps: Vec<bool>) -> Result<Self> {
        if steps.len() != self.steps() {
            return Err(Error::shape(vec![self.steps()], vec![steps.len()]));
        }
        self.mask = Some(Mask(steps));
        Ok(self)
    }

    /// Number of time steps (rows).
    pub fn steps(&self) -> usize {
        self.features.shape()[0]
    }

    /// Feature channels per step (columns), ≥ 1.
    pub fn channels(&self) -> usize {
        self.features.shape()[1]
    }

    /// The backing `[steps, channels]` tensor.
    pub fn features(&self) -> &Tensor {
        &self.features
    }

    /// The mask, if one was attached. Its length equals
    /// [`steps`](Self::steps) by construction.
    pub fn mask(&self) -> Option<&Mask> {
        self.mask.as_ref()
    }
}

/// Wire form routed back through the validating constructors.
#[derive(Deserialize)]
struct TemporalRepr {
    features: Tensor,
    mask: Option<Vec<bool>>,
}

impl TryFrom<TemporalRepr> for TemporalLatent {
    type Error = Error;

    fn try_from(repr: TemporalRepr) -> Result<Self> {
        let latent = TemporalLatent::new(repr.features)?;
        match repr.mask {
            Some(mask) => latent.with_mask(mask),
            None => Ok(latent),
        }
    }
}

/// The conditioning bundle a [`SourceEncoder`](crate::seams::SourceEncoder)
/// produces and a [`LatentGenerator`](crate::seams::LatentGenerator)
/// consumes.
///
/// Slots are private and populated through validating adders; the reserved
/// PRD slots (`z_acoustic`, `z_event`, `z_noise`) carry raw tensors until
/// their first real producer earns them a typed contract — reserving a name
/// is cheap, freezing wrong semantics is not.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Conditioning {
    global: Option<GlobalLatent>,
    temporal: Option<TemporalLatent>,
    acoustic: Option<Tensor>,
    event: Option<Tensor>,
    noise: Option<Tensor>,
    #[serde(default)]
    metadata: Metadata,
}

impl Conditioning {
    /// An empty bundle.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the clip-level slot.
    pub fn with_global(mut self, global: GlobalLatent) -> Self {
        self.global = Some(global);
        self
    }

    /// Set the time-varying slot.
    pub fn with_temporal(mut self, temporal: TemporalLatent) -> Self {
        self.temporal = Some(temporal);
        self
    }

    /// Set the reserved timbre slot (untyped until a producer exists).
    pub fn with_acoustic(mut self, acoustic: Tensor) -> Self {
        self.acoustic = Some(acoustic);
        self
    }

    /// Set the reserved event slot (untyped until a producer exists).
    pub fn with_event(mut self, event: Tensor) -> Self {
        self.event = Some(event);
        self
    }

    /// Set the reserved stochastic-microvariation slot.
    pub fn with_noise(mut self, noise: Tensor) -> Self {
        self.noise = Some(noise);
        self
    }

    /// Attach free-form provenance and labels.
    pub fn with_metadata(mut self, metadata: Metadata) -> Self {
        self.metadata = metadata;
        self
    }

    /// The clip-level slot.
    pub fn global(&self) -> Option<&GlobalLatent> {
        self.global.as_ref()
    }

    /// The time-varying slot.
    pub fn temporal(&self) -> Option<&TemporalLatent> {
        self.temporal.as_ref()
    }

    /// The reserved timbre slot.
    pub fn acoustic(&self) -> Option<&Tensor> {
        self.acoustic.as_ref()
    }

    /// The reserved event slot.
    pub fn event(&self) -> Option<&Tensor> {
        self.event.as_ref()
    }

    /// The reserved stochastic slot.
    pub fn noise(&self) -> Option<&Tensor> {
        self.noise.as_ref()
    }

    /// Free-form provenance and labels.
    pub fn metadata(&self) -> &Metadata {
        &self.metadata
    }

    /// Whether every latent slot is empty. Total over the slots by
    /// construction — a new slot extends this match, and the mask cannot be
    /// forgotten because it lives inside [`TemporalLatent`].
    pub fn is_empty(&self) -> bool {
        let Self {
            global,
            temporal,
            acoustic,
            event,
            noise,
            metadata: _,
        } = self;
        global.is_none()
            && temporal.is_none()
            && acoustic.is_none()
            && event.is_none()
            && noise.is_none()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn rank2(rows: usize, cols: usize) -> Tensor {
        Tensor::zeros([rows, cols]).unwrap()
    }

    #[test]
    fn slots_prove_their_invariants() {
        assert!(GlobalLatent::new(Tensor::vector([1.0])).is_err());
        assert!(GlobalLatent::new(rank2(2, 0)).is_err());
        assert!(GlobalLatent::new(Tensor::new([1, 1], vec![f32::NAN]).unwrap()).is_err());
        let g = GlobalLatent::new(rank2(3, 4)).unwrap();
        assert_eq!((g.items(), g.dim()), (3, 4));
    }

    #[test]
    fn mask_length_is_tied_to_steps() {
        let t = TemporalLatent::new(rank2(3, 2)).unwrap();
        assert!(t.clone().with_mask(vec![true, false]).is_err());
        let masked = t.with_mask(vec![true, false, true]).unwrap();
        assert_eq!(masked.mask().unwrap().len(), masked.steps());
    }

    #[test]
    fn deserialization_re_proves_the_mask_tie() {
        let good = r#"{"features":{"bytes":[0,0,0,0,0,0,0,0],"shape":[2,1],"dtype":"F32"},"mask":[true,false]}"#;
        // The exact Tensor wire form is TensorData's; build the good case
        // programmatically instead of pinning that format here.
        let t = TemporalLatent::new(rank2(2, 1))
            .unwrap()
            .with_mask(vec![true, false])
            .unwrap();
        let json = serde_json::to_string(&t).unwrap();
        assert_eq!(serde_json::from_str::<TemporalLatent>(&json).unwrap(), t);

        // A hand-edited record with a wrong-length mask fails closed.
        let bad = json.replace("[true,false]", "[true]");
        assert!(serde_json::from_str::<TemporalLatent>(&bad).is_err());
        let _ = good; // exact wire form is TensorData's own; not pinned here
    }

    #[test]
    fn emptiness_is_total_over_slots() {
        assert!(Conditioning::default().is_empty());
        let with_temporal =
            Conditioning::new().with_temporal(TemporalLatent::new(rank2(1, 1)).unwrap());
        assert!(!with_temporal.is_empty());
        let with_reserved = Conditioning::new().with_event(Tensor::vector([1.0]));
        assert!(!with_reserved.is_empty());
        // Metadata alone is not conditioning.
        let mut metadata = Metadata::default();
        metadata.insert("k".into(), serde_json::json!("v"));
        assert!(Conditioning::new().with_metadata(metadata).is_empty());
    }
}