speakrs 0.6.0

#[cfg(any(feature = "coreml", feature = "_metrics"))]
use crate::inference::CoreMlComputeUnits;
#[cfg(feature = "_cuda-libraries")]
use crate::inference::CudaLstmAlgorithm;
use crate::inference::ExecutionMode;
#[cfg(feature = "_cuda")]
use crate::inference::{CudaGraphs, CudaMath};
#[cfg(feature = "_metrics")]
use crate::pipeline::SphereVbxPfConfig;
use crate::pipeline::{ActivityCleanup, AhcConfig, VbxConfig};

/// Speaker clustering model and its valid configuration
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub enum ClusteringBackend {
    /// PLDA-transformed Gaussian variational Bayes clustering
    GaussianVbx(VbxConfig),
    /// Parameter-free spherical variational Bayes clustering
    #[cfg(feature = "_metrics")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
    SphereVbxPf(SphereVbxPfConfig),
}

/// Invalid clustering configuration
#[derive(Debug, Clone, Copy, PartialEq, thiserror::Error)]
pub enum ClusteringConfigError {
    /// Speaker-keep threshold was negative or non-finite
    #[error("speaker keep threshold must be finite and non-negative, got {0}")]
    InvalidKeepThreshold(f64),
}

/// Checked clustering settings owned by the clustering domain
///
/// Replaces the former split `PipelineConfig` fields `vbx`, `experimental_clustering`,
/// `ahc`, and `speaker_keep_threshold`. Negative VBx smoothing maps to
/// [`crate::pipeline::ResponsibilityInitialization::Hard`], zero to
/// [`crate::pipeline::ResponsibilityInitialization::Uniform`], and a positive
/// finite value to [`crate::pipeline::ResponsibilityInitialization::Smoothed`].
#[derive(Debug, Clone, Copy)]
pub struct ClusteringConfig {
    ahc: AhcConfig,
    speaker_keep_threshold: f64,
    backend: ClusteringBackend,
}

impl Default for ClusteringConfig {
    fn default() -> Self {
        Self {
            ahc: AhcConfig::default(),
            speaker_keep_threshold: 1e-7,
            backend: ClusteringBackend::GaussianVbx(VbxConfig::default()),
        }
    }
}

impl ClusteringConfig {
    /// Create a checked clustering configuration
    pub fn new(
        ahc: AhcConfig,
        speaker_keep_threshold: f64,
        backend: ClusteringBackend,
    ) -> Result<Self, ClusteringConfigError> {
        if !(speaker_keep_threshold.is_finite() && speaker_keep_threshold >= 0.0) {
            return Err(ClusteringConfigError::InvalidKeepThreshold(
                speaker_keep_threshold,
            ));
        }
        Ok(Self {
            ahc,
            speaker_keep_threshold,
            backend,
        })
    }

    /// Agglomerative hierarchical clustering settings
    pub const fn ahc(self) -> AhcConfig {
        self.ahc
    }

    /// Minimum speaker activity weight used to keep a speaker
    pub const fn speaker_keep_threshold(self) -> f64 {
        self.speaker_keep_threshold
    }

    /// Clustering backend that will execute
    pub const fn backend(self) -> ClusteringBackend {
        self.backend
    }

    /// Replace the clustering backend
    pub const fn with_backend(mut self, backend: ClusteringBackend) -> Self {
        self.backend = backend;
        self
    }

    /// Replace the AHC settings
    pub const fn with_ahc(mut self, ahc: AhcConfig) -> Self {
        self.ahc = ahc;
        self
    }

    /// Replace the speaker-keep threshold
    pub fn with_speaker_keep_threshold(
        self,
        speaker_keep_threshold: f64,
    ) -> Result<Self, ClusteringConfigError> {
        Self::new(self.ahc, speaker_keep_threshold, self.backend)
    }
}

/// How to map cluster assignments back to per-frame speaker activations
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
pub enum ReconstructMethod {
    /// Standard top-K selection (pyannote-compatible)
    Standard,
    /// Temporal smoothing. If scores are within epsilon, keep the previous speaker.
    Smoothed {
        /// Score difference below which the previous frame's speaker is preferred
        epsilon: f32,
    },
}

/// Tunable parameters for the diarization pipeline
///
/// # Configuration mapping
///
/// - [`Self::activity`] replaces `BinarizeConfig`
/// - [`Self::clustering`] replaces `vbx`, `experimental_clustering`, `ahc`, and
///   `speaker_keep_threshold`
/// - Negative VBx smoothing maps to [`crate::pipeline::ResponsibilityInitialization::Hard`],
///   zero to [`crate::pipeline::ResponsibilityInitialization::Uniform`], and a
///   positive finite value to [`crate::pipeline::ResponsibilityInitialization::Smoothed`]
#[derive(Debug, Clone)]
pub struct PipelineConfig {
    /// Frame-count activity cleanup applied after reconstruction
    pub activity: ActivityCleanup,
    /// Clustering backend, AHC, and speaker-keep settings
    pub clustering: ClusteringConfig,
    /// Minimum single-speaker activity used to select clustering embeddings in experiments
    #[cfg(feature = "_metrics")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
    pub clean_frame_duration: CleanFrameDuration,
    /// Maximum gap in seconds between segments to merge into one
    pub merge_gap: f64,
    /// Strategy for mapping clusters back to frame activations
    pub reconstruct_method: ReconstructMethod,
}

