scirs2-sparse 0.4.2

//! Reinforcement learning components for adaptive sparse matrix optimization
//!
//! This module implements RL agents that learn optimal strategies for sparse
//! matrix operations based on performance feedback and matrix characteristics.

use super::neural_network::NeuralNetwork;
use super::pattern_memory::OptimizationStrategy;
use crate::error::SparseResult;
use scirs2_core::random::{Rng, RngExt};
use std::collections::VecDeque;

/// Reinforcement learning algorithms
#[derive(Debug, Clone, Copy)]
pub enum RLAlgorithm {
    /// Q-Learning with experience replay
    DQN,
    /// Policy gradient methods
    PolicyGradient,
    /// Actor-Critic methods
    ActorCritic,
    /// Proximal Policy Optimization
    PPO,
    /// Soft Actor-Critic
    SAC,
}

/// Reinforcement learning agent
#[derive(Debug)]
#[allow(dead_code)]
pub(crate) struct RLAgent {
    pub q_network: NeuralNetwork,
    pub target_network: Option<NeuralNetwork>,
    pub policy_network: Option<NeuralNetwork>,
    pub value_network: Option<NeuralNetwork>,
    pub algorithm: RLAlgorithm,
    pub epsilon: f64,
    pub learningrate: f64,
}

/// Experience for reinforcement learning
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub(crate) struct Experience {
    pub state: Vec<f64>,
    pub action: OptimizationStrategy,
    pub reward: f64,
    pub next_state: Vec<f64>,
    pub done: bool,
    pub timestamp: u64,
}

/// Experience replay buffer
#[derive(Debug)]
pub(crate) struct ExperienceBuffer {
    pub buffer: VecDeque<Experience>,
    pub capacity: usize,
    pub priority_weights: Vec<f64>,
}

/// Performance metrics for reinforcement learning
#[derive(Debug, Clone)]
pub struct PerformanceMetrics {
    #[allow(dead_code)]
    pub executiontime: f64,
    #[allow(dead_code)]
    pub cache_efficiency: f64,
    #[allow(dead_code)]
    pub simd_utilization: f64,
    #[allow(dead_code)]
    pub parallel_efficiency: f64,
    #[allow(dead_code)]
    pub memory_bandwidth: f64,
    pub strategy_used: OptimizationStrategy,
}

impl RLAgent {
    /// Create a new RL agent
    pub fn new(
        state_size: usize,
        action_size: usize,
        algorithm: RLAlgorithm,
        learning_rate: f64,
        epsilon: f64,
    ) -> Self {
        let q_network = NeuralNetwork::new(state_size, 3, 64, action_size, 4);

        let target_network = match algorithm {
            RLAlgorithm::DQN => Some(q_network.clone()),
            _ => None,
        };

        let (policy_network, value_network) = match algorithm {
            RLAlgorithm::ActorCritic | RLAlgorithm::PPO | RLAlgorithm::SAC => {
                let policy = NeuralNetwork::new(state_size, 2, 32, action_size, 4);
                let value = NeuralNetwork::new(state_size, 2, 32, 1, 4);
                (Some(policy), Some(value))
            }
            _ => (None, None),
        };

        Self {
            q_network,
            target_network,
            policy_network,
            value_network,
            algorithm,
            epsilon,
            learningrate: learning_rate,
        }
    }

    /// Select action using current policy
    pub fn select_action(&self, state: &[f64]) -> OptimizationStrategy {
        let mut rng = scirs2_core::random::thread_rng();

        // Epsilon-greedy action selection for DQN
        if matches!(self.algorithm, RLAlgorithm::DQN) && rng.random::<f64>() < self.epsilon {
            // Random action
            self.random_action()
        } else {
            // Greedy action
            self.greedy_action(state)
        }
    }

