torsh-optim 0.1.3

//! Rectified Adam optimizer

use crate::{
    Optimizer, OptimizerError, OptimizerResult, OptimizerState, ParamGroup, ParamGroupState,
};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::ops::Add;
use std::sync::Arc;
use torsh_core::error::{Result, TorshError};
use torsh_tensor::Tensor;

/// Rectified Adam (RAdam) optimizer
///
/// RAdam addresses the bad convergence problem of Adam during the early stage of training.
/// It provides a dynamic rectification term to stabilize the variance of adaptive learning rate.
pub struct RAdam {
    param_groups: Vec<ParamGroup>,
    state: HashMap<String, HashMap<String, Tensor>>,
    step_count: usize,
    beta1: f32,
    beta2: f32,
    eps: f32,
}

impl RAdam {
    /// Create a new RAdam optimizer
    ///
    /// # Arguments
    /// * `params` - Parameters to optimize
    /// * `lr` - Learning rate (default: 1e-3)
    /// * `beta1` - First moment decay rate (default: 0.9)
    /// * `beta2` - Second moment decay rate (default: 0.999)
    /// * `eps` - Small constant for numerical stability (default: 1e-8)
    /// * `weight_decay` - Weight decay (L2 penalty) (default: 0.0)
    pub fn new(
        params: Vec<Arc<RwLock<Tensor>>>,
        lr: Option<f32>,
        beta1: Option<f32>,
        beta2: Option<f32>,
        eps: Option<f32>,
        weight_decay: Option<f32>,
    ) -> Self {
        let lr = lr.unwrap_or(1e-3);
        let beta1 = beta1.unwrap_or(0.9);
        let beta2 = beta2.unwrap_or(0.999);
        let eps = eps.unwrap_or(1e-8);
        let weight_decay = weight_decay.unwrap_or(0.0);

        let mut options = HashMap::new();
        options.insert("weight_decay".to_string(), weight_decay);

        let param_group = ParamGroup::new(params, lr).with_options(options);

        Self {
            param_groups: vec![param_group],
            state: HashMap::new(),
            step_count: 0,
            beta1,
            beta2,
            eps,
        }
    }

    fn get_param_id(param: &Arc<RwLock<Tensor>>) -> String {
        format!("{:p}", Arc::as_ptr(param))
    }

    fn variance_rectification(&self, step: usize) -> Option<f32> {
        let rho_inf = 2.0 / (1.0 - self.beta2) - 1.0;
        let rho_t = rho_inf
            - 2.0 * step as f32 * self.beta2.powi(step as i32)
                / (1.0 - self.beta2.powi(step as i32));

        if rho_t > 5.0 {
            let r_t = ((rho_t - 4.0) * (rho_t - 2.0) * rho_inf
                / ((rho_inf - 4.0) * (rho_inf - 2.0) * rho_t))
                .sqrt();
            Some(r_t)
        } else {
            None
        }
    }

    /// Get beta1 parameter
    pub fn beta1(&self) -> f32 {
        self.beta1
    }

    /// Get beta2 parameter
    pub fn beta2(&self) -> f32 {
        self.beta2
    }

    /// Get epsilon parameter
    pub fn eps(&self) -> f32 {
        self.eps
    }

    /// Get weight decay parameter
    pub fn weight_decay(&self) -> f32 {
        self.param_groups
            .first()
            .and_then(|group| group.options.get("weight_decay"))
            .copied()
            .unwrap_or(0.0)
    }

    /// Set beta1 parameter
    pub fn set_beta1(&mut self, beta1: f32) {
        self.beta1 = beta1;
    }

    /// Set beta2 parameter
    pub fn set_beta2(&mut self, beta2: f32) {
        self.beta2 = beta2;
    }

    /// Set epsilon parameter
    pub fn set_eps(&mut self, eps: f32) {
        self.eps = eps;
    }

    /// Set weight decay parameter
    pub fn set_weight_decay(&mut self, weight_decay: f32) {
        for group in &mut self.param_groups {
            group
                .options
                .insert("weight_decay".to_string(), weight_decay);
        }
    }

    /// Check if rectification is enabled (always true for RAdam)
    pub fn is_rectification_enabled(&self) -> bool {
        true
    }

    /// Set rectification (no-op for RAdam as it's always enabled)
    pub fn set_rectification(&mut self, _enabled: bool) {
        // RAdam always uses rectification, so this is a no-op
    }

    /// Get the current rectification coefficient
    pub fn rectification_coefficient(&self) -> Option<f32> {
        self.variance_rectification(self.step_count)
    }
}

