helena 0.1.0

//! The swappable component contracts (PRD §13, NFR-020).
//!
//! These are the stable seams of the platform. Concrete implementations live
//! in the outer crates and in user plugins:
//!
//! - [`SourceEncoder`] — `helena-data` baseline encoders and user modules.
//! - [`AudioEncoder`] / [`AudioDecoder`] / [`AudioCodec`] — the waveform
//!   interface: the trained causal VAE (`helena-infer::VaeCodec`, the MVP
//!   engine) and analytic/frozen codecs (`helena-data`, the §16.4 baseline)
//!   implement the *same* seam — the proof that both architecture postures
//!   collapse onto one vocabulary.
//! - [`LatentGenerator`] — the learned mapping `f: Z_data → Z_audio`
//!   (PRD §8): a thin control→latent map on the surface posture, a diffusion
//!   prior on the baseline posture.
//!
//! The seams are generic over the latent kind, so a token generator wired to
//! a continuous decoder is a type error, not a runtime discovery. There is
//! deliberately no differentiable-decode method on [`AudioCodec`]: keeping
//! gradients out of the frozen decode path is a structural rule, not a
//! discipline.

use serde::{Deserialize, Serialize};

use crate::error::{Error, Result};
use crate::latent::{Conditioning, Latent, LatentKind, Metadata};
use crate::signal::Pcm;
use crate::tensor::Tensor;
use crate::time::TimeBase;

/// A batch of source items presented to a [`SourceEncoder`].
///
/// Each item holds one source's features as a raw [`Tensor`]; the layout (a
/// flat vector for tabular, `[steps, channels]` for time series, …) is the
/// encoder's contract — source data is genuinely heterogeneous, so the raw
/// carrier is the honest type here (A2). Items are *ordered evidence*:
/// manifest order is preserved so training, evaluation, and report rows join
/// without re-sorting. Fields are private so an encoder can trust the batch
/// it validated.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct SourceBatch {
    items: Vec<Tensor>,
    metadata: Metadata,
}

impl SourceBatch {
    /// Build a batch from per-item feature tensors, in manifest order.
    pub fn new(items: Vec<Tensor>) -> Self {
        Self {
            items,
            metadata: Metadata::default(),
        }
    }

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

    /// Per-item feature tensors, in manifest order.
    pub fn items(&self) -> &[Tensor] {
        &self.items
    }

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

    /// Number of items in the batch.
    pub fn len(&self) -> usize {
        self.items.len()
    }

    /// Whether the batch is empty.
    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }
}

/// Converts source data into conditioning latents (PRD §13.1).
///
/// Encoders are *ready at construction* (fitted parameters are frozen
/// artifacts loaded beforehand); the runtime contract is pure.
pub trait SourceEncoder {
    /// Encode a batch of source items into [`Conditioning`].
    fn encode(&self, batch: &SourceBatch) -> Result<Conditioning>;
}

impl<E: SourceEncoder + ?Sized> SourceEncoder for Box<E> {
    fn encode(&self, batch: &SourceBatch) -> Result<Conditioning> {
        (**self).encode(batch)
    }
}

/// A component with a constructor-fixed latent clock.
///
/// One method shared by both codec directions, so an encoder and decoder on
/// the same type cannot disagree about the clock they tick on.
pub trait Clocked {
    /// The latent clock: sample rate over stride, fixed at construction.
    fn time_base(&self) -> TimeBase;
}

/// Analysis direction: waveform → latent (PRD §13.2).
pub trait AudioEncoder<K: LatentKind>: Clocked {
    /// Encode a waveform into a latent on [`Clocked::time_base`].
    fn encode(&self, pcm: &Pcm) -> Result<Latent<K>>;
}

/// Synthesis direction: latent → waveform (PRD §13.2).
pub trait AudioDecoder<K: LatentKind>: Clocked {
    /// Decode a latent back into a waveform.
    fn decode(&self, latent: &Latent<K>) -> Result<Pcm>;
}

/// A codec is both directions sharing one clock. Blanket-implemented: any
/// type that can encode and decode the same kind *is* a codec — there is no
/// separate conformance to forget.
pub trait AudioCodec<K: LatentKind>: AudioEncoder<K> + AudioDecoder<K> {}

impl<K: LatentKind, C: AudioEncoder<K> + AudioDecoder<K> + ?Sized> AudioCodec<K> for C {}

