//! Distributional Reinforcement Learning Implementation
//!
//! Supports two distributional RL algorithms:
//! 1. **C51 (Categorical DQN)**: Fixed categorical bins (Bellemare et al., 2017)
//! 2. **QR-DQN (Quantile Regression DQN)**: Adaptive quantiles (Dabney et al., 2018)
//!
//! Instead of learning scalar Q-values, we learn the full return distribution.
//!
//! **For Trading, Use QR-DQN:**
//! - Better tail risk modeling (CVaR/VaR)
//! - No need to tune `v_min/v_max`
//! - More stable for asymmetric returns

use std::sync::Arc;

use cudarc::driver::CudaStream;
use serde::{Deserialize, Serialize};

use ml_core::cuda_autograd::GpuTensor;
use ml_core::MLError;

/// Distributional type enum (C51 vs QR-DQN)
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum DistributionalType {
    /// Categorical DQN (C51) - fixed bins
    C51,
    /// Quantile Regression DQN - adaptive quantiles (RECOMMENDED for trading)
    QRDQN,
}

impl Default for DistributionalType {
    fn default() -> Self {
        // Default to QR-DQN for better trading risk modeling
        Self::QRDQN
    }
}

/// Configuration for distributional RL
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributionalConfig {
    pub num_atoms: usize,
    pub v_min: f64,
    pub v_max: f64,
}

impl Default for DistributionalConfig {
    fn default() -> Self {
        Self {
            num_atoms: 51,
            v_min: -25.0,  // DSR Q-values: rewards +/-2 with gamma=0.92 -> Q approx +/-25
            v_max: 25.0,   // DSR Q-values: rewards +/-2 with gamma=0.92 -> Q approx +/-25
        }
    }
}

/// Categorical distribution for value function approximation
#[derive(Debug)]
pub struct CategoricalDistribution {
    config: DistributionalConfig,
    /// Support values on GPU, shape `[num_atoms]`.
    support: GpuTensor,
    /// Support values cached on host for CPU-side computations.
    support_host: Vec<f32>,
    delta_z: f32,
    stream: Arc<CudaStream>,
}

impl CategoricalDistribution {
    pub fn new(config: &DistributionalConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let delta_z = ((config.v_max - config.v_min) / (config.num_atoms - 1) as f64) as f32;

        let support_host: Vec<f32> = (0..config.num_atoms)
            .map(|i| config.v_min as f32 + i as f32 * delta_z)
            .collect();

        let support = GpuTensor::from_host(
            &support_host,
            vec![config.num_atoms],
            stream,
        )?;

        Ok(Self {
            config: config.clone(),
            support,
            support_host,
            delta_z,
            stream: stream.clone(),
        })
    }

    /// Reinitialize distribution with new `v_min/v_max` bounds
    pub fn reinit(&mut self, v_min: f64, v_max: f64) -> Result<(), MLError> {
        self.config.v_min = v_min;
        self.config.v_max = v_max;

        self.delta_z = ((v_max - v_min) / (self.config.num_atoms - 1) as f64) as f32;

        self.support_host = (0..self.config.num_atoms)
            .map(|i| v_min as f32 + i as f32 * self.delta_z)
            .collect();

        self.support = GpuTensor::from_host(
            &self.support_host,
            vec![self.config.num_atoms],
            &self.stream,
        )?;

        Ok(())
    }