    /// Select greedy action based on current Q-values or policy
    fn greedy_action(&self, state: &[f64]) -> OptimizationStrategy {
        match self.algorithm {
            RLAlgorithm::DQN => {
                let q_values = self.q_network.forward(state);
                let best_action_idx = q_values
                    .iter()
                    .enumerate()
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).expect("Operation failed"))
                    .map(|(idx, _)| idx)
                    .unwrap_or(0);
                self.idx_to_strategy(best_action_idx)
            }
            RLAlgorithm::PolicyGradient
            | RLAlgorithm::ActorCritic
            | RLAlgorithm::PPO
            | RLAlgorithm::SAC => {
                if let Some(ref policy_network) = self.policy_network {
                    let action_probs = policy_network.forward(state);
                    let best_action_idx = action_probs
                        .iter()
                        .enumerate()
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).expect("Operation failed"))
                        .map(|(idx, _)| idx)
                        .unwrap_or(0);
                    self.idx_to_strategy(best_action_idx)
                } else {
                    self.random_action()
                }
            }
        }
    }

    /// Select random action
    fn random_action(&self) -> OptimizationStrategy {
        let mut rng = scirs2_core::random::thread_rng();
        let strategies = [
            OptimizationStrategy::RowWiseCache,
            OptimizationStrategy::ColumnWiseLocality,
            OptimizationStrategy::BlockStructured,
            OptimizationStrategy::DiagonalOptimized,
            OptimizationStrategy::Hierarchical,
            OptimizationStrategy::StreamingCompute,
            OptimizationStrategy::SIMDVectorized,
            OptimizationStrategy::ParallelWorkStealing,
            OptimizationStrategy::AdaptiveHybrid,
        ];

        strategies[rng.random_range(0..strategies.len())]
    }

    /// Convert action index to optimization strategy
    fn idx_to_strategy(&self, idx: usize) -> OptimizationStrategy {
        match idx % 9 {
            0 => OptimizationStrategy::RowWiseCache,
            1 => OptimizationStrategy::ColumnWiseLocality,
            2 => OptimizationStrategy::BlockStructured,
            3 => OptimizationStrategy::DiagonalOptimized,
            4 => OptimizationStrategy::Hierarchical,
            5 => OptimizationStrategy::StreamingCompute,
            6 => OptimizationStrategy::SIMDVectorized,
            7 => OptimizationStrategy::ParallelWorkStealing,
            _ => OptimizationStrategy::AdaptiveHybrid,
        }
    }

    /// Convert optimization strategy to action index
    fn strategy_to_idx(&self, strategy: OptimizationStrategy) -> usize {
        Self::strategy_to_idx_static(strategy)
    }

    /// Static version of strategy_to_idx to avoid borrowing issues
    fn strategy_to_idx_static(strategy: OptimizationStrategy) -> usize {
        match strategy {
            OptimizationStrategy::RowWiseCache => 0,
            OptimizationStrategy::ColumnWiseLocality => 1,
            OptimizationStrategy::BlockStructured => 2,
            OptimizationStrategy::DiagonalOptimized => 3,
            OptimizationStrategy::Hierarchical => 4,
            OptimizationStrategy::StreamingCompute => 5,
            OptimizationStrategy::SIMDVectorized => 6,
            OptimizationStrategy::ParallelWorkStealing => 7,
            OptimizationStrategy::AdaptiveHybrid => 8,
        }
    }

    /// Train the agent on a batch of experiences
    pub fn train(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        if experiences.is_empty() {
            return Ok(());
        }

        match self.algorithm {
            RLAlgorithm::DQN => self.train_dqn(experiences),
            RLAlgorithm::PolicyGradient => self.train_policy_gradient(experiences),
            RLAlgorithm::ActorCritic => self.train_actor_critic(experiences),
            RLAlgorithm::PPO => self.train_ppo(experiences),
            RLAlgorithm::SAC => self.train_sac(experiences),
        }
    }

    /// Train DQN algorithm
    fn train_dqn(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        for experience in experiences {
            let current_q_values = self.q_network.forward(&experience.state);
            let action_idx = self.strategy_to_idx(experience.action);

            let target = if experience.done {
                experience.reward
            } else if let Some(ref target_network) = self.target_network {
                let next_q_values = target_network.forward(&experience.next_state);
                let max_next_q = next_q_values
                    .iter()
                    .fold(f64::NEG_INFINITY, |a, &b| a.max(b));
                experience.reward + 0.99 * max_next_q // gamma = 0.99
            } else {
                experience.reward
            };

            // Simplified Q-learning update
            // In practice, you'd compute proper gradients and update weights
            let mut target_q_values = current_q_values;
            if action_idx < target_q_values.len() {
                target_q_values[action_idx] = target;
            }

            // Update Q-network (simplified)
            let (_, cache) = self.q_network.forward_with_cache(&experience.state);
            let gradients =
                self.q_network
                    .compute_gradients(&experience.state, &target_q_values, &cache);
            self.q_network.update_weights(&gradients, self.learningrate);
        }

        Ok(())
    }

    /// Train policy gradient algorithm
    fn train_policy_gradient(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        // Simplified policy gradient implementation
        let learning_rate = self.learningrate;
        if let Some(ref mut policy_network) = self.policy_network {
            for experience in experiences {
                let action_probs = policy_network.forward(&experience.state);
                let action_idx = Self::strategy_to_idx_static(experience.action);

                // Compute policy gradient (simplified)
                let mut target_probs = action_probs;
                if action_idx < target_probs.len() {
                    target_probs[action_idx] += learning_rate * experience.reward;
                }

                // Update policy network
                let (_, cache) = policy_network.forward_with_cache(&experience.state);
                let gradients =
                    policy_network.compute_gradients(&experience.state, &target_probs, &cache);
                policy_network.update_weights(&gradients, learning_rate);
            }
        }

        Ok(())
    }

    /// Train Actor-Critic algorithm
    fn train_actor_critic(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        // Simplified Actor-Critic implementation
        let learning_rate = self.learningrate;
        for experience in experiences {
            // Update critic (value network)
            if let Some(ref mut value_network) = self.value_network {
                let current_value = value_network.forward(&experience.state)[0];
                let target_value = if experience.done {
                    experience.reward
                } else {
                    let next_value = value_network.forward(&experience.next_state)[0];
                    experience.reward + 0.99 * next_value
                };

                let (_, cache) = value_network.forward_with_cache(&experience.state);
                let gradients =
                    value_network.compute_gradients(&experience.state, &[target_value], &cache);
                value_network.update_weights(&gradients, learning_rate);

                // Update actor (policy network)
                if let Some(ref mut policy_network) = self.policy_network {
                    let advantage = target_value - current_value;
                    let action_probs = policy_network.forward(&experience.state);
                    let action_idx = Self::strategy_to_idx_static(experience.action);

                    let mut target_probs = action_probs;
                    if action_idx < target_probs.len() {
                        target_probs[action_idx] += learning_rate * advantage;
                    }

                    let (_, cache) = policy_network.forward_with_cache(&experience.state);
                    let gradients =
                        policy_network.compute_gradients(&experience.state, &target_probs, &cache);
                    policy_network.update_weights(&gradients, learning_rate);
                }
            }
        }

        Ok(())
    }

    /// Train PPO algorithm (simplified)
    fn train_ppo(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        // Simplified PPO implementation
        self.train_actor_critic(experiences) // Using Actor-Critic as base
    }

    /// Train SAC algorithm (simplified)
    fn train_sac(&mut self, experiences: &[Experience]) -> SparseResult<()> {
        // Simplified SAC implementation
        self.train_actor_critic(experiences) // Using Actor-Critic as base
    }

    /// Update target network (for DQN)
    pub fn update_target_network(&mut self) {
        if let Some(ref mut target_network) = self.target_network {
            let params = self.q_network.get_parameters();
            target_network.set_parameters(&params);
        }
    }

    /// Decay exploration rate
    pub fn decay_epsilon(&mut self, decay_rate: f64) {
        self.epsilon *= decay_rate;
        self.epsilon = self.epsilon.max(0.01); // Minimum epsilon
    }

    /// Compute value estimate for a state
    pub fn estimate_value(&self, state: &[f64]) -> f64 {
        match self.algorithm {
            RLAlgorithm::DQN => {
                let q_values = self.q_network.forward(state);
                q_values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b))
            }
            _ => {
                if let Some(ref value_network) = self.value_network {
                    value_network.forward(state)[0]
                } else {
                    0.0
                }
            }
        }
    }
}