/// Maps conditioning latents to audio latents (PRD §13.3): the learned
/// mapping `f: Z_data → Z_audio`, whatever its posture.
pub trait LatentGenerator<K: LatentKind> {
    /// Generate a latent from conditioning and per-call parameters.
    ///
    /// Determinism is a *value*: [`GenerationParams::seed`] is a
    /// [`SeedPolicy`], and an implementation must replay bit-identically for
    /// [`SeedPolicy::replayable`] policies (PRD FR-031, NFR-012) and may
    /// yield distinct valid outputs otherwise (FR-032).
    fn generate(&self, cond: &Conditioning, params: &GenerationParams) -> Result<Latent<K>>;
}

impl<K: LatentKind, G: LatentGenerator<K> + ?Sized> LatentGenerator<K> for Box<G> {
    fn generate(&self, cond: &Conditioning, params: &GenerationParams) -> Result<Latent<K>> {
        (**self).generate(cond, params)
    }
}

/// A random seed for replayable generation.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Seed(pub u64);

/// A strictly positive, finite sampling temperature.
///
/// Zero temperature is not a `Temperature` — determinism is
/// [`SeedPolicy::Deterministic`], so "seeded but temperature zero" (the old
/// ambiguous state) is unrepresentable.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "f32", into = "f32")]
pub struct Temperature(f32);

impl Temperature {
    /// Wrap a temperature. Returns [`Error::Validation`] unless
    /// `0 < value` and the value is finite.
    pub fn new(value: f32) -> Result<Self> {
        if value.is_finite() && value > 0.0 {
            Ok(Self(value))
        } else {
            Err(Error::validation(format!(
                "temperature must be finite and positive, got {value}"
            )))
        }
    }

    /// The raw value.
    pub fn get(self) -> f32 {
        self.0
    }
}

impl TryFrom<f32> for Temperature {
    type Error = Error;

    fn try_from(value: f32) -> Result<Self> {
        Self::new(value)
    }
}

impl From<Temperature> for f32 {
    fn from(t: Temperature) -> f32 {
        t.get()
    }
}

/// The conditioning-strength control (PRD FR-033): `1.0` applies the
/// conditioning as-is, lower weakens, higher amplifies. Finite, non-negative.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "f32", into = "f32")]
pub struct GuidanceScale(f32);

impl GuidanceScale {
    /// The neutral scale: conditioning applied as-is.
    pub const NEUTRAL: GuidanceScale = GuidanceScale(1.0);

    /// Wrap a guidance scale. Returns [`Error::Validation`] unless the value
    /// is finite and non-negative.
    pub fn new(value: f32) -> Result<Self> {
        if value.is_finite() && value >= 0.0 {
            Ok(Self(value))
        } else {
            Err(Error::validation(format!(
                "guidance scale must be finite and non-negative, got {value}"
            )))
        }
    }

    /// The raw value.
    pub fn get(self) -> f32 {
        self.0
    }

    /// Whether this is exactly the neutral scale (no separate null pass).
    pub fn is_neutral(self) -> bool {
        self.0 == 1.0
    }
}

impl TryFrom<f32> for GuidanceScale {
    type Error = Error;

    fn try_from(value: f32) -> Result<Self> {
        Self::new(value)
    }
}

impl From<GuidanceScale> for f32 {
    fn from(g: GuidanceScale) -> f32 {
        g.get()
    }
}

/// A strictly positive, finite clip duration in seconds.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "f32", into = "f32")]
pub struct ClipDuration(f32);

impl ClipDuration {
    /// Wrap a duration. Returns [`Error::Validation`] unless
    /// `0 < seconds` and the value is finite.
    pub fn new(seconds: f32) -> Result<Self> {
        if seconds.is_finite() && seconds > 0.0 {
            Ok(Self(seconds))
        } else {
            Err(Error::validation(format!(
                "clip duration must be finite and positive, got {seconds}"
            )))
        }
    }

    /// The duration in seconds.
    pub fn seconds(self) -> f32 {
        self.0
    }
}

impl TryFrom<f32> for ClipDuration {
    type Error = Error;

    fn try_from(seconds: f32) -> Result<Self> {
        Self::new(seconds)
    }
}

impl From<ClipDuration> for f32 {
    fn from(d: ClipDuration) -> f32 {
        d.seconds()
    }
}

