kaccy-core 0.2.0

//! Deep Learning Framework for Trading
//!
//! This module provides deep learning abstractions including:
//! - LSTM for time series prediction
//! - Transformer models for price prediction
//! - GANs for synthetic data generation
//! - Attention mechanisms for feature selection

use anyhow::Result;
use serde::{Deserialize, Serialize};

/// Deep learning model type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DLModelType {
    /// Long Short-Term Memory
    LSTM,
    /// Gated Recurrent Unit
    GRU,
    /// Transformer
    Transformer,
    /// Generative Adversarial Network
    GAN,
    /// Convolutional Neural Network
    CNN,
}

/// LSTM configuration for time series
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LSTMConfig {
    /// Number of input features per time step
    pub input_size: usize,
    /// Number of hidden units per LSTM layer
    pub hidden_size: usize,
    /// Number of stacked LSTM layers
    pub num_layers: usize,
    /// Number of output values to predict
    pub output_size: usize,
    /// Dropout rate applied between layers (0.0–1.0)
    pub dropout: f64,
    /// Whether to use a bidirectional LSTM
    pub bidirectional: bool,
}

impl Default for LSTMConfig {
    fn default() -> Self {
        Self {
            input_size: 10,
            hidden_size: 128,
            num_layers: 2,
            output_size: 1,
            dropout: 0.2,
            bidirectional: false,
        }
    }
}

/// LSTM model for time series prediction
pub struct LSTMModel {
    /// Configuration parameters for this model
    pub config: LSTMConfig,
    /// Whether the model has been trained
    pub is_trained: bool,
    /// Training loss recorded at each epoch
    pub training_loss: Vec<f64>,
}

impl LSTMModel {
    /// Create a new untrained LSTM model with the given configuration
    pub fn new(config: LSTMConfig) -> Self {
        Self {
            config,
            is_trained: false,
            training_loss: Vec::new(),
        }
    }

    /// Train the LSTM model
    pub fn train(
        &mut self,
        _sequences: &[Vec<f64>],
        _targets: &[f64],
        epochs: usize,
    ) -> Result<()> {
        // NOTE: In production, this would use actual LSTM implementation
        // For now, simulate training by storing dummy loss values
        for epoch in 0..epochs {
            let loss = 1.0 / (epoch + 1) as f64; // Simulated decreasing loss
            self.training_loss.push(loss);
        }

        self.is_trained = true;
        Ok(())
    }

    /// Predict next values
    pub fn predict(&self, sequence: &[f64]) -> Result<Vec<f64>> {
        if !self.is_trained {
            anyhow::bail!("Model must be trained before prediction");
        }

        // NOTE: Would use actual LSTM forward pass in production
        // Placeholder: return last value (naive prediction)
        Ok(vec![sequence.last().copied().unwrap_or(0.0)])
    }

    /// Get attention weights (for interpretability)
    pub fn get_attention_weights(&self, sequence: &[f64]) -> Vec<f64> {
        // Simplified attention: give more weight to recent values
        let len = sequence.len();
        (0..len).map(|i| (i + 1) as f64 / len as f64).collect()
    }
}

/// Transformer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransformerConfig {
    /// Model embedding dimension
    pub d_model: usize,
    /// Number of parallel attention heads
    pub nhead: usize,
    /// Number of encoder layers
    pub num_encoder_layers: usize,
    /// Number of decoder layers
    pub num_decoder_layers: usize,
    /// Dimension of the feed-forward sub-layer
    pub dim_feedforward: usize,
    /// Dropout rate applied within the transformer (0.0–1.0)
    pub dropout: f64,
    /// Maximum supported input sequence length
    pub max_seq_length: usize,
}

impl Default for TransformerConfig {
    fn default() -> Self {
        Self {
            d_model: 256,
            nhead: 8,
            num_encoder_layers: 6,
            num_decoder_layers: 6,
            dim_feedforward: 1024,
            dropout: 0.1,
            max_seq_length: 100,
        }
    }
}

/// Transformer model for sequence-to-sequence tasks
pub struct TransformerModel {
    /// Configuration parameters for this model
    pub config: TransformerConfig,
    /// Whether the model has been trained
    pub is_trained: bool,
}

impl TransformerModel {
    /// Create a new untrained transformer with the given configuration
    pub fn new(config: TransformerConfig) -> Self {
        Self {
            config,
            is_trained: false,
        }
    }

    /// Train transformer
    pub fn train(
        &mut self,
        _src_sequences: &[Vec<f64>],
        _tgt_sequences: &[Vec<f64>],
    ) -> Result<()> {
        // NOTE: Would implement actual transformer training in production
        self.is_trained = true;
        Ok(())
    }

    /// Predict with transformer
    pub fn predict(&self, src_sequence: &[f64]) -> Result<Vec<f64>> {
        if !self.is_trained {
            anyhow::bail!("Model must be trained first");
        }

        // Placeholder prediction
        Ok(vec![src_sequence.last().copied().unwrap_or(0.0)])
    }
}

