helena 0.1.0

//! Causality settled at construction (A4).
//!
//! A [`CausalPlan`] describes a causal topology as host-side data: an ordered
//! stack of layer descriptors whose receptive-field, stride, and lookahead
//! arithmetic are total functions. The latency promise is therefore proven
//! *before a single weight exists*: [`CausalPlan::check`] against a
//! [`LatencyBudget`] mints the [`CheckedPlan`] that a model interpreter
//! (`helena-model`) accepts — an unchecked plan is not buildable.
//!
//! `CheckedPlan` is deliberately *not* serializable: a proof is re-derived on
//! load, never trusted from disk. The plan itself serializes and
//! fingerprints, so the architecture that produced a checkpoint is recorded
//! structurally (NFR-010).

use std::num::NonZeroU32;

use serde::{Deserialize, Serialize};

use crate::error::Error;
use crate::stream::StreamSpec;
use crate::time::{Frames, SampleRate, TimeBase};

/// The morph band's upper edge (PRD §4.2, §12.1): gestural timbre
/// transitions tolerate up to 60 ms of algorithmic latency.
pub const MORPH_BAND_MAX_SECONDS: f64 = 0.060;

/// One layer of a causal stack, described by its temporal geometry only —
/// channel widths are value-level model config, not causality.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CausalLayer {
    /// A causal (left-context) convolution: consumes `(kernel−1)·dilation`
    /// past samples, emits one output per `stride` inputs.
    Conv {
        /// Kernel width.
        kernel: NonZeroU32,
        /// Dilation factor.
        dilation: NonZeroU32,
        /// Temporal downsampling factor.
        stride: NonZeroU32,
    },
    /// A causal sub-pixel upsampler: emits `factor` outputs per input.
    Upsample {
        /// Temporal upsampling factor.
        factor: NonZeroU32,
        /// Kernel width of the underlying stride-1 causal conv.
        kernel: NonZeroU32,
    },
}

impl CausalLayer {
    /// Past context this layer needs, in its own input ticks.
    fn left_context(self) -> u64 {
        match self {
            CausalLayer::Conv {
                kernel, dilation, ..
            } => u64::from(kernel.get() - 1) * u64::from(dilation.get()),
            CausalLayer::Upsample { kernel, .. } => u64::from(kernel.get() - 1),
        }
    }

    /// Future context this layer needs, in its own input ticks. Zero for
    /// every causal layer; a future non-causal descriptor (a centered
    /// window, a PQMF synthesis tail) carries its lookahead here and the
    /// plan arithmetic already accounts for it.
    fn lookahead(self) -> u64 {
        0
    }
}

/// A plan error: why a topology cannot stream within its budget.
#[derive(Debug, thiserror::Error, PartialEq)]
#[non_exhaustive]
pub enum PlanError {
    /// A plan needs at least one layer.
    #[error("a causal plan needs at least one layer")]
    Empty,
    /// Stride or receptive-field arithmetic overflowed `u64` — a topology
    /// this deep is a configuration error, not a model.
    #[error("plan arithmetic overflows: {0}")]
    Overflow(&'static str),
    /// The plan's worst-case latency exceeds the budget.
    #[error(
        "latency {latency_seconds:.4}s exceeds the {max_seconds:.4}s budget \
         (block {block} + lookahead {lookahead})"
    )]
    BudgetExceeded {
        /// The derived worst-case latency.
        latency_seconds: f64,
        /// The budget's ceiling.
        max_seconds: f64,
        /// Latent frames per streaming step.
        block: Frames,
        /// Future context in latent frames.
        lookahead: Frames,
    },
}

impl From<PlanError> for Error {
    fn from(err: PlanError) -> Self {
        Error::validation(err)
    }
}

/// The latency budget a streaming deployment must honor: the clock the
/// latent ticks on, the block size the runtime will push, and the ceiling.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct LatencyBudget {
    /// Sample rate of the synthesized audio.
    pub sample_rate: SampleRate,
    /// Latent frames pushed per streaming step.
    pub block: Frames,
    /// Ceiling in seconds.
    pub max_seconds: f64,
}

