sevensense-embedding 0.1.0

//! Domain entities for the embedding bounded context.
//!
//! This module defines the core domain entities:
//! - `Embedding`: A 1536-dimensional vector representation of an audio segment
//! - `EmbeddingModel`: Configuration and metadata for Perch 2.0 ONNX model
//! - `EmbeddingBatch`: Collection of embeddings processed together

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::EMBEDDING_DIM;

/// Unique identifier for an embedding
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct EmbeddingId(Uuid);

impl EmbeddingId {
    /// Create a new unique embedding ID
    #[must_use]
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }

    /// Create from an existing UUID
    #[must_use]
    pub const fn from_uuid(uuid: Uuid) -> Self {
        Self(uuid)
    }

    /// Get the inner UUID
    #[must_use]
    pub const fn as_uuid(&self) -> &Uuid {
        &self.0
    }
}

impl Default for EmbeddingId {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for EmbeddingId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Unique identifier for an audio segment
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SegmentId(Uuid);

impl SegmentId {
    /// Create a new unique segment ID
    #[must_use]
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }

    /// Create from an existing UUID
    #[must_use]
    pub const fn from_uuid(uuid: Uuid) -> Self {
        Self(uuid)
    }

    /// Get the inner UUID
    #[must_use]
    pub const fn as_uuid(&self) -> &Uuid {
        &self.0
    }
}

impl Default for SegmentId {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for SegmentId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Storage tier for embeddings based on access patterns
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum StorageTier {
    /// Hot storage: frequently accessed, lowest latency
    /// Full f32 precision, in-memory or SSD
    Hot,

    /// Warm storage: occasional access, moderate latency
    /// F16 quantized, SSD storage
    Warm,

    /// Cold storage: rare access, higher latency acceptable
    /// INT8 quantized, archive storage
    Cold,
}

impl Default for StorageTier {
    fn default() -> Self {
        Self::Hot
    }
}

impl StorageTier {
    /// Get the bytes per dimension for this tier
    #[must_use]
    pub const fn bytes_per_dim(&self) -> usize {
        match self {
            Self::Hot => 4,   // f32
            Self::Warm => 2,  // f16
            Self::Cold => 1,  // i8
        }
    }

    /// Get the total bytes for an embedding in this tier
    #[must_use]
    pub const fn embedding_bytes(&self) -> usize {
        self.bytes_per_dim() * EMBEDDING_DIM
    }
}

/// Timestamp type alias for consistency
pub type Timestamp = DateTime<Utc>;

/// A 1536-dimensional embedding vector representing an audio segment.
///
/// This is the aggregate root for the embedding context. Each embedding
/// is generated by the Perch 2.0 model from a preprocessed audio segment.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Embedding {
    /// Unique identifier for this embedding
    pub id: EmbeddingId,

    /// Reference to the source audio segment
    pub segment_id: SegmentId,

    /// The 1536-dimensional embedding vector (L2 normalized)
    pub vector: Vec<f32>,

    /// Version of the model used to generate this embedding
    pub model_version: String,

    /// When this embedding was created
    pub created_at: Timestamp,

    /// Storage tier for this embedding
    pub tier: StorageTier,

    /// Additional metadata about the embedding
    pub metadata: EmbeddingMetadata,
}

impl Embedding {
    /// Create a new embedding with the given parameters
    ///
    /// # Errors
    ///
    /// Returns an error if the vector dimension is not 1536
    pub fn new(
        segment_id: SegmentId,
        vector: Vec<f32>,
        model_version: String,
    ) -> Result<Self, crate::EmbeddingError> {
        if vector.len() != EMBEDDING_DIM {
            return Err(crate::EmbeddingError::InvalidDimensions {
                expected: EMBEDDING_DIM,
                actual: vector.len(),
            });
        }

        Ok(Self {
            id: EmbeddingId::new(),
            segment_id,
            vector,
            model_version,
            created_at: Utc::now(),
            tier: StorageTier::default(),
            metadata: EmbeddingMetadata::default(),
        })
    }