impl Default for PipelineConfig {
    fn default() -> Self {
        Self {
            activity: ActivityCleanup::default(),
            clustering: ClusteringConfig::default(),
            #[cfg(feature = "_metrics")]
            clean_frame_duration: CleanFrameDuration::default(),
            merge_gap: 0.0,
            reconstruct_method: ReconstructMethod::Smoothed { epsilon: 0.1 },
        }
    }
}

/// Minimum single-speaker activity required for a clustering embedding
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct CleanFrameDuration(f64);

impl CleanFrameDuration {
    /// Create a positive finite duration in seconds
    pub fn new(seconds: f64) -> Result<Self, CleanFrameDurationError> {
        if seconds.is_finite() && seconds > 0.0 {
            Ok(Self(seconds))
        } else {
            Err(CleanFrameDurationError(seconds))
        }
    }

    /// Return the duration in seconds
    pub const fn seconds(self) -> f64 {
        self.0
    }

    pub(crate) fn minimum_frames(self) -> f32 {
        (self.0 / FRAME_STEP_SECONDS).floor() as f32
    }
}

impl Default for CleanFrameDuration {
    fn default() -> Self {
        Self(2.0)
    }
}

/// Error returned for a non-positive or non-finite clean-frame duration
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CleanFrameDurationError(f64);

impl std::fmt::Display for CleanFrameDurationError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            formatter,
            "clean-frame duration must be finite and greater than zero, got {}",
            self.0
        )
    }
}

impl std::error::Error for CleanFrameDurationError {}

impl PipelineConfig {
    pub(crate) fn effective_clean_frame_duration(&self) -> CleanFrameDuration {
        #[cfg(feature = "_metrics")]
        {
            self.clean_frame_duration
        }
        #[cfg(not(feature = "_metrics"))]
        {
            CleanFrameDuration::default()
        }
    }

    /// Mode-specific defaults. Fast modes use min-duration filtering to remove
    /// single-frame speaker flicker from the larger step size.
    pub fn for_mode(mode: ExecutionMode) -> Self {
        match mode {
            ExecutionMode::CoreMlFast | ExecutionMode::CudaFast => Self {
                activity: ActivityCleanup::new(3, 3, 0, 0),
                clustering: ClusteringConfig::default().with_backend(
                    ClusteringBackend::GaussianVbx(
                        VbxConfig::default()
                            .with_max_iters(3)
                            .expect("fast-mode iteration count is valid"),
                    ),
                ),
                ..Self::default()
            },
            _ => Self::default(),
        }
    }

    /// Select the clustering backend used by reconstruction
    pub fn with_clustering(mut self, backend: ClusteringBackend) -> Self {
        self.clustering = self.clustering.with_backend(backend);
        self
    }

    /// Return the clustering backend that will execute
    pub const fn clustering_backend(&self) -> ClusteringBackend {
        self.clustering.backend()
    }
}

/// CoreML chunk or per-window inference layout used by metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoreMlChunkLayout {
    /// Two phased chunk models with an exact 1 second window step
    OneSecondPhased,
    /// One-pass chunk models aligned to 25 ResNet frames at a 2 second step
    FastS25,
    /// Per-window embedding with an exact 1 second window step
    PerWindow,
}

/// Fixed-shape CoreML chunk-model set used by metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum CoreMlShapeLadder {
    /// Use every production fixed-shape model
    #[default]
    Full,
    /// Keep the 21, 51, and 111-window one-second models
    Reduced,
}

/// Segmentation worker count used by metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum CoreMlSegmentationWorkers {
    /// Use the production host-derived worker count
    #[default]
    Automatic,
    /// Use four workers
    Four,
    /// Use six workers
    Six,
    /// Use eight workers
    Eight,
}

#[cfg(feature = "_metrics")]
impl CoreMlSegmentationWorkers {
    #[cfg(feature = "coreml")]
    pub(crate) fn resolve(self) -> usize {
        match self {
            Self::Automatic => default_coreml_segmentation_worker_count(),
            Self::Four => 4,
            Self::Six => 6,
            Self::Eight => 8,
        }
    }
}

/// Filterbank preparation worker count used by metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum CoreMlFbankPreparationWorkers {
    /// Use one worker
    One,
    /// Use the production count of two workers
    #[default]
    Two,
    /// Use four workers
    Four,
}

/// Filterbank normalization scope used by metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum CoreMlFbankNormalizationScope {
    /// Normalize all samples in each selected fixed-shape chunk together
    #[default]
    Chunk,
    /// Normalize independent 10-second filterbank segments before stitching them
    TenSecondSegments,
}

#[cfg(feature = "_metrics")]
impl CoreMlFbankPreparationWorkers {
    #[cfg(feature = "coreml")]
    pub(crate) const fn get(self) -> usize {
        match self {
            Self::One => 1,
            Self::Two => 2,
            Self::Four => 4,
        }
    }
}

