foxhunt/adaptive-strategy/src/models/mod.rs

//! ML model interfaces and implementations
//!
//! This module provides a unified interface for different types of machine learning
//! models used in adaptive trading strategies, including deep learning models,
//! traditional ML models, and custom algorithmic models.

use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tracing::{debug, info, warn};

// Add missing core types
// ML types removed due to compilation issues - using local definitions instead
// use ml::prelude::{MarketRegime, TensorSpec};  // REMOVED

pub mod deep_learning;
pub mod ensemble_models;
pub mod tlob_model;
pub mod traditional;

/// Unified interface for all ML models
#[async_trait]
pub trait ModelTrait: std::fmt::Debug + Send + Sync {
    /// Get model name/identifier
    fn name(&self) -> &str;

    /// Get model type (lstm, transformer, random_forest, etc.)
    fn model_type(&self) -> &str;

    /// Make a prediction based on input features
    ///
    /// # Arguments
    ///
    /// * `features` - Input feature vector
    ///
    /// # Returns
    ///
    /// Model prediction with confidence score
    async fn predict(&self, features: &[f64]) -> Result<ModelPrediction>;

    /// Train/update the model with new data
    ///
    /// # Arguments
    ///
    /// * `training_data` - Training dataset
    ///
    /// # Returns
    ///
    /// Training metrics and model performance
    async fn train(&mut self, training_data: &TrainingData) -> Result<TrainingMetrics>;

    /// Get model metadata and configuration
    fn get_metadata(&self) -> ModelMetadata;

    /// Get current model performance metrics
    async fn get_performance(&self) -> Result<ModelPerformance>;

    /// Update model parameters/configuration
    async fn update_config(&mut self, config: ModelConfig) -> Result<()>;

    /// Check if model is ready for predictions
    fn is_ready(&self) -> bool;

    /// Get memory usage of the model
    fn memory_usage(&self) -> usize;

    /// Save model state to storage
    async fn save(&self, path: &str) -> Result<()>;

    /// Load model state from storage
    async fn load(&mut self, path: &str) -> Result<()>;
}

/// Model prediction with confidence and metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPrediction {
    /// Predicted value (e.g., price movement, return)
    pub value: f64,
    /// Confidence score (0.0 to 1.0)
    pub confidence: f64,
    /// Features used for this prediction
    pub features_used: Vec<String>,
    /// Additional metadata about the prediction
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

/// Training data structure
#[derive(Debug, Clone)]
pub struct TrainingData {
    /// Input features matrix
    pub features: Vec<Vec<f64>>,
    /// Target values
    pub targets: Vec<f64>,
    /// Feature names
    pub feature_names: Vec<String>,
    /// Timestamps for each sample
    pub timestamps: Vec<chrono::DateTime<chrono::Utc>>,
    /// Sample weights (optional)
    pub weights: Option<Vec<f64>>,
}

/// Training metrics and results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingMetrics {
    /// Training loss
    pub training_loss: f64,
    /// Validation loss
    pub validation_loss: f64,
    /// Training accuracy
    pub training_accuracy: f64,
    /// Validation accuracy
    pub validation_accuracy: f64,
    /// Number of epochs/iterations
    pub epochs: u32,
    /// Training time in seconds
    pub training_time_seconds: f64,
    /// Additional metrics
    pub additional_metrics: HashMap<String, f64>,
}

/// Model metadata and configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
    /// Model name
    pub name: String,
    /// Model type/architecture
    pub model_type: String,
    /// Model version
    pub version: String,
    /// Creation timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Last updated timestamp
    pub updated_at: chrono::DateTime<chrono::Utc>,
    /// Model parameters
    pub parameters: HashMap<String, serde_json::Value>,
    /// Expected input dimensions
    pub input_dimensions: usize,
    /// Model description
    pub description: Option<String>,
}

/// Model performance metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPerformance {
    /// Accuracy on validation set
    pub accuracy: f64,
    /// Precision score
    pub precision: f64,
    /// Recall score
    pub recall: f64,
    /// F1 score
    pub f1_score: f64,
    /// Sharpe ratio (for trading models)
    pub sharpe_ratio: f64,
    /// Maximum drawdown
    pub max_drawdown: f64,
    /// Total number of predictions made
    pub prediction_count: u64,
    /// Last evaluation timestamp
    pub last_evaluated: chrono::DateTime<chrono::Utc>,
}

