#![allow(unsafe_code)] // Required for CUDA kernel launches

//! GPU-accelerated portfolio simulation for DQN training.
//!
//! Compiles CUDA kernels at runtime via NVRTC and launches them
//! to replace the CPU-bound inner loop in the DQN trainer.
//!
//! Uses cudarc 0.19 API:
//!   - CudaContext / CudaStream (not CudaDevice — that's the older API)
//!   - stream.launch_builder(&func).arg(&buf).launch(config)
//!   - stream.alloc_zeros / memcpy_htod / memcpy_dtoh

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use tracing::{debug, info};

use crate::MLError;

/// Maximum batch size for GPU portfolio simulation
const MAX_BATCH_SIZE: usize = 256;

/// Number of f32 elements in portfolio state
const PORTFOLIO_STATE_SIZE: usize = 8;

/// GPU-accelerated portfolio simulator for DQN training.
///
/// Replaces per-bar `PortfolioTracker::execute_action()` + reward
/// calculation with a CUDA kernel that processes an entire batch
/// sequentially on GPU.
#[allow(missing_debug_implementations)] // CudaFunction/CudaSlice don't impl Debug
pub struct GpuPortfolioSimulator {
    stream: Arc<CudaStream>,
    kernel_func: CudaFunction,

    // GPU buffers (persistent across batches)
    actions_buf: CudaSlice<i32>,
    portfolio_state_buf: CudaSlice<f32>,
    portfolio_out_buf: CudaSlice<f32>,
    rewards_out_buf: CudaSlice<f32>,
    done_out_buf: CudaSlice<i32>,

    // Config
    max_position: f32,
    episode_length: i32,
    total_bars: i32,
}

/// GPU-resident output from portfolio simulation (zero CPU download).
///
/// All buffers stay on GPU as `CudaSlice` for direct consumption by
/// downstream GPU stages (state builder, forward pass, loss computation).
#[allow(missing_debug_implementations)]
pub struct GpuPortfolioSimResult {
    /// Normalized portfolio features [batch_size * 3] on GPU
    pub portfolio_features: CudaSlice<f32>,
    /// Raw PnL rewards [batch_size] on GPU
    pub rewards: CudaSlice<f32>,
    /// Episode boundary flags [batch_size] on GPU (1=done, 0=continue)
    pub done_flags: CudaSlice<i32>,
    /// Number of bars processed
    pub batch_size: usize,
}

fn compile_experience_kernels(context: &Arc<CudaContext>) -> Result<CudaFunction, MLError> {
    let kernel_src = include_str!("experience_kernels.cu");
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context).map_err(|e| {
        MLError::ModelError(format!(
            "CUDA experience kernel compilation failed (is NVRTC available?): {e}"
        ))
    })?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("Failed to load experience kernel module: {e}"))
    })?;
    module.load_function("portfolio_sim_kernel").map_err(|e| {
        MLError::ModelError(format!("Failed to load portfolio_sim_kernel: {e}"))
    })
}

impl GpuPortfolioSimulator {
    /// Create a new GPU portfolio simulator.
    ///
    /// Compiles CUDA kernels via NVRTC and allocates persistent GPU buffers.
    pub fn new(
        stream: Arc<CudaStream>,
        initial_capital: f32,
        avg_spread: f32,
        cash_reserve_pct: f32,
        max_position: f32,
        episode_length: usize,
        total_bars: usize,
    ) -> Result<Self, MLError> {
        // Compile and load CUDA kernel
        let kernel_func = compile_experience_kernels(stream.context())?;
        info!("GPU portfolio simulator: CUDA kernels compiled and loaded");

        // Allocate persistent buffers
        let actions_buf = stream
            .alloc_zeros::<i32>(MAX_BATCH_SIZE)
            .map_err(|e| MLError::ModelError(format!("Failed to alloc actions buffer: {e}")))?;
        let mut portfolio_state_buf = stream
            .alloc_zeros::<f32>(PORTFOLIO_STATE_SIZE)
            .map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio state: {e}")))?;
        let portfolio_out_buf = stream
            .alloc_zeros::<f32>(MAX_BATCH_SIZE * 3)
            .map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio output: {e}")))?;
        let rewards_out_buf = stream
            .alloc_zeros::<f32>(MAX_BATCH_SIZE)
            .map_err(|e| MLError::ModelError(format!("Failed to alloc rewards output: {e}")))?;
        let done_out_buf = stream
            .alloc_zeros::<i32>(MAX_BATCH_SIZE)
            .map_err(|e| MLError::ModelError(format!("Failed to alloc done output: {e}")))?;

        // Initialize portfolio state: [cash, position, entry_price, initial_cap, spread, last_price, reserve_pct, cum_costs]
        let init_state = [
            initial_capital,
            0.0_f32,
            0.0,
            initial_capital,
            avg_spread,
            0.0,
            cash_reserve_pct,
            0.0,
        ];
        stream
            .memcpy_htod(&init_state, &mut portfolio_state_buf)
            .map_err(|e| MLError::ModelError(format!("Failed to upload portfolio state: {e}")))?;

        debug!(
            "GPU portfolio sim initialized: capital={}, spread={}, reserve={}%, max_pos={}, episode_len={}, total_bars={}",
            initial_capital, avg_spread, cash_reserve_pct, max_position, episode_length, total_bars
        );

        Ok(Self {
            stream,
            kernel_func,
            actions_buf,
            portfolio_state_buf,
            portfolio_out_buf,
            rewards_out_buf,
            done_out_buf,
            max_position,
            episode_length: episode_length as i32,
            total_bars: total_bars as i32,
        })
    }

