//! Sophia optimizer implementation
//!
//! Sophia: Second-order Clipped Stochastic Optimization (Stanford, 2023)
//!
//! Key features:
//! - Second-order information via Hessian diagonal estimation
//! - Adaptive learning with gradient clipping
//! - Parameters: lr, betas, rho (Hessian EMA), weight_decay
//! - Update rule uses Hessian for per-parameter learning rates
//! - Numerical stability with clipping
//!
//! # Mathematical Foundation
//!
//! Sophia approximates the diagonal Hessian and uses it for adaptive learning:
//! - m_t = β₁ * m_{t-1} + (1 - β₁) * g_t           (momentum)
//! - h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)     (Hessian diagonal approximation)  
//! - u_t = m_t / max(h_t, ε)                        (adaptive update)
//! - u_t = clip(u_t, γ)                             (gradient clipping)
//! - θ_t = θ_{t-1} - η * u_t                        (parameter update)
//!
//! # Performance Characteristics  
//! - Uses second-order curvature information for better conditioning
//! - Adaptive per-parameter learning rates
//! - Numerically stable with clipping mechanisms

use crate::{Result, TransformerError};
use crate::optimizers::{Optimizer, BaseOptimizer};
use rtx_tensor::{Tensor, TensorError, DType};
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use tracing::{debug, trace};

/// Sophia optimizer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SophiaConfig {
    /// Learning rate
    pub learning_rate: f64,
    /// Beta1 parameter (momentum coefficient)
    pub beta1: f64,
    /// Beta2 parameter (second moment coefficient)
    pub beta2: f64,
    /// Rho parameter (Hessian EMA coefficient)
    pub rho: f64,
    /// Weight decay coefficient
    pub weight_decay: f64,
    /// Clipping threshold for gradient clipping
    pub clip_threshold: f64,
    /// Epsilon for numerical stability
    pub eps: f64,
}

impl Default for SophiaConfig {
    fn default() -> Self {
        Self {
            learning_rate: 1e-3,
            beta1: 0.965,
            beta2: 0.99,
            rho: 0.04,
            weight_decay: 0.0,
            clip_threshold: 1.0,
            eps: 1e-8,
        }
    }
}

/// Sophia optimizer state for a single parameter
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SophiaState {
    /// Exponential moving average of gradient (momentum)
    pub momentum: Tensor,
    /// Exponential moving average of Hessian diagonal approximation
    pub hessian_diagonal: Tensor,
    /// Step count
    pub step: i64,
}

/// Sophia optimizer with second-order curvature information
///
/// Sophia uses an approximation of the diagonal Hessian to provide adaptive
/// per-parameter learning rates, leading to better convergence properties
/// especially for ill-conditioned optimization problems.
///
/// # Mathematical Details
///
/// The Sophia update rule consists of several phases:
/// 1. Update momentum: m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
/// 2. Update Hessian diagonal: h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
/// 3. Compute adaptive update: u_t = m_t / max(h_t, ε)
/// 4. Apply clipping: u_t = clip(u_t, γ)  
/// 5. Update parameter: θ_t = θ_{t-1} - η * u_t
///
/// # Adaptive Learning
/// The Hessian diagonal approximation provides per-parameter scaling,
/// allowing the optimizer to adapt to the local curvature of the loss surface.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SophiaOptimizer {
    /// Base optimizer functionality
    base: BaseOptimizer,
    /// Beta1 parameter (momentum coefficient)
    beta1: f64,
    /// Beta2 parameter (second moment coefficient)
    beta2: f64,
    /// Rho parameter (Hessian EMA coefficient)
    rho: f64,
    /// Weight decay (L2 regularization)
    weight_decay: f64,
    /// Clipping threshold for gradient clipping
    clip_threshold: f64,
    /// Epsilon for numerical stability
    eps: f64,
    /// Per-parameter state
    state: HashMap<String, SophiaState>,
}

