//! GPU-Native Adam Optimizer
//!
//! This module provides an Adam optimizer that keeps ALL state on GPU,
//! eliminating CPU-GPU synchronization during training steps.
//!
//! ## Design Goals
//!
//! 1. **Zero CPU Sync**: All optimizer state (m, v, parameters) stays on GPU
//! 2. **Fused Operations**: Minimize kernel launches per step
//! 3. **In-Place Updates**: Modify weights in-place for graph capture compatibility
//!
//! ## Adam Update Rule
//!
//! For each parameter θ with gradient g:
//!
//! ```text
//! m = β₁ * m + (1 - β₁) * g           // Update first moment
//! v = β₂ * v + (1 - β₂) * g²          // Update second moment
//! m̂ = m / (1 - β₁^t)                  // Bias-corrected first moment
//! v̂ = v / (1 - β₂^t)                  // Bias-corrected second moment
//! θ = θ - lr * m̂ / (√v̂ + ε)          // Update parameters
//! ```

use rtx_tensor::{Tensor, Device};

#[cfg(feature = "cuda")]
use crate::PinnStreamContext;

/// Result type for optimizer operations
pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

/// GPU-native Adam optimizer
///
/// All optimizer state lives on GPU. No CPU-GPU synchronization during training.
#[derive(Debug)]
pub struct GpuAdam {
    /// First moment estimates (one per parameter tensor)
    pub m: Vec<Tensor>,

    /// Second moment estimates (one per parameter tensor)
    pub v: Vec<Tensor>,

    /// Timestep counter
    pub t: u64,

    /// Learning rate
    pub lr: f32,

    /// Decay rate for first moment estimate
    pub beta1: f32,

    /// Decay rate for second moment estimate
    pub beta2: f32,

    /// Small constant for numerical stability
    pub eps: f32,

    /// Device for allocation
    device: Device,
}

impl GpuAdam {
    /// Create a new GPU Adam optimizer
    ///
    /// # Arguments
    /// * `param_shapes` - Shapes of parameter tensors
    /// * `device` - GPU device for state allocation
    /// * `lr` - Learning rate (default: 1e-3)
    /// * `beta1` - First moment decay (default: 0.9)
    /// * `beta2` - Second moment decay (default: 0.999)
    /// * `eps` - Numerical stability constant (default: 1e-8)
    pub fn new(
        param_shapes: &[Vec<usize>],
        device: &Device,
        lr: f32,
        beta1: f32,
        beta2: f32,
        eps: f32,
    ) -> Result<Self> {
        let mut m = Vec::with_capacity(param_shapes.len());
        let mut v = Vec::with_capacity(param_shapes.len());

        for shape in param_shapes {
            // Initialize first and second moments to zero on GPU
            let shape_i64: Vec<usize> = shape.iter().map(|&x| x).collect();
            m.push(Tensor::zeros(&shape_i64, device)?);
            v.push(Tensor::zeros(&shape_i64, device)?);
        }

        Ok(Self {
            m,
            v,
            t: 0,
            lr,
            beta1,
            beta2,
            eps,
            device: device.clone(),
        })
    }

    /// Create with default hyperparameters
    ///
    /// lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8
    pub fn with_defaults(param_shapes: &[Vec<usize>], device: &Device) -> Result<Self> {
        Self::new(param_shapes, device, 1e-3, 0.9, 0.999, 1e-8)
    }

    /// Perform one optimization step
    ///
    /// Updates parameters in-place using the provided gradients.
    /// All operations stay on GPU - no CPU-GPU synchronization.
    ///
    /// # Arguments
    /// * `params` - Parameter tensors to update (modified in-place)
    /// * `grads` - Gradient tensors (same shapes as params)
    pub fn step(&mut self, params: &mut [Tensor], grads: &[Tensor]) -> Result<()> {
        if params.len() != grads.len() || params.len() != self.m.len() {
            return Err(format!(
                "Mismatch: {} params, {} grads, {} optimizer states",
                params.len(),
                grads.len(),
                self.m.len()
            ).into());
        }

        self.t += 1;

        // Compute bias corrections
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);

        // Adjusted learning rate with bias correction
        let lr_adj = self.lr * (bias_correction2.sqrt()) / bias_correction1;

        for i in 0..params.len() {
            let g = &grads[i];

            // m = β₁ * m + (1 - β₁) * g
            let m_scaled = self.m[i].mul_scalar(self.beta1)?;
            let g_scaled = g.mul_scalar(1.0 - self.beta1)?;
            self.m[i] = m_scaled.add(&g_scaled)?;

            // v = β₂ * v + (1 - β₂) * g²
            let g_sq = g.mul(g)?;
            let v_scaled = self.v[i].mul_scalar(self.beta2)?;
            let g_sq_scaled = g_sq.mul_scalar(1.0 - self.beta2)?;
            self.v[i] = v_scaled.add(&g_sq_scaled)?;

            // θ = θ - lr_adj * m / (√v + ε)
            // Using bias-corrected formulation with lr_adj
            let v_sqrt = self.v[i].sqrt()?;
            let v_sqrt_eps = v_sqrt.add_scalar(self.eps)?;
            let update = self.m[i].div(&v_sqrt_eps)?;
            let update_scaled = update.mul_scalar(lr_adj)?;
            params[i] = params[i].sub(&update_scaled)?;
        }