#[cfg(feature = "_metrics")]
impl CoreMlChunkLayout {
    /// Segmentation step in seconds required by this layout
    pub const fn step_seconds(self) -> f64 {
        match self {
            Self::OneSecondPhased | Self::PerWindow => 1.0,
            Self::FastS25 => 2.0,
        }
    }

    /// Whether this layout uses the native full-audio chunk session ladder
    pub const fn uses_native_chunk_sessions(self) -> bool {
        !matches!(self, Self::PerWindow)
    }

    /// Validate that this layout belongs to the requested execution mode
    pub const fn supports_mode(self, mode: ExecutionMode) -> bool {
        matches!(
            (mode, self),
            (
                ExecutionMode::CoreMl,
                Self::OneSecondPhased | Self::PerWindow
            ) | (ExecutionMode::CoreMlFast, Self::FastS25)
        )
    }
}

/// Layout and shape ladder that can execute together
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ExperimentInferenceLayout {
    OneSecondPhased { shape_ladder: CoreMlShapeLadder },
    PerWindow,
    FastS25,
}

#[cfg(feature = "_metrics")]
impl ExperimentInferenceLayout {
    const fn from_full_layout(layout: CoreMlChunkLayout) -> Self {
        match layout {
            CoreMlChunkLayout::OneSecondPhased => Self::OneSecondPhased {
                shape_ladder: CoreMlShapeLadder::Full,
            },
            CoreMlChunkLayout::PerWindow => Self::PerWindow,
            CoreMlChunkLayout::FastS25 => Self::FastS25,
        }
    }

    fn try_from_parts(
        layout: CoreMlChunkLayout,
        shape_ladder: CoreMlShapeLadder,
    ) -> Result<Self, ExperimentInferenceConfigError> {
        match (layout, shape_ladder) {
            (CoreMlChunkLayout::OneSecondPhased, shape_ladder) => {
                Ok(Self::OneSecondPhased { shape_ladder })
            }
            (CoreMlChunkLayout::PerWindow, CoreMlShapeLadder::Full) => Ok(Self::PerWindow),
            (CoreMlChunkLayout::FastS25, CoreMlShapeLadder::Full) => Ok(Self::FastS25),
            (layout, shape_ladder) => {
                Err(ExperimentInferenceConfigError::IncompatibleShapeLadder {
                    layout,
                    shape_ladder,
                })
            }
        }
    }

    const fn chunk_layout(self) -> CoreMlChunkLayout {
        match self {
            Self::OneSecondPhased { .. } => CoreMlChunkLayout::OneSecondPhased,
            Self::PerWindow => CoreMlChunkLayout::PerWindow,
            Self::FastS25 => CoreMlChunkLayout::FastS25,
        }
    }

    const fn shape_ladder(self) -> CoreMlShapeLadder {
        match self {
            Self::OneSecondPhased { shape_ladder } => shape_ladder,
            Self::PerWindow | Self::FastS25 => CoreMlShapeLadder::Full,
        }
    }
}

/// Typed inference configuration for metrics experiments
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ExperimentInferenceConfig {
    layout: ExperimentInferenceLayout,
    /// Parallel segmentation worker policy
    pub segmentation_workers: CoreMlSegmentationWorkers,
    /// Parallel filterbank preparation worker policy
    pub fbank_preparation_workers: CoreMlFbankPreparationWorkers,
    /// Filterbank normalization scope
    pub fbank_normalization_scope: CoreMlFbankNormalizationScope,
    /// Compute units for native embedding models; filterbank preparation has separate placement
    pub embedding_compute_units: CoreMlComputeUnits,
}

#[cfg(feature = "_metrics")]
impl ExperimentInferenceConfig {
    /// Create an experiment configuration for a fixed CoreML layout
    pub const fn new(coreml_chunk_layout: CoreMlChunkLayout) -> Self {
        Self {
            layout: ExperimentInferenceLayout::from_full_layout(coreml_chunk_layout),
            segmentation_workers: CoreMlSegmentationWorkers::Automatic,
            fbank_preparation_workers: CoreMlFbankPreparationWorkers::Two,
            fbank_normalization_scope: CoreMlFbankNormalizationScope::Chunk,
            embedding_compute_units: CoreMlComputeUnits::All,
        }
    }

    /// Create an experiment configuration with an explicit execution policy
    pub fn with_execution_policy(
        coreml_chunk_layout: CoreMlChunkLayout,
        shape_ladder: CoreMlShapeLadder,
        segmentation_workers: CoreMlSegmentationWorkers,
        fbank_preparation_workers: CoreMlFbankPreparationWorkers,
    ) -> Result<Self, ExperimentInferenceConfigError> {
        Ok(Self {
            layout: ExperimentInferenceLayout::try_from_parts(coreml_chunk_layout, shape_ladder)?,
            segmentation_workers,
            fbank_preparation_workers,
            fbank_normalization_scope: CoreMlFbankNormalizationScope::Chunk,
            embedding_compute_units: CoreMlComputeUnits::All,
        })
    }

    /// Return the CoreML chunk or per-window layout
    pub const fn coreml_chunk_layout(self) -> CoreMlChunkLayout {
        self.layout.chunk_layout()
    }

