//! Optimizers for transformer training
//!
//! This module provides high-performance optimizers including Adam and `AdamW`
//! with support for distributed training and revolutionary enhancements.

use crate::Result;
use rtx_tensor::Tensor;
use std::collections::HashMap;
use tracing::debug;

// ============================================================================
// PHASE 1: CORE OPTIMIZERS (Essential functionality)
// ============================================================================

pub mod adam;
pub mod adamw;

#[cfg(all(test, feature = "disabled_tests"))]
pub mod adam_test;

// Re-export main optimizer types
pub use adam::AdamState;

// ============================================================================
// PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
// ============================================================================
// These will be re-enabled once core optimizers are stable

// pub mod ademamix;          // TODO: Re-enable in Phase 2
// pub mod adabound;          // TODO: Re-enable in Phase 2
// pub mod lion;              // TODO: Re-enable in Phase 2
// pub mod sophia;            // TODO: Re-enable in Phase 2
// pub mod shampoo;           // TODO: Re-enable in Phase 2
// pub mod kfac;              // TODO: Re-enable in Phase 2
// pub mod kfac_advanced;     // TODO: Re-enable in Phase 2
// pub mod kfac_example;      // TODO: Re-enable in Phase 2
// pub mod ranger;            // TODO: Re-enable in Phase 2
// pub mod novograd;          // TODO: Re-enable in Phase 2
// pub mod natural_gradient;  // TODO: Re-enable in Phase 2
// pub mod trust_region;      // TODO: Re-enable in Phase 2
// pub mod lbfgs;             // TODO: Re-enable in Phase 2

pub mod matrix_utils;

// Re-export tensor bridge traits for all optimizer modules
pub use crate::tensor_bridge::{TensorBridge, TensorBridgeStatic, TensorCompat};

// Test modules temporarily disabled due to missing dependencies
// #[cfg(test)]
// pub mod parameter_update_tests;

// All test modules temporarily disabled
// Will be re-enabled after implementing missing optimizer types

// ============================================================================
// PHASE 1: CORE OPTIMIZER EXPORTS (Only enabled optimizers)
// ============================================================================

pub use adam::AdamOptimizer;
pub use adamw::AdamWOptimizer;

// ============================================================================
// PHASE 2: ADVANCED OPTIMIZER EXPORTS (Temporarily disabled)
// ============================================================================
// These will be re-enabled when the modules are uncommented

// pub use ademamix::{AdEMAMixOptimizer, AdEMAMixConfig};
// pub use adabound::{AdaBoundOptimizer, AdaBoundConfig};
// pub use lion::{LionOptimizer, LionConfig};
// pub use sophia::{SophiaOptimizer, SophiaConfig};
// pub use shampoo::{ShampooOptimizer, ShampooConfig};
// pub use kfac::{KFacOptimizer, KFacConfig, LayerType};
// pub use kfac_advanced::{AdvancedKFacOptimizer, AdvancedKFacConfig, ConvergenceMetrics};
// pub use ranger::{RangerOptimizer, RangerConfig};
// pub use novograd::{NovoGradOptimizer, NovoGradConfig};
// pub use natural_gradient::{NaturalGradientOptimizer, NaturalGradientConfig, FisherApproximation};
// pub use trust_region::{TrustRegionOptimizer, TrustRegionConfig, SubproblemSolver};
// pub use lbfgs::{LBFGSOptimizer, LBFGSConfig};

/// Base optimizer implementation with gradient storage
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct BaseOptimizer {
    /// Stored gradients for batch processing
    #[serde(skip)]
    pub stored_gradients: HashMap<String, Tensor>,
    /// Current learning rate
    pub learning_rate: f64,
}

impl BaseOptimizer {
    /// Create a new base optimizer
    #[must_use]
    pub fn new(learning_rate: f64) -> Self {
        Self {
            stored_gradients: HashMap::new(),
            learning_rate,
        }
    }

