#![allow(unsafe_code)]

//! GPU monitoring reduction — aggregates per-experience rewards/actions
//! into a compact summary without downloading full arrays.

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

/// Compact monitoring summary from GPU reduction (48-byte GPU transfer, 12 × f32).
#[derive(Debug, Clone, Default)]
pub struct MonitoringSummary {
    pub mean_reward: f32,
    pub reward_std: f32,
    pub min_reward: f32,
    pub max_reward: f32,
    pub sharpe_estimate: f32,
    pub action_counts: [usize; 9],
    pub total_experiences: usize,
}

/// GPU monitoring reducer.
#[allow(missing_debug_implementations)]
pub struct GpuMonitoringReducer {
    stream: Arc<CudaStream>,
    kernel_func: CudaFunction,
    summary_buf: CudaSlice<f32>,  // [12]
}

// ── PTX cache ─────────────────────────────────────────────────────────────

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

fn compile_monitoring_ptx(context: &CudaContext) -> Result<Ptx, String> {
    let defines = "\
        #define STATE_DIM 72\n\
        #define MARKET_DIM 42\n\
        #define PORTFOLIO_DIM 8\n";
    let common_src = include_str!("common_device_functions.cuh");
    let kernel_src = include_str!("monitoring_kernel.cu");
    let full_source = format!("{defines}{common_src}\n{kernel_src}");
    crate::cuda_pipeline::compile_ptx_for_device(&full_source, context)
}

impl GpuMonitoringReducer {
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx_result = MONITORING_PTX.get_or_init(|| compile_monitoring_ptx(&context));
        let ptx = ptx_result.as_ref()
            .map_err(|e| MLError::ModelError(format!("monitoring PTX: {e}")))?;
        let module = context.load_module(ptx.clone())
            .map_err(|e| MLError::ModelError(format!("monitoring module load: {e}")))?;
        let kernel_func = module.load_function("monitoring_reduce")
            .map_err(|e| MLError::ModelError(format!("monitoring_reduce load: {e}")))?;
        let summary_buf = stream.alloc_zeros::<f32>(16) // 5 stats + 9 action counts + 1 total + 1 pad
            .map_err(|e| MLError::ModelError(format!("monitoring summary alloc: {e}")))?;

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

    /// Launch reduction over rewards/actions buffers already on GPU.
    /// Does NOT synchronize — caller must sync before reading result.
    pub fn reduce(
        &mut self,
        rewards: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        n: usize,
    ) -> Result<(), MLError> {
        let _nvtx = NvtxRange::new("gpu_monitoring");
        if n == 0 {
            return Ok(());  // nothing to reduce; summary_buf retains zeros from alloc
        }
        if n > i32::MAX as usize {
            return Err(MLError::ModelError(format!("monitoring n={n} exceeds i32::MAX")));
        }
        let config = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,  // kernel uses static __shared__ only — no dynamic shmem needed
        };
        // Safety: rewards and actions are valid GPU slices with at least n elements.
        // summary_buf is pre-allocated 12-float device buffer on the same stream.
        unsafe {
            self.stream
                .launch_builder(&self.kernel_func)
                .arg(rewards)
                .arg(actions)
                .arg(&self.summary_buf)
                .arg(&(n as i32))
                .launch(config)
                .map_err(|e| MLError::ModelError(format!("monitoring_reduce launch: {e}")))?;
        }
        Ok(())
    }

    /// Download summary from GPU (single 48-byte transfer).
    pub fn download_summary(&self) -> Result<MonitoringSummary, MLError> {
        let mut raw = vec![0.0_f32; 16];
        self.stream.memcpy_dtoh(&self.summary_buf, &mut raw) // gpu-exit: 64B monitoring summary (5 stats + 9 actions + 1 total + 1 pad)
            .map_err(|e| MLError::ModelError(format!("monitoring download: {e}")))?;
        Ok(MonitoringSummary {
            mean_reward: raw[0],
            reward_std: raw[1],
            min_reward: raw[2],
            max_reward: raw[3],
            sharpe_estimate: raw[4],
            action_counts: [
                raw[5] as usize, raw[6] as usize, raw[7] as usize,
                raw[8] as usize, raw[9] as usize, raw[10] as usize,
                raw[11] as usize, raw[12] as usize, raw[13] as usize,
            ],
            total_experiences: raw[14] as usize,
        })
    }
}

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

    #[test]
    fn test_monitoring_summary_default() {
        let s = MonitoringSummary::default();
        assert_eq!(s.total_experiences, 0);
        assert_eq!(s.action_counts, [0; 9]);
    }

    #[test]

    fn test_monitoring_ptx_compilation() {
        let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        let stream = context.default_stream();
        let context = stream.context();
        let result = compile_monitoring_ptx(&context);
        if let Err(ref e) = result {
            panic!("monitoring kernel PTX compilation failed: {e}");
        }
    }
}
