helena 0.1.0

//! The clocked latent value and its dynamic edge representation.

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};
use crate::time::{Frames, TimeBase};

use super::kind::{Framed, KindTag, LatentKind};

/// An audio-side latent of a statically known kind on a known clock: the
/// generator's output and the decoder's input.
///
/// The kind is the type parameter; the clock is carried alongside the
/// payload so every latent knows its own frame rate — "50 Hz frames treated
/// as 44.1 kHz samples" is not expressible.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(bound(serialize = "", deserialize = ""))]
pub struct Latent<K: LatentKind> {
    repr: K::Repr,
    time_base: TimeBase,
}

impl<K: LatentKind> Latent<K> {
    /// Bind a payload to its clock.
    pub fn new(repr: K::Repr, time_base: TimeBase) -> Self {
        Self { repr, time_base }
    }

    /// The payload.
    pub fn repr(&self) -> &K::Repr {
        &self.repr
    }

    /// The clock this latent ticks on.
    pub fn time_base(&self) -> TimeBase {
        self.time_base
    }

    /// Number of frames.
    pub fn frames(&self) -> Frames {
        Frames(self.repr.frames() as u64)
    }

    /// Duration in seconds on this latent's clock.
    pub fn duration_seconds(&self) -> f64 {
        self.time_base.seconds(self.frames())
    }

    /// Split into payload and clock.
    pub fn into_parts(self) -> (K::Repr, TimeBase) {
        (self.repr, self.time_base)
    }

    /// Cut the latent into consecutive [`LatentBlock`]s of `block` frames.
    ///
    /// Returns [`Error::Validation`] if `block` is zero or does not evenly
    /// divide the frame count — partial blocks are a caller decision (pad or
    /// reject), never a silent tail-drop. This is the cut the streaming law
    /// feeds through a session.
    pub fn blocks(&self, block: Frames) -> Result<Vec<LatentBlock<K>>> {
        let frames = self.repr.frames();
        let size = usize::try_from(block.get())
            .map_err(|_| Error::validation("block size overflows usize"))?;
        if size == 0 {
            return Err(Error::validation("block size must be nonzero"));
        }
        if !frames.is_multiple_of(size) {
            return Err(Error::validation(format!(
                "{frames} frames is not a whole number of {size}-frame blocks"
            )));
        }
        Ok((0..frames / size)
            .map(|i| LatentBlock::new(self.repr.slice_frames(i * size, size)))
            .collect())
    }
}

/// One streaming input: a block of latent payload pushed into a
/// [`StreamSession`](crate::stream::StreamSession). Deliberately unclocked —
/// the session's decoder owns the clock — and deliberately distinct from
/// [`Latent`], so a block is never confused with a whole sequence.
#[derive(Clone, Debug, PartialEq)]
pub struct LatentBlock<K: LatentKind>(K::Repr);

impl<K: LatentKind> LatentBlock<K> {
    /// Wrap a block of latent payload.
    pub fn new(repr: K::Repr) -> Self {
        Self(repr)
    }

    /// The payload.
    pub fn repr(&self) -> &K::Repr {
        &self.0
    }

    /// Number of frames in this block.
    pub fn frames(&self) -> Frames {
        Frames(self.0.frames() as u64)
    }

    /// Consume into the payload.
    pub fn into_repr(self) -> K::Repr {
        self.0
    }
}

/// The *edge* representation of a latent: what caches, artifacts, and CLI
/// reports see. Inside the pipeline the kind is static; this enum exists
/// only where data is genuinely dynamic (a sidecar that may hold either
/// kind), with [`downcast`](Self::downcast) the one checked gate back.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AnyLatent {
    /// A continuous latent.
    Continuous(Latent<super::Continuous>),
    /// A token latent.
    Tokens(Latent<super::Tokens>),
}

impl AnyLatent {
    /// The stored kind tag.
    pub fn tag(&self) -> KindTag {
        match self {
            AnyLatent::Continuous(_) => KindTag::Continuous,
            AnyLatent::Tokens(_) => KindTag::Tokens,
        }
    }