    /// Run portfolio simulation keeping all results on GPU (zero CPU download).
    ///
    /// Returns DtoD copies of the kernel output buffers so they can be consumed
    /// directly by the next GPU stage (state builder, forward pass, etc.).
    pub fn simulate_batch_gpu(
        &mut self,
        targets_buf: &CudaSlice<f32>,
        actions: &[i32],
        batch_start: usize,
    ) -> Result<GpuPortfolioSimResult, MLError> {
        let batch_size = actions.len();
        if batch_size == 0 {
            return Ok(GpuPortfolioSimResult {
                portfolio_features: self.stream.alloc_zeros::<f32>(0)
                    .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
                rewards: self.stream.alloc_zeros::<f32>(0)
                    .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
                done_flags: self.stream.alloc_zeros::<i32>(0)
                    .map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
                batch_size: 0,
            });
        }
        if batch_size > MAX_BATCH_SIZE {
            return Err(MLError::ModelError(format!(
                "Batch size {} exceeds max {}",
                batch_size, MAX_BATCH_SIZE
            )));
        }

        // Upload actions to GPU
        self.stream
            .memcpy_htod(actions, &mut self.actions_buf)
            .map_err(|e| MLError::ModelError(format!("Failed to upload actions: {e}")))?;

        let config = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (1, 1, 1),
            shared_mem_bytes: 0,
        };

        let batch_start_i32 = batch_start as i32;
        let batch_size_i32 = batch_size as i32;

        // Safety: kernel parameters match the CUDA function signature exactly.
        unsafe {
            self.stream
                .launch_builder(&self.kernel_func)
                .arg(targets_buf)
                .arg(&self.actions_buf)
                .arg(&mut self.portfolio_state_buf)
                .arg(&mut self.portfolio_out_buf)
                .arg(&mut self.rewards_out_buf)
                .arg(&mut self.done_out_buf)
                .arg(&batch_start_i32)
                .arg(&batch_size_i32)
                .arg(&self.max_position)
                .arg(&self.episode_length)
                .arg(&self.total_bars)
                .launch(config)
                .map_err(|e| MLError::ModelError(format!("Kernel launch failed: {e}")))?;
        }

        // DtoD copy result slices into independent buffers (no CPU download)
        let portfolio_features = Self::dtod_clone_range_f32(
            &self.portfolio_out_buf, batch_size * 3, &self.stream,
        )?;
        let rewards = Self::dtod_clone_range_f32(
            &self.rewards_out_buf, batch_size, &self.stream,
        )?;
        let done_flags = Self::dtod_clone_range_i32(
            &self.done_out_buf, batch_size, &self.stream,
        )?;

        Ok(GpuPortfolioSimResult {
            portfolio_features,
            rewards,
            done_flags,
            batch_size,
        })
    }

    /// DtoD clone the first `count` f32 elements of `src`.
    #[allow(unsafe_code)]
    fn dtod_clone_range_f32(
        src: &CudaSlice<f32>,
        count: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let mut dst = stream.alloc_zeros::<f32>(count)
            .map_err(|e| MLError::ModelError(format!("dtod_clone alloc: {e}")))?;
        {
            let nbytes = count * std::mem::size_of::<f32>();
            let src_view = src.slice(..count);
            let (src_ptr, _sg) = src_view.device_ptr(stream);
            let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("dtod_clone copy: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// DtoD clone the first `count` i32 elements of `src`.
    #[allow(unsafe_code)]
    fn dtod_clone_range_i32(
        src: &CudaSlice<i32>,
        count: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<CudaSlice<i32>, MLError> {
        let mut dst = stream.alloc_zeros::<i32>(count)
            .map_err(|e| MLError::ModelError(format!("dtod_clone_i32 alloc: {e}")))?;
        {
            let nbytes = count * std::mem::size_of::<i32>();
            let src_view = src.slice(..count);
            let (src_ptr, _sg) = src_view.device_ptr(stream);
            let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("dtod_clone_i32 copy: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// Reset portfolio state to initial capital.
    pub fn reset(&mut self, initial_capital: f32, avg_spread: f32, cash_reserve_pct: f32) -> Result<(), MLError> {
        let init_state = [initial_capital, 0.0, 0.0, initial_capital, avg_spread, 0.0, cash_reserve_pct, 0.0_f32];
        self.stream
            .memcpy_htod(&init_state, &mut self.portfolio_state_buf)
            .map_err(|e| MLError::ModelError(format!("Failed to reset portfolio state: {e}")))?;
        Ok(())
    }

    /// Get reference to the GPU-resident portfolio state buffer (for diagnostics).
    ///
    /// Returns the raw `CudaSlice<f32>` of shape `[PORTFOLIO_STATE_SIZE]` on GPU.
    /// Use `memcpy_dtoh` explicitly at checkpoint boundaries only.
    pub fn portfolio_state_buf(&self) -> &CudaSlice<f32> {
        &self.portfolio_state_buf
    }
}
