//! Quantile Regression DQN (QR-DQN) Implementation
//!
//! Implementation of quantile regression for distributional value function approximation
//! as described in "Distributional Reinforcement Learning with Quantile Regression" (Dabney et al., 2018)
//!
//! **Why QR-DQN is Better for Trading:**
//! - C51 uses fixed categorical bins → poor tail risk modeling
//! - QR-DQN learns quantiles directly → better CVaR/VaR estimation
//! - More stable for asymmetric return distributions (common in trading)
//! - Adaptive to distribution shape without manual `v_min/v_max` tuning
//!
//! **Key Advantages:**
//! 1. **Risk Modeling**: Direct `CVaR` computation for risk-averse policies
//! 2. **Tail Sensitivity**: Better captures extreme events (crashes, rallies)
//! 3. **Flexibility**: No need to specify value ranges (`v_min/v_max`)
//! 4. **Stability**: Quantile Huber loss is more robust than cross-entropy

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore};
use serde::{Deserialize, Serialize};
use std::f32::consts::PI;

use ml_core::MLError;

/// Configuration for Quantile Regression DQN
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantileConfig {
    /// Number of quantile samples (typically 200 for trading)
    /// More quantiles = better tail resolution
    pub num_quantiles: usize,
    /// Embedding dimension for quantile τ (typically 64)
    pub quantile_embedding_dim: usize,
    /// Kappa parameter for quantile Huber loss (typically 1.0)
    /// Controls transition between L1 (robust) and L2 (smooth)
    pub kappa: f32,
    /// Number of actions in the DQN action space
    pub num_actions: usize,
}

impl Default for QuantileConfig {
    fn default() -> Self {
        Self {
            num_quantiles: 200,  // High resolution for risk modeling
            quantile_embedding_dim: 64,
            kappa: 1.0,  // Standard quantile Huber parameter
            num_actions: 9,  // 9 exposure levels
        }
    }
}

/// Quantile Network for distributional value function approximation
///
/// Architecture:
/// 1. State embedding: φ(s) via standard Q-network layers
/// 2. Quantile embedding: ψ(τ) via cosine embedding
/// 3. Element-wise product: φ(s) ⊙ ψ(τ)
/// 4. Output layer: Projects to quantile values
pub struct QuantileNetwork {
    config: QuantileConfig,
    /// Cosine embedding layer for quantiles
    quantile_embedding: GpuLinear,
    /// Output layer after element-wise product
    output_layer: GpuLinear,
    /// Network variables for optimizer access
    vars: GpuVarStore,
    /// CUDA stream for GPU operations
    stream: Arc<CudaStream>,
}

impl std::fmt::Debug for QuantileNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("QuantileNetwork")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

impl QuantileNetwork {
    /// Create new quantile network
    ///
    /// # Arguments
    /// * `config` - Quantile configuration
    /// * `state_dim` - State embedding dimension (from base Q-network)
    /// * `stream` - CUDA stream for GPU operations
    pub fn new(
        config: &QuantileConfig,
        state_dim: usize,
        stream: Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut vars = GpuVarStore::new(stream.clone());

        let quantile_embedding = vars.linear_xavier(
            "quantile_embedding",
            config.quantile_embedding_dim,
            state_dim,
        )?;

        let output_layer = vars.linear_xavier(
            "quantile_output",
            state_dim,
            config.num_actions,
        )?;

        Ok(Self {
            config: config.clone(),
            quantile_embedding,
            output_layer,
            vars,
            stream,
        })
    }

    /// Get network variables for optimizer
    pub const fn vars(&self) -> &GpuVarStore {
        &self.vars
    }

    /// Copy weights from another `QuantileNetwork` (for target network sync)
    pub fn copy_weights_from(&mut self, other: &QuantileNetwork) -> Result<(), MLError> {
        let host_data = other.vars.export_to_host()?; // gpu-exit: checkpoint weight export
        self.vars.import_from_host(&host_data)
    }