/// Model configuration parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Learning rate
    pub learning_rate: f64,
    /// Batch size
    pub batch_size: usize,
    /// Regularization parameters
    pub regularization: f64,
    /// Dropout rate
    pub dropout_rate: f64,
    /// Hidden layer dimensions
    pub hidden_dimensions: Vec<usize>,
    /// Maximum training epochs
    pub max_epochs: u32,
    /// Early stopping patience
    pub early_stopping_patience: u32,
    /// Additional model-specific parameters
    pub custom_parameters: HashMap<String, serde_json::Value>,
}

/// Model factory for creating different types of models
pub struct ModelFactory;

impl ModelFactory {
    /// Create a new model instance based on configuration
    ///
    /// # Arguments
    ///
    /// * `model_type` - Type of model to create
    /// * `name` - Name for the model instance
    /// * `config` - Model configuration
    ///
    /// # Returns
    ///
    /// Boxed model instance implementing ModelTrait
    pub async fn create_model(
        model_type: &str,
        name: String,
        config: ModelConfig,
    ) -> Result<Box<dyn ModelTrait>> {
        info!("Creating model: {} of type: {}", name, model_type);

        match model_type.to_lowercase().as_str() {
            "lstm" => Ok(Box::new(deep_learning::LSTMModel::new(name, config).await?)),
            "gru" => Ok(Box::new(deep_learning::GRUModel::new(name, config).await?)),
            "transformer" => Ok(Box::new(
                deep_learning::TransformerModel::new(name, config).await?,
            )),
            "cnn" => Ok(Box::new(deep_learning::CNNModel::new(name, config).await?)),
            "random_forest" => Ok(Box::new(
                traditional::RandomForestModel::new(name, config).await?,
            )),
            "xgboost" => Ok(Box::new(
                traditional::XGBoostModel::new(name, config).await?,
            )),
            "svm" => Ok(Box::new(traditional::SVMModel::new(name, config).await?)),
            "linear_regression" => Ok(Box::new(
                traditional::LinearRegressionModel::new(name, config).await?,
            )),
            "ensemble" => Ok(Box::new(
                ensemble_models::EnsembleModel::new(name, config).await?,
            )),
            "tlob" => {
                // Create TLOB model if available, otherwise use mock
                match tlob_model::TLOBModel::new(name.clone(), config).await {
                    Ok(model) => {
                        info!("Created TLOB model with sub-50μs inference capability");
                        Ok(Box::new(model))
                    },
                    Err(_) => {
                        warn!("TLOB model creation failed, using mock model for {}", name);
                        Ok(Box::new(MockModel::new(name)))
                    },
                }
            },
            _ => {
                warn!("Unknown model type: {}, creating mock model", model_type);
                Ok(Box::new(MockModel::new(name)))
            },
        }
    }

    /// Get available model types
    pub fn available_models() -> Vec<&'static str> {
        vec![
            "lstm",
            "gru",
            "transformer",
            "cnn",
            "random_forest",
            "xgboost",
            "svm",
            "linear_regression",
            "ensemble",
            "tlob",
        ]
    }
}

/// Mock model implementation for testing and development
#[derive(Debug, Clone)]
pub struct MockModel {
    name: String,
    ready: bool,
    prediction_count: u64,
}

#[async_trait]
impl ModelTrait for MockModel {
    fn name(&self) -> &str {
        &self.name
    }

    fn model_type(&self) -> &str {
        "mock"
    }

    async fn predict(&self, features: &[f64]) -> Result<ModelPrediction> {
        debug!(
            "MockModel {} predicting with {} features",
            self.name,
            features.len()
        );

        if !self.ready {
            anyhow::bail!("Model {} is not ready for predictions", self.name);
        }

        // Simple mock prediction based on feature sum
        let feature_sum: f64 = features.iter().sum();
        let prediction_value = (feature_sum % 2.0) - 1.0; // Range [-1, 1]
        let confidence = 0.5 + (feature_sum % 0.5); // Range [0.5, 1.0]

        Ok(ModelPrediction {
            value: prediction_value,
            confidence,
            features_used: (0..features.len())
                .map(|i| format!("feature_{}", i))
                .collect(),
            metadata: Some(HashMap::from([
                (
                    "model_type".to_owned(),
                    serde_json::Value::String("mock".to_owned()),
                ),
                (
                    "feature_sum".to_owned(),
                    serde_json::Value::Number(
                        serde_json::Number::from_f64(feature_sum)
                            .unwrap_or_else(|| serde_json::Number::from(0)),
                    ),
                ),
            ])),
        })
    }