impl SophiaOptimizer {
    /// Create a new Sophia optimizer from configuration
    pub fn new(config: SophiaConfig) -> Result<Self> {
        // Validate parameters
        Self::validate_parameters(
            config.learning_rate,
            config.beta1,
            config.beta2,
            config.rho,
            config.weight_decay,
            config.clip_threshold,
            config.eps,
        )?;

        debug!(
            "Creating Sophia optimizer: lr={}, β₁={}, β₂={}, ρ={}, weight_decay={}, clip={}, ε={}",
            config.learning_rate, config.beta1, config.beta2, config.rho,
            config.weight_decay, config.clip_threshold, config.eps
        );

        Ok(Self {
            base: BaseOptimizer::new(config.learning_rate),
            beta1: config.beta1,
            beta2: config.beta2,
            rho: config.rho,
            weight_decay: config.weight_decay,
            clip_threshold: config.clip_threshold,
            eps: config.eps,
            state: HashMap::new(),
        })
    }

    /// Validate optimizer parameters
    fn validate_parameters(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        rho: f64,
        weight_decay: f64,
        clip_threshold: f64,
        eps: f64,
    ) -> Result<()> {
        if learning_rate <= 0.0 {
            return Err(TransformerError::generic(
                format!("learning_rate {} must be positive", learning_rate)
            ));
        }

        if !(0.0..1.0).contains(&beta1) {
            return Err(TransformerError::generic(
                format!("beta1 {} must be in [0, 1)", beta1)
            ));
        }

        if !(0.0..1.0).contains(&beta2) {
            return Err(TransformerError::generic(
                format!("beta2 {} must be in [0, 1)", beta2)
            ));
        }

        if !(0.0..=1.0).contains(&rho) {
            return Err(TransformerError::generic(
                format!("rho {} must be in [0, 1]", rho)
            ));
        }

        if weight_decay < 0.0 {
            return Err(TransformerError::generic(
                format!("weight_decay {} must be non-negative", weight_decay)
            ));
        }

        if clip_threshold <= 0.0 {
            return Err(TransformerError::generic(
                format!("clip_threshold {} must be positive", clip_threshold)
            ));
        }

        if eps < 0.0 {
            return Err(TransformerError::generic(
                format!("eps {} must be non-negative", eps)
            ));
        }

        Ok(())
    }

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

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

    /// Get rho parameter
    pub fn rho(&self) -> f64 {
        self.rho
    }

    /// Get weight decay parameter
    pub fn weight_decay(&self) -> f64 {
        self.weight_decay
    }

    /// Get clip threshold parameter
    pub fn clip_threshold(&self) -> f64 {
        self.clip_threshold
    }

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

    /// Initialize state for a parameter if it doesn't exist
    fn ensure_state(&mut self, param_name: &str, param: &Tensor) -> Result<()> {
        if !self.state.contains_key(param_name) {
            let shape = param.shape();
            let device = param.device();

            let momentum = Tensor::zeros(shape.clone(), device)?;
            let hessian_diagonal = Tensor::zeros(shape.clone(), device)?;

            let state = SophiaState {
                momentum,
                hessian_diagonal,
                step: 0,
            };

            self.state.insert(param_name.to_string(), state);
            trace!("Initialized Sophia state for parameter: {}", param_name);
        }

        Ok(())
    }

    /// Apply gradient clipping element-wise
    fn apply_clipping(update: &Tensor, threshold: f64) -> Result<Tensor> {
        // Element-wise clipping: clip(x, threshold) = max(-threshold, min(threshold, x))
        let threshold_tensor = Tensor::scalar(threshold as f32, DType::F32, update.device())?;
        let neg_threshold_tensor = Tensor::scalar(-threshold as f32, DType::F32, update.device())?;
        
        // Clamp to [-threshold, threshold]
        let clipped = update.clamp(&neg_threshold_tensor, &threshold_tensor)?;
        Ok(clipped)
    }

