//! Continuous Policy Network for PPO with Gaussian Action Distribution
//!
//! Provides a continuous policy network for position sizing in [0.0, 1.0].
//! Forward pass and sampling are delegated to `cuda_nn` primitives.

use std::f32::consts::PI;

use rand::thread_rng;
use rand::Rng;
use serde::{Deserialize, Serialize};
use statrs::distribution::{ContinuousCDF, Normal};
use tracing::warn;

use ml_core::MLError;

use crate::cuda_nn::{GpuContext, CudaLinear, cuda_relu, cuda_sigmoid, cuda_from_slice};

/// Configuration for continuous policy network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContinuousPolicyConfig {
    /// State dimension
    pub state_dim: usize,
    /// Policy network hidden dimensions
    pub hidden_dims: Vec<usize>,
    /// Minimum log standard deviation
    pub min_log_std: f32,
    /// Maximum log standard deviation
    pub max_log_std: f32,
    /// Initial log standard deviation
    pub init_log_std: f32,
    /// Whether to use learnable log std
    pub learnable_std: bool,
    /// Action bounds [min, max]
    pub action_bounds: (f32, f32),
}

impl Default for ContinuousPolicyConfig {
    fn default() -> Self {
        Self {
            state_dim: 64,
            hidden_dims: vec![128, 64],
            min_log_std: -5.0,
            max_log_std: 2.0,
            init_log_std: -1.0,
            learnable_std: true,
            action_bounds: (0.0, 1.0),
        }
    }
}

/// Continuous policy network using Gaussian distributions (GPU-native).
///
/// All layers backed by `CudaLinear` with cuBLAS forward pass.
pub struct ContinuousPolicyNetwork {
    /// Feature layers
    feature_layers: Vec<CudaLinear>,
    /// Mean head
    mean_head: CudaLinear,
    /// Log std head (if learnable)
    log_std_head: Option<CudaLinear>,
    /// Fixed log std value (if not learnable)
    fixed_log_std: f32,
    /// Configuration
    config: ContinuousPolicyConfig,
    /// GPU context
    ctx: GpuContext,
}

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

impl ContinuousPolicyNetwork {
    /// Create new continuous policy network
    pub fn new(config: ContinuousPolicyConfig) -> Result<Self, MLError> {
        let ctx = GpuContext::new()?;

        let mut feature_layers = Vec::new();
        let mut current_dim = config.state_dim;

        for &hidden_dim in &config.hidden_dims {
            let layer = CudaLinear::new(ctx.clone(), current_dim, hidden_dim)?;
            feature_layers.push(layer);
            current_dim = hidden_dim;
        }

        let mean_head = CudaLinear::new(ctx.clone(), current_dim, 1)?;

        let log_std_head = config.learnable_std
            .then(|| CudaLinear::new(ctx.clone(), current_dim, 1))
            .transpose()?;

        Ok(Self {
            feature_layers,
            mean_head,
            log_std_head,
            fixed_log_std: config.init_log_std,
            config,
            ctx,
        })
    }

    /// Forward pass returning mean and `log_std` as host values (single sample).
    pub fn forward_host(&self, state: &[f32]) -> Result<(f32, f32), MLError> {
        let batch_size = 1;
        let mut x = cuda_from_slice(&self.ctx.stream, state)?;

        for layer in &self.feature_layers {
            let out = layer.forward(&x.data, batch_size)?;
            x = cuda_relu(&self.ctx.stream, &out.data, out.len)?;
        }

        // Mean head
        let mean_raw = self.mean_head.forward(&x.data, batch_size)?;

        // Sigmoid to bound mean to [0, 1]
        let mean_sigmoid = cuda_sigmoid(&self.ctx.stream, &mean_raw.data, mean_raw.len)?;
        let mean_host = mean_sigmoid.to_vec(&self.ctx.stream)?;
        let mean_val = mean_host.first().copied().unwrap_or(0.5);

        // Scale to action bounds
        let action_range = self.config.action_bounds.1 - self.config.action_bounds.0;
        let mean_scaled = self.config.action_bounds.0 + mean_val * action_range;

        // Log std
        let log_std_val = if let Some(ref log_std_head) = self.log_std_head {
            let log_std_out = log_std_head.forward(&x.data, batch_size)?;
            let log_std_host = log_std_out.to_vec(&self.ctx.stream)?;
            let raw = log_std_host.first().copied().unwrap_or(self.fixed_log_std);
            let min_ls = self.config.min_log_std.max(-20.0);
            let max_ls = self.config.max_log_std.min(2.0);
            raw.clamp(min_ls, max_ls)
        } else {
            self.fixed_log_std
        };

        Ok((mean_scaled, log_std_val))
    }