    async fn train(&mut self, training_data: &TrainingData) -> Result<TrainingMetrics> {
        info!(
            "MockModel {} training with {} samples",
            self.name,
            training_data.features.len()
        );

        // Simulate training time
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        self.ready = true;

        Ok(TrainingMetrics {
            training_loss: 0.1 + (rand::random::<f64>() * 0.05),
            validation_loss: 0.12 + (rand::random::<f64>() * 0.05),
            training_accuracy: 0.85 + (rand::random::<f64>() * 0.1),
            validation_accuracy: 0.83 + (rand::random::<f64>() * 0.1),
            epochs: 10,
            training_time_seconds: 0.1,
            additional_metrics: HashMap::from([(
                "samples_processed".to_owned(),
                training_data.features.len() as f64,
            )]),
        })
    }

    fn get_metadata(&self) -> ModelMetadata {
        ModelMetadata {
            name: self.name.clone(),
            model_type: "mock".to_owned(),
            version: "1.0.0".to_owned(),
            created_at: chrono::Utc::now(),
            updated_at: chrono::Utc::now(),
            parameters: HashMap::new(),
            input_dimensions: 0,
            description: Some("Mock model for testing".to_owned()),
        }
    }

    async fn get_performance(&self) -> Result<ModelPerformance> {
        Ok(ModelPerformance {
            accuracy: 0.85,
            precision: 0.82,
            recall: 0.88,
            f1_score: 0.85,
            sharpe_ratio: 1.5,
            max_drawdown: 0.05,
            prediction_count: self.prediction_count,
            last_evaluated: chrono::Utc::now(),
        })
    }

    async fn update_config(&mut self, _config: ModelConfig) -> Result<()> {
        info!("MockModel {} config updated", self.name);
        Ok(())
    }

    fn is_ready(&self) -> bool {
        self.ready
    }

    fn memory_usage(&self) -> usize {
        1024 // 1KB mock usage
    }

    async fn save(&self, path: &str) -> Result<()> {
        info!("MockModel {} saved to {}", self.name, path);
        Ok(())
    }

    async fn load(&mut self, path: &str) -> Result<()> {
        info!("MockModel {} loaded from {}", self.name, path);
        self.ready = true;
        Ok(())
    }
}

impl MockModel {
    /// Create a new mock model
    pub fn new(name: String) -> Self {
        Self {
            name,
            ready: true, // Set to true for testing - models are ready by default
            prediction_count: 0,
        }
    }
}

/// Model registry for managing model instances
pub struct ModelRegistry {
    models: HashMap<String, Box<dyn ModelTrait>>,
}

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

impl ModelRegistry {
    /// Create a new model registry
    pub fn new() -> Self {
        Self {
            models: HashMap::new(),
        }
    }

    /// Register a model in the registry
    pub fn register(&mut self, model: Box<dyn ModelTrait>) {
        let name = model.name().to_string();
        info!("Registering model: {}", name);
        self.models.insert(name, model);
    }

    /// Get a model by name
    pub fn get(&self, name: &str) -> Option<&dyn ModelTrait> {
        self.models.get(name).map(|m| m.as_ref())
    }

    // TODO: Fix lifetime issues with get_mut method
    // /// Get a mutable reference to a model by name
    // pub fn get_mut<'a>(&'a mut self, name: &str) -> Option<&'a mut (dyn ModelTrait + 'a)> {
    //     self.models.get_mut(name).map(move |m| m.as_mut())
    // }

    /// Remove a model from the registry
    pub fn remove(&mut self, name: &str) -> Option<Box<dyn ModelTrait>> {
        info!("Removing model: {}", name);
        self.models.remove(name)
    }

    /// List all registered model names
    pub fn list_models(&self) -> Vec<&str> {
        self.models.keys().map(|s| s.as_str()).collect()
    }

    /// Get total memory usage of all models
    pub fn total_memory_usage(&self) -> usize {
        self.models.values().map(|m| m.memory_usage()).sum()
    }
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            learning_rate: 0.001,
            batch_size: 32,
            regularization: 0.01,
            dropout_rate: 0.1,
            hidden_dimensions: vec![128, 64, 32],
            max_epochs: 100,
            early_stopping_patience: 10,
            custom_parameters: HashMap::new(),
        }
    }
}