    /// Return the fixed-shape chunk-model ladder
    pub const fn shape_ladder(self) -> CoreMlShapeLadder {
        self.layout.shape_ladder()
    }

    /// Select the filterbank normalization scope for a controlled comparison
    pub const fn with_fbank_normalization_scope(
        mut self,
        scope: CoreMlFbankNormalizationScope,
    ) -> Self {
        self.fbank_normalization_scope = scope;
        self
    }

    /// Select compute units for native embedding models
    pub const fn with_embedding_compute_units(mut self, units: CoreMlComputeUnits) -> Self {
        self.embedding_compute_units = units;
        self
    }

    /// Return the fixed segmentation step required by this layout
    pub const fn step_seconds(self) -> f64 {
        self.coreml_chunk_layout().step_seconds()
    }

    /// Return the fixed segmentation step required by this layout
    pub const fn segmentation_step_seconds(self) -> f64 {
        self.step_seconds()
    }

    /// Validate this configuration for an execution mode
    pub const fn validate(self, mode: ExecutionMode) -> Result<(), ExperimentInferenceConfigError> {
        let layout = self.coreml_chunk_layout();
        if !layout.supports_mode(mode) {
            return Err(ExperimentInferenceConfigError::IncompatibleMode { mode, layout });
        }

        Ok(())
    }

    /// Validate this configuration for an execution mode and segmentation step
    pub fn validate_for_step(
        self,
        mode: ExecutionMode,
        step_seconds: f64,
    ) -> Result<(), ExperimentInferenceConfigError> {
        self.validate(mode)?;

        let expected = self.step_seconds();
        if !step_seconds.is_finite() || (step_seconds - expected).abs() > 1e-9 {
            return Err(ExperimentInferenceConfigError::IncompatibleStep {
                layout: self.coreml_chunk_layout(),
                expected,
                actual: step_seconds,
            });
        }

        Ok(())
    }
}

#[cfg(feature = "_metrics")]
impl Default for ExperimentInferenceConfig {
    fn default() -> Self {
        Self::new(CoreMlChunkLayout::OneSecondPhased)
    }
}

/// Error returned when a metrics inference configuration is incompatible
#[cfg(feature = "_metrics")]
#[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
pub enum ExperimentInferenceConfigError {
    /// The layout is not supported by the requested execution mode
    IncompatibleMode {
        /// Requested execution mode
        mode: ExecutionMode,
        /// Layout selected for the experiment
        layout: CoreMlChunkLayout,
    },
    /// The step does not match the selected fixed layout
    IncompatibleStep {
        /// Layout selected for the experiment
        layout: CoreMlChunkLayout,
        /// Required step in seconds
        expected: f64,
        /// Requested step in seconds
        actual: f64,
    },
    /// The shape ladder is not defined for the selected layout
    IncompatibleShapeLadder {
        /// Selected inference layout
        layout: CoreMlChunkLayout,
        /// Selected shape ladder
        shape_ladder: CoreMlShapeLadder,
    },
}

#[cfg(feature = "_metrics")]
impl std::fmt::Display for ExperimentInferenceConfigError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::IncompatibleMode { mode, layout } => {
                write!(
                    formatter,
                    "CoreML experiment layout {layout:?} is incompatible with execution mode {mode}"
                )
            }
            Self::IncompatibleStep {
                layout,
                expected,
                actual,
            } => write!(
                formatter,
                "CoreML experiment layout {layout:?} requires a {expected} second step, got {actual}"
            ),
            Self::IncompatibleShapeLadder {
                layout,
                shape_ladder,
            } => write!(
                formatter,
                "CoreML experiment shape ladder {shape_ladder:?} is incompatible with layout {layout:?}"
            ),
        }
    }
}

#[cfg(feature = "_metrics")]
impl std::error::Error for ExperimentInferenceConfigError {}

/// Number of intra-operation threads for one ONNX Runtime session
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OrtThreadCount(std::num::NonZeroI32);

impl OrtThreadCount {
    /// Create a nonzero ONNX Runtime thread count
    pub const fn new(threads: usize) -> Result<Self, OrtThreadCountError> {
        if threads == 0 {
            return Err(OrtThreadCountError::Zero);
        }
        if threads > i32::MAX as usize {
            return Err(OrtThreadCountError::TooLarge(threads));
        }

        Ok(Self(
            std::num::NonZeroI32::new(threads as i32).expect("thread count is nonzero"),
        ))
    }

    /// Return the configured thread count
    pub const fn get(self) -> usize {
        self.0.get() as usize
    }
}

impl Default for OrtThreadCount {
    fn default() -> Self {
        let threads = std::thread::available_parallelism()
            .map(std::num::NonZeroUsize::get)
            .unwrap_or(1)
            .min(4);

        Self(std::num::NonZeroI32::new(threads as i32).expect("thread count is at least one"))
    }
}

/// Invalid ONNX Runtime thread count
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum OrtThreadCountError {
    /// A zero thread count was requested
    #[error("ONNX Runtime thread count must be greater than zero")]
    Zero,
    /// The thread count cannot be represented by the ONNX Runtime C API
    #[error("ONNX Runtime thread count must fit in a positive 32-bit signed integer, got {0}")]
    TooLarge(usize),
}