    /// Convert distribution to expected value (scalar Q-value) (cold path, CPU).
    ///
    /// **Hot-path C51 distributional forward is fused into `dqn_forward_only_kernel`
    /// and `dqn_forward_loss_kernel` via `warp_expected_q()` in CUDA. This CPU-based
    /// method exists for unit tests and the backward pass.**
    ///
    /// # Arguments
    /// * `distribution_host` - Probabilities `[batch * num_actions * num_atoms]` (host, row-major)
    /// * `batch` - Batch size
    /// * `num_actions` - Number of actions
    ///
    /// # Returns
    /// Expected Q-values `[batch * num_actions]` (host)
    #[cold]
    pub fn to_scalar_host(
        &self,
        distribution_host: &[f32],
        batch: usize,
        num_actions: usize,
    ) -> Result<Vec<f32>, MLError> {
        let na = self.config.num_atoms;
        let expected_len = batch * num_actions * na;
        if distribution_host.len() != expected_len {
            return Err(MLError::DimensionMismatch {
                expected: expected_len,
                actual: distribution_host.len(),
            });
        }

        let mut result = Vec::with_capacity(batch * num_actions);
        for b in 0..batch {
            for a in 0..num_actions {
                let base = (b * num_actions + a) * na;
                let mut expected = 0.0_f32;
                for i in 0..na {
                    let prob = distribution_host.get(base + i).copied().unwrap_or(0.0);
                    let support_val = self.support_host.get(i).copied().unwrap_or(0.0);
                    expected += prob * support_val;
                }
                result.push(expected);
            }
        }
        Ok(result)
    }

    /// Project target distribution onto current support (CPU, cold path).
    ///
    /// Implements the distributional Bellman operator:
    /// `T_z` = r + gamma * z for each support atom z
    /// Projects this onto the fixed support using linear interpolation.
    pub fn project_distribution_host(
        &self,
        target_support_host: &[f32],
        probabilities_host: &[f32],
        batch_size: usize,
    ) -> Result<Vec<f32>, MLError> {
        let num_atoms = self.config.num_atoms;
        let expected_len = batch_size * num_atoms;
        if target_support_host.len() != expected_len || probabilities_host.len() != expected_len {
            return Err(MLError::DimensionMismatch {
                expected: expected_len,
                actual: target_support_host.len().min(probabilities_host.len()),
            });
        }

        let v_min = self.config.v_min as f32;
        let v_max = self.config.v_max as f32;

        let mut projected = vec![0.0_f32; batch_size * num_atoms];

        for b in 0..batch_size {
            let base = b * num_atoms;
            for j in 0..num_atoms {
                let tz = target_support_host.get(base + j).copied().unwrap_or(0.0);
                let tz_clamped = tz.clamp(v_min, v_max);
                let atom_idx = (tz_clamped - v_min) / self.delta_z;

                let lower = atom_idx.floor() as usize;
                let upper = atom_idx.ceil() as usize;
                let lower = lower.min(num_atoms - 1);
                let upper = upper.min(num_atoms - 1);

                let frac = atom_idx - atom_idx.floor();
                let prob = probabilities_host.get(base + j).copied().unwrap_or(0.0);

                if let Some(p) = projected.get_mut(b * num_atoms + lower) {
                    *p += prob * (1.0 - frac);
                }
                if let Some(p) = projected.get_mut(b * num_atoms + upper) {
                    *p += prob * frac;
                }
            }
        }

        Ok(projected)
    }

    pub fn support_host(&self) -> &[f32] {
        &self.support_host
    }

    pub const fn support(&self) -> &GpuTensor {
        &self.support
    }

    pub const fn num_atoms(&self) -> usize {
        self.config.num_atoms
    }

    /// Compute categorical cross-entropy loss (CPU, cold path).
    ///
    /// `Loss = -sum_i target_i * log(pred_i)`
    pub fn categorical_loss_host(
        &self,
        predicted_probs: &[f32],
        target_probs: &[f32],
    ) -> Result<f32, MLError> {
        if predicted_probs.len() != target_probs.len() {
            return Err(MLError::DimensionMismatch {
                expected: target_probs.len(),
                actual: predicted_probs.len(),
            });
        }
        if predicted_probs.is_empty() {
            return Err(MLError::InvalidInput("empty input".into()));
        }

        let mut loss = 0.0_f32;
        for (p, t) in predicted_probs.iter().zip(target_probs.iter()) {
            loss -= t * (p + 1e-8).ln();
        }
        Ok(loss / predicted_probs.len() as f32 * self.config.num_atoms as f32)
    }