impl TrainingData {
    /// Create new training data
    pub fn new(features: Vec<Vec<f64>>, targets: Vec<f64>, feature_names: Vec<String>) -> Self {
        let timestamps = vec![chrono::Utc::now(); features.len()];

        Self {
            features,
            targets,
            feature_names,
            timestamps,
            weights: None,
        }
    }

    /// Add sample weights
    pub fn with_weights(mut self, weights: Vec<f64>) -> Self {
        self.weights = Some(weights);
        self
    }

    /// Validate training data consistency
    pub fn validate(&self) -> Result<()> {
        if self.features.len() != self.targets.len() {
            anyhow::bail!("Features and targets length mismatch");
        }

        if self.features.len() != self.timestamps.len() {
            anyhow::bail!("Features and timestamps length mismatch");
        }

        if let Some(ref weights) = self.weights {
            if weights.len() != self.features.len() {
                anyhow::bail!("Weights and features length mismatch");
            }
        }

        if !self.features.is_empty()
            && self.features.first().map(|f| f.len()).unwrap_or(0) != self.feature_names.len()
        {
            anyhow::bail!("Feature dimensions and feature names length mismatch");
        }

        Ok(())
    }

    /// Get number of samples
    pub fn len(&self) -> usize {
        self.features.len()
    }

    /// Check if dataset is empty
    pub fn is_empty(&self) -> bool {
        self.features.is_empty()
    }
}

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

    #[tokio::test]
    async fn test_mock_model_creation() {
        let model = MockModel::new("test_model".to_owned());
        assert_eq!(model.name(), "test_model");
        assert_eq!(model.model_type(), "mock");
        // MockModel is ready by default for testing (Wave 122 change)
        assert!(model.is_ready());
    }

    #[tokio::test]
    async fn test_mock_model_training() {
        let mut model = MockModel::new("test_model".to_owned());

        let training_data = TrainingData::new(
            vec![vec![1.0, 2.0, 3.0]; 100],
            vec![0.5; 100],
            vec!["f1".to_owned(), "f2".to_owned(), "f3".to_owned()],
        );

        let metrics = model.train(&training_data).await.unwrap();
        assert!(metrics.training_accuracy > 0.0);
        assert!(model.is_ready());
    }

    #[tokio::test]
    async fn test_mock_model_prediction() {
        let mut model = MockModel::new("test_model".to_owned());

        // Train first
        let training_data = TrainingData::new(
            vec![vec![1.0, 2.0]; 10],
            vec![0.5; 10],
            vec!["f1".to_owned(), "f2".to_owned()],
        );
        model.train(&training_data).await.unwrap();

        // Test prediction
        let features = vec![1.0, 2.0, 3.0];
        let prediction = model.predict(&features).await.unwrap();

        assert!(prediction.confidence >= 0.0 && prediction.confidence <= 1.0);
        assert!(!prediction.features_used.is_empty());
    }

    #[test]
    fn test_model_factory_available_models() {
        let models = ModelFactory::available_models();
        assert!(models.contains(&"lstm"));
        assert!(models.contains(&"transformer"));
        assert!(models.contains(&"random_forest"));
    }

    #[test]
    fn test_model_registry() {
        let mut registry = ModelRegistry::new();
        let model = Box::new(MockModel::new("test_model".to_owned()));

        registry.register(model);
        assert!(registry.get("test_model").is_some());
        assert_eq!(registry.list_models(), vec!["test_model"]);

        let removed = registry.remove("test_model");
        assert!(removed.is_some());
        assert!(registry.get("test_model").is_none());
    }

    #[test]
    fn test_training_data_validation() {
        let data = TrainingData::new(
            vec![vec![1.0, 2.0]; 3],
            vec![0.5; 3],
            vec!["f1".to_owned(), "f2".to_owned()],
        );

        assert!(data.validate().is_ok());
        assert_eq!(data.len(), 3);
        assert!(!data.is_empty());
    }

    #[test]
    fn test_training_data_invalid() {
        let data = TrainingData::new(
            vec![vec![1.0, 2.0]; 3],
            vec![0.5; 2], // Wrong length
            vec!["f1".to_owned(), "f2".to_owned()],
        );

        assert!(data.validate().is_err());
    }
}