impl ExperienceBuffer {
    /// Create a new experience buffer
    pub fn new(capacity: usize) -> Self {
        Self {
            buffer: VecDeque::new(),
            capacity,
            priority_weights: Vec::new(),
        }
    }

    /// Add experience to buffer
    pub fn add(&mut self, experience: Experience) {
        if self.buffer.len() >= self.capacity {
            self.buffer.pop_front();
            if !self.priority_weights.is_empty() {
                self.priority_weights.remove(0);
            }
        }

        self.buffer.push_back(experience);
        self.priority_weights.push(1.0); // Default priority
    }

    /// Sample a batch of experiences
    pub fn sample(&self, batch_size: usize) -> Vec<Experience> {
        let mut rng = scirs2_core::random::thread_rng();
        let mut batch = Vec::new();

        for _ in 0..batch_size.min(self.buffer.len()) {
            let idx = rng.random_range(0..self.buffer.len());
            if let Some(experience) = self.buffer.get(idx) {
                batch.push(experience.clone());
            }
        }

        batch
    }

    /// Sample with priority weights
    pub fn sample_prioritized(&self, batch_size: usize) -> Vec<Experience> {
        if self.priority_weights.is_empty() {
            return self.sample(batch_size);
        }

        let mut rng = scirs2_core::random::thread_rng();
        let mut batch = Vec::new();
        let total_weight: f64 = self.priority_weights.iter().sum();

        for _ in 0..batch_size.min(self.buffer.len()) {
            let mut weight_sum = 0.0;
            let target = rng.random::<f64>() * total_weight;

            for (idx, &weight) in self.priority_weights.iter().enumerate() {
                weight_sum += weight;
                if weight_sum >= target {
                    if let Some(experience) = self.buffer.get(idx) {
                        batch.push(experience.clone());
                        break;
                    }
                }
            }
        }

        batch
    }