impl LatencyBudget {
    /// The morph-band budget (PRD §12.1): the product's default ceiling.
    pub fn morph(sample_rate: SampleRate, block: Frames) -> Self {
        Self {
            sample_rate,
            block,
            max_seconds: MORPH_BAND_MAX_SECONDS,
        }
    }
}

/// An ordered causal topology with total geometry arithmetic.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CausalPlan {
    layers: Vec<CausalLayer>,
}

impl CausalPlan {
    /// Wrap an ordered layer stack. Returns [`PlanError::Empty`] for zero
    /// layers.
    pub fn new(layers: Vec<CausalLayer>) -> Result<Self, PlanError> {
        if layers.is_empty() {
            return Err(PlanError::Empty);
        }
        Ok(Self { layers })
    }

    /// The ordered layers.
    pub fn layers(&self) -> &[CausalLayer] {
        &self.layers
    }

    /// Product of conv strides: input samples per output frame (before
    /// upsampling), or an overflow error.
    pub fn total_stride(&self) -> Result<u64, PlanError> {
        self.layers
            .iter()
            .filter_map(|layer| match layer {
                CausalLayer::Conv { stride, .. } => Some(u64::from(stride.get())),
                CausalLayer::Upsample { .. } => None,
            })
            .try_fold(1u64, |acc, s| acc.checked_mul(s))
            .ok_or(PlanError::Overflow("total stride"))
    }

    /// Product of upsample factors: output samples per input frame, or an
    /// overflow error.
    pub fn total_upsample(&self) -> Result<u64, PlanError> {
        self.layers
            .iter()
            .filter_map(|layer| match layer {
                CausalLayer::Upsample { factor, .. } => Some(u64::from(factor.get())),
                CausalLayer::Conv { .. } => None,
            })
            .try_fold(1u64, |acc, f| acc.checked_mul(f))
            .ok_or(PlanError::Overflow("total upsample"))
    }

    /// Receptive field in *input* ticks: how much past input can influence
    /// one output tick. Computed by walking the stack with an exact rational
    /// jump (stride product / upsample product so far).
    pub fn receptive_field(&self) -> Result<u64, PlanError> {
        // jump = num/den input ticks per current-layer input tick.
        let (mut num, mut den) = (1u64, 1u64);
        let mut field = 1u64;
        for layer in &self.layers {
            let reach = layer.left_context();
            // reach layer-ticks × (num/den) input ticks per layer-tick,
            // rounded up: a partial input tick still requires that input.
            let scaled = reach
                .checked_mul(num)
                .ok_or(PlanError::Overflow("receptive field"))?
                .div_ceil(den);
            field = field
                .checked_add(scaled)
                .ok_or(PlanError::Overflow("receptive field"))?;
            match layer {
                CausalLayer::Conv { stride, .. } => {
                    num = num
                        .checked_mul(u64::from(stride.get()))
                        .ok_or(PlanError::Overflow("receptive field"))?;
                }
                CausalLayer::Upsample { factor, .. } => {
                    den = den
                        .checked_mul(u64::from(factor.get()))
                        .ok_or(PlanError::Overflow("receptive field"))?;
                }
            }
            let g = gcd(num, den);
            num /= g;
            den /= g;
        }
        Ok(field)
    }

    /// Future context in output frames. Zero for a stack of causal layers;
    /// the arithmetic is kept total so a future non-causal descriptor slots
    /// in without a signature change.
    pub fn lookahead(&self) -> Frames {
        Frames(self.layers.iter().map(|l| l.lookahead()).sum())
    }

    /// Prove this plan fits a latency budget, minting the [`CheckedPlan`] a
    /// model interpreter accepts. The stream's clock is derived from the
    /// budget's sample rate and the plan's total stride.
    pub fn check(self, budget: &LatencyBudget) -> Result<CheckedPlan, PlanError> {
        let stride = self.total_stride()?;
        let stride = u32::try_from(stride)
            .ok()
            .and_then(NonZeroU32::new)
            .ok_or(PlanError::Overflow("stride exceeds u32"))?;
        let time_base = TimeBase::new(budget.sample_rate, stride);
        let lookahead = self.lookahead();
        let spec = StreamSpec::new(budget.block, lookahead, time_base)
            .map_err(|_| PlanError::Overflow("zero-frame block"))?;
        let latency = spec.latency_seconds();
        if latency > budget.max_seconds {
            return Err(PlanError::BudgetExceeded {
                latency_seconds: latency,
                max_seconds: budget.max_seconds,
                block: budget.block,
                lookahead,
            });
        }
        Ok(CheckedPlan { plan: self, spec })
    }
}

