//! Adaptive entropy coefficient for PPO (SAC-style alpha tuning)
//!
//! Auto-tunes exploration by learning an entropy coefficient that maintains
//! a target entropy level. Implemented as a simple gradient-free controller
//! since the Candle autograd backend has been removed.

use serde::{Deserialize, Serialize};

use ml_core::MLError;

/// Configuration for adaptive entropy coefficient tuning.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdaptiveEntropyConfig {
    /// Initial entropy coefficient. Default: 0.05
    pub initial_alpha: f64,
    /// Target entropy as fraction of max entropy. Default: 0.5
    pub target_ratio: f64,
    /// Learning rate for alpha adjustment. Default: 3e-4
    pub alpha_lr: f64,
    /// Number of discrete actions. Default: 45
    pub num_actions: usize,
}

impl Default for AdaptiveEntropyConfig {
    fn default() -> Self {
        Self {
            initial_alpha: 0.05,
            target_ratio: 0.5,
            alpha_lr: 3e-4,
            num_actions: 45,
        }
    }
}

/// Adaptive entropy coefficient using simple proportional control.
///
/// Replaces the Candle-based dual gradient descent implementation.
/// Uses a lightweight proportional controller that adjusts `log(alpha)`
/// based on the difference between current and target entropy.
pub struct AdaptiveEntropyCoeff {
    /// Current log(alpha) value
    log_alpha: f64,
    /// Target entropy (positive for discrete actions)
    target_entropy: f64,
    /// Learning rate for alpha adjustment
    alpha_lr: f64,
}

impl std::fmt::Debug for AdaptiveEntropyCoeff {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AdaptiveEntropyCoeff")
            .field("alpha", &self.log_alpha.exp())
            .field("target_entropy", &self.target_entropy)
            .finish()
    }
}

impl AdaptiveEntropyCoeff {
    /// Create a new adaptive entropy coefficient.
    pub fn new(config: &AdaptiveEntropyConfig) -> Result<Self, MLError> {
        if config.initial_alpha <= 0.0 {
            return Err(MLError::ConfigError(format!(
                "initial_alpha must be positive, got {}",
                config.initial_alpha
            )));
        }
        if config.num_actions == 0 {
            return Err(MLError::ConfigError("num_actions must be > 0".to_owned()));
        }

        let target_entropy = config.target_ratio * (config.num_actions as f64).ln();
        let log_alpha = config.initial_alpha.ln();

        Ok(Self {
            log_alpha,
            target_entropy,
            alpha_lr: config.alpha_lr,
        })
    }

    /// Return the current entropy coefficient `alpha = exp(log_alpha)`.
    pub fn alpha(&self) -> Result<f64, MLError> {
        Ok(self.log_alpha.exp())
    }

    /// Update the entropy coefficient given the mean log-probability.
    ///
    /// When entropy (`-mean_log_pi`) is below target, alpha increases.
    /// When entropy is above target, alpha decreases.
    ///
    /// Returns the new alpha value.
    pub fn update(&mut self, mean_log_pi: f64) -> Result<f64, MLError> {
        let entropy_estimate = -mean_log_pi;
        let entropy_error = entropy_estimate - self.target_entropy;

        // Proportional update: when entropy < target, error < 0, alpha should increase
        // When entropy > target, error > 0, alpha should decrease
        // gradient of alpha_loss w.r.t. log_alpha: alpha * (entropy - target)
        let alpha = self.log_alpha.exp();
        let grad = alpha * entropy_error;
        self.log_alpha -= self.alpha_lr * grad;

        // Clamp log_alpha to prevent extreme values
        self.log_alpha = self.log_alpha.clamp(-10.0, 5.0);

        self.alpha()
    }

    /// Return the target entropy value.
    pub const fn target_entropy(&self) -> f64 {
        self.target_entropy
    }
}

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

    fn default_config() -> AdaptiveEntropyConfig {
        AdaptiveEntropyConfig::default()
    }

    #[test]
    fn test_adaptive_entropy_initial_alpha() -> Result<(), MLError> {
        let config = default_config();
        let ae = AdaptiveEntropyCoeff::new(&config)?;

        let alpha = ae.alpha()?;
        assert!(
            (alpha - config.initial_alpha).abs() < 1e-5,
            "Expected alpha ~ {}, got {}",
            config.initial_alpha,
            alpha,
        );
        Ok(())
    }

    #[test]
    fn test_adaptive_entropy_target_entropy() -> Result<(), MLError> {
        let config = default_config();
        let ae = AdaptiveEntropyCoeff::new(&config)?;

        let expected = 0.5 * (45.0_f64).ln();
        let actual = ae.target_entropy();
        assert!(
            (actual - expected).abs() < 1e-6,
            "Expected target_entropy ~ {expected:.6}, got {actual:.6}",
        );
        Ok(())
    }

    #[test]
    fn test_adaptive_entropy_update_increases_alpha() -> Result<(), MLError> {
        let config = AdaptiveEntropyConfig {
            initial_alpha: 0.05,
            target_ratio: 0.5,
            alpha_lr: 0.01,
            num_actions: 45,
        };
        let mut ae = AdaptiveEntropyCoeff::new(&config)?;
        let initial_alpha = ae.alpha()?;

        // Low entropy (deterministic policy, mean_log_pi ~ 0)
        let mut last_alpha = initial_alpha;
        for _ in 0..20 {
            last_alpha = ae.update(-0.1)?;
        }

        assert!(
            last_alpha > initial_alpha,
            "Alpha should increase when entropy < target; initial={initial_alpha}, final={last_alpha}",
        );

        Ok(())
    }

    #[test]
    fn test_adaptive_entropy_update_decreases_alpha() -> Result<(), MLError> {
        let config = AdaptiveEntropyConfig {
            initial_alpha: 0.05,
            target_ratio: 0.5,
            alpha_lr: 0.01,
            num_actions: 45,
        };
        let mut ae = AdaptiveEntropyCoeff::new(&config)?;
        let initial_alpha = ae.alpha()?;

        // High entropy (near uniform, mean_log_pi ~ -ln(45) ~ -3.81)
        let mut last_alpha = initial_alpha;
        for _ in 0..20 {
            last_alpha = ae.update(-3.81)?;
        }

        assert!(
            last_alpha < initial_alpha,
            "Alpha should decrease when entropy > target; initial={initial_alpha}, final={last_alpha}",
        );

        Ok(())
    }

    #[test]
    fn test_adaptive_entropy_invalid_config() {
        let config = AdaptiveEntropyConfig {
            initial_alpha: 0.0,
            ..default_config()
        };
        let result = AdaptiveEntropyCoeff::new(&config);
        assert!(result.is_err(), "Should reject initial_alpha = 0");

        let config = AdaptiveEntropyConfig {
            num_actions: 0,
            ..default_config()
        };
        let result = AdaptiveEntropyCoeff::new(&config);
        assert!(result.is_err(), "Should reject num_actions = 0");
    }

    #[test]
    fn test_adaptive_entropy_custom_num_actions() -> Result<(), MLError> {
        let config = AdaptiveEntropyConfig {
            num_actions: 10,
            target_ratio: 0.5,
            ..default_config()
        };
        let ae = AdaptiveEntropyCoeff::new(&config)?;

        let expected_target = 0.5 * (10.0_f64).ln();
        assert!(
            (ae.target_entropy() - expected_target).abs() < 1e-10,
            "target_entropy for 10 actions: expected {expected_target}, got {}",
            ae.target_entropy(),
        );
        Ok(())
    }
}