/// ONNX Runtime CPU filterbank session pool policy for MIGraphX
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum FbankSessionPool {
    /// Size the pool from host parallelism and the per-session thread count
    #[default]
    Automatic,
    /// Do not create a filterbank session pool
    Disabled,
    /// Create the requested nonzero number of sessions
    Fixed(FbankSessionPoolSize),
}

const MAX_FBANK_SESSIONS: usize = 8;

/// Validated fixed filterbank session pool size
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FbankSessionPoolSize(std::num::NonZeroUsize);

impl FbankSessionPoolSize {
    /// Create a session pool size between one and eight
    pub const fn new(sessions: usize) -> Result<Self, FbankSessionPoolSizeError> {
        match std::num::NonZeroUsize::new(sessions) {
            Some(sessions) if sessions.get() <= MAX_FBANK_SESSIONS => Ok(Self(sessions)),
            Some(_) => Err(FbankSessionPoolSizeError::TooLarge(sessions)),
            None => Err(FbankSessionPoolSizeError::Zero),
        }
    }

    /// Return the validated session count
    pub const fn get(self) -> usize {
        self.0.get()
    }
}

impl FbankSessionPool {
    /// Create a fixed session pool with between one and eight sessions
    pub const fn fixed(sessions: usize) -> Result<Self, FbankSessionPoolSizeError> {
        match FbankSessionPoolSize::new(sessions) {
            Ok(sessions) => Ok(Self::Fixed(sessions)),
            Err(error) => Err(error),
        }
    }

    #[cfg(feature = "migraphx")]
    pub(crate) fn resolve(self, threads: OrtThreadCount) -> usize {
        match self {
            Self::Automatic => {
                let available = std::thread::available_parallelism()
                    .map(std::num::NonZeroUsize::get)
                    .unwrap_or(1);

                (available / threads.get()).clamp(1, MAX_FBANK_SESSIONS)
            }
            Self::Disabled => 0,
            Self::Fixed(sessions) => sessions.get(),
        }
    }
}

/// Invalid fixed filterbank session pool size
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum FbankSessionPoolSizeError {
    /// A zero session count was requested
    #[error("fixed filterbank session pool size must be greater than zero")]
    Zero,
    /// The session count exceeds the supported maximum
    #[error("fixed filterbank session pool size must not exceed eight, got {0}")]
    TooLarge(usize),
}

/// Runtime configuration for the diarization pipeline
///
/// Controls execution parameters that can affect numerical output and performance.
#[derive(Debug, Clone)]
// the CUDA modes need per-stage defaults that differ from their field types' defaults
#[cfg_attr(not(feature = "_cuda"), derive(Default))]
pub struct RuntimeConfig {
    /// ONNX Runtime CPU filterbank pool for MIGraphX split inference
    ///
    /// Native CPU inference ignores this setting
    pub fbank_pool: FbankSessionPool,
    /// Intra-operation threads for each MIGraphX ONNX Runtime CPU filterbank session
    ///
    /// Native CPU inference ignores this setting
    pub fbank_threads: OrtThreadCount,
    /// CoreML compute units for native embedding models (CoreML modes only)
    ///
    /// This also applies to [`EmbeddingModel::embed`](crate::inference::EmbeddingModel::embed)
    /// and `embed_masked` in CoreML modes. `CpuOnly` gives the closest match to the CPU backend
    #[cfg(feature = "coreml")]
    #[cfg_attr(docsrs, doc(cfg(feature = "coreml")))]
    pub chunk_emb_compute_units: CoreMlComputeUnits,
    /// Precision of the segmentation model's cuBLAS GEMMs, cuDNN convolutions and cuDNN
    /// LSTM (CUDA modes only)
    ///
    /// [`CudaMath::Fp32`] by default. [`CudaMath::Tf32`] is faster on Ampere and newer
    /// GPUs, but it moves the segmentation logits by up to about 0.23 and, on
    /// VoxConverse-dev, made one file's DER 4.5 points worse
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub cuda_segmentation_math: CudaMath,
    /// Precision of the embedding model's cuDNN convolutions and cuBLAS embedding layer
    /// (CUDA modes only)
    ///
    /// [`CudaMath::Tf32`] by default, as the ONNX Runtime CUDA path ran it; embeddings
    /// stay within a cosine distance of about 1e-6 of FP32 and DER matched FP32 on all
    /// 216 VoxConverse-dev files. [`CudaMath::Fp32`] matches
    /// the references most closely, and is the choice to compare against if clustering
    /// looks off, since embedding drift can change PLDA/VBx clustering. The filterbank
    /// always runs in FP32, because TF32 moves its log-mel values by up to 2.7
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub cuda_embedding_math: CudaMath,
    /// cuDNN RNN algorithm for the segmentation model's LSTM layers (CUDA modes only)
    #[cfg(feature = "_cuda-libraries")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda-libraries")))]
    pub cuda_lstm_algorithm: CudaLstmAlgorithm,
    /// Whether the CUDA modes capture and replay CUDA graphs (CUDA modes only)
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub cuda_graphs: CudaGraphs,
    /// Optional typed inference layout for metrics experiments
    #[cfg(feature = "_metrics")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
    pub experiment: Option<ExperimentInferenceConfig>,
}