    /// Sample action from the Gaussian policy (returns (action, `log_prob`))
    pub fn sample_action_host(&self, state: &[f32]) -> Result<(f32, f32), MLError> {
        let (mean, log_std) = self.forward_host(state)?;
        let log_std_clamped = log_std.clamp(-20.0, 2.0);
        let std = log_std_clamped.exp();

        let mut rng = thread_rng();
        let normal = Normal::new(mean as f64, std as f64).map_err(|e| {
            MLError::ModelError(format!("Failed to create normal distribution: {}", e))
        })?;

        let uniform_sample = rng.gen::<f64>();
        let action_raw = normal.inverse_cdf(uniform_sample);

        let action = action_raw.clamp(
            self.config.action_bounds.0 as f64,
            self.config.action_bounds.1 as f64,
        );

        let log_prob = Self::compute_log_prob(action as f32, mean, log_std_clamped)?;

        Ok((action as f32, log_prob))
    }

    /// Compute Gaussian log probability
    fn compute_log_prob(action: f32, mean: f32, log_std: f32) -> Result<f32, MLError> {
        let std = log_std.exp();
        let normalized_diff = (action - mean) / std;
        let log_prob = -0.5 * (2.0 * PI).ln() - log_std - 0.5 * normalized_diff * normalized_diff;

        if !log_prob.is_finite() {
            warn!(
                "Non-finite log probability: action={}, mean={}, std={}, log_std={}",
                action, mean, std, log_std
            );
            return Err(MLError::ModelError(
                "Non-finite log probability computed".to_owned(),
            ));
        }

        Ok(log_prob)
    }

    /// Get configuration
    pub const fn config(&self) -> &ContinuousPolicyConfig {
        &self.config
    }
}

/// Continuous action for position sizing
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct ContinuousAction {
    /// Position size as a fraction (0.0 to 1.0)
    pub position_size: f32,
}

impl ContinuousAction {
    /// Create new continuous action
    pub const fn new(position_size: f32) -> Self {
        Self {
            position_size: position_size.clamp(0.0, 1.0),
        }
    }

    /// Get position size
    pub const fn position_size(&self) -> f32 {
        self.position_size
    }

    /// Validate action is within bounds
    pub fn is_valid(&self) -> bool {
        self.position_size >= 0.0 && self.position_size <= 1.0 && self.position_size.is_finite()
    }
}

impl Default for ContinuousAction {
    fn default() -> Self {
        Self { position_size: 0.0 }
    }
}

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

    #[test]
    fn test_continuous_policy_creation() {
        let config = ContinuousPolicyConfig::default();
        let policy = ContinuousPolicyNetwork::new(config);
        assert!(policy.is_ok());
    }

    #[test]
    fn test_continuous_action() {
        let action = ContinuousAction::new(0.5);
        assert_eq!(action.position_size(), 0.5);
        assert!(action.is_valid());

        let action_high = ContinuousAction::new(1.5);
        assert_eq!(action_high.position_size(), 1.0);

        let action_low = ContinuousAction::new(-0.5);
        assert_eq!(action_low.position_size(), 0.0);
    }

    #[test]
    fn test_forward_and_sample() {
        let config = ContinuousPolicyConfig {
            state_dim: 10,
            hidden_dims: vec![16, 8],
            ..ContinuousPolicyConfig::default()
        };
        let policy = ContinuousPolicyNetwork::new(config).unwrap();

        let state = vec![0.1_f32; 10];
        let (mean, log_std) = policy.forward_host(&state).unwrap();

        assert!(mean >= 0.0 && mean <= 1.0, "Mean {} out of bounds", mean);
        assert!(log_std.is_finite());

        let (action, log_prob) = policy.sample_action_host(&state).unwrap();
        assert!(action >= 0.0 && action <= 1.0, "Action {} out of bounds", action);
        assert!(log_prob.is_finite());
    }
}
