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

//! GPU-resident curiosity forward model training.
//!
//! Trains the curiosity forward model (2-layer MLP) entirely on GPU with zero
//! CPU involvement. Experience data (states, actions, next_states) is already
//! on GPU from the experience collector. Weights in [`CuriosityWeightSet`] are
//! updated in-place via Adam optimizer -- no CPU roundtrips.
//!
//! Architecture: `[MARKET_DIM+3] -> [128] LeakyReLU -> [MARKET_DIM]` (MARKET_DIM=42: 11_954 params)
//!
//! Kernels (fused path -- 2 launches per step):
//! - `curiosity_shift_states`: builds shifted next_states from states buffer
//! - `curiosity_fused_zero_fwd_bwd_adam`: zeros grads + forward/backward + Adam in one launch
//!
//! Legacy kernels (kept for fallback, not used in hot path):
//! - `curiosity_forward_backward`: standalone forward + backward pass
//! - `curiosity_adam_step` / `curiosity_adam_step_fused`: standalone Adam optimizers

use std::sync::{Arc, OnceLock};

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

use crate::MLError;
use super::gpu_weights::CuriosityWeightSet;

// ---------------------------------------------------------------------------
// Constants (must match CUDA kernel defines)
// ---------------------------------------------------------------------------

const CUR_MARKET_DIM: usize = 42;                         // Must match MARKET_DIM in CUDA
const CUR_HIDDEN: usize = 128;                            // Must match CUR_HIDDEN in CUDA
const CUR_INPUT: usize = CUR_MARKET_DIM + 3;              // 45
const CUR_OUTPUT: usize = CUR_MARKET_DIM;                  // 42

const CUR_W1_LEN: usize = CUR_HIDDEN * CUR_INPUT;         // [128, 45] = 5760
const CUR_B1_LEN: usize = CUR_HIDDEN;                     // [128]
const CUR_W2_LEN: usize = CUR_OUTPUT * CUR_HIDDEN;        // [42, 128] = 5376
const CUR_B2_LEN: usize = CUR_OUTPUT;                     // [42]

/// Adam optimizer hyperparameters.
const ADAM_LR: f32 = 0.001;
const ADAM_BETA1: f32 = 0.9;
const ADAM_BETA2: f32 = 0.999;
const ADAM_EPS: f32 = 1e-8;

// ---------------------------------------------------------------------------
// PTX cache
// ---------------------------------------------------------------------------

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

fn compile_curiosity_training_ptx(context: &CudaContext) -> Result<Ptx, String> {
    let defines = "\
        #define STATE_DIM 48\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!("curiosity_training_kernel.cu");
    let full_source = format!("{defines}{common_src}\n{kernel_src}");
    crate::cuda_pipeline::compile_ptx_for_device(&full_source, context)
}

// ---------------------------------------------------------------------------
// GpuCuriosityTrainer
// ---------------------------------------------------------------------------

/// GPU-resident curiosity forward model trainer.
///
/// Trains the curiosity MLP entirely on GPU using experience data that is
/// already device-resident. Maintains gradient buffers and Adam optimizer
/// state. Modifies [`CuriosityWeightSet`] in-place -- zero CPU traffic.
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct GpuCuriosityTrainer {
    stream: Arc<CudaStream>,

    // Kernel functions
    shift_func: CudaFunction,
    fwd_bwd_func: CudaFunction,
    adam_func: CudaFunction,
    adam_fused_func: CudaFunction,
    /// Fully fused kernel: zero grads + fwd/bwd + Adam in one launch.
    fused_zero_fwd_bwd_adam_func: CudaFunction,

    // Gradient buffers
    grad_w1: CudaSlice<f32>,  // [CUR_W1_LEN]
    grad_b1: CudaSlice<f32>,  // [CUR_B1_LEN]
    grad_w2: CudaSlice<f32>,  // [CUR_W2_LEN]
    grad_b2: CudaSlice<f32>,  // [CUR_B2_LEN]

    // Adam first moment (per-param-group)
    adam_m_w1: CudaSlice<f32>, // [CUR_W1_LEN]
    adam_m_b1: CudaSlice<f32>, // [CUR_B1_LEN]
    adam_m_w2: CudaSlice<f32>, // [CUR_W2_LEN]
    adam_m_b2: CudaSlice<f32>, // [CUR_B2_LEN]

    // Adam second moment (per-param-group)
    adam_v_w1: CudaSlice<f32>, // [CUR_W1_LEN]
    adam_v_b1: CudaSlice<f32>, // [CUR_B1_LEN]
    adam_v_w2: CudaSlice<f32>, // [CUR_W2_LEN]
    adam_v_b2: CudaSlice<f32>, // [CUR_B2_LEN]

    // Shifted next_states buffer
    next_states_buf: CudaSlice<f32>,

    /// Atomic block-arrival counter for fused kernel grid-wide sync (single i32 on GPU).
    block_counter: CudaSlice<i32>,

    // Adam step counter (1-based)
    step: i32,

    // State dimension of the experience data
    state_dim: usize,

    // Maximum number of samples the next_states buffer was allocated for
    buf_capacity: usize,
}