    /// Apply distributional Bellman operator (CPU, cold path).
    pub fn apply_bellman_operator_host(
        &self,
        rewards: &[f32],
        next_probs: &[f32],
        dones: &[f32],
        gamma: f32,
    ) -> Result<Vec<f32>, MLError> {
        let batch_size = rewards.len();
        let num_atoms = self.config.num_atoms;

        if next_probs.len() != batch_size * num_atoms {
            return Err(MLError::DimensionMismatch {
                expected: batch_size * num_atoms,
                actual: next_probs.len(),
            });
        }

        // Compute target support: T_z = r + gamma * z * (1 - done)
        let mut target_support = Vec::with_capacity(batch_size * num_atoms);
        for b in 0..batch_size {
            let r = rewards.get(b).copied().unwrap_or(0.0);
            let d = dones.get(b).copied().unwrap_or(0.0);
            for j in 0..num_atoms {
                let z = self.support_host.get(j).copied().unwrap_or(0.0);
                target_support.push(r + gamma * z * (1.0 - d));
            }
        }

        self.project_distribution_host(&target_support, next_probs, batch_size)
    }

    /// Get `V_min` value
    pub const fn v_min(&self) -> f64 {
        self.config.v_min
    }

    /// Get `V_max` value
    pub const fn v_max(&self) -> f64 {
        self.config.v_max
    }

    /// Get `delta_z` (atom spacing)
    pub const fn delta_z(&self) -> f32 {
        self.delta_z
    }
}

#[cfg(test)]
#[allow(
    clippy::doc_markdown,
    clippy::needless_borrow
)]
mod tests {
    use super::*;

    fn make_stream() -> Arc<CudaStream> {
        let device = ml_core::device::MlDevice::cuda(0).expect("CUDA required");
        device.cuda_stream().expect("stream").clone()
    }

    #[test]
    fn test_categorical_distribution_creation() -> Result<(), MLError> {
        let config = DistributionalConfig::default();
        let stream = make_stream();
        let _dist = CategoricalDistribution::new(&config, &stream)?;
        Ok(())
    }

    #[test]
    fn test_support_creation() -> Result<(), MLError> {
        let config = DistributionalConfig {
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
        };

        let stream = make_stream();
        let dist = CategoricalDistribution::new(&config, &stream)?;
        let support = dist.support_host();

        assert_eq!(support.len(), 51);

        // Check first and last values
        let first_val = support.first().copied().unwrap_or(f32::NAN);
        let last_val = support.last().copied().unwrap_or(f32::NAN);

        assert!((first_val - (-10.0)).abs() < 1e-6);
        assert!((last_val - 10.0).abs() < 1e-6);

        Ok(())
    }

    #[test]
    fn test_basic_functionality() -> Result<(), MLError> {
        let config = DistributionalConfig::default();
        let stream = make_stream();
        let dist = CategoricalDistribution::new(&config, &stream)?;

        assert_eq!(dist.num_atoms(), config.num_atoms);

        Ok(())
    }

    /// Verify categorical_loss produces correct cross-entropy.
    #[test]
    fn test_categorical_loss_values() -> Result<(), MLError> {
        let stream = make_stream();
        let config = DistributionalConfig {
            num_atoms: 5,
            v_min: -1.0,
            v_max: 1.0,
        };
        let dist = CategoricalDistribution::new(&config, &stream)?;

        // Uniform prediction vs peaked target
        let pred = vec![0.2_f32; 10]; // batch=2, atoms=5
        let mut target_data = vec![0.0_f32; 10];
        target_data[2] = 1.0; // peak at atom 2 for sample 0
        target_data[7] = 1.0; // peak at atom 2 for sample 1

        let loss_val = dist.categorical_loss_host(&pred, &target_data)?;

        // -log(0.2) * 5 / 10 * 5 ≈ positive value
        assert!(loss_val > 1.0, "Cross-entropy of uniform vs peaked should be > 1, got {loss_val}");

        Ok(())
    }
}