    /// Update priority for experience
    pub fn update_priority(&mut self, idx: usize, priority: f64) {
        if idx < self.priority_weights.len() {
            self.priority_weights[idx] = priority.max(0.01); // Minimum priority
        }
    }

    /// Get buffer size
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

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

    /// Clear the buffer
    pub fn clear(&mut self) {
        self.buffer.clear();
        self.priority_weights.clear();
    }
}

impl PerformanceMetrics {
    /// Create new performance metrics
    pub fn new(
        execution_time: f64,
        cache_efficiency: f64,
        simd_utilization: f64,
        parallel_efficiency: f64,
        memory_bandwidth: f64,
        strategy_used: OptimizationStrategy,
    ) -> Self {
        Self {
            executiontime: execution_time,
            cache_efficiency,
            simd_utilization,
            parallel_efficiency,
            memory_bandwidth,
            strategy_used,
        }
    }

    /// Compute reward for reinforcement learning
    pub fn compute_reward(&self, baseline_time: f64) -> f64 {
        // Reward based on performance improvement
        let time_improvement = (baseline_time - self.executiontime) / baseline_time;
        let efficiency_score =
            (self.cache_efficiency + self.simd_utilization + self.parallel_efficiency) / 3.0;

        // Combined reward considering both time improvement and efficiency
        time_improvement * 10.0 + efficiency_score * 5.0
    }

    /// Get overall performance score
    pub fn performance_score(&self) -> f64 {
        let time_score = 1.0 / (1.0 + self.executiontime); // Lower time is better
        let efficiency_score = (self.cache_efficiency
            + self.simd_utilization
            + self.parallel_efficiency
            + self.memory_bandwidth)
            / 4.0;

        (time_score + efficiency_score) / 2.0
    }
}