/// Attention mechanism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AttentionMechanism {
    /// Which variant of attention to compute
    pub attention_type: AttentionType,
    /// Number of parallel attention heads (for multi-head variants)
    pub num_heads: usize,
}

/// Types of attention
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AttentionType {
    /// Scaled dot-product attention
    ScaledDotProduct,
    /// Multi-head attention
    MultiHead,
    /// Self-attention
    SelfAttention,
    /// Cross-attention
    CrossAttention,
}

impl AttentionMechanism {
    /// Create a new attention mechanism of the given type
    pub fn new(attention_type: AttentionType, num_heads: usize) -> Self {
        Self {
            attention_type,
            num_heads,
        }
    }

    /// Compute attention scores
    pub fn compute_attention(
        &self,
        query: &[f64],
        keys: &[Vec<f64>],
        values: &[Vec<f64>],
    ) -> Result<Vec<f64>> {
        // NOTE: Would implement actual attention computation in production
        // Simplified: return weighted average based on similarity to query

        let scores: Vec<f64> = keys
            .iter()
            .map(|key| {
                // Simplified dot product
                query
                    .iter()
                    .zip(key.iter())
                    .map(|(q, k)| q * k)
                    .sum::<f64>()
            })
            .collect();

        // Softmax (simplified)
        let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let exp_scores: Vec<f64> = scores.iter().map(|s| (s - max_score).exp()).collect();
        let sum_exp: f64 = exp_scores.iter().sum();
        let attention_weights: Vec<f64> = exp_scores.iter().map(|e| e / sum_exp).collect();

        // Weighted sum of values
        let output_dim = values[0].len();
        let mut output = vec![0.0; output_dim];

        for (weight, value) in attention_weights.iter().zip(values.iter()) {
            for (i, &v) in value.iter().enumerate() {
                output[i] += weight * v;
            }
        }

        Ok(output)
    }
}

/// GAN configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GANConfig {
    /// Dimensionality of the random latent input to the generator
    pub latent_dim: usize,
    /// Hidden layer sizes for the generator network
    pub generator_layers: Vec<usize>,
    /// Hidden layer sizes for the discriminator network
    pub discriminator_layers: Vec<usize>,
    /// Learning rate for the generator optimizer
    pub learning_rate_g: f64,
    /// Learning rate for the discriminator optimizer
    pub learning_rate_d: f64,
}

impl Default for GANConfig {
    fn default() -> Self {
        Self {
            latent_dim: 100,
            generator_layers: vec![128, 256, 512],
            discriminator_layers: vec![512, 256, 128],
            learning_rate_g: 0.0002,
            learning_rate_d: 0.0002,
        }
    }
}

/// GAN for synthetic data generation
pub struct GANModel {
    /// Configuration parameters for this GAN
    pub config: GANConfig,
    /// Whether the GAN has been trained
    pub is_trained: bool,
    /// Generator loss at each training epoch
    pub generator_loss: Vec<f64>,
    /// Discriminator loss at each training epoch
    pub discriminator_loss: Vec<f64>,
}

impl GANModel {
    /// Create a new untrained GAN with the given configuration
    pub fn new(config: GANConfig) -> Self {
        Self {
            config,
            is_trained: false,
            generator_loss: Vec::new(),
            discriminator_loss: Vec::new(),
        }
    }

    /// Train GAN
    pub fn train(&mut self, _real_data: &[Vec<f64>], epochs: usize) -> Result<()> {
        // NOTE: Would implement actual GAN training in production
        // Simulated adversarial training
        for epoch in 0..epochs {
            let g_loss = 1.0 / (epoch + 1) as f64;
            let d_loss = 0.5 + (0.5 / (epoch + 1) as f64);

            self.generator_loss.push(g_loss);
            self.discriminator_loss.push(d_loss);
        }

        self.is_trained = true;
        Ok(())
    }

    /// Generate synthetic data
    pub fn generate(&self, num_samples: usize) -> Result<Vec<Vec<f64>>> {
        if !self.is_trained {
            anyhow::bail!("GAN must be trained first");
        }

        // NOTE: Would use actual generator network in production
        // Placeholder: generate random data
        let samples = (0..num_samples)
            .map(|_| vec![0.0; self.config.latent_dim])
            .collect();

        Ok(samples)
    }

    /// Discriminate real vs fake
    pub fn discriminate(&self, _data: &[f64]) -> Result<f64> {
        // NOTE: Would use actual discriminator network
        // Placeholder: return probability
        Ok(0.5)
    }
}

/// Feature importance selector backed by an attention mechanism
pub struct AttentionFeatureSelector {
    /// Underlying attention mechanism used to score features
    pub attention: AttentionMechanism,
    /// Most recently computed importance scores, one per feature
    pub importance_scores: Vec<f64>,
}

