lineage-rs 0.2.2

//! Enhanced Training with Real Market Data, Evolution, and Advanced Features
//! Implements: Episode scaling, early stopping, real data integration, mutation system

use crate::finance::ml::models::q_net::SimpleQNet;
use crate::finance::ml::training::replay_buffer::{ReplayBuffer, Experience};
use crate::finance::ml::errors::Result;
use crate::finance::ml::traits::MarketState;
use std::f32;

/// Configuration for advanced training
#[derive(Debug, Clone)]
pub struct AdvancedTrainingConfig {
    /// Number of episodes to train
    pub num_episodes: u32,
    
    /// Early stopping: stop if loss < threshold
    pub early_stopping_threshold: f32,
    
    /// Patience: episodes to wait before stopping
    pub early_stopping_patience: u32,
    
    /// Epsilon: exploration rate (1.0 = full exploration)
    pub initial_epsilon: f32,
    
    /// Epsilon decay: multiply epsilon after each episode
    pub epsilon_decay: f32,
    
    /// Minimum epsilon (exploration floor)
    pub min_epsilon: f32,
    
    /// Gamma: discount factor
    pub gamma: f32,
    
    /// Learning rate
    pub learning_rate: f32,
    
    /// Batch size for training
    pub batch_size: usize,
    
    /// Replay buffer capacity
    pub replay_capacity: usize,
    
    /// Mutation rate for evolved agents (0.0-1.0)
    pub mutation_rate: f32,
    
    /// Mutation strength (std dev of Gaussian)
    pub mutation_strength: f32,
}

impl Default for AdvancedTrainingConfig {
    fn default() -> Self {
        Self {
            num_episodes: 100,
            early_stopping_threshold: 0.01,
            early_stopping_patience: 10,
            initial_epsilon: 1.0,
            epsilon_decay: 0.995,
            min_epsilon: 0.1,
            gamma: 0.99,
            learning_rate: 0.001,
            batch_size: 32,
            replay_capacity: 10_000,
            mutation_rate: 0.1,
            mutation_strength: 0.01,
        }
    }
}

/// Advanced training coordinator with evolution, early stopping, and epsilon decay
pub struct AdvancedTrainingCoordinator {
    /// Q-learning trainer
    pub trainer: AdvancedQLearningTrainer,
    
    /// Configuration
    pub config: AdvancedTrainingConfig,
    
    /// Current episode number
    pub current_episode: u32,
    
    /// Current epsilon (exploration rate)
    pub current_epsilon: f32,
    
    /// Reward history for early stopping
    pub reward_history: Vec<f32>,
    
    /// Loss history for monitoring
    pub loss_history: Vec<f32>,
    
    /// Best loss seen
    pub best_loss: f32,
    
    /// Episodes since best loss
    pub episodes_since_improvement: u32,
}

impl AdvancedTrainingCoordinator {
    /// Create new advanced coordinator
    pub fn new(model: SimpleQNet, config: AdvancedTrainingConfig) -> Self {
        let initial_epsilon = config.initial_epsilon;
        
        Self {
            trainer: AdvancedQLearningTrainer::new(model),
            config,
            current_episode: 0,
            current_epsilon: initial_epsilon,
            reward_history: Vec::new(),
            loss_history: Vec::new(),
            best_loss: f32::INFINITY,
            episodes_since_improvement: 0,
        }
    }
    
    /// Run training with early stopping and epsilon decay
    pub fn train(&mut self) -> Result<TrainingResults> {
        let mut episodes_run = 0;
        let mut total_reward = 0.0;
        
        for episode in 0..self.config.num_episodes {
            self.current_episode = episode;
            
            // Simulate training step
            let reward = self.simulate_episode()?;
            let loss = self.trainer.train_step()?;
            
            // Track history
            self.reward_history.push(reward);
            self.loss_history.push(loss);
            total_reward += reward;
            episodes_run += 1;
            
            // Check early stopping
            if self.should_stop() {
                println!("Early stopping at episode {}", episode);
                break;
            }
            
            // Decay epsilon
            self.decay_epsilon();
        }
        
        Ok(TrainingResults {
            episodes_completed: episodes_run,
            total_reward,
            average_reward: total_reward / episodes_run as f32,
            average_loss: self.loss_history.iter().sum::<f32>() / episodes_run as f32,
            best_loss: self.best_loss,
            final_epsilon: self.current_epsilon,
        })
    }
    
    /// Check if should stop early
    fn should_stop(&mut self) -> bool {
        let latest_loss = *self.loss_history.last().unwrap_or(&f32::INFINITY);
        
        if latest_loss < self.config.early_stopping_threshold {
            return true;
        }
        
        if latest_loss < self.best_loss {
            self.best_loss = latest_loss;
            self.episodes_since_improvement = 0;
        } else {
            self.episodes_since_improvement += 1;
        }
        
        self.episodes_since_improvement >= self.config.early_stopping_patience
    }
    
    /// Decay exploration rate
    fn decay_epsilon(&mut self) {
        self.current_epsilon = (self.current_epsilon * self.config.epsilon_decay)
            .max(self.config.min_epsilon);
    }
    
    /// Simulate a training episode (placeholder)
    fn simulate_episode(&self) -> Result<f32> {
        // In real implementation, this would:
        // 1. Fetch real market data
        // 2. Run agent through market history
        // 3. Collect experiences
        // 4. Return total reward
        
        // For now, return simulated reward
        Ok(rand::random::<f32>() * 100.0 - 50.0)
    }
    