#[cfg(feature = "_cuda")]
impl Default for RuntimeConfig {
    fn default() -> Self {
        Self {
            fbank_pool: FbankSessionPool::default(),
            fbank_threads: OrtThreadCount::default(),
            #[cfg(feature = "coreml")]
            chunk_emb_compute_units: CoreMlComputeUnits::default(),
            // segmentation TF32 failed the per-file DER gate on VoxConverse-dev while
            // embedding TF32 and the persistent LSTM passed, see the CUDA DER report
            cuda_segmentation_math: CudaMath::Fp32,
            cuda_embedding_math: CudaMath::Tf32,
            #[cfg(feature = "_cuda-libraries")]
            cuda_lstm_algorithm: CudaLstmAlgorithm::PersistStaticSmallH,
            cuda_graphs: CudaGraphs::Enabled,
            #[cfg(feature = "_metrics")]
            experiment: None,
        }
    }
}

impl RuntimeConfig {
    /// Select the ONNX Runtime CPU filterbank session pool for MIGraphX
    pub const fn with_fbank_pool(mut self, pool: FbankSessionPool) -> Self {
        self.fbank_pool = pool;
        self
    }

    /// Select the thread count for each MIGraphX ONNX Runtime CPU filterbank session
    pub const fn with_fbank_threads(mut self, threads: OrtThreadCount) -> Self {
        self.fbank_threads = threads;
        self
    }

    /// Select the segmentation precision of the CUDA modes
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub const fn with_cuda_segmentation_math(mut self, math: CudaMath) -> Self {
        self.cuda_segmentation_math = math;
        self
    }

    /// Select the embedding precision of the CUDA modes
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub const fn with_cuda_embedding_math(mut self, math: CudaMath) -> Self {
        self.cuda_embedding_math = math;
        self
    }

    /// Select the segmentation LSTM algorithm of the CUDA modes
    #[cfg(feature = "_cuda-libraries")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda-libraries")))]
    pub const fn with_cuda_lstm_algorithm(mut self, algorithm: CudaLstmAlgorithm) -> Self {
        self.cuda_lstm_algorithm = algorithm;
        self
    }

    /// Turn CUDA graph capture on or off in the CUDA modes
    #[cfg(feature = "_cuda")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_cuda")))]
    pub const fn with_cuda_graphs(mut self, graphs: CudaGraphs) -> Self {
        self.cuda_graphs = graphs;
        self
    }

    /// Set a typed inference layout for metrics experiments
    #[cfg(feature = "_metrics")]
    #[cfg_attr(docsrs, doc(cfg(feature = "_metrics")))]
    pub fn with_experiment(mut self, experiment: ExperimentInferenceConfig) -> Self {
        self.experiment = Some(experiment);
        self
    }

    #[cfg(feature = "coreml")]
    pub(crate) fn coreml_embedding_compute_units(&self) -> CoreMlComputeUnits {
        #[cfg(feature = "_metrics")]
        if let Some(experiment) = self.experiment {
            return experiment.embedding_compute_units;
        }

        self.chunk_emb_compute_units
    }

    #[cfg(feature = "coreml")]
    pub(crate) fn coreml_chunk_execution_policy(&self) -> CoreMlChunkExecutionPolicy {
        #[cfg(feature = "_metrics")]
        if let Some(experiment) = self.experiment {
            return CoreMlChunkExecutionPolicy {
                segmentation_workers: experiment.segmentation_workers.resolve(),
                fbank_preparation_workers: experiment.fbank_preparation_workers.get(),
                fbank_normalization_scope: match experiment.fbank_normalization_scope {
                    CoreMlFbankNormalizationScope::Chunk => ChunkFbankNormalizationScope::Chunk,
                    CoreMlFbankNormalizationScope::TenSecondSegments => {
                        ChunkFbankNormalizationScope::TenSecondSegments
                    }
                },
            };
        }

        CoreMlChunkExecutionPolicy::default()
    }
}

#[derive(Clone, Copy)]
#[cfg(feature = "coreml")]
pub(crate) struct CoreMlChunkExecutionPolicy {
    pub(crate) segmentation_workers: usize,
    pub(crate) fbank_preparation_workers: usize,
    pub(crate) fbank_normalization_scope: ChunkFbankNormalizationScope,
}

#[derive(Clone, Copy, Eq, PartialEq)]
#[cfg(feature = "coreml")]
pub(crate) enum ChunkFbankNormalizationScope {
    Chunk,
    #[cfg(feature = "_metrics")]
    TenSecondSegments,
}

#[cfg(feature = "coreml")]
impl ChunkFbankNormalizationScope {
    pub(crate) const fn uses_chunk_scope(self) -> bool {
        match self {
            Self::Chunk => true,
            #[cfg(feature = "_metrics")]
            Self::TenSecondSegments => false,
        }
    }
}