/// Launch Adam optimizer step for one parameter group (legacy, kept for fallback).
///
/// Free function to avoid borrow conflicts when calling from
/// `train_on_collector_buffers` (which mutably borrows multiple fields).
#[allow(dead_code)]
fn launch_adam_step(
    stream: &CudaStream,
    adam_func: &CudaFunction,
    params: &mut CudaSlice<f32>,
    grads: &CudaSlice<f32>,
    num_params: usize,
    m: &mut CudaSlice<f32>,
    v: &mut CudaSlice<f32>,
    batch_size: usize,
    step: i32,
) -> Result<(), MLError> {
    let cfg = LaunchConfig {
        grid_dim: (((num_params as u32) + 255) / 256, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };
    let np_i32 = num_params as i32;
    let bs_i32 = batch_size as i32;
    unsafe {
        stream
            .launch_builder(adam_func)
            .arg(params)
            .arg(grads)
            .arg(m)
            .arg(v)
            .arg(&np_i32)
            .arg(&bs_i32)
            .arg(&ADAM_LR)
            .arg(&ADAM_BETA1)
            .arg(&ADAM_BETA2)
            .arg(&ADAM_EPS)
            .arg(&step)
            .launch(cfg)
            .map_err(|e| {
                MLError::ModelError(format!("curiosity_adam_step launch: {e}"))
            })?;
    }
    Ok(())
}

impl GpuCuriosityTrainer {
    /// Create a new GPU curiosity trainer.
    ///
    /// Compiles the CUDA training kernel via NVRTC, allocates gradient and
    /// Adam optimizer state buffers as zeros on GPU.
    ///
    /// # Arguments
    /// * `stream` - CUDA stream for all GPU operations
    /// * `state_dim` - Dimensionality of state vectors in the experience buffer
    /// * `max_samples` - Maximum number of training samples (for next_states buffer sizing)
    pub fn new(
        stream: Arc<CudaStream>,
        state_dim: usize,
        max_samples: usize,
    ) -> Result<Self, MLError> {
        // The fused curiosity kernel allocates ~3KB/thread on stack (6 arrays of
        // 42-128 floats: input[45], pre_act[128], hidden[128], pred[42], d_pred[42],
        // Stack sizing: curiosity kernel needs ~3KB/thread (513 floats for
        // input/hidden/pred arrays). The fused training kernel (GpuDqnTrainer)
        // sets 64KB stack when it inits — that covers all kernels including
        // curiosity. Don't set stack here: cuCtxSetLimit reserves stack_bytes ×
        // max_threads of VRAM upfront, causing OOM on 4GB GPUs.

        // ---- Compile and load kernels ----
        let ptx_result = CURIOSITY_TRAINING_PTX.get_or_init(|| compile_curiosity_training_ptx(&stream.context()));
        let ptx = ptx_result
            .as_ref()
            .map_err(|e| MLError::ModelError(format!("curiosity training PTX: {e}")))?;

        let context = stream.context();
        let module = context.load_module(ptx.clone()).map_err(|e| {
            MLError::ModelError(format!("curiosity training module load: {e}"))
        })?;

        let shift_func = module.load_function("curiosity_shift_states").map_err(|e| {
            MLError::ModelError(format!("curiosity_shift_states load: {e}"))
        })?;
        let fwd_bwd_func = module.load_function("curiosity_forward_backward").map_err(|e| {
            MLError::ModelError(format!("curiosity_forward_backward load: {e}"))
        })?;
        let adam_func = module.load_function("curiosity_adam_step").map_err(|e| {
            MLError::ModelError(format!("curiosity_adam_step load: {e}"))
        })?;
        let adam_fused_func = module.load_function("curiosity_adam_step_fused").map_err(|e| {
            MLError::ModelError(format!("curiosity_adam_step_fused load: {e}"))
        })?;
        let fused_zero_fwd_bwd_adam_func = module
            .load_function("curiosity_fused_zero_fwd_bwd_adam")
            .map_err(|e| {
                MLError::ModelError(format!("curiosity_fused_zero_fwd_bwd_adam load: {e}"))
            })?;

        // ---- Allocate gradient buffers ----
        let grad_w1 = stream.alloc_zeros::<f32>(CUR_W1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc grad_w1: {e}"))
        })?;
        let grad_b1 = stream.alloc_zeros::<f32>(CUR_B1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc grad_b1: {e}"))
        })?;
        let grad_w2 = stream.alloc_zeros::<f32>(CUR_W2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc grad_w2: {e}"))
        })?;
        let grad_b2 = stream.alloc_zeros::<f32>(CUR_B2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc grad_b2: {e}"))
        })?;

        // ---- Allocate Adam first moment buffers ----
        let adam_m_w1 = stream.alloc_zeros::<f32>(CUR_W1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_m_w1: {e}"))
        })?;
        let adam_m_b1 = stream.alloc_zeros::<f32>(CUR_B1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_m_b1: {e}"))
        })?;
        let adam_m_w2 = stream.alloc_zeros::<f32>(CUR_W2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_m_w2: {e}"))
        })?;
        let adam_m_b2 = stream.alloc_zeros::<f32>(CUR_B2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_m_b2: {e}"))
        })?;

        // ---- Allocate Adam second moment buffers ----
        let adam_v_w1 = stream.alloc_zeros::<f32>(CUR_W1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_v_w1: {e}"))
        })?;
        let adam_v_b1 = stream.alloc_zeros::<f32>(CUR_B1_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_v_b1: {e}"))
        })?;
        let adam_v_w2 = stream.alloc_zeros::<f32>(CUR_W2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_v_w2: {e}"))
        })?;
        let adam_v_b2 = stream.alloc_zeros::<f32>(CUR_B2_LEN).map_err(|e| {
            MLError::ModelError(format!("alloc adam_v_b2: {e}"))
        })?;

        // ---- Allocate shifted next_states buffer ----
        let next_states_buf = stream
            .alloc_zeros::<f32>(max_samples * state_dim)
            .map_err(|e| {
                MLError::ModelError(format!("alloc next_states_buf: {e}"))
            })?;

        // ---- Allocate atomic block-arrival counter for fused kernel ----
        let block_counter = stream.alloc_zeros::<i32>(1).map_err(|e| {
            MLError::ModelError(format!("alloc block_counter: {e}"))
        })?;

        debug!(
            state_dim,
            max_samples,
            total_params = CUR_W1_LEN + CUR_B1_LEN + CUR_W2_LEN + CUR_B2_LEN,
            "GPU curiosity trainer initialized (Adam optimizer)"
        );

        Ok(Self {
            stream,
            shift_func,
            fwd_bwd_func,
            adam_func,
            adam_fused_func,
            fused_zero_fwd_bwd_adam_func,
            grad_w1,
            grad_b1,
            grad_w2,
            grad_b2,
            adam_m_w1,
            adam_m_b1,
            adam_m_w2,
            adam_m_b2,
            adam_v_w1,
            adam_v_b1,
            adam_v_w2,
            adam_v_b2,
            next_states_buf,
            block_counter,
            step: 0,
            state_dim,
            buf_capacity: max_samples,
        })
    }

    /// Train the curiosity forward model on GPU-resident experience data.
    ///
    /// Performs one training step (forward + backward + Adam update) entirely
    /// on GPU. The states buffer is shifted by one timestep to produce
    /// next_states -- episode boundary noise is negligible for this tiny
    /// auxiliary model.
    ///
    /// # Arguments
    /// * `weights` - Curiosity model weights to update in-place on GPU
    /// * `states` - State observations `[n_samples * state_dim]` on GPU
    /// * `actions` - Action indices `[n_samples]` on GPU (i32, 0-4 for DQN)
    /// * `n_samples` - Number of experience samples (will use n_samples-1 for training,
    ///   since the last sample has no valid next_state)
    ///
    /// # Errors
    /// Returns `MLError::ModelError` on kernel launch or buffer size mismatch.
    pub fn train_on_collector_buffers(
        &mut self,
        weights: &mut CuriosityWeightSet,
        states: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        n_samples: usize,
    ) -> Result<(), MLError> {
        // Need at least 2 samples (one for state, one for next_state via shift)
        if n_samples < 2 {
            return Ok(());
        }

        // Effective training samples: n_samples - 1 (last has no valid next_state)
        let n_train = n_samples - 1;

        if n_train > self.buf_capacity {
            return Err(MLError::ModelError(format!(
                "curiosity trainer: n_train={n_train} exceeds buf_capacity={}",
                self.buf_capacity
            )));
        }

        let sd = self.state_dim;
        let sd_i32 = sd as i32;
        let n_i32 = n_train as i32;

        self.stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("curiosity PRE-SHIFT sync FAILED: {e}")))?;

        // ---- Launch 1/2: Build shifted next_states buffer ----
        // next_states[i] = states[i + state_dim] (shift by one timestep)
        // Separate launch because grid dimensions differ from the fwd/bwd grid.
        let shift_total = n_train * sd;
        let shift_cfg = LaunchConfig {
            grid_dim: (((shift_total as u32) + 255) / 256, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            self.stream
                .launch_builder(&self.shift_func)
                .arg(states)
                .arg(&mut self.next_states_buf)
                .arg(&n_i32)
                .arg(&sd_i32)
                .launch(shift_cfg)
                .map_err(|e| {
                    MLError::ModelError(format!("curiosity_shift_states launch: {e}"))
                })?;
        }

        // Sync after shift kernel to catch crashes
        self.stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("curiosity shift kernel CRASHED: {e}")))?;

        // ---- Launch 2/2: Fused zero + forward/backward + Adam ----
        self.step += 1;
        let step = self.step;
        let _bs_i32 = n_train as i32;

        // Zero gradient buffers + block-arrival counter before fused kernel launch.
        // Gradients are accumulated via atomicAdd in the kernel — they MUST start at zero.
        // Previously this was done inside the kernel (Phase 1) but had an inter-block race.
        // memset_zeros is GPU-side cuMemsetD8Async, ordered on the same stream.
        self.stream.memset_zeros(&mut self.grad_w1)
            .map_err(|e| MLError::ModelError(format!("memset grad_w1: {e}")))?;
        self.stream.memset_zeros(&mut self.grad_b1)
            .map_err(|e| MLError::ModelError(format!("memset grad_b1: {e}")))?;
        self.stream.memset_zeros(&mut self.grad_w2)
            .map_err(|e| MLError::ModelError(format!("memset grad_w2: {e}")))?;
        self.stream.memset_zeros(&mut self.grad_b2)
            .map_err(|e| MLError::ModelError(format!("memset grad_b2: {e}")))?;
        self.stream.memset_zeros(&mut self.block_counter)
            .map_err(|e| MLError::ModelError(format!("memset block_counter: {e}")))?;

        // Step 2a: Forward + backward (separate kernel — no inter-block sync needed)
        let fwd_bwd_cfg = LaunchConfig {
            grid_dim: (((n_train as u32) + 255) / 256, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            self.stream
                .launch_builder(&self.fwd_bwd_func)
                .arg(states)
                .arg(actions)
                .arg(&self.next_states_buf)
                .arg(&weights.w1)
                .arg(&weights.b1)
                .arg(&weights.w2)
                .arg(&weights.b2)
                .arg(&mut self.grad_w1)
                .arg(&mut self.grad_b1)
                .arg(&mut self.grad_w2)
                .arg(&mut self.grad_b2)
                .arg(&n_i32)
                .arg(&sd_i32)
                .launch(fwd_bwd_cfg)
                .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd launch: {e}")))?;
        }
        self.stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd CRASHED: {e}")))?;

        // Step 2b: Adam optimizer update (4 separate launches, one per param group)
        launch_adam_step(
            &self.stream, &self.adam_func,
            &mut weights.w1, &self.grad_w1, CUR_W1_LEN,
            &mut self.adam_m_w1, &mut self.adam_v_w1,
            n_train, step,
        )?;
        launch_adam_step(
            &self.stream, &self.adam_func,
            &mut weights.b1, &self.grad_b1, CUR_B1_LEN,
            &mut self.adam_m_b1, &mut self.adam_v_b1,
            n_train, step,
        )?;
        launch_adam_step(
            &self.stream, &self.adam_func,
            &mut weights.w2, &self.grad_w2, CUR_W2_LEN,
            &mut self.adam_m_w2, &mut self.adam_v_w2,
            n_train, step,
        )?;
        launch_adam_step(
            &self.stream, &self.adam_func,
            &mut weights.b2, &self.grad_b2, CUR_B2_LEN,
            &mut self.adam_m_b2, &mut self.adam_v_b2,
            n_train, step,
        )?;

        debug!(step, n_train, "curiosity GPU training step complete (fused)");

        Ok(())
    }

    /// Get the current Adam step count.
    pub fn step_count(&self) -> i32 {
        self.step
    }
}