    /// Get the L2 norm of the embedding vector
    #[must_use]
    pub fn norm(&self) -> f32 {
        self.vector
            .iter()
            .map(|x| x * x)
            .sum::<f32>()
            .sqrt()
    }

    /// Check if the embedding is properly L2 normalized (norm close to 1.0)
    #[must_use]
    pub fn is_normalized(&self) -> bool {
        let norm = self.norm();
        (0.99..=1.01).contains(&norm)
    }

    /// Check if the embedding contains any invalid values (NaN or Inf)
    #[must_use]
    pub fn is_valid(&self) -> bool {
        !self.vector.iter().any(|x| x.is_nan() || x.is_infinite())
    }

    /// Compute cosine similarity with another embedding
    ///
    /// For L2-normalized embeddings, this is equivalent to dot product
    #[must_use]
    pub fn cosine_similarity(&self, other: &Self) -> f32 {
        self.vector
            .iter()
            .zip(other.vector.iter())
            .map(|(a, b)| a * b)
            .sum()
    }

    /// Change the storage tier for this embedding
    pub fn set_tier(&mut self, tier: StorageTier) {
        self.tier = tier;
    }
}

/// Metadata about an embedding's generation
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct EmbeddingMetadata {
    /// Time taken for inference in milliseconds
    pub inference_latency_ms: Option<f32>,

    /// Batch ID if processed in a batch
    pub batch_id: Option<String>,

    /// Whether GPU was used for inference
    pub gpu_used: bool,

    /// Original norm before L2 normalization
    pub original_norm: Option<f32>,

    /// Sparsity of the embedding (fraction of near-zero values)
    pub sparsity: Option<f32>,

    /// Quality score based on various metrics
    pub quality_score: Option<f32>,
}

/// Model version identifier
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ModelVersion {
    /// Model name (e.g., "perch-v2")
    pub name: String,

    /// Semantic version string (e.g., "2.0.0")
    pub version: String,

    /// Model variant (e.g., "base", "quantized", "pruned")
    pub variant: String,
}

impl ModelVersion {
    /// Create a new model version
    #[must_use]
    pub fn new(name: impl Into<String>, version: impl Into<String>, variant: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            version: version.into(),
            variant: variant.into(),
        }
    }

    /// Get the default Perch 2.0 base model version
    #[must_use]
    pub fn perch_v2_base() -> Self {
        Self::new("perch-v2", "2.0.0", "base")
    }

    /// Get the Perch 2.0 quantized model version
    #[must_use]
    pub fn perch_v2_quantized() -> Self {
        Self::new("perch-v2", "2.0.0", "quantized")
    }

    /// Get the full version string
    #[must_use]
    pub fn full_version(&self) -> String {
        format!("{}-{}-{}", self.name, self.version, self.variant)
    }
}

impl Default for ModelVersion {
    fn default() -> Self {
        Self::perch_v2_base()
    }
}

impl std::fmt::Display for ModelVersion {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.full_version())
    }
}

/// Input specification for the embedding model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InputSpecification {
    /// Expected sample rate in Hz
    pub sample_rate: u32,

    /// Window duration in seconds
    pub window_duration: f32,

    /// Number of samples per window
    pub window_samples: usize,

    /// Number of mel frequency bins
    pub mel_bins: usize,

    /// Number of time frames
    pub mel_frames: usize,

    /// Frequency range (low, high) in Hz
    pub frequency_range: (f32, f32),
}

impl Default for InputSpecification {
    fn default() -> Self {
        Self {
            sample_rate: crate::TARGET_SAMPLE_RATE,
            window_duration: crate::TARGET_WINDOW_SECONDS,
            window_samples: crate::TARGET_WINDOW_SAMPLES,
            mel_bins: crate::MEL_BINS,
            mel_frames: crate::MEL_FRAMES,
            frequency_range: (60.0, 16000.0),
        }
    }
}