impl Optimizer for RAdam {
    fn step(&mut self) -> OptimizerResult<()> {
        self.step_count += 1;

        for group in &self.param_groups {
            let lr = group.lr;
            let weight_decay = group.options.get("weight_decay").copied().unwrap_or(0.0);

            for param in &group.params {
                let param_id = Self::get_param_id(param);
                let param_read = param.read();

                let grad = param_read.grad().ok_or_else(|| {
                    TorshError::invalid_argument_with_context(
                        "Parameter has no gradient",
                        "radam_step",
                    )
                })?;

                // Get or initialize state
                let param_state = self.state.entry(param_id.clone()).or_default();

                let exp_avg = if !param_state.contains_key("exp_avg") {
                    let exp_avg = Tensor::zeros_like(&param_read)?;
                    param_state.insert("exp_avg".to_string(), exp_avg.clone());
                    exp_avg
                } else {
                    param_state
                        .get("exp_avg")
                        .expect("exp_avg state should exist")
                        .clone()
                };

                let exp_avg_sq = if !param_state.contains_key("exp_avg_sq") {
                    let exp_avg_sq = Tensor::zeros_like(&param_read)?;
                    param_state.insert("exp_avg_sq".to_string(), exp_avg_sq.clone());
                    exp_avg_sq
                } else {
                    param_state
                        .get("exp_avg_sq")
                        .expect("exp_avg_sq state should exist")
                        .clone()
                };

                // Apply weight decay
                let mut grad_to_use = grad.clone();
                if weight_decay != 0.0 {
                    grad_to_use = grad_to_use.add(&param_read.mul_scalar(weight_decay)?)?;
                }

                // Update biased first moment estimate
                let new_exp_avg = exp_avg
                    .mul_scalar(self.beta1)?
                    .add(&grad_to_use.mul_scalar(1.0 - self.beta1)?)?;
                param_state.insert("exp_avg".to_string(), new_exp_avg.clone());

                // Update biased second raw moment estimate
                let grad_sq = grad_to_use.mul_op(&grad_to_use)?;
                let new_exp_avg_sq = exp_avg_sq
                    .mul_scalar(self.beta2)?
                    .add(&grad_sq.mul_scalar(1.0 - self.beta2)?)?;
                param_state.insert("exp_avg_sq".to_string(), new_exp_avg_sq.clone());

                // Bias correction
                let bias_correction1 = 1.0 - self.beta1.powi(self.step_count as i32);
                let bias_correction2 = 1.0 - self.beta2.powi(self.step_count as i32);

                // Apply variance rectification
                drop(param_read);
                let mut param_write = param.write();

                if let Some(r_t) = self.variance_rectification(self.step_count) {
                    // Adaptive update with variance rectification
                    let corrected_exp_avg = new_exp_avg.div_scalar(bias_correction1)?;
                    let corrected_exp_avg_sq = new_exp_avg_sq.div_scalar(bias_correction2)?;
                    let denom = corrected_exp_avg_sq.sqrt()?.add_scalar(self.eps)?;

                    let update = corrected_exp_avg.div(&denom)?.mul_scalar(r_t * lr)?;
                    *param_write = param_write.sub(&update)?;
                } else {
                    // Simple momentum update without variance rectification
                    let corrected_exp_avg = new_exp_avg.div_scalar(bias_correction1)?;
                    let update = corrected_exp_avg.mul_scalar(lr)?;
                    *param_write = param_write.sub(&update)?;
                }
            }
        }

        Ok(())
    }

    fn zero_grad(&mut self) {
        for group in &self.param_groups {
            for param in &group.params {
                param.write().zero_grad();
            }
        }
    }

    fn get_lr(&self) -> Vec<f32> {
        self.param_groups.iter().map(|g| g.lr).collect()
    }

    fn set_lr(&mut self, lr: f32) {
        for group in &mut self.param_groups {
            group.lr = lr;
        }
    }

    fn add_param_group(&mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>) {
        let lr = options.get("lr").copied().unwrap_or(1e-3);
        let group = ParamGroup::new(params, lr).with_options(options);
        self.param_groups.push(group);
    }

    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
        crate::optimizer::collect_parameters(&self.param_groups)
    }

    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
        let param_groups = self
            .param_groups
            .iter()
            .map(|g| ParamGroupState {
                lr: g.lr,
                options: g.options.clone(),
                param_count: g.params.len(),
            })
            .collect();

        Ok(OptimizerState {
            param_groups,
            state: self.state.clone(),
            optimizer_type: "RAdam".to_string(),
            version: "1.0".to_string(),
            global_state: HashMap::new(),
        })
    }

    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
        if state.param_groups.len() != self.param_groups.len() {
            return Err(OptimizerError::InvalidParameter(
                "Parameter group count mismatch".to_string(),
            ));
        }

        for (i, group_state) in state.param_groups.iter().enumerate() {
            self.param_groups[i].lr = group_state.lr;
            self.param_groups[i].options = group_state.options.clone();
        }

        self.state = state.state;
        Ok(())
    }
}

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

    #[test]
    fn test_radam_creation() -> OptimizerResult<()> {
        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
        let optimizer = RAdam::new(params, None, None, None, None, None);
        assert_eq!(optimizer.get_lr()[0], 1e-3);
        Ok(())
    }

    #[test]
    #[ignore = "Temporarily disabled due to potential deadlock"]
    fn test_radam_step() -> OptimizerResult<()> {
        let param = Arc::new(RwLock::new(randn::<f32>(&[2, 2])?));
        let mut param_write = param.write();
        param_write.set_grad(Some(randn::<f32>(&[2, 2])?));
        drop(param_write);

        let params = vec![param];
        let mut optimizer = RAdam::new(params, Some(0.1), None, None, None, None);

        optimizer.step()?;
        // step succeeded if we reached this point
        Ok(())
    }

    #[test]
    fn test_radam_basic() -> OptimizerResult<()> {
        // Simplified test that just checks the optimizer can be created and configured
        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
        let mut optimizer = RAdam::new(
            params,
            Some(0.01),
            Some(0.9),
            Some(0.999),
            Some(1e-8),
            Some(0.01),
        );

        assert_eq!(optimizer.get_lr()[0], 0.01);
        optimizer.set_lr(0.001);
        assert_eq!(optimizer.get_lr()[0], 0.001);

        // Test variance rectification calculation
        assert!(optimizer.variance_rectification(1).is_none());
        assert!(optimizer.variance_rectification(100).is_some());
        Ok(())
    }
}