    /// The clock, kind-independently.
    pub fn time_base(&self) -> TimeBase {
        match self {
            AnyLatent::Continuous(latent) => latent.time_base(),
            AnyLatent::Tokens(latent) => latent.time_base(),
        }
    }

    /// Number of frames, kind-independently.
    pub fn frames(&self) -> Frames {
        match self {
            AnyLatent::Continuous(latent) => latent.frames(),
            AnyLatent::Tokens(latent) => latent.frames(),
        }
    }

    /// The single checked gate from the dynamic edge into the typed
    /// pipeline. Returns [`Error::KindMismatch`] naming both kinds if the
    /// value holds the other one.
    pub fn downcast<K: LatentKind>(self) -> Result<Latent<K>> {
        let actual = self.tag();
        K::unwrap_any(self).ok_or(Error::KindMismatch {
            expected: K::TAG,
            actual,
        })
    }
}

impl<K: LatentKind> From<Latent<K>> for AnyLatent {
    fn from(latent: Latent<K>) -> Self {
        K::wrap_any(latent)
    }
}

#[cfg(test)]
mod tests {
    use std::num::NonZeroU32;

    use super::*;
    use crate::latent::{Continuous, FrameSeq, TokenGrid, Tokens};
    use crate::time::SampleRate;

    fn tb() -> TimeBase {
        TimeBase::new(
            SampleRate::try_from(48_000).unwrap(),
            NonZeroU32::new(480).unwrap(),
        )
    }

    fn seq(frames: usize, dim: usize) -> FrameSeq {
        let values: Vec<f32> = (0..frames * dim).map(|i| i as f32).collect();
        FrameSeq::new(frames, dim, values).unwrap()
    }

    #[test]
    fn latent_knows_its_clock() {
        let latent = Latent::<Continuous>::new(seq(100, 4), tb());
        assert_eq!(latent.frames(), Frames(100));
        assert!((latent.duration_seconds() - 1.0).abs() < 1e-12);
    }

    #[test]
    fn blocks_cut_exactly_or_reject() {
        let latent = Latent::<Continuous>::new(seq(6, 2), tb());
        let blocks = latent.blocks(Frames(2)).unwrap();
        assert_eq!(blocks.len(), 3);
        assert!(blocks.iter().all(|b| b.frames() == Frames(2)));
        // The middle block is rows 2..4.
        assert_eq!(blocks[1].repr().values(), &[4.0, 5.0, 6.0, 7.0]);

        assert!(latent.blocks(Frames(0)).is_err());
        assert!(latent.blocks(Frames(4)).is_err()); // 6 % 4 != 0
    }

    #[test]
    fn downcast_is_the_one_checked_gate() {
        let any: AnyLatent = Latent::<Continuous>::new(seq(2, 2), tb()).into();
        assert_eq!(any.tag(), KindTag::Continuous);
        assert!(any.clone().downcast::<Continuous>().is_ok());
        let err = any.downcast::<Tokens>().unwrap_err();
        assert!(matches!(
            err,
            Error::KindMismatch {
                expected: KindTag::Tokens,
                actual: KindTag::Continuous
            }
        ));
    }

    #[test]
    fn token_latents_flow_the_same_way() {
        let grid = TokenGrid::new(vec![vec![1, 2, 3], vec![4, 5, 6]]).unwrap();
        let any: AnyLatent = Latent::<Tokens>::new(grid, tb()).into();
        assert_eq!(any.frames(), Frames(3));
        let back = any.downcast::<Tokens>().unwrap();
        assert_eq!(back.repr().codebooks(), 2);
    }

    #[test]
    fn serde_round_trips_typed_and_dynamic() {
        let latent = Latent::<Continuous>::new(seq(2, 3), tb());
        let json = serde_json::to_string(&latent).unwrap();
        assert_eq!(
            serde_json::from_str::<Latent<Continuous>>(&json).unwrap(),
            latent
        );

        let any = AnyLatent::from(latent);
        let json = serde_json::to_string(&any).unwrap();
        assert_eq!(serde_json::from_str::<AnyLatent>(&json).unwrap(), any);
    }
}