/// Status of an embedding model
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelStatus {
    /// Model is available and ready for inference
    Active,

    /// Model is being loaded
    Loading,

    /// Model failed to load
    Failed,

    /// Model is deprecated and should not be used
    Deprecated,
}

/// Configuration and metadata for an embedding model.
///
/// Represents the Perch 2.0 ONNX model used for generating embeddings.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingModel {
    /// Model name
    pub name: String,

    /// Model version
    pub version: ModelVersion,

    /// Output embedding dimensions
    pub dimensions: usize,

    /// SHA-256 checksum of the model file
    pub checksum: String,

    /// Input specification
    pub input_spec: InputSpecification,

    /// Current model status
    pub status: ModelStatus,

    /// When the model was last loaded
    pub loaded_at: Option<Timestamp>,

    /// Path to the model file
    pub model_path: Option<String>,
}

impl EmbeddingModel {
    /// Create a new embedding model configuration
    #[must_use]
    pub fn new(name: String, version: ModelVersion, checksum: String) -> Self {
        Self {
            name,
            version,
            dimensions: EMBEDDING_DIM,
            checksum,
            input_spec: InputSpecification::default(),
            status: ModelStatus::Loading,
            loaded_at: None,
            model_path: None,
        }
    }

    /// Create a default Perch 2.0 model configuration
    #[must_use]
    pub fn perch_v2_default() -> Self {
        Self::new(
            "perch-v2".to_string(),
            ModelVersion::perch_v2_base(),
            String::new(), // Checksum will be computed on load
        )
    }

    /// Check if the model is ready for inference
    #[must_use]
    pub const fn is_ready(&self) -> bool {
        matches!(self.status, ModelStatus::Active)
    }

    /// Mark the model as active
    pub fn mark_active(&mut self) {
        self.status = ModelStatus::Active;
        self.loaded_at = Some(Utc::now());
    }

    /// Mark the model as failed
    pub fn mark_failed(&mut self) {
        self.status = ModelStatus::Failed;
    }
}

/// A batch of embeddings processed together
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingBatch {
    /// Unique batch identifier
    pub id: String,

    /// Segment IDs in this batch
    pub segment_ids: Vec<SegmentId>,

    /// Batch processing status
    pub status: BatchStatus,

    /// When batch processing started
    pub started_at: Timestamp,

    /// When batch processing completed
    pub completed_at: Option<Timestamp>,

    /// Batch processing metrics
    pub metrics: BatchMetrics,
}

/// Status of a batch operation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum BatchStatus {
    /// Batch is queued for processing
    Pending,

    /// Batch is currently being processed
    Processing,

    /// Batch processing completed successfully
    Completed,

    /// Batch processing failed
    Failed,
}

/// Metrics for batch processing
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct BatchMetrics {
    /// Total segments in batch
    pub total_segments: u32,

    /// Successfully processed segments
    pub success_count: u32,

    /// Failed segments
    pub failure_count: u32,

    /// Average inference latency in milliseconds
    pub avg_latency_ms: f32,

    /// Throughput in segments per second
    pub throughput: f32,
}

impl EmbeddingBatch {
    /// Create a new embedding batch
    #[must_use]
    pub fn new(segment_ids: Vec<SegmentId>) -> Self {
        let total = segment_ids.len() as u32;
        Self {
            id: Uuid::new_v4().to_string(),
            segment_ids,
            status: BatchStatus::Pending,
            started_at: Utc::now(),
            completed_at: None,
            metrics: BatchMetrics {
                total_segments: total,
                ..Default::default()
            },
        }
    }

    /// Mark batch as processing
    pub fn mark_processing(&mut self) {
        self.status = BatchStatus::Processing;
    }