#[cfg(feature = "coreml")]
impl Default for CoreMlChunkExecutionPolicy {
    fn default() -> Self {
        Self {
            segmentation_workers: default_coreml_segmentation_worker_count(),
            fbank_preparation_workers: 2,
            fbank_normalization_scope: ChunkFbankNormalizationScope::Chunk,
        }
    }
}

#[cfg(feature = "coreml")]
fn default_coreml_segmentation_worker_count() -> usize {
    std::thread::available_parallelism()
        .map(usize::from)
        .unwrap_or(4)
        .min(8)
}

/// Segmentation step size in seconds for the selected execution mode
pub const fn segmentation_step_seconds(mode: ExecutionMode) -> f64 {
    match mode {
        ExecutionMode::CoreMlFast | ExecutionMode::CudaFast => FAST_SEGMENTATION_STEP_SECONDS,
        ExecutionMode::CoreMl => COREML_SEGMENTATION_STEP_SECONDS,
        ExecutionMode::Cuda => CUDA_SEGMENTATION_STEP_SECONDS,
        ExecutionMode::MiGraphX => CUDA_SEGMENTATION_STEP_SECONDS,
        ExecutionMode::Cpu => SEGMENTATION_STEP_SECONDS,
    }
}

/// Sliding window length for segmentation model input, in seconds
pub const SEGMENTATION_WINDOW_SECONDS: f64 = 10.0;
/// Default sliding window step for segmentation, in seconds
pub const SEGMENTATION_STEP_SECONDS: f64 = 1.0;
/// CoreML step in seconds
///
/// The chunk embedding model uses two aligned phases to support this exact step
pub const COREML_SEGMENTATION_STEP_SECONDS: f64 = 1.0;
/// CUDA segmentation step, in seconds
pub const CUDA_SEGMENTATION_STEP_SECONDS: f64 = 1.0;
/// Step size for fast modes, in seconds
pub const FAST_SEGMENTATION_STEP_SECONDS: f64 = 2.0;
/// Duration of each output frame from the segmentation model, in seconds
pub const FRAME_DURATION_SECONDS: f64 = 0.0619375;
/// Hop between consecutive output frames from the segmentation model, in seconds
pub const FRAME_STEP_SECONDS: f64 = 0.016875;

/// Minimum speaker activity (sum of weights) to run embedding inference.
/// Speakers below this threshold are skipped because their NaN embedding is filtered out later
pub(crate) const MIN_SPEAKER_ACTIVITY: f32 = 10.0;

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

    #[test]
    fn ort_thread_count_rejects_zero() {
        assert_eq!(OrtThreadCount::new(0), Err(OrtThreadCountError::Zero));
        assert_eq!(OrtThreadCount::new(3).unwrap().get(), 3);
    }

    #[test]
    fn ort_thread_count_stays_within_the_c_api_range() {
        assert_eq!(
            OrtThreadCount::new(i32::MAX as usize).unwrap().get(),
            i32::MAX as usize
        );
        assert_eq!(
            OrtThreadCount::new(i32::MAX as usize + 1),
            Err(OrtThreadCountError::TooLarge(i32::MAX as usize + 1))
        );
    }

    #[cfg(feature = "migraphx")]
    #[test]
    fn fbank_pool_models_disabled_automatic_and_fixed_policies() {
        let threads = OrtThreadCount::new(i32::MAX as usize).unwrap();

        assert_eq!(FbankSessionPool::Disabled.resolve(threads), 0);
        assert_eq!(FbankSessionPool::Automatic.resolve(threads), 1);
        assert_eq!(
            FbankSessionPool::fixed(0),
            Err(FbankSessionPoolSizeError::Zero)
        );
        assert_eq!(
            FbankSessionPool::fixed(MAX_FBANK_SESSIONS + 1),
            Err(FbankSessionPoolSizeError::TooLarge(MAX_FBANK_SESSIONS + 1))
        );
        assert_eq!(
            FbankSessionPoolSize::new(MAX_FBANK_SESSIONS).unwrap().get(),
            MAX_FBANK_SESSIONS
        );
        assert_eq!(FbankSessionPool::fixed(3).unwrap().resolve(threads), 3);
    }

    #[test]
    fn clustering_config_rejects_invalid_keep_threshold() {
        for threshold in [f64::NAN, f64::INFINITY, -1.0] {
            assert!(
                ClusteringConfig::new(
                    AhcConfig::default(),
                    threshold,
                    ClusteringBackend::GaussianVbx(VbxConfig::default())
                )
                .is_err()
            );
        }
        assert_eq!(ClusteringConfig::default().speaker_keep_threshold(), 1e-7);
    }

    fn gaussian_vbx(backend: ClusteringBackend) -> VbxConfig {
        match backend {
            ClusteringBackend::GaussianVbx(vbx) => vbx,
            #[cfg(feature = "_metrics")]
            _ => panic!("clustering must be gaussian"),
        }
    }

    #[test]
    fn pipeline_defaults_keep_gaussian_vbx_and_fixed_mode_steps() {
        let standard = gaussian_vbx(PipelineConfig::default().clustering_backend());
        let fast =
            gaussian_vbx(PipelineConfig::for_mode(ExecutionMode::CoreMlFast).clustering_backend());

        assert_eq!(standard.max_iters(), 20);
        assert_eq!(fast.max_iters(), 3);
        assert_eq!(segmentation_step_seconds(ExecutionMode::CoreMl), 1.0);
        assert_eq!(segmentation_step_seconds(ExecutionMode::CoreMlFast), 2.0);
    }

    #[test]
    fn default_clean_frame_duration_preserves_the_previous_frame_threshold() {
        assert_eq!(CleanFrameDuration::default().minimum_frames(), 118.0);
    }

    #[test]
    fn clean_frame_duration_rejects_invalid_values() {
        for seconds in [0.0, -1.0, f64::NAN, f64::INFINITY] {
            assert!(CleanFrameDuration::new(seconds).is_err());
        }
    }
}