/// A plan whose latency promise has been proven. Not serializable by design:
/// proofs are re-derived on load ([`CausalPlan::check`]), never trusted from
/// disk.
#[derive(Clone, Debug, PartialEq)]
pub struct CheckedPlan {
    plan: CausalPlan,
    spec: StreamSpec,
}

impl CheckedPlan {
    /// The proven plan.
    pub fn plan(&self) -> &CausalPlan {
        &self.plan
    }

    /// The stream facts the proof derived (block, lookahead, clock).
    pub fn spec(&self) -> StreamSpec {
        self.spec
    }
}

fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b != 0 {
        (a, b) = (b, a % b);
    }
    a.max(1)
}

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

    fn nz(v: u32) -> NonZeroU32 {
        NonZeroU32::new(v).unwrap()
    }

    fn conv(kernel: u32, dilation: u32, stride: u32) -> CausalLayer {
        CausalLayer::Conv {
            kernel: nz(kernel),
            dilation: nz(dilation),
            stride: nz(stride),
        }
    }

    #[test]
    fn stride_and_upsample_products() {
        let plan = CausalPlan::new(vec![
            conv(3, 1, 2),
            conv(3, 2, 4),
            CausalLayer::Upsample {
                factor: nz(8),
                kernel: nz(3),
            },
        ])
        .unwrap();
        assert_eq!(plan.total_stride().unwrap(), 8);
        assert_eq!(plan.total_upsample().unwrap(), 8);
        assert_eq!(plan.lookahead(), Frames(0));
    }

    #[test]
    fn receptive_field_composes_with_stride() {
        // conv(k=3,d=1,s=1): rf 3. Then conv(k=3,d=1,s=2) at jump 1: +2.
        // Then conv(k=3,d=1,s=1) at jump 2: +4.
        let plan = CausalPlan::new(vec![conv(3, 1, 1), conv(3, 1, 2), conv(3, 1, 1)]).unwrap();
        assert_eq!(plan.receptive_field().unwrap(), 1 + 2 + 2 + 4);
    }

    #[test]
    fn empty_and_overflow_fail_closed() {
        assert_eq!(CausalPlan::new(vec![]).unwrap_err(), PlanError::Empty);
        let deep = CausalPlan::new(vec![conv(3, 1, u32::MAX); 3]).unwrap();
        assert!(matches!(deep.total_stride(), Err(PlanError::Overflow(_))));
    }

    #[test]
    fn budget_check_mints_or_rejects() {
        let rate = SampleRate::try_from(48_000).unwrap();
        // Total stride 512 ⇒ 93.75 frames/s ⇒ 10.67 ms per frame.
        let plan = CausalPlan::new(vec![conv(3, 1, 8), conv(3, 1, 8), conv(3, 1, 8)]).unwrap();

        // 3-frame blocks: 32 ms — inside the morph band.
        let checked = plan
            .clone()
            .check(&LatencyBudget::morph(rate, Frames(3)))
            .unwrap();
        assert_eq!(checked.spec().time_base().stride().get(), 512);
        assert!(checked.spec().latency_seconds() < MORPH_BAND_MAX_SECONDS);

        // 8-frame blocks: 85 ms — over budget, named precisely.
        let err = plan
            .check(&LatencyBudget::morph(rate, Frames(8)))
            .unwrap_err();
        assert!(matches!(
            err,
            PlanError::BudgetExceeded {
                block: Frames(8),
                ..
            }
        ));
    }

    #[test]
    fn plans_serialize_proofs_do_not() {
        let plan = CausalPlan::new(vec![conv(3, 2, 2)]).unwrap();
        let json = serde_json::to_string(&plan).unwrap();
        assert_eq!(serde_json::from_str::<CausalPlan>(&json).unwrap(), plan);
        // CheckedPlan has no Serialize impl — enforced by the type system;
        // this test documents the intent.
    }
}