    /// Mark batch as completed with metrics
    pub fn mark_completed(&mut self, success_count: u32, failure_count: u32, avg_latency_ms: f32) {
        self.status = BatchStatus::Completed;
        self.completed_at = Some(Utc::now());
        self.metrics.success_count = success_count;
        self.metrics.failure_count = failure_count;
        self.metrics.avg_latency_ms = avg_latency_ms;

        if let Some(completed) = self.completed_at {
            let duration = (completed - self.started_at).num_milliseconds() as f32 / 1000.0;
            if duration > 0.0 {
                self.metrics.throughput = self.metrics.total_segments as f32 / duration;
            }
        }
    }

    /// Mark batch as failed
    pub fn mark_failed(&mut self) {
        self.status = BatchStatus::Failed;
        self.completed_at = Some(Utc::now());
    }
}

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

    #[test]
    fn test_embedding_id_generation() {
        let id1 = EmbeddingId::new();
        let id2 = EmbeddingId::new();
        assert_ne!(id1, id2);
    }

    #[test]
    fn test_embedding_creation() {
        let segment_id = SegmentId::new();
        let vector = vec![0.0; EMBEDDING_DIM];
        let embedding = Embedding::new(segment_id, vector, "perch-v2-2.0.0-base".to_string());
        assert!(embedding.is_ok());
    }

    #[test]
    fn test_embedding_invalid_dimensions() {
        let segment_id = SegmentId::new();
        let vector = vec![0.0; 100]; // Wrong dimension
        let result = Embedding::new(segment_id, vector, "perch-v2-2.0.0-base".to_string());
        assert!(result.is_err());
    }

    #[test]
    fn test_embedding_norm() {
        let segment_id = SegmentId::new();
        let mut vector = vec![0.0; EMBEDDING_DIM];
        vector[0] = 1.0; // Unit vector
        let embedding = Embedding::new(segment_id, vector, "perch-v2-2.0.0-base".to_string()).unwrap();
        assert!((embedding.norm() - 1.0).abs() < 1e-6);
        assert!(embedding.is_normalized());
    }

    #[test]
    fn test_cosine_similarity() {
        let segment_id = SegmentId::new();

        // Create two identical normalized vectors
        let mut vector = vec![0.0; EMBEDDING_DIM];
        vector[0] = 1.0;

        let emb1 = Embedding::new(segment_id, vector.clone(), "test".to_string()).unwrap();
        let emb2 = Embedding::new(segment_id, vector, "test".to_string()).unwrap();

        let similarity = emb1.cosine_similarity(&emb2);
        assert!((similarity - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_storage_tier_bytes() {
        assert_eq!(StorageTier::Hot.bytes_per_dim(), 4);
        assert_eq!(StorageTier::Warm.bytes_per_dim(), 2);
        assert_eq!(StorageTier::Cold.bytes_per_dim(), 1);

        assert_eq!(StorageTier::Hot.embedding_bytes(), EMBEDDING_DIM * 4);
    }

    #[test]
    fn test_model_version() {
        let version = ModelVersion::perch_v2_base();
        assert_eq!(version.name, "perch-v2");
        assert_eq!(version.version, "2.0.0");
        assert_eq!(version.variant, "base");
        assert_eq!(version.full_version(), "perch-v2-2.0.0-base");
    }

    #[test]
    fn test_input_specification_defaults() {
        let spec = InputSpecification::default();
        assert_eq!(spec.sample_rate, 32000);
        assert_eq!(spec.window_samples, 160_000);
        assert_eq!(spec.mel_bins, 128);
        assert_eq!(spec.mel_frames, 500);
    }

    #[test]
    fn test_embedding_batch_lifecycle() {
        let segment_ids = vec![SegmentId::new(), SegmentId::new()];
        let mut batch = EmbeddingBatch::new(segment_ids);

        assert_eq!(batch.status, BatchStatus::Pending);
        assert_eq!(batch.metrics.total_segments, 2);

        batch.mark_processing();
        assert_eq!(batch.status, BatchStatus::Processing);

        batch.mark_completed(2, 0, 50.0);
        assert_eq!(batch.status, BatchStatus::Completed);
        assert_eq!(batch.metrics.success_count, 2);
        assert!(batch.completed_at.is_some());
    }
}