    /// Perform Sophia parameter update (static version)
    fn update_parameter_static(
        param: &Tensor,
        grad: &Tensor,
        state: &mut SophiaState,
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        rho: f64,
        weight_decay: f64,
        clip_threshold: f64,
        eps: f64,
    ) -> Result<Tensor> {
        // Increment step count
        state.step += 1;
        
        trace!("Sophia step {} for parameter", state.step);

        // Step 1: Update momentum m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
        let momentum_term = (&state.momentum * beta1)?;
        let grad_term = (grad * (1.0 - beta1))?;
        state.momentum = (momentum_term + grad_term)?;

        // Step 2: Update Hessian diagonal approximation h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
        let hessian_term = (&state.hessian_diagonal * rho)?;
        let grad_squared = (grad * grad)?;
        let hessian_update = (&grad_squared * (1.0 - rho))?;
        state.hessian_diagonal = (hessian_term + hessian_update)?;

        // Step 3: Compute adaptive update u_t = m_t / max(h_t, ε)
        let eps_tensor = Tensor::scalar(eps as f32, DType::F32, state.hessian_diagonal.device())?;
        let hessian_with_eps = state.hessian_diagonal.maximum(&eps_tensor)?;
        let hessian_sqrt = hessian_with_eps.sqrt()?;
        let adaptive_update = state.momentum.div(&hessian_sqrt)?;

        // Step 4: Apply gradient clipping u_t = clip(u_t, γ)
        let clipped_update = Self::apply_clipping(&adaptive_update, clip_threshold)?;

        // Step 5: Apply weight decay and parameter update
        let scaled_update = (&clipped_update * learning_rate)?;
        let final_param = if weight_decay > 0.0 {
            // L2 regularization: θ_t = θ_{t-1} - α * λ * θ_{t-1} - update
            (param * (1.0 - learning_rate * weight_decay))? - &scaled_update
        } else {
            // Standard update: θ_t = θ_{t-1} - update
            param - &scaled_update
        }?;

        Ok(final_param)
    }
}

impl Optimizer for SophiaOptimizer {
    fn step_param(&mut self, param_name: &str, param: &Tensor, grad: &Tensor) -> Result<Tensor> {
        // Validate input shapes match
        if param.shape() != grad.shape() {
            return Err(TransformerError::shape_mismatch(
                format!("Parameter shape {:?} doesn't match gradient shape {:?}",
                       param.shape().dims().to_vec(),
                       grad.shape().dims().to_vec())
            ));
        }

        // Ensure state exists
        self.ensure_state(param_name, param)?;

        // Get the state and perform update
        let state = self.state.get_mut(param_name).unwrap();
        let learning_rate = self.base.learning_rate;
        let beta1 = self.beta1;
        let beta2 = self.beta2;
        let rho = self.rho;
        let weight_decay = self.weight_decay;
        let clip_threshold = self.clip_threshold;
        let eps = self.eps;

        // Perform update using static method
        Self::update_parameter_static(
            param, grad, state, learning_rate, beta1, beta2, 
            rho, weight_decay, clip_threshold, eps
        )
    }

    fn learning_rate(&self) -> f64 {
        self.base.learning_rate
    }

    fn set_learning_rate(&mut self, lr: f64) -> Result<()> {
        if lr <= 0.0 {
            return Err(TransformerError::generic(
                format!("learning_rate {} must be positive", lr)
            ));
        }

        self.base.set_learning_rate(lr);
        debug!("Updated Sophia learning rate to: {}", lr);
        Ok(())
    }

    fn has_state(&self, param_name: &str) -> bool {
        self.state.contains_key(param_name)
    }

    fn reset_state(&mut self, param_name: &str) -> Result<()> {
        if self.state.remove(param_name).is_some() {
            debug!("Reset Sophia state for parameter: {}", param_name);
        }
        Ok(())
    }

    fn reset_all_state(&mut self) {
        let count = self.state.len();
        self.state.clear();
        debug!("Reset all Sophia state ({} parameters)", count);
    }

    fn get_step_count(&self, param_name: &str) -> Result<i64> {
        self.state
            .get(param_name)
            .map(|state| state.step)
            .ok_or_else(|| {
                TransformerError::optimizer(format!("No state found for parameter: {}", param_name))
            })
    }

    fn optimizer_type(&self) -> &'static str {
        "Sophia"
    }

    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }

    fn store_gradients_internal(&mut self, gradients: HashMap<String, Tensor>) -> Result<()> {
        debug!("Storing {} gradients for Sophia optimizer", gradients.len());
        self.base.store_gradients(gradients);
        Ok(())
    }

    fn process_stored_gradients(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>> {
        debug!("Processing {} stored gradients with lr={} using Sophia algorithm", 
               self.base.stored_gradients.len(), learning_rate);

        // Update learning rate if different
        if (self.learning_rate() - learning_rate).abs() > f64::EPSILON {
            self.set_learning_rate(learning_rate)?;
        }

        let mut parameter_updates = HashMap::new();

        // Process each stored gradient with Sophia algorithm
        for (param_name, grad) in &self.base.stored_gradients {
            // Initialize state if it doesn't exist
            if !self.state.contains_key(param_name) {
                debug!("Creating Sophia state for parameter: {}", param_name);
                let shape = grad.shape();
                let device = grad.device();

                let momentum = Tensor::zeros(shape.clone(), device)?;
                let hessian_diagonal = Tensor::zeros(shape.clone(), device)?;

                let state = SophiaState {
                    momentum,
                    hessian_diagonal,
                    step: 0,
                };

                self.state.insert(param_name.clone(), state);
            }

            // Get mutable reference to state
            let state = self.state.get_mut(param_name).unwrap();

            // Increment step count
            state.step += 1;
            
            trace!("Sophia step {} for parameter {}", state.step, param_name);

            // Step 1: Update momentum m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
            let momentum_term = (&state.momentum * self.beta1)?;
            let grad_term = (grad * (1.0 - self.beta1))?;
            state.momentum = (momentum_term + grad_term)?;

            // Step 2: Update Hessian diagonal approximation h_t = ρ * h_{t-1} + (1 - ρ) * (g_t ⊙ g_t)
            let hessian_term = (&state.hessian_diagonal * self.rho)?;
            let grad_squared = (grad * grad)?;
            let hessian_update = (&grad_squared * (1.0 - self.rho))?;
            state.hessian_diagonal = (hessian_term + hessian_update)?;

            // Step 3: Compute adaptive update u_t = m_t / max(h_t, ε)
            let eps_tensor = Tensor::scalar(self.eps as f32, DType::F32, state.hessian_diagonal.device())?;
            let hessian_with_eps = state.hessian_diagonal.maximum(&eps_tensor)?;
            let hessian_sqrt = hessian_with_eps.sqrt()?;
            let adaptive_update = state.momentum.div(&hessian_sqrt)?;

            // Step 4: Apply gradient clipping u_t = clip(u_t, γ)
            let clipped_update = Self::apply_clipping(&adaptive_update, self.clip_threshold)?;

            // Step 5: Scale by learning rate and store negative update
            let scaled_update = (&clipped_update * learning_rate)?;
            let final_update = if self.weight_decay > 0.0 {
                // Note: Weight decay requires parameter values, which we don't have here
                debug!("Weight decay configured but parameter values not available in process_stored_gradients");
                (scaled_update * -1.0)? // Negative because we want to subtract the update
            } else {
                (scaled_update * -1.0)? // Negative because we want to subtract the update
            };

            parameter_updates.insert(param_name.clone(), final_update);
            debug!("Created Sophia update for parameter: {} (step {})", param_name, state.step);
        }

        // Clear stored gradients after processing
        self.base.clear_gradients();

        debug!("Generated {} parameter updates using Sophia algorithm", parameter_updates.len());
        Ok(parameter_updates)
    }
}

#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
    use super::*;
    use rtx_tensor::{Device, Shape};

    #[test]
    fn test_sophia_parameter_validation() {
        let config = SophiaConfig::default();

        // Test valid parameters
        assert!(SophiaOptimizer::new(config).is_ok());

        // Test invalid learning rate
        let mut invalid_config = SophiaConfig::default();
        invalid_config.learning_rate = -0.1;
        assert!(SophiaOptimizer::new(invalid_config).is_err());
    }

    #[test] 
    fn test_sophia_basic_functionality() -> Result<()> {
        let config = SophiaConfig::default();
        let mut optimizer = SophiaOptimizer::new(config)?;
        let device = Device::cpu();

        let param = Tensor::from_data(vec![1.0, 2.0], &[2], &device)?;
        let grad = Tensor::from_data(vec![0.1, 0.2], &[2], &device)?;

        // First step
        let updated_param = optimizer.step_param("test_param", &param, &grad)?;
        
        assert!(optimizer.has_state("test_param"));
        assert_eq!(optimizer.get_step_count("test_param")?, 1);

        // Verify parameter changed
        let param_data = param.to_cpu()?;
        let updated_data = updated_param.to_cpu()?;
        assert_ne!(param_data, updated_data);

        Ok(())
    }
}

// Include tests if compiled with test features
#[cfg(all(test, feature = "disabled_tests"))]
#[path = "sophia_tests.rs"]
mod sophia_tests;