#![allow(unsafe_code)]

//! GPU batch statistics computation via parallel reduction.
//!
//! Replaces per-sample CPU accumulation of reward mean, variance,
//! drawdown, PnL, and win rate with a single CUDA kernel launch
//! and 40-byte readback. Derived statistics (mean, variance) are
//! computed on the host from raw sums for numerical correctness.

use std::sync::Arc;
use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use std::sync::OnceLock;

use ml_core::nvtx::NvtxRange;
use crate::MLError;

static STATISTICS_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

fn compile_statistics_ptx(context: &CudaContext) -> Result<Ptx, String> {
    let kernel_src = include_str!("statistics_kernel.cu");
    crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context)
}

/// Aggregated batch statistics from GPU parallel reduction.
#[derive(Debug, Clone)]
pub struct BatchStatistics {
    pub mean_reward: f32,
    pub reward_variance: f32,
    pub max_drawdown: f32,
    pub total_pnl: f32,
    pub win_rate: f32,
    pub total_count: f32,
    pub reward_sum: f32,
    pub reward_sq_sum: f32,
    pub mean_position_size: f32,
}

/// GPU batch statistics computer.
#[allow(missing_debug_implementations)]
pub struct GpuStatistics {
    kernel_func: CudaFunction,
    output_buf: CudaSlice<f32>,
    stream: Arc<CudaStream>,
}

impl GpuStatistics {
    /// Create a new GPU statistics computer from a CUDA stream.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        let ptx_result = STATISTICS_PTX.get_or_init(|| compile_statistics_ptx(&context));
        let ptx = ptx_result.as_ref().map_err(|e| {
            MLError::ModelError(format!("statistics PTX: {e}"))
        })?;
        let module = context.load_module(ptx.clone()).map_err(|e| {
            MLError::ModelError(format!("statistics module load: {e}"))
        })?;
        let kernel_func = module.load_function("batch_statistics").map_err(|e| {
            MLError::ModelError(format!("batch_statistics function load: {e}"))
        })?;

        let output_buf = stream.alloc_zeros::<f32>(10).map_err(|e| {
            MLError::ModelError(format!("alloc output_buf: {e}"))
        })?;

        Ok(Self { kernel_func, output_buf, stream: Arc::clone(stream) })
    }

    /// Compute batch statistics from GPU-resident buffers.
    ///
    /// Single kernel launch + 40-byte readback. Mean, variance, and
    /// drawdown are computed on the host from raw sums for correctness.
    pub fn compute(
        &mut self,
        rewards: &CudaSlice<f32>,
        portfolio_values: &CudaSlice<f32>,
        positions: &CudaSlice<f32>,
        n: usize,
    ) -> Result<BatchStatistics, MLError> {
        let _nvtx = NvtxRange::new("gpu_statistics");
        if n == 0 {
            return Ok(BatchStatistics {
                mean_reward: 0.0,
                reward_variance: 0.0,
                max_drawdown: 0.0,
                total_pnl: 0.0,
                win_rate: 0.0,
                total_count: 0.0,
                reward_sum: 0.0,
                reward_sq_sum: 0.0,
                mean_position_size: 0.0,
            });
        }

        let stream = &self.stream;

        // Zero output buffer
        stream.memset_zeros(&mut self.output_buf).map_err(|e| {
            MLError::ModelError(format!("memset output_buf: {e}"))
        })?;

        let n_i32 = n as i32;
        // Single block — grid-stride handles N > 256
        let cfg = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            stream
                .launch_builder(&self.kernel_func)
                .arg(rewards)
                .arg(portfolio_values)
                .arg(positions)
                .arg(&mut self.output_buf)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("batch_statistics launch: {e}")))?;
        }

        // Single 40-byte readback
        let mut host = [0.0_f32; 10];
        stream.memcpy_dtoh(&self.output_buf, &mut host).map_err(|e| { // gpu-exit: 40B batch stats
            MLError::ModelError(format!("statistics readback: {e}"))
        })?;

        let reward_sum = host[0];
        let reward_sq_sum = host[1];
        let global_min_pv = host[2];
        let global_max_pv = host[3];
        let win_count = host[4];
        let total_count = host[5].max(1.0);
        let position_sum = host[7];

        // Compute derived statistics on host (mathematically correct)
        let mean = reward_sum / total_count;
        let variance = (reward_sq_sum / total_count) - mean * mean;
        let max_drawdown = (global_max_pv - global_min_pv).max(0.0);
        let win_rate = win_count / total_count;
        let mean_position = position_sum / total_count;

        Ok(BatchStatistics {
            mean_reward: mean,
            reward_variance: variance.max(0.0),
            max_drawdown,
            total_pnl: reward_sum,
            win_rate,
            total_count,
            reward_sum,
            reward_sq_sum,
            mean_position_size: mean_position,
        })
    }
}
