foxhunt/ml/src/model.rs

//! Real Candle-based ML model implementations to replace mocks
//!
//! This module provides actual neural network implementations using the Candle framework
//! for production-ready HFT models.

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

    #[tokio::test]
    async fn test_real_model_creation() {
        let config = ModelConfig::default();
        let model =
            RealMLModel::new(config).map_err(|e| anyhow!("Failed to create model: {:?}", e))?;

        assert!(!model.is_trained);
        assert_eq!(
            model.metadata.model_type,
            MLModelType::Custom("MLP".to_string())
        );
        assert_eq!(model.network.config.input_dim, 16);
    }

    #[tokio::test]
    async fn test_synthetic_data_generation() {
        let device = Device::Cpu;
        let (inputs, targets) = create_synthetic_training_data(&device, 100, 16, 1)
            .map_err(|e| anyhow!("Failed to create synthetic data: {:?}", e))?;

        assert_eq!(inputs.dims(), &[100, 16]);
        assert_eq!(targets.dims(), &[1, 100]);
    }

    #[tokio::test]
    async fn test_model_training() {
        let config = ModelConfig {
            input_dim: 8,
            hidden_dims: vec![16, 8],
            output_dim: 1,
            learning_rate: 0.01,
            batch_size: 32,
        };

        let mut model =
            RealMLModel::new(config).map_err(|e| anyhow!("Failed to create model: {:?}", e))?;
        let device = model.device.clone();

        // Create small dataset for quick test
        let (train_x, train_y) = create_synthetic_training_data(&device, 64, 8, 1)
            .map_err(|e| anyhow!("Failed to create training data: {:?}", e))?;

        // Train for few epochs
        let metrics = model
            .train(&train_x, &train_y, 10)
            .await
            .map_err(|e| anyhow!("Training failed: {:?}", e))?;

        assert!(model.is_trained);
        assert_eq!(metrics.epochs_trained, 10);
        assert!(metrics.train_loss >= 0.0);
        assert!(metrics.training_time_seconds > 0.0);
    }

    #[tokio::test]
    async fn test_real_training_pipeline() {
        let training_config = TrainingConfig {
            epochs: 5,
            learning_rate: 0.01,
            ..Default::default()
        };

        let mut pipeline = RealTrainingPipeline::new(training_config);

        // Add two models
        let model1 = RealMLModel::new(ModelConfig {
            input_dim: 4,
            hidden_dims: vec![8],
            output_dim: 1,
            ..Default::default()
        })
        .map_err(|e| anyhow!("Failed to create model 1: {:?}", e))?;

        let model2 = RealMLModel::new(ModelConfig {
            input_dim: 4,
            hidden_dims: vec![8, 4],
            output_dim: 1,
            ..Default::default()
        })
        .map_err(|e| anyhow!("Failed to create model 2: {:?}", e))?;

        pipeline.add_model("model1".to_string(), model1);
        pipeline.add_model("model2".to_string(), model2);

        // Create training data
        let device = Device::Cpu;
        let (train_x, train_y) = create_synthetic_training_data(&device, 32, 4, 1)
            .map_err(|e| anyhow!("Failed to create training data: {:?}", e))?;

        // Train all models
        let results = pipeline
            .train_all(&train_x, &train_y)
            .await
            .map_err(|e| anyhow!("Pipeline training failed: {:?}", e))?;

        assert_eq!(results.len(), 2);
        assert!(results.contains_key("model1"));
        assert!(results.contains_key("model2"));

        // Test predictions
        let predictions = pipeline
            .predict_all(&train_x)
            .map_err(|e| anyhow!("Prediction failed: {:?}", e))?;

        assert_eq!(predictions.len(), 2);
    }
}