    /// Forward pass: Compute quantile values Z(s, a, tau) for all actions (cold path).
    ///
    /// **Hot-path IQN forward is handled by `gpu_iqn_head::GpuIqnHead` which uses
    /// the fused CUDA kernel `iqn_dual_head_kernel`.**
    #[cold]
    pub fn forward(&self, state_embed: &GpuTensor, taus: &GpuTensor) -> Result<GpuTensor, MLError> {
        // IQN forward pass (cold path):
        // 1. Cosine embedding: psi(tau) = ReLU(W_embed @ cos(pi * i * tau) + b_embed)
        // 2. Element-wise product: phi(s) * psi(tau)
        // 3. Output: Q(s, a, tau) = W_out @ (phi * psi) + b_out
        let cublas = cudarc::cublas::CudaBlas::new(self.stream.clone()).map_err(|e| {
            MLError::ModelError(format!("IQN forward cublas init: {e}"))
        })?;
        let activations = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?;

        let (batch_size, num_quantiles) = taus.dims2().map_err(|e| {
            MLError::ModelError(format!("IQN taus dims2: {e}"))
        })?;
        let embed_dim = self.config.quantile_embedding_dim;

        // Download taus to host for cosine embedding computation
        let taus_host = taus.to_host(&self.stream)?;
        let state_dim = state_embed.shape().get(1).copied().unwrap_or(
            state_embed.numel() / batch_size.max(1)
        );

        // Compute cosine features: cos(pi * i * tau) for i in 0..embed_dim
        // Shape: [batch * num_quantiles, embed_dim]
        let total_taus = batch_size * num_quantiles;
        let mut cos_features = Vec::with_capacity(total_taus * embed_dim);
        for t_idx in 0..total_taus {
            let tau = taus_host.get(t_idx).copied().unwrap_or(0.5);
            for i in 0..embed_dim {
                let val = (PI * (i + 1) as f32 * tau).cos();
                cos_features.push(val);
            }
        }
        let cos_tensor = GpuTensor::from_host(
            &cos_features,
            vec![total_taus, embed_dim],
            &self.stream,
        )?;

        // Quantile embedding: psi = ReLU(cos_features @ W_embed^T + b_embed)
        // quantile_embedding: [embed_dim -> state_dim]
        let (psi_raw, _acts) = self.quantile_embedding.forward(
            &cos_tensor, &self.vars, &cublas, &self.stream,
        )?;
        let (psi, _mask) = activations.relu_fwd(&psi_raw, &self.stream)?;
        // psi: [total_taus, state_dim]

        // Download state_embed and psi, broadcast multiply
        let state_host = state_embed.to_host(&self.stream)?;
        let psi_host = psi.to_host(&self.stream)?;

        // Element-wise product: phi(s) * psi(tau) for each (batch, quantile)
        // state_embed: [batch, state_dim], psi: [batch * num_quantiles, state_dim]
        let mut product = Vec::with_capacity(total_taus * state_dim);
        for b in 0..batch_size {
            for q in 0..num_quantiles {
                let tau_idx = b * num_quantiles + q;
                for d in 0..state_dim {
                    let s_val = state_host.get(b * state_dim + d).copied().unwrap_or(0.0);
                    let p_val = psi_host.get(tau_idx * state_dim + d).copied().unwrap_or(0.0);
                    product.push(s_val * p_val);
                }
            }
        }
        let product_tensor = GpuTensor::from_host(
            &product,
            vec![total_taus, state_dim],
            &self.stream,
        )?;

        // Output: Q(s, a, tau) = output_layer(product)
        // output_layer: [state_dim -> num_actions]
        let (output, _acts) = self.output_layer.forward(
            &product_tensor, &self.vars, &cublas, &self.stream,
        )?;
        // output: [total_taus, num_actions]

        // Reshape to [batch, num_quantiles, num_actions]
        output.reshape(vec![batch_size, num_quantiles, self.config.num_actions])
    }

    /// Sample fixed quantiles uniformly in [0, 1]
    ///
    /// `tau_i` = (i + 0.5) / N for i in 0..N
    pub fn sample_uniform_quantiles(&self, batch_size: usize) -> Result<GpuTensor, MLError> {
        let num_quantiles = self.config.num_quantiles;
        let host: Vec<f32> = (0..batch_size)
            .flat_map(|_| {
                (0..num_quantiles).map(|i| (i as f32 + 0.5) / num_quantiles as f32)
            })
            .collect();
        GpuTensor::from_host(&host, vec![batch_size, num_quantiles], &self.stream)
    }