/// How randomness enters one generation call. The three legal states are the
/// three variants; illegal combinations (a seed that a zero temperature
/// would silently ignore) are unrepresentable.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SeedPolicy {
    /// Fully deterministic: same conditioning ⇒ identical output (FR-031).
    Deterministic,
    /// Seeded stochastic: replayable variety — same conditioning and seed ⇒
    /// identical output (NFR-012).
    Seeded {
        /// The replay seed.
        seed: Seed,
        /// Sampling temperature (strictly positive).
        temperature: Temperature,
    },
    /// OS-entropy stochastic: distinct valid outputs across calls (FR-032).
    Entropy {
        /// Sampling temperature (strictly positive).
        temperature: Temperature,
    },
}

impl SeedPolicy {
    /// Whether this policy promises bit-identical replay.
    pub fn replayable(&self) -> bool {
        !matches!(self, SeedPolicy::Entropy { .. })
    }

    /// The temperature, or `None` for the deterministic policy.
    pub fn temperature(&self) -> Option<Temperature> {
        match self {
            SeedPolicy::Deterministic => None,
            SeedPolicy::Seeded { temperature, .. } | SeedPolicy::Entropy { temperature } => {
                Some(*temperature)
            }
        }
    }

    /// The seed, if one was fixed.
    pub fn seed(&self) -> Option<Seed> {
        match self {
            SeedPolicy::Seeded { seed, .. } => Some(*seed),
            _ => None,
        }
    }
}

/// Per-call generation parameters (PRD §13.3). Explicit by design: there is
/// no `Default` — a duration and a randomness policy are decisions the
/// calling boundary must make, not defaults it can forget.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
pub struct GenerationParams {
    /// Requested clip duration.
    pub duration: ClipDuration,
    /// How randomness enters this call.
    pub seed: SeedPolicy,
    /// Conditioning strength (FR-033).
    pub guidance: GuidanceScale,
}

impl GenerationParams {
    /// Deterministic parameters at neutral guidance — the replayable
    /// baseline call.
    pub fn deterministic(duration: ClipDuration) -> Self {
        Self {
            duration,
            seed: SeedPolicy::Deterministic,
            guidance: GuidanceScale::NEUTRAL,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::latent::{Continuous, FrameSeq};
    use crate::time::SampleRate;
    use std::num::NonZeroU32;

    #[test]
    fn newtypes_reject_illegal_values() {
        assert!(Temperature::new(0.0).is_err());
        assert!(Temperature::new(-1.0).is_err());
        assert!(Temperature::new(f32::NAN).is_err());
        assert!(Temperature::new(0.8).is_ok());

        assert!(GuidanceScale::new(-0.1).is_err());
        assert!(GuidanceScale::new(f32::INFINITY).is_err());
        assert!(GuidanceScale::new(0.0).is_ok());
        assert!(GuidanceScale::NEUTRAL.is_neutral());

        assert!(ClipDuration::new(0.0).is_err());
        assert!(ClipDuration::new(12.5).is_ok());
    }

    #[test]
    fn seed_policy_states_are_exactly_the_legal_ones() {
        let det = SeedPolicy::Deterministic;
        assert!(det.replayable());
        assert_eq!(det.temperature(), None);

        let seeded = SeedPolicy::Seeded {
            seed: Seed(7),
            temperature: Temperature::new(1.0).unwrap(),
        };
        assert!(seeded.replayable());
        assert_eq!(seeded.seed(), Some(Seed(7)));

        let entropy = SeedPolicy::Entropy {
            temperature: Temperature::new(0.5).unwrap(),
        };
        assert!(!entropy.replayable());

        // Serde gates hold on the nested newtypes.
        assert!(
            serde_json::from_str::<SeedPolicy>(r#"{"seeded":{"seed":1,"temperature":0.0}}"#)
                .is_err()
        );
    }

    #[test]
    fn boxed_dyn_generator_forwards_to_inner() {
        struct Echo(Latent<Continuous>);
        impl LatentGenerator<Continuous> for Echo {
            fn generate(
                &self,
                _cond: &Conditioning,
                _params: &GenerationParams,
            ) -> Result<Latent<Continuous>> {
                Ok(self.0.clone())
            }
        }

        let tb = crate::time::TimeBase::new(
            SampleRate::try_from(48_000).unwrap(),
            NonZeroU32::new(480).unwrap(),
        );
        let latent = Latent::new(FrameSeq::new(1, 2, vec![0.25, -0.5]).unwrap(), tb);
        let boxed: Box<dyn LatentGenerator<Continuous>> = Box::new(Echo(latent.clone()));
        let params = GenerationParams::deterministic(ClipDuration::new(1.0).unwrap());
        let out = boxed.generate(&Conditioning::default(), &params).unwrap();
        assert_eq!(out, latent);
    }
}