impl AttentionFeatureSelector {
    /// Create a new feature selector with the given number of attention heads
    pub fn new(num_heads: usize) -> Self {
        Self {
            attention: AttentionMechanism::new(AttentionType::MultiHead, num_heads),
            importance_scores: Vec::new(),
        }
    }

    /// Compute feature importance
    pub fn compute_importance(&mut self, features: &[Vec<f64>]) -> Result<Vec<f64>> {
        // Use attention to determine feature importance
        let query = features.last().unwrap(); // Use latest as query
        let attention_output = self
            .attention
            .compute_attention(query, features, features)?;

        // Compute importance as magnitude of attention output
        self.importance_scores = attention_output.iter().map(|x| x.abs()).collect();

        Ok(self.importance_scores.clone())
    }

    /// Select top-k features
    pub fn select_top_features(&self, k: usize) -> Vec<usize> {
        let mut indexed_scores: Vec<_> = self.importance_scores.iter().enumerate().collect();

        indexed_scores.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap());

        indexed_scores.iter().take(k).map(|(i, _)| *i).collect()
    }
}

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

    #[test]
    fn test_lstm_config_default() {
        let config = LSTMConfig::default();
        assert_eq!(config.hidden_size, 128);
        assert_eq!(config.num_layers, 2);
        assert!(!config.bidirectional);
    }

    #[test]
    fn test_lstm_training() {
        let mut model = LSTMModel::new(LSTMConfig::default());
        let sequences = vec![vec![1.0, 2.0, 3.0], vec![2.0, 3.0, 4.0]];
        let targets = vec![4.0, 5.0];

        model.train(&sequences, &targets, 10).unwrap();
        assert!(model.is_trained);
        assert_eq!(model.training_loss.len(), 10);
    }

    #[test]
    fn test_lstm_prediction() {
        let mut model = LSTMModel::new(LSTMConfig::default());
        let sequences = vec![vec![1.0, 2.0, 3.0]];
        let targets = vec![4.0];

        model.train(&sequences, &targets, 5).unwrap();

        let prediction = model.predict(&[1.0, 2.0, 3.0]).unwrap();
        assert!(!prediction.is_empty());
    }

    #[test]
    fn test_lstm_attention_weights() {
        let model = LSTMModel::new(LSTMConfig::default());
        let sequence = vec![1.0, 2.0, 3.0, 4.0, 5.0];

        let weights = model.get_attention_weights(&sequence);
        assert_eq!(weights.len(), 5);
        assert!(weights.last().unwrap() > weights.first().unwrap()); // Recent values have more weight
    }

    #[test]
    fn test_transformer_config() {
        let config = TransformerConfig::default();
        assert_eq!(config.d_model, 256);
        assert_eq!(config.nhead, 8);
    }

    #[test]
    fn test_transformer_model() {
        let mut model = TransformerModel::new(TransformerConfig::default());
        let src = vec![vec![1.0, 2.0, 3.0]];
        let tgt = vec![vec![4.0]];

        model.train(&src, &tgt).unwrap();
        assert!(model.is_trained);
    }

    #[test]
    fn test_attention_mechanism() {
        let attention = AttentionMechanism::new(AttentionType::ScaledDotProduct, 1);
        let query = vec![1.0, 0.0];
        let keys = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
        let values = vec![vec![1.0], vec![2.0], vec![3.0]];

        let output = attention.compute_attention(&query, &keys, &values).unwrap();
        assert_eq!(output.len(), 1);
    }

    #[test]
    fn test_gan_config() {
        let config = GANConfig::default();
        assert_eq!(config.latent_dim, 100);
        assert_eq!(config.generator_layers.len(), 3);
    }

    #[test]
    fn test_gan_training() {
        let mut gan = GANModel::new(GANConfig::default());
        let real_data = vec![vec![1.0, 2.0, 3.0]; 100];

        gan.train(&real_data, 10).unwrap();
        assert!(gan.is_trained);
        assert_eq!(gan.generator_loss.len(), 10);
        assert_eq!(gan.discriminator_loss.len(), 10);
    }

    #[test]
    fn test_gan_generation() {
        let mut gan = GANModel::new(GANConfig::default());
        let real_data = vec![vec![1.0, 2.0, 3.0]; 10];

        gan.train(&real_data, 5).unwrap();

        let synthetic = gan.generate(5).unwrap();
        assert_eq!(synthetic.len(), 5);
    }

    #[test]
    fn test_attention_feature_selector() {
        let mut selector = AttentionFeatureSelector::new(4);
        let features = vec![
            vec![1.0, 2.0, 3.0],
            vec![2.0, 3.0, 4.0],
            vec![3.0, 4.0, 5.0],
        ];

        let importance = selector.compute_importance(&features).unwrap();
        assert_eq!(importance.len(), 3);

        let top_features = selector.select_top_features(2);
        assert_eq!(top_features.len(), 2);
    }
}