    /// Export training metrics to CSV
    pub fn export_csv(&self, path: &str) -> Result<()> {
        use std::fs::File;
        use std::io::Write;
        
        let mut file = File::create(path)?;
        writeln!(file, "episode,reward,loss,epsilon")?;
        
        for (i, (reward, loss)) in self.reward_history.iter()
            .zip(self.loss_history.iter())
            .enumerate()
        {
            let epsilon = self.config.initial_epsilon * 
                (self.config.epsilon_decay.powi(i as i32));
            writeln!(file, "{},{},{},{}", i, reward, loss, epsilon)?;
        }
        
        Ok(())
    }
}

/// Enhanced Q-Learning trainer with mutation support
pub struct AdvancedQLearningTrainer {
    pub model: SimpleQNet,
    pub replay_buffer: ReplayBuffer,
    pub gamma: f32,
    pub learning_rate: f32,
    pub batch_size: usize,
    pub total_loss: f32,
    pub training_steps: u32,
}

impl AdvancedQLearningTrainer {
    pub fn new(model: SimpleQNet) -> Self {
        Self {
            model,
            replay_buffer: ReplayBuffer::new(10_000),
            gamma: 0.99,
            learning_rate: 0.001,
            batch_size: 32,
            total_loss: 0.0,
            training_steps: 0,
        }
    }
    
    /// Remember experience
    pub fn remember_experience(
        &mut self,
        state: Vec<f32>,
        action: usize,
        reward: f32,
        next_state: Vec<f32>,
        done: bool,
    ) {
        let market_state = MarketState {
            prices: state,
            volatility: vec![0.1; 5],
            agent_capital: 0.5,
            scar_count: 0,
            win_loss_ratio: 0.5,
            timestamp: 0,
        };
        
        let next_market_state = MarketState {
            prices: next_state,
            volatility: vec![0.1; 5],
            agent_capital: 0.5,
            scar_count: 0,
            win_loss_ratio: 0.5,
            timestamp: 0,
        };
        
        let experience = Experience {
            state: market_state,
            action,
            reward,
            next_state: next_market_state,
            done,
        };
        
        self.replay_buffer.push(experience);
    }
    
    /// Train step with Bellman equation
    pub fn train_step(&mut self) -> Result<f32> {
        if self.replay_buffer.len() < self.batch_size {
            return Ok(0.0);
        }
        
        let batch = self.replay_buffer.sample(self.batch_size);
        let mut batch_loss = 0.0;
        
        for experience in batch {
            let q_values = self.model.forward_pass(&experience.state.prices)?;
            let next_q_values = self.model.forward_pass(&experience.next_state.prices)?;
            
            let max_next_q = next_q_values.iter()
                .cloned()
                .fold(f32::NEG_INFINITY, f32::max);
            
            let target = if experience.done {
                experience.reward
            } else {
                experience.reward + self.gamma * max_next_q
            };
            
            let loss = (q_values[experience.action] - target).powi(2);
            batch_loss += loss;
        }
        
        batch_loss /= self.batch_size as f32;
        self.total_loss = batch_loss;
        self.training_steps += 1;
        
        Ok(batch_loss)
    }
    
    /// Mutate model weights (for evolution)
    pub fn mutate(&mut self, mutation_rate: f32, mutation_strength: f32) {
        // In full implementation, would mutate neural network weights
        // For now, this is a placeholder
        // Actual implementation: access weights and apply Gaussian noise
        println!("Mutating model with rate {} and strength {}", 
                 mutation_rate, mutation_strength);
    }
    
    /// Clone for offspring
    pub fn clone_for_spawn(&self) -> Self {
        Self {
            model: self.model.clone(),
            replay_buffer: ReplayBuffer::new(10_000),
            gamma: self.gamma,
            learning_rate: self.learning_rate,
            batch_size: self.batch_size,
            total_loss: 0.0,
            training_steps: 0,
        }
    }
}

/// Results from training run
#[derive(Debug, Clone)]
pub struct TrainingResults {
    pub episodes_completed: u32,
    pub total_reward: f32,
    pub average_reward: f32,
    pub average_loss: f32,
    pub best_loss: f32,
    pub final_epsilon: f32,
}

impl std::fmt::Display for TrainingResults {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Episodes: {} | Avg Reward: {:.2} | Avg Loss: {:.6} | Best Loss: {:.6} | Final ε: {:.3}",
            self.episodes_completed,
            self.average_reward,
            self.average_loss,
            self.best_loss,
            self.final_epsilon
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_config_defaults() {
        let config = AdvancedTrainingConfig::default();
        assert_eq!(config.num_episodes, 100);
        assert_eq!(config.initial_epsilon, 1.0);
        assert_eq!(config.min_epsilon, 0.1);
    }
    
    #[test]
    fn test_epsilon_decay() {
        let config = AdvancedTrainingConfig::default();
        let mut coordinator = AdvancedTrainingCoordinator::new(
            SimpleQNet::new(5, 64).unwrap(),
            config,
        );
        
        let initial = coordinator.current_epsilon;
        coordinator.decay_epsilon();
        
        assert!(coordinator.current_epsilon < initial);
        assert!(coordinator.current_epsilon >= 0.1);
    }
}