    /// Sample random quantiles from Uniform(0, 1) -- IQN training mode
    pub fn sample_random_quantiles(&self, batch_size: usize) -> Result<GpuTensor, MLError> {
        let num_quantiles = self.config.num_quantiles;
        let host: Vec<f32> = (0..batch_size * num_quantiles)
            .map(|_| rand::random::<f32>())
            .collect();
        GpuTensor::from_host(&host, vec![batch_size, num_quantiles], &self.stream)
    }

    /// Compute expected Q-values from quantile distributions (mean over quantiles)
    pub fn to_expected_q(&self, quantiles: &GpuTensor) -> Result<GpuTensor, MLError> {
        // quantiles: [batch, num_quantiles, num_actions]
        // Expected Q = mean over quantiles dim -> [batch, num_actions]
        let (batch, num_q, num_a) = quantiles.dims3().map_err(|e| {
            MLError::ModelError(format!("to_expected_q dims3: {e}"))
        })?;

        let host = quantiles.to_host(&self.stream)?;
        let mut result = Vec::with_capacity(batch * num_a);
        for b in 0..batch {
            for a in 0..num_a {
                let mut sum = 0.0_f32;
                for q in 0..num_q {
                    let idx = b * num_q * num_a + q * num_a + a;
                    sum += host.get(idx).copied().unwrap_or(0.0);
                }
                result.push(sum / num_q as f32);
            }
        }

        GpuTensor::from_host(&result, vec![batch, num_a], &self.stream)
    }

    /// Extract `CVaR` (Conditional Value at Risk) for each action
    pub fn compute_cvar(&self, quantiles: &GpuTensor, alpha: f32) -> Result<GpuTensor, MLError> {
        // CVaR_alpha = mean of the lowest alpha-fraction of quantiles per action.
        // quantiles: [batch, num_quantiles, num_actions]
        let (batch, num_q, num_a) = quantiles.dims3().map_err(|e| {
            MLError::ModelError(format!("compute_cvar dims3: {e}"))
        })?;

        // Number of quantiles in the lower alpha fraction
        let k = ((alpha * num_q as f32).ceil() as usize).max(1).min(num_q);

        let host = quantiles.to_host(&self.stream)?;
        let mut result = Vec::with_capacity(batch * num_a);

        for b in 0..batch {
            for a in 0..num_a {
                // Gather quantile values for this (batch, action)
                let mut vals: Vec<f32> = (0..num_q)
                    .map(|q| {
                        let idx = b * num_q * num_a + q * num_a + a;
                        host.get(idx).copied().unwrap_or(0.0)
                    })
                    .collect();

                // Sort ascending
                vals.sort_by(|x, y| x.partial_cmp(y).unwrap_or(std::cmp::Ordering::Equal));

                // Mean of the lowest k values
                let cvar_sum: f32 = vals.iter().take(k).sum();
                result.push(cvar_sum / k as f32);
            }
        }

        GpuTensor::from_host(&result, vec![batch, num_a], &self.stream)
    }
}