    /// Get stored gradients
    #[must_use]
    pub fn stored_gradients(&self) -> &HashMap<String, Tensor> {
        &self.stored_gradients
    }

    /// Clear stored gradients
    pub fn clear_gradients(&mut self) {
        self.stored_gradients.clear();
    }

    /// Set learning rate
    pub fn set_learning_rate(&mut self, lr: f64) {
        self.learning_rate = lr;
    }

    /// Store gradients
    pub fn store_gradients(&mut self, gradients: HashMap<String, Tensor>) {
        self.stored_gradients.clear();
        self.stored_gradients.extend(gradients);
    }
}

/// Configuration for different optimizer types (Phase 1: Core optimizers only)
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum OptimizerConfig {
    /// Adam optimizer configuration
    Adam(AdamConfig),
    /// `AdamW` optimizer configuration
    AdamW(AdamWConfig),
    // ========================================================================
    // PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
    // ========================================================================
    // These will be re-enabled once the corresponding modules are uncommented

    // /// AdEMAMix optimizer configuration
    // AdEMAMix(AdEMAMixConfig),
    // /// AdaBound optimizer configuration
    // AdaBound(AdaBoundConfig),
    // /// Lion optimizer configuration
    // Lion(LionConfig),
    // /// Sophia optimizer configuration
    // Sophia(SophiaConfig),
    // /// Shampoo optimizer configuration
    // Shampoo(ShampooConfig),
    // /// K-FAC optimizer configuration
    // KFac(KFacConfig),
    // /// Advanced K-FAC optimizer configuration
    // AdvancedKFac(AdvancedKFacConfig),
    // /// Ranger optimizer configuration (RAdam + Lookahead)
    // Ranger(RangerConfig),
    // /// NovoGrad optimizer configuration (layer-wise gradient normalization)
    // NovoGrad(NovoGradConfig),
    // /// Natural Gradient optimizer configuration (Fisher Information Matrix preconditioning)
    // NaturalGradient(NaturalGradientConfig),
    // /// Trust Region optimizer configuration (second-order with adaptive radius)
    // TrustRegion(TrustRegionConfig),
    // /// L-BFGS optimizer configuration (quasi-Newton with limited memory)
    // LBFGS(LBFGSConfig),
}

/// Adam optimizer configuration
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct AdamConfig {
    /// Learning rate
    pub learning_rate: f64,
    /// Beta1 parameter (momentum)
    pub beta1: f64,
    /// Beta2 parameter (`RMSprop`)
    pub beta2: f64,
    /// Epsilon for numerical stability
    pub epsilon: f64,
    /// Weight decay coefficient
    pub weight_decay: f64,
    /// Whether to use `AMSGrad` variant
    pub amsgrad: bool,
}

impl Default for AdamConfig {
    fn default() -> Self {
        Self {
            learning_rate: 1e-3,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 0.0,
            amsgrad: false,
        }
    }
}

/// `AdamW` optimizer configuration
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct AdamWConfig {
    /// Learning rate
    pub learning_rate: f64,
    /// Beta1 parameter (momentum)
    pub beta1: f64,
    /// Beta2 parameter (`RMSprop`)
    pub beta2: f64,
    /// Epsilon for numerical stability
    pub epsilon: f64,
    /// Weight decay coefficient
    pub weight_decay: f64,
    /// Whether to use bias correction
    pub amsgrad: bool,
}

impl Default for AdamWConfig {
    fn default() -> Self {
        Self {
            learning_rate: 1e-3,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-2,
            amsgrad: false,
        }
    }
}