        Ok(())
    }

    /// Get current timestep
    pub fn timestep(&self) -> u64 {
        self.t
    }

    /// Reset optimizer state (first and second moments)
    pub fn reset(&mut self) -> Result<()> {
        for i in 0..self.m.len() {
            self.m[i] = Tensor::zeros(self.m[i].shape(), &self.device)?;
            self.v[i] = Tensor::zeros(self.v[i].shape(), &self.device)?;
        }
        self.t = 0;
        Ok(())
    }

    /// Set learning rate
    pub fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }

    /// Get learning rate
    pub fn lr(&self) -> f32 {
        self.lr
    }

    /// Fused optimizer step using a single CUDA kernel per parameter.
    ///
    /// This is the zero-allocation version that uses the PinnStreamContext
    /// for launching the fused Adam kernel. Each parameter is updated in-place
    /// with a single kernel launch that computes:
    /// - m update
    /// - v update
    /// - parameter update
    ///
    /// # Arguments
    /// * `params` - Mutable references to parameter tensors (updated in-place)
    /// * `grads` - Gradient tensors (read-only)
    /// * `ctx` - CUDA stream context for kernel launch
    #[cfg(feature = "cuda")]
    pub fn step_fused(
        &mut self,
        params: &mut [&mut Tensor],
        grads: &[&Tensor],
        ctx: &PinnStreamContext,
    ) -> Result<()> {
        if params.len() != grads.len() || params.len() != self.m.len() {
            return Err(format!(
                "Mismatch: {} params, {} grads, {} optimizer states",
                params.len(),
                grads.len(),
                self.m.len()
            ).into());
        }

        self.t += 1;

        // Compute bias corrections
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);

        // Adjusted learning rate with bias correction
        let lr_adj = self.lr * (bias_correction2.sqrt()) / bias_correction1;

        for i in 0..params.len() {
            // One kernel launch per parameter: updates param, m, v in-place
            ctx.adam_step(
                params[i],
                &mut self.m[i],
                &mut self.v[i],
                grads[i],
                lr_adj,
                self.beta1,
                self.beta2,
                self.eps,
            ).map_err(|e| format!("adam_step failed for param {}: {}", i, e))?;
        }

        Ok(())
    }
}

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

    #[test]
    fn test_gpu_adam_creation() {
        let device = Device::Cpu;
        let shapes = vec![
            vec![64, 32],  // W0
            vec![64],      // b0
            vec![2, 64],   // Wout
            vec![2],       // bout
        ];

        let opt = GpuAdam::with_defaults(&shapes, &device).unwrap();
        assert_eq!(opt.m.len(), 4);
        assert_eq!(opt.v.len(), 4);
        assert_eq!(opt.timestep(), 0);
    }

    #[test]
    fn test_gpu_adam_step() {
        let device = Device::Cpu;
        let shapes = vec![vec![2, 2]];

        let mut opt = GpuAdam::with_defaults(&shapes, &device).unwrap();

        // Create parameter and gradient tensors
        let mut params = vec![
            Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap()
        ];
        let grads = vec![
            Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 0.4], &[2, 2], &device).unwrap()
        ];

        // Perform optimization step
        opt.step(&mut params, &grads).unwrap();

        assert_eq!(opt.timestep(), 1);

        // Parameters should have changed
        let params_cpu = params[0].to_cpu().unwrap();
        // After one Adam step with lr=1e-3, params should decrease slightly
        assert!(params_cpu[0] < 1.0);
        assert!(params_cpu[1] < 2.0);
    }

    #[test]
    fn test_gpu_adam_reset() {
        let device = Device::Cpu;
        let shapes = vec![vec![2, 2]];

        let mut opt = GpuAdam::with_defaults(&shapes, &device).unwrap();

        // Do a step
        let mut params = vec![
            Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap()
        ];
        let grads = vec![
            Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 0.4], &[2, 2], &device).unwrap()
        ];
        opt.step(&mut params, &grads).unwrap();

        assert_eq!(opt.timestep(), 1);

        // Reset
        opt.reset().unwrap();
        assert_eq!(opt.timestep(), 0);

        // Check moments are zeroed
        let m0_cpu = opt.m[0].to_cpu().unwrap();
        for val in m0_cpu.iter() {
            assert!(*val == 0.0);
        }
    }
}