/// Quantile Huber loss for stable quantile regression (GPU-resident).
///
/// Computes the asymmetric Huber loss used in QR-DQN/IQN training.
///
/// # Arguments
/// * `predicted` - Predicted quantile values [batch, `num_quantiles`]
/// * `target` - Target quantile values [batch, `num_quantiles`]
/// * `taus` - Quantile fractions [batch, `num_quantiles`]
/// * `kappa` - Huber threshold
///
/// # Returns
/// Mean quantile Huber loss (scalar `GpuTensor`)
pub fn quantile_huber_loss(
    predicted: &GpuTensor,
    target: &GpuTensor,
    taus: &GpuTensor,
    kappa: f32,
) -> Result<GpuTensor, MLError> {
    // Compute quantile Huber loss via host-side calculation (cold path).
    // predicted, target, taus: [batch, num_quantiles]
    // Returns: scalar mean loss
    let per_sample = quantile_huber_loss_per_sample(predicted, target, taus, kappa)?;
    // per_sample: [batch]
    // Mean over batch
    let batch = per_sample.numel();
    if batch == 0 {
        let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
            MLError::DeviceError(format!("CUDA context: {e}"))
        })?;
        let stream = ctx.new_stream().map_err(|e| {
            MLError::DeviceError(format!("CUDA stream: {e}"))
        })?;
        return GpuTensor::scalar(0.0, &stream);
    }
    let host = per_sample.to_host(
        // We need a stream -- extract from tensor data context
        // The tensor's CudaSlice carries its context
        &{
            let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
                MLError::DeviceError(format!("CUDA context: {e}"))
            })?;
            ctx.new_stream().map_err(|e| {
                MLError::DeviceError(format!("CUDA stream: {e}"))
            })?
        },
    )?;
    let mean_loss: f32 = host.iter().sum::<f32>() / batch as f32;
    let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
        MLError::DeviceError(format!("CUDA context: {e}"))
    })?;
    let stream = ctx.new_stream().map_err(|e| {
        MLError::DeviceError(format!("CUDA stream: {e}"))
    })?;
    GpuTensor::scalar(mean_loss, &stream)
}

/// Per-sample quantile Huber loss for PER importance-sampling weight correction.
///
/// Identical to [`quantile_huber_loss`] except the batch dimension is preserved.
///
/// # Returns
/// Per-sample loss tensor `[batch]` (mean over quantiles, **not** over batch)
pub fn quantile_huber_loss_per_sample(
    predicted: &GpuTensor,
    target: &GpuTensor,
    taus: &GpuTensor,
    kappa: f32,
) -> Result<GpuTensor, MLError> {
    // Per-sample quantile Huber loss (cold path).
    // predicted, target, taus: [batch, num_quantiles]
    // Returns: [batch] (mean over quantiles, NOT over batch)
    let (batch, num_q) = predicted.dims2().map_err(|e| {
        MLError::ModelError(format!("quantile_huber_loss_per_sample predicted dims2: {e}"))
    })?;

    let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
        MLError::DeviceError(format!("CUDA context: {e}"))
    })?;
    let stream = ctx.new_stream().map_err(|e| {
        MLError::DeviceError(format!("CUDA stream: {e}"))
    })?;

    let pred_host = predicted.to_host(&stream)?;
    let tgt_host = target.to_host(&stream)?;
    let tau_host = taus.to_host(&stream)?;

    let mut result = Vec::with_capacity(batch);
    for b in 0..batch {
        let mut sample_loss = 0.0_f32;
        for q in 0..num_q {
            let idx = b * num_q + q;
            let pred_val = pred_host.get(idx).copied().unwrap_or(0.0);
            let tgt_val = tgt_host.get(idx).copied().unwrap_or(0.0);
            let tau_val = tau_host.get(idx).copied().unwrap_or(0.5);

            let diff = tgt_val - pred_val;
            let abs_diff = diff.abs();

            // Huber loss
            let huber = if abs_diff <= kappa {
                0.5 * diff * diff
            } else {
                kappa * (abs_diff - 0.5 * kappa)
            };

            // Quantile weighting: |tau - I(diff < 0)|
            let indicator = if diff < 0.0 { 1.0_f32 } else { 0.0_f32 };
            let weight = (tau_val - indicator).abs();

            sample_loss += weight * huber;
        }
        result.push(sample_loss / num_q as f32);
    }

    GpuTensor::from_host(&result, vec![batch], &stream)
}

// Re-export from distributional module (single source of truth)
pub use super::distributional::DistributionalType;

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

    #[test]
    fn test_quantile_config_default() {
        let config = QuantileConfig::default();
        assert_eq!(config.num_quantiles, 200);
        assert_eq!(config.quantile_embedding_dim, 64);
        assert_eq!(config.kappa, 1.0);
        assert_eq!(config.num_actions, 9); // 9 exposure levels (default)
    }
}
