foxhunt/ml-data/src/lib.rs

#![allow(missing_docs)] // Internal implementation details don't require documentation
//! ML Data Repository
//!
//! Production-ready ML data management for HFT trading systems.
//! Provides repositories for training data, model artifacts, performance tracking,
//! and feature engineering with PostgreSQL integration.

use database::Database;
use std::sync::Arc;

pub mod features;
pub mod models;
pub mod performance;
pub mod training;

/// Common error types for ML data operations
#[derive(thiserror::Error, Debug)]
pub enum MlDataError {
    #[error("Database error: {0}")]
    Database(#[from] database::DatabaseError),

    #[error("SQL error: {0}")]
    Sql(#[from] sqlx::Error),

    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),

    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("Validation error: {message}")]
    Validation { message: String },

    #[error("Version conflict: {message}")]
    VersionConflict { message: String },

    #[error("Not found: {resource_type} with id {id}")]
    NotFound { resource_type: String, id: String },

    #[error("Configuration error: {0}")]
    Configuration(#[from] config::ConfigError),
}

pub type Result<T> = std::result::Result<T, MlDataError>;

/// Configuration for ML data repositories
#[derive(Debug, Clone, serde::Deserialize)]
pub struct MlDataConfig {
    /// Database connection configuration
    pub database: config::DatabaseConfig,

    /// Model artifact storage path
    pub model_storage_path: String,

    /// Feature store configuration
    pub feature_store: FeatureStoreConfig,

    /// Performance tracking settings
    pub performance: PerformanceConfig,

    /// Training data management
    pub training: TrainingConfig,
}

#[derive(Debug, Clone, serde::Deserialize)]
pub struct FeatureStoreConfig {
    /// Maximum number of features to cache in memory
    pub cache_size: usize,

    /// Feature TTL in seconds
    pub ttl_seconds: u64,

    /// Batch size for feature computation
    pub batch_size: usize,
}

#[derive(Debug, Clone, serde::Deserialize)]
pub struct PerformanceConfig {
    /// Metrics retention period in days
    pub retention_days: u32,

    /// Performance degradation threshold
    pub degradation_threshold: f64,

    /// A/B test confidence level
    pub confidence_level: f64,
}

#[derive(Debug, Clone, serde::Deserialize)]
pub struct TrainingConfig {
    /// Maximum dataset size in bytes
    pub max_dataset_size: u64,

    /// Data validation rules
    pub validation_rules: ValidationRules,

    /// Default train/validation/test split ratios
    pub default_splits: SplitRatios,
}

#[derive(Debug, Clone, serde::Deserialize)]
pub struct ValidationRules {
    /// Minimum number of samples required
    pub min_samples: usize,

    /// Maximum missing value ratio allowed
    pub max_missing_ratio: f64,

    /// Required features list
    pub required_features: Vec<String>,
}

#[derive(Debug, Clone, serde::Deserialize)]
pub struct SplitRatios {
    pub train: f64,
    pub validation: f64,
    pub test: f64,
}

impl Default for SplitRatios {
    fn default() -> Self {
        Self {
            train: 0.7,
            validation: 0.2,
            test: 0.1,
        }
    }
}

/// ML Data Repository Manager
///
/// Central coordinator for all ML data operations
#[derive(Clone)]
pub struct MlDataManager {
    pub training: Arc<training::TrainingDataRepository>,
    pub models: Arc<models::ModelRepository>,
    pub performance: Arc<performance::PerformanceRepository>,
    pub features: Arc<features::FeatureRepository>,
    config: MlDataConfig,
}

impl MlDataManager {
    /// Create a new ML data manager with the provided configuration
    pub async fn new(config: MlDataConfig) -> Result<Self> {
        let db = Database::new(config.database.clone()).await?;

        let training = Arc::new(
            training::TrainingDataRepository::new(db.clone(), config.training.clone()).await?,
        );

        let models = Arc::new(
            models::ModelRepository::new(db.clone(), config.model_storage_path.clone()).await?,
        );

        let performance = Arc::new(
            performance::PerformanceRepository::new(db.clone(), config.performance.clone()).await?,
        );

        let features = Arc::new(
            features::FeatureRepository::new(db.clone(), config.feature_store.clone()).await?,
        );

        Ok(Self {
            training,
            models,
            performance,
            features,
            config,
        })
    }

    /// Get the configuration
    pub fn config(&self) -> &MlDataConfig {
        &self.config
    }

    /// Initialize database schema for all repositories
    pub async fn initialize_schema(&self) -> Result<()> {
        tracing::info!("Initializing ML data repository schemas");

        self.training.initialize_schema().await?;
        self.models.initialize_schema().await?;
        self.performance.initialize_schema().await?;
        self.features.initialize_schema().await?;

        tracing::info!("ML data repository schemas initialized successfully");
        Ok(())
    }

    /// Health check for all repositories
    pub async fn health_check(&self) -> Result<HealthStatus> {
        let training_health = self.training.health_check().await?;
        let models_health = self.models.health_check().await?;
        let performance_health = self.performance.health_check().await?;
        let features_health = self.features.health_check().await?;

        Ok(HealthStatus {
            training: training_health,
            models: models_health,
            performance: performance_health,
            features: features_health,
            overall: training_health && models_health && performance_health && features_health,
        })
    }
}

/// Health status for all ML data repositories
#[derive(Debug, serde::Serialize)]
pub struct HealthStatus {
    pub training: bool,
    pub models: bool,
    pub performance: bool,
    pub features: bool,
    pub overall: bool,
}