foxhunt/ml-data/src/training.rs

//! Training Data Repository
//! 
//! Manages training datasets, versioning, validation, and data splits
//! for machine learning model training in HFT environments.

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

use crate::{MlDataError, Result, TrainingConfig};
use database::{Database, DatabaseTransaction};

/// Training data repository for ML workflows
#[derive(Clone)]
pub struct TrainingDataRepository {
    db: Database,
    config: TrainingConfig,
}

impl TrainingDataRepository {
    pub async fn new(db: Database, config: TrainingConfig) -> Result<Self> {
        let repo = Self { db, config };
        repo.initialize_schema().await?;
        Ok(repo)
    }
    
    /// Initialize database schema for training data
    pub async fn initialize_schema(&self) -> Result<()> {
        // Training datasets table
        self.db.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_training_datasets (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                name VARCHAR NOT NULL,
                version INTEGER NOT NULL,
                description TEXT,
                created_at TIMESTAMPTZ DEFAULT NOW(),
                updated_at TIMESTAMPTZ DEFAULT NOW(),
                created_by VARCHAR NOT NULL,
                status dataset_status DEFAULT 'draft',
                validation_results JSONB,
                metadata JSONB DEFAULT '{}',
                UNIQUE(name, version)
            )
        "#).await?;
        
        // Create enum if not exists
        self.db.execute(r#"
            DO $$ BEGIN
                CREATE TYPE dataset_status AS ENUM ('draft', 'validated', 'active', 'deprecated');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;
        
        // Data splits table
        self.db.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_data_splits (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                dataset_id UUID NOT NULL REFERENCES ml_training_datasets(id) ON DELETE CASCADE,
                split_type split_type_enum NOT NULL,
                sample_count BIGINT NOT NULL,
                start_time TIMESTAMPTZ,
                end_time TIMESTAMPTZ,
                metadata JSONB DEFAULT '{}',
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        "#).await?;
        
        // Create split type enum
        self.db.execute(r#"
            DO $$ BEGIN
                CREATE TYPE split_type_enum AS ENUM ('train', 'validation', 'test');
            EXCEPTION
                WHEN duplicate_object THEN null;
            END $$;
        "#).await?;
        
        // Dataset samples table for large datasets
        self.db.execute(r#"
            CREATE TABLE IF NOT EXISTS ml_dataset_samples (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                dataset_id UUID NOT NULL REFERENCES ml_training_datasets(id) ON DELETE CASCADE,
                split_id UUID REFERENCES ml_data_splits(id) ON DELETE SET NULL,
                timestamp TIMESTAMPTZ NOT NULL,
                features JSONB NOT NULL,
                labels JSONB NOT NULL,
                weight DOUBLE PRECISION DEFAULT 1.0,
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        "#).await?;
        
        // Indexes for performance
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_datasets_name_version ON ml_training_datasets(name, version)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_samples_dataset_timestamp ON ml_dataset_samples(dataset_id, timestamp)").await?;
        self.db.execute("CREATE INDEX IF NOT EXISTS idx_samples_split ON ml_dataset_samples(split_id)").await?;
        
        Ok(())
    }
    /// Create a new training dataset
    pub async fn create_dataset(&self, request: CreateDatasetRequest) -> Result<TrainingDataset> {
        // Validate dataset configuration
        self.validate_dataset_request(&request).await?;
    
        // Check for version conflicts
        if self.dataset_version_exists(&request.name, request.version).await? {
            return Err(MlDataError::VersionConflict {
                message: format!("Dataset {} version {} already exists", request.name, request.version)
            });
        }
    
        let dataset_id = Uuid::new_v4();

        // Insert dataset record using transaction
        let query = format!(
            r#"INSERT INTO ml_training_datasets
               (id, name, version, description, created_by, metadata)
               VALUES ('{}', '{}', {}, {}, '{}', '{}')"#,
            dataset_id,
            request.name.replace("'", "''"),
            request.version,
            match &request.description {
                Some(desc) => format!("'{}'", desc.replace("'", "''")),
                None => "NULL".to_string()
            },
            request.created_by.replace("'", "''"),
            request.metadata.to_string().replace("'", "''")
        );

        let mut tx = self.db.begin_transaction().await?;
        tx.execute(&query).await.map_err(|e| database::DatabaseError::Unknown { message: e.to_string() })?;
        tx.commit().await?;
        
        let dataset = TrainingDataset {
            id: dataset_id,
            name: request.name,
            version: request.version,
            description: request.description,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            created_by: request.created_by,
            status: DatasetStatus::Draft,
            validation_results: None,
            metadata: request.metadata,
            splits: HashMap::new(),
            sample_count: 0,
        };
        
        tracing::info!("Created training dataset: {} v{}", dataset.name, dataset.version);
        Ok(dataset)
    }
    
    /// Add training samples to a dataset
    pub async fn add_samples(
        &self,
        dataset_id: Uuid,
        samples: Vec<TrainingSample>
    ) -> Result<()> {
        let sample_count = samples.len();

        let mut tx = self.db.begin_transaction().await?;

        for sample in samples {
            let sample_id = Uuid::new_v4();
            let query = format!(
                r#"INSERT INTO ml_dataset_samples
                   (id, dataset_id, timestamp, features, labels, weight)
                   VALUES ('{}', '{}', '{}', '{}', '{}', {})"#,
                sample_id,
                dataset_id,
                sample.timestamp.to_rfc3339(),
                sample.features.to_string().replace("'", "''"),
                sample.labels.to_string().replace("'", "''"),
                sample.weight
            );
            tx.execute(&query).await.map_err(|e| database::DatabaseError::Unknown { message: e.to_string() })?;
        }

        tx.commit().await?;
    
        tracing::info!("Added {} samples to dataset {}", sample_count, dataset_id);
        Ok(())
    }
    
    /// Create data splits for a dataset
    pub async fn create_splits(
        &self, 
        dataset_id: Uuid, 
        split_config: SplitConfiguration
    ) -> Result<HashMap<DataSplit, DataSplitInfo>> {
        // Get total sample count
        let mut conn = self.db.acquire().await?;
        let total_samples: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM ml_dataset_samples WHERE dataset_id = $1"
        )
        .bind(dataset_id)
        .fetch_one(conn.as_mut())
        .await?;
        
        if total_samples < self.config.validation_rules.min_samples as i64 {
            return Err(MlDataError::Validation {
                message: format!("Insufficient samples: {} < {}", 
                    total_samples, self.config.validation_rules.min_samples)
            });
        }
        
        // Calculate split sizes
        let train_size = (total_samples as f64 * split_config.ratios.train) as i64;
        let val_size = (total_samples as f64 * split_config.ratios.validation) as i64;
        let test_size = total_samples - train_size - val_size;
        
        // Create splits in transaction
        let splits = self.db.with_transaction(|mut tx| async move {
            let mut splits = HashMap::new();
            let mut offset = 0i64;
        
            // Create training split
            let train_split_id = self.create_split_record(
                &mut tx, dataset_id, DataSplit::Train, train_size, offset
            ).await.map_err(|e| database::DatabaseError::Unknown { message: e.to_string() })?;
            splits.insert(DataSplit::Train, DataSplitInfo {
                id: train_split_id,
                sample_count: train_size as usize,
                start_offset: offset as usize,
                end_offset: (offset + train_size) as usize,
            });
            offset += train_size;
        
            // Create validation split
            let val_split_id = self.create_split_record(
                &mut tx, dataset_id, DataSplit::Validation, val_size, offset
            ).await.map_err(|e| database::DatabaseError::Unknown { message: e.to_string() })?;
            splits.insert(DataSplit::Validation, DataSplitInfo {
                id: val_split_id,
                sample_count: val_size as usize,
                start_offset: offset as usize,
                end_offset: (offset + val_size) as usize,
            });
            offset += val_size;
        
            // Create test split
            let test_split_id = self.create_split_record(
                &mut tx, dataset_id, DataSplit::Test, test_size, offset
            ).await.map_err(|e| database::DatabaseError::Unknown { message: e.to_string() })?;
            splits.insert(DataSplit::Test, DataSplitInfo {
                id: test_split_id,
                sample_count: test_size as usize,
                start_offset: offset as usize,
                end_offset: (offset + test_size) as usize,
            });
        
            Ok((splits, tx))
        }).await?;
        
        tracing::info!("Created splits for dataset {}: train={}, val={}, test={}", 
            dataset_id, train_size, val_size, test_size);
        
        Ok(splits)
    }
    
    /// Get training data for a specific split
    pub async fn get_split_data(
        &self,
        dataset_id: Uuid,
        split: DataSplit,
        batch_size: Option<usize>
    ) -> Result<TrainingDataStream> {
        let split_info = self.get_split_info(dataset_id, split.clone()).await?;
        
        Ok(TrainingDataStream {
            dataset_id,
            split_id: split_info.id,
            split_type: split,
            batch_size: batch_size.unwrap_or(1000),
            current_offset: 0,
            total_samples: split_info.sample_count,
            db: self.db.clone(),
        })
    }
    
    /// Validate dataset against configuration rules
    async fn validate_dataset_request(&self, request: &CreateDatasetRequest) -> Result<()> {
        if request.name.trim().is_empty() {
            return Err(MlDataError::Validation {
                message: "Dataset name cannot be empty".to_string()
            });
        }
        
        if request.version < 1 {
            return Err(MlDataError::Validation {
                message: "Dataset version must be >= 1".to_string()
            });
        }
        
        Ok(())
    }
    
    /// Check if dataset version already exists
    async fn dataset_version_exists(&self, name: &str, version: i32) -> Result<bool> {
        let mut conn = self.db.acquire().await?;
        let count: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM ml_training_datasets WHERE name = $1 AND version = $2"
        )
        .bind(name)
        .bind(version)
        .fetch_one(conn.as_mut())
        .await?;
        
        Ok(count > 0)
    }
    
    /// Create a split record in the database
    async fn create_split_record(
        &self,
        tx: &mut DatabaseTransaction,
        dataset_id: Uuid,
        split_type: DataSplit,
        sample_count: i64,
        offset: i64
    ) -> Result<Uuid> {
        let split_id = Uuid::new_v4();
    
        let query = format!(
            r#"INSERT INTO ml_data_splits 
               (id, dataset_id, split_type, sample_count, metadata)
               VALUES ('{}', '{}', '{}', {}, '{}')"#,
            split_id,
            dataset_id,
            split_type.to_string(),
            sample_count,
            serde_json::json!({"offset": offset})
        );
    
        tx.execute(&query).await.map_err(|e| MlDataError::Database(e))?;
    
        Ok(split_id)
    }
    
    /// Get split information
    async fn get_split_info(&self, dataset_id: Uuid, split: DataSplit) -> Result<DataSplitInfo> {
        let mut conn = self.db.acquire().await?;
        
        let row = sqlx::query_as::<_, (Uuid, i64, serde_json::Value)>(
            "SELECT id, sample_count, metadata FROM ml_data_splits WHERE dataset_id = $1 AND split_type = $2"
        )
        .bind(&dataset_id)
        .bind(&split.to_string())
        .fetch_one(conn.as_mut())
        .await
        .map_err(|_| MlDataError::NotFound {
            resource_type: "DataSplit".to_string(),
            id: format!("{}:{:?}", dataset_id, split),
        })?;
        
        let (id, sample_count, metadata) = row;
        let offset = metadata.get("offset").and_then(|v| v.as_i64()).unwrap_or(0) as usize;
        
        Ok(DataSplitInfo {
            id,
            sample_count: sample_count as usize,
            start_offset: offset,
            end_offset: offset + sample_count as usize,
        })
    }
    
    /// Health check for training repository
    pub async fn health_check(&self) -> Result<bool> {
        let mut conn = self.db.acquire().await?;
        let _: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM ml_training_datasets")
            .fetch_one(conn.as_mut())
            .await?;
        Ok(true)
    }
}

/// Request to create a new training dataset
#[derive(Debug, Serialize, Deserialize)]
pub struct CreateDatasetRequest {
    pub name: String,
    pub version: i32,
    pub description: Option<String>,
    pub created_by: String,
    pub metadata: serde_json::Value,
}

/// Training dataset representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingDataset {
    pub id: Uuid,
    pub name: String,
    pub version: i32,
    pub description: Option<String>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub created_by: String,
    pub status: DatasetStatus,
    pub validation_results: Option<ValidationResults>,
    pub metadata: serde_json::Value,
    pub splits: HashMap<DataSplit, DataSplitInfo>,
    pub sample_count: usize,
}

/// Dataset status enumeration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DatasetStatus {
    Draft,
    Validated,
    Active,
    Deprecated,
}

/// Data split types
#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
pub enum DataSplit {
    Train,
    Validation,
    Test,
}

impl ToString for DataSplit {
    fn to_string(&self) -> String {
        match self {
            DataSplit::Train => "train".to_string(),
            DataSplit::Validation => "validation".to_string(),
            DataSplit::Test => "test".to_string(),
        }
    }
}

/// Information about a data split
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataSplitInfo {
    pub id: Uuid,
    pub sample_count: usize,
    pub start_offset: usize,
    pub end_offset: usize,
}

/// Configuration for creating data splits
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SplitConfiguration {
    pub ratios: SplitRatios,
    pub stratify_by: Option<String>,
    pub shuffle: bool,
    pub random_seed: Option<u64>,
}

/// Split ratios for train/validation/test
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SplitRatios {
    pub train: f64,
    pub validation: f64,
    pub test: f64,
}

/// Individual training sample
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingSample {
    pub timestamp: DateTime<Utc>,
    pub features: serde_json::Value,
    pub labels: serde_json::Value,
    pub weight: f64,
}

/// Validation results for a dataset
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationResults {
    pub is_valid: bool,
    pub errors: Vec<ValidationError>,
    pub warnings: Vec<ValidationWarning>,
    pub statistics: DataStatistics,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationError {
    pub code: String,
    pub message: String,
    pub severity: ErrorSeverity,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationWarning {
    pub code: String,
    pub message: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ErrorSeverity {
    Critical,
    Major,
    Minor,
}

/// Dataset statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataStatistics {
    pub total_samples: usize,
    pub feature_count: usize,
    pub missing_ratio: f64,
    pub class_distribution: HashMap<String, usize>,
    pub time_range: (DateTime<Utc>, DateTime<Utc>),
}

/// Async stream for loading training data in batches
pub struct TrainingDataStream {
    dataset_id: Uuid,
    split_id: Uuid,
    split_type: DataSplit,
    batch_size: usize,
    current_offset: usize,
    total_samples: usize,
    db: Database,
}

impl TrainingDataStream {
    /// Get the next batch of training data
    pub async fn next_batch(&mut self) -> Result<Option<TrainingBatch>> {
        if self.current_offset >= self.total_samples {
            return Ok(None);
        }
        
        let mut conn = self.db.acquire().await?;
        let limit = std::cmp::min(self.batch_size, self.total_samples - self.current_offset);
        
        let rows = sqlx::query_as::<_, (DateTime<Utc>, serde_json::Value, serde_json::Value, Option<f64>)>(
            r#"SELECT timestamp, features, labels, weight 
               FROM ml_dataset_samples 
               WHERE split_id = $1 
               ORDER BY timestamp 
               LIMIT $2 OFFSET $3"#
        )
        .bind(&self.split_id)
        .bind(limit as i64)
        .bind(self.current_offset as i64)
        .fetch_all(conn.as_mut())
        .await?;
        
        let mut samples = Vec::with_capacity(rows.len());
        for (timestamp, features, labels, weight) in rows {
            samples.push(TrainingSample {
                timestamp,
                features,
                labels,
                weight: weight.unwrap_or(1.0),
            });
        }
        
        self.current_offset += samples.len();
        
        Ok(Some(TrainingBatch {
            samples,
            split_type: self.split_type.clone(),
            batch_index: self.current_offset / self.batch_size,
            is_last: self.current_offset >= self.total_samples,
        }))
    }
}

/// A batch of training data
#[derive(Debug, Clone)]
pub struct TrainingBatch {
    pub samples: Vec<TrainingSample>,
    pub split_type: DataSplit,
    pub batch_index: usize,
    pub is_last: bool,
}

/// Dataset version information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetVersion {
    pub name: String,
    pub version: i32,
    pub created_at: DateTime<Utc>,
    pub status: DatasetStatus,
    pub sample_count: usize,
}