#[cfg(all(test, feature = "_metrics"))]
mod tests {
    use super::*;

    #[test]
    fn runtime_defaults_to_mode_owned_inference() {
        assert_eq!(RuntimeConfig::default().experiment, None);
    }

    #[test]
    fn experiment_layouts_have_fixed_steps_and_modes() {
        let cases = [
            (
                CoreMlChunkLayout::OneSecondPhased,
                ExecutionMode::CoreMl,
                1.0,
            ),
            (CoreMlChunkLayout::FastS25, ExecutionMode::CoreMlFast, 2.0),
            (CoreMlChunkLayout::PerWindow, ExecutionMode::CoreMl, 1.0),
        ];

        for (layout, mode, expected_step) in cases {
            let config = ExperimentInferenceConfig::new(layout);
            assert_eq!(config.step_seconds(), expected_step);
            assert!(config.validate(mode).is_ok());
        }
    }

    #[test]
    fn experiment_execution_policy_preserves_production_defaults() {
        let config = ExperimentInferenceConfig::new(CoreMlChunkLayout::OneSecondPhased);

        assert_eq!(config.shape_ladder(), CoreMlShapeLadder::Full);
        assert_eq!(
            config.segmentation_workers,
            CoreMlSegmentationWorkers::Automatic
        );
        assert_eq!(
            config.fbank_preparation_workers,
            CoreMlFbankPreparationWorkers::Two
        );
        assert_eq!(
            config.fbank_normalization_scope,
            CoreMlFbankNormalizationScope::Chunk
        );
        assert_eq!(config.embedding_compute_units, CoreMlComputeUnits::All);

        #[cfg(feature = "coreml")]
        {
            let resolved = RuntimeConfig::default()
                .with_experiment(config)
                .coreml_chunk_execution_policy();
            assert_eq!(resolved.fbank_preparation_workers, 2);
            assert!(resolved.segmentation_workers > 0);
            assert!(matches!(
                resolved.fbank_normalization_scope,
                ChunkFbankNormalizationScope::Chunk
            ));
        }
    }

    #[cfg(feature = "coreml")]
    #[test]
    fn experiment_embedding_compute_units_override_runtime_fallback() {
        let experiment = ExperimentInferenceConfig::new(CoreMlChunkLayout::PerWindow)
            .with_embedding_compute_units(CoreMlComputeUnits::CpuOnly);
        let runtime = RuntimeConfig {
            chunk_emb_compute_units: CoreMlComputeUnits::CpuAndNeuralEngine,
            experiment: Some(experiment),
            ..RuntimeConfig::default()
        };

        assert_eq!(
            runtime.coreml_embedding_compute_units(),
            CoreMlComputeUnits::CpuOnly
        );

        let runtime = RuntimeConfig {
            chunk_emb_compute_units: CoreMlComputeUnits::CpuAndNeuralEngine,
            experiment: None,
            ..RuntimeConfig::default()
        };
        assert_eq!(
            runtime.coreml_embedding_compute_units(),
            CoreMlComputeUnits::CpuAndNeuralEngine
        );
    }

    #[test]
    fn reduced_shape_ladder_rejects_unsupported_layouts() {
        let config = ExperimentInferenceConfig::with_execution_policy(
            CoreMlChunkLayout::FastS25,
            CoreMlShapeLadder::Reduced,
            CoreMlSegmentationWorkers::Four,
            CoreMlFbankPreparationWorkers::One,
        );

        assert!(matches!(
            config,
            Err(ExperimentInferenceConfigError::IncompatibleShapeLadder { .. })
        ));
    }

    #[test]
    fn experiment_validation_rejects_incompatible_mode_and_step() {
        let fast = ExperimentInferenceConfig::new(CoreMlChunkLayout::FastS25);
        assert!(matches!(
            fast.validate(ExecutionMode::CoreMl),
            Err(ExperimentInferenceConfigError::IncompatibleMode { .. })
        ));

        let aligned = ExperimentInferenceConfig::new(CoreMlChunkLayout::OneSecondPhased);
        assert!(matches!(
            aligned.validate_for_step(ExecutionMode::CoreMl, 1.04),
            Err(ExperimentInferenceConfigError::IncompatibleStep { .. })
        ));
    }

    #[test]
    fn per_window_layouts_disable_native_chunk_sessions() {
        assert!(!CoreMlChunkLayout::PerWindow.uses_native_chunk_sessions());
        assert!(CoreMlChunkLayout::OneSecondPhased.uses_native_chunk_sessions());
    }
}