/// Create an optimizer from configuration (Phase 1: Core optimizers only)
pub fn create_optimizer(
    config: OptimizerConfig,
    _parameters: HashMap<String, Tensor>,
) -> Result<Box<dyn Optimizer>> {
    match config {
        OptimizerConfig::Adam(adam_config) => {
            debug!("Creating Adam optimizer with config: {:?}", adam_config);
            Ok(Box::new(AdamOptimizer::new(adam_config)?))
        }
        OptimizerConfig::AdamW(adamw_config) => {
            debug!("Creating AdamW optimizer with config: {:?}", adamw_config);
            Ok(Box::new(AdamWOptimizer::new(
                adamw_config.learning_rate,
                adamw_config.beta1,
                adamw_config.beta2,
                adamw_config.epsilon,
                adamw_config.weight_decay,
                adamw_config.amsgrad, // Using amsgrad field as decoupled flag for now
            )?))
        } // ====================================================================
          // PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
          // ====================================================================
          // These cases will be re-enabled when the optimizer modules are uncommented

          // OptimizerConfig::AdEMAMix(ademamix_config) => {
          //     debug!("Creating AdEMAMix optimizer with config: {:?}", ademamix_config);
          //     Ok(Box::new(AdEMAMixOptimizer::new(ademamix_config)?))
          // }

          // All other optimizer cases are temporarily disabled since the modules are commented out
    }
}

/// Trait for all optimizers in the transformer training infrastructure
pub trait Optimizer: Send + Sync {
    /// Perform a single optimization step for a parameter
    fn step_param(&mut self, param_name: &str, param: &Tensor, grad: &Tensor) -> Result<Tensor>;

    /// Get the current learning rate
    fn learning_rate(&self) -> f64;

    /// Set the learning rate
    fn set_learning_rate(&mut self, lr: f64) -> Result<()>;

    /// Check if optimizer has state for a parameter
    fn has_state(&self, param_name: &str) -> bool;

    /// Reset state for a specific parameter
    fn reset_state(&mut self, param_name: &str) -> Result<()>;

    /// Reset all optimizer state
    fn reset_all_state(&mut self);

    /// Get step count for a parameter
    fn get_step_count(&self, param_name: &str) -> Result<i64>;

    /// Get optimizer type name
    fn optimizer_type(&self) -> &'static str;

    /// Get mutable reference as Any for downcasting
    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;

    /// Zero all gradients (clear stored gradients)
    fn zero_grad(&mut self) {
        // Default implementation clears stored gradients
        if let Ok(()) = self.store_gradients_internal(HashMap::new()) {
            // Gradients cleared successfully
        }
    }

    /// Store gradients for batch processing (used by trainer)
    fn set_gradients(&mut self, gradients: HashMap<String, Tensor>) -> Result<()> {
        // Default implementation stores gradients in internal map
        self.store_gradients_internal(gradients)
    }

    /// Perform optimization step on all stored gradients with given learning rate
    fn step(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>> {
        // Default implementation processes stored gradients
        self.process_stored_gradients(learning_rate)
    }

    /// Internal method to store gradients (default implementation)
    fn store_gradients_internal(&mut self, gradients: HashMap<String, Tensor>) -> Result<()>;

    /// Internal method to process stored gradients (default implementation)
    fn process_stored_gradients(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>>;
}

/// Optimizer type enumeration
#[derive(Debug, Clone)]
pub enum OptimizerType {
    /// Adam optimizer
    Adam,
    /// `AdamW` optimizer with decoupled weight decay
    AdamW,
    /// `AdEMAMix` optimizer with adaptive EMA mixture
    AdEMAMix,
    /// `AdaBound` optimizer with dynamic learning rate bounds
    AdaBound,
    /// Lion optimizer with evolved sign momentum
    Lion,
    /// Sophia optimizer with second-order clipped stochastic optimization
    Sophia,
    /// Shampoo optimizer with full matrix preconditioning
    Shampoo,
    /// K-FAC optimizer with Kronecker-factored approximation
    KFac,
    /// Ranger optimizer combining `RAdam` with Lookahead
    Ranger,
    /// `NovoGrad` optimizer with layer-wise gradient normalization
    NovoGrad,
    /// Natural Gradient optimizer with Fisher Information Matrix preconditioning
    NaturalGradient,
    /// Trust Region optimizer with adaptive radius and second-order approximation
    TrustRegion,
    /// L-BFGS optimizer with limited-memory quasi-Newton approximation
    LBFGS,
}