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

//! GPU-accelerated Hindsight Experience Replay (HER).
//!
//! Operates on `GpuReplayBuffer` tensor storage via a fused CUDA kernel that
//! relabels experiences entirely on GPU. Produces relabeled experience batches
//! that are concatenated with normal PER samples before feeding to the training
//! kernel.
//!
//! ## Algorithm (Andrychowicz et al., 2017)
//!
//! For each experience to relabel:
//! 1. Copy original (state, next_state, action, done) from a PER-sampled source
//! 2. Read achieved goal from a donor experience (strategy-dependent)
//! 3. Replace goal portion (first `goal_dim` elements) in state and next_state
//! 4. Recompute sparse reward (+1 if goal achieved, -0.01 otherwise)
//!
//! ## Kernel launch
//!
//! `grid=(her_batch_size, 1, 1), block=(32, 1, 1)` -- one warp per sample.
//! Warp-cooperative coalesced memory copy for state vectors, lane 0 scalar
//! work for reward recomputation.
//!
//! ## Strategy
//!
//! Currently implements Random strategy (donor = random experience from buffer).
//! Final and Future strategies require episode boundary tracking on GPU and
//! will be added in a follow-up.

use std::sync::Arc;

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

use crate::MLError;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// GPU HER relabeling strategy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HerGpuStrategy {
    /// Donor = last experience in the same episode (requires episode tracking).
    Final,
    /// Donor = random future experience from same episode (requires episode tracking).
    Future,
    /// Donor = random experience from entire buffer (no episode tracking needed).
    Random,
}

/// Configuration for GPU HER.
#[derive(Debug, Clone)]
pub struct GpuHerConfig {
    /// Dimension of the goal portion at the front of the state vector.
    pub goal_dim: usize,
    /// Fraction of batch from HER relabeling (0.0 = disabled).
    pub her_ratio: f64,
    /// Relabeling strategy.
    pub strategy: HerGpuStrategy,
    /// L2 distance threshold for goal achievement (default: 0.01).
    pub goal_threshold: f32,
    /// Total training batch size (normal + HER).
    pub batch_size: usize,
    /// State dimension (tensor-core-aligned, e.g. 48 or 56).
    pub state_dim: usize,
}

impl Default for GpuHerConfig {
    fn default() -> Self {
        Self {
            goal_dim: 1,
            her_ratio: 0.0,
            strategy: HerGpuStrategy::Random,
            goal_threshold: 0.01,
            batch_size: 256,
            state_dim: 72,
        }
    }
}

impl GpuHerConfig {
    /// Compute the HER batch size (number of relabeled samples per batch).
    pub fn her_batch_size(&self) -> usize {
        (self.batch_size as f64 * self.her_ratio) as usize
    }

    /// Compute the normal (non-HER) batch size.
    pub fn normal_batch_size(&self) -> usize {
        self.batch_size.saturating_sub(self.her_batch_size())
    }
}

// ---------------------------------------------------------------------------
// GPU-resident relabeled batch
// ---------------------------------------------------------------------------

/// GPU-resident HER relabeled batch -- staging buffer slices ready for
/// concatenation with the normal PER batch.
#[allow(missing_debug_implementations)]
pub struct HerBatch {
    /// Relabeled states `[her_batch_size, state_dim]` -- f32 on device.
    pub states: CudaSlice<f32>,
    /// Relabeled next_states `[her_batch_size, state_dim]` -- f32 on device.
    pub next_states: CudaSlice<f32>,
    /// Actions (copied unchanged) `[her_batch_size]` -- i32 on device.
    pub actions: CudaSlice<i32>,
    /// Recomputed rewards `[her_batch_size]` -- f32 on device.
    pub rewards: CudaSlice<f32>,
    /// Dones (copied unchanged) `[her_batch_size]` -- f32 on device.
    pub dones: CudaSlice<f32>,
    /// Number of relabeled samples in this batch.
    pub batch_size: usize,
}

// ---------------------------------------------------------------------------
// GpuHer
// ---------------------------------------------------------------------------

/// GPU-accelerated Hindsight Experience Replay.
///
/// Owns pre-allocated staging buffers for relabeled experiences and the compiled
/// relabel kernel. Call [`relabel_batch`] to produce a [`HerBatch`] from the
/// GPU replay buffer's tensor storage.
#[allow(missing_debug_implementations)]
pub struct GpuHer {
    /// Configuration.
    pub config: GpuHerConfig,

    // Staging buffers for relabeled experiences (pre-allocated at max HER batch size)
    out_states: CudaSlice<f32>,
    out_next_states: CudaSlice<f32>,
    out_actions: CudaSlice<i32>,
    out_rewards: CudaSlice<f32>,
    out_dones: CudaSlice<f32>,

    // Index buffers (uploaded each step)
    source_indices: CudaSlice<i32>,
    /// GPU-resident donor indices -- written by `relabel_batch_with_strategy()`.
    /// Readable by the training loop for strategy-aware relabeling.
    pub(crate) donor_indices: CudaSlice<i32>,

    // Compiled kernels
    relabel_func: CudaFunction,
    /// GPU kernel: sample a future donor in the same episode (HER Future).
    future_donors_func: CudaFunction,
    /// GPU kernel: find the last transition in the same episode (HER Final).
    episode_end_func: CudaFunction,

    // Per-sample LCG RNG state for Future strategy donor selection.
    // Seeded once at construction; survives across relabel_batch_with_strategy calls.
    rng_states: CudaSlice<u32>,

    // Stream for kernel launches
    stream: Arc<CudaStream>,
}

impl GpuHer {
    /// Create a new GPU HER instance.
    ///
    /// Compiles the relabel kernel via NVRTC and pre-allocates all staging
    /// buffers at the maximum HER batch size.
    pub fn new(
        stream: Arc<CudaStream>,
        config: GpuHerConfig,
    ) -> Result<Self, MLError> {
        let her_batch = config.her_batch_size();
        if her_batch == 0 {
            return Err(MLError::ConfigError(
                "GpuHer: her_ratio results in zero HER samples".to_owned(),
            ));
        }

        // Compile the relabel kernel
        let relabel_func = compile_her_kernel(&stream, &config)?;

        // Compile HER episode boundary kernels (Future + Final strategies)
        let (future_donors_func, episode_end_func) = compile_her_episode_kernels(&stream)?;

        // Pre-allocate staging buffers
        let out_states = alloc_f32(&stream, her_batch * config.state_dim, "her_out_states")?;
        let out_next_states = alloc_f32(&stream, her_batch * config.state_dim, "her_out_next_states")?;
        let out_actions = alloc_i32(&stream, her_batch, "her_out_actions")?;
        let out_rewards = alloc_f32(&stream, her_batch, "her_out_rewards")?;
        let out_dones = alloc_f32(&stream, her_batch, "her_out_dones")?;

        // Index buffers
        let source_indices = alloc_i32(&stream, her_batch, "her_source_indices")?;
        let donor_indices = alloc_i32(&stream, her_batch, "her_donor_indices")?;

        // Per-sample RNG state for Future strategy LCG — seeded with index.
        // Uploaded once at construction; subsequent calls advance the state in-place on GPU.
        let rng_init: Vec<u32> = (0..her_batch as u32)
            .map(|i| i.wrapping_mul(2654435761u32).wrapping_add(1013904223u32))
            .collect();
        let mut rng_states = stream.alloc_zeros::<u32>(her_batch)
            .map_err(|e| MLError::ModelError(format!("GpuHer alloc rng_states: {e}")))?;
        stream.memcpy_htod(&rng_init, &mut rng_states)
            .map_err(|e| MLError::ModelError(format!("GpuHer rng_states upload: {e}")))?;

        let vram_bytes = (her_batch * config.state_dim * 2 + her_batch * 3) * 4
            + her_batch * 3 * 4; // output bufs + index bufs + rng_states
        info!(
            her_batch_size = her_batch,
            state_dim = config.state_dim,
            goal_dim = config.goal_dim,
            strategy = ?config.strategy,
            vram_kb = vram_bytes / 1024,
            "GpuHer initialized: relabel + episode kernels compiled, staging buffers allocated"
        );

        Ok(Self {
            config,
            out_states,
            out_next_states,
            out_actions,
            out_rewards,
            out_dones,
            source_indices,
            donor_indices,
            relabel_func,
            future_donors_func,
            episode_end_func,
            rng_states,
            stream,
        })
    }

    /// Generate a relabeled batch from GPU replay buffer storage.
    ///
    /// 1. Accepts pre-generated source indices (PER-weighted) and random donor
    ///    indices (for Random strategy).
    /// 2. Uploads index arrays to GPU.
    /// 3. Launches the relabel kernel.
    /// 4. Returns staging buffer references as a `HerBatch`.
    ///
    /// # Arguments
    ///
    /// * `buffer_states` - `[capacity, state_dim]` f32 on device (replay buffer storage)
    /// * `buffer_next_states` - `[capacity, state_dim]` f32 on device
    /// * `buffer_actions` - `[capacity]` i32 on device
    /// * `buffer_rewards` - `[capacity]` f32 on device
    /// * `buffer_dones` - `[capacity]` f32 on device
    /// * `source_idx_host` - PER-sampled source indices (CPU)
    /// * `donor_idx_host` - Donor indices (CPU, strategy-dependent)
    /// * `buffer_size` - Current replay buffer occupancy
    /// * `her_batch_size` - Number of experiences to relabel
    pub fn relabel_batch(
        &mut self,
        buffer_states: &CudaSlice<f32>,
        buffer_next_states: &CudaSlice<f32>,
        buffer_actions: &CudaSlice<i32>,
        buffer_rewards: &CudaSlice<f32>,
        buffer_dones: &CudaSlice<f32>,
        source_idx_host: &[i32],
        donor_idx_host: &[i32],
        buffer_size: usize,
        her_batch_size: usize,
    ) -> Result<HerBatch, MLError> {
        if her_batch_size == 0 {
            return Err(MLError::InvalidInput(
                "HER batch size must be > 0".to_owned(),
            ));
        }

        let max_her = self.config.her_batch_size();
        if her_batch_size > max_her {
            return Err(MLError::InvalidInput(format!(
                "HER batch size {her_batch_size} exceeds pre-allocated max {max_her}"
            )));
        }

        // Upload source and donor indices to GPU
        self.stream
            .memcpy_htod(&source_idx_host[..her_batch_size], &mut self.source_indices)
            .map_err(|e| MLError::ModelError(format!("HER source_indices upload: {e}")))?;
        self.stream
            .memcpy_htod(&donor_idx_host[..her_batch_size], &mut self.donor_indices)
            .map_err(|e| MLError::ModelError(format!("HER donor_indices upload: {e}")))?;

        // Launch relabel kernel: one warp (32 threads) per sample
        let launch_config = LaunchConfig {
            grid_dim: (her_batch_size as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };
        let buffer_size_i32 = buffer_size as i32;

        // Safety: all buffer pointers are valid GPU allocations on the same context.
        // source_indices and donor_indices contain valid buffer indices in [0, buffer_size).
        // Output buffers are pre-allocated with sufficient capacity.
        unsafe {
            self.stream
                .launch_builder(&self.relabel_func)
                .arg(buffer_states)
                .arg(buffer_next_states)
                .arg(buffer_actions)
                .arg(buffer_rewards)
                .arg(buffer_dones)
                .arg(&self.source_indices)
                .arg(&self.donor_indices)
                .arg(&buffer_size_i32)
                .arg(&mut self.out_states)
                .arg(&mut self.out_next_states)
                .arg(&mut self.out_actions)
                .arg(&mut self.out_rewards)
                .arg(&mut self.out_dones)
                .launch(launch_config)
                .map_err(|e| MLError::ModelError(format!("HER relabel kernel launch: {e}")))?;
        }

        // Return references to the staging buffers (data is on GPU, ready for concat)
        // Clone the CudaSlice handles — they point to the same device memory.
        // The caller will use these for D2D copies into the merged training batch.
        let states_slice = slice_clone_f32(&self.stream, &self.out_states, her_batch_size * self.config.state_dim)?;
        let next_states_slice = slice_clone_f32(&self.stream, &self.out_next_states, her_batch_size * self.config.state_dim)?;
        let actions_slice = slice_clone_i32(&self.stream, &self.out_actions, her_batch_size)?;
        let rewards_slice = slice_clone_f32(&self.stream, &self.out_rewards, her_batch_size)?;
        let dones_slice = slice_clone_f32(&self.stream, &self.out_dones, her_batch_size)?;

        Ok(HerBatch {
            states: states_slice,
            next_states: next_states_slice,
            actions: actions_slice,
            rewards: rewards_slice,
            dones: dones_slice,
            batch_size: her_batch_size,
        })
    }

    /// Generate random donor indices for the Random strategy.
    ///
    /// Each donor is a uniform random index in `[0, buffer_size)`.
    /// This is done on CPU since it is O(her_batch_size) with trivial
    /// per-element cost and avoids cuRAND setup.
    pub fn generate_random_donors(
        buffer_size: usize,
        her_batch_size: usize,
    ) -> Vec<i32> {
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let mut donors = Vec::with_capacity(her_batch_size);
        donors.resize_with(her_batch_size, || rng.gen_range(0..buffer_size as i32));
        donors
    }

    /// GPU-native donor selection for Future and Final strategies.
    ///
    /// Dispatches to the appropriate episode-aware kernel based on the configured
    /// strategy. Requires `episode_ids` from `GpuExperienceBatch` (filled by
    /// `fill_episode_ids_kernel`) so there is zero CPU work in the hot path.
    ///
    /// After this call, `self.donor_indices` contains the GPU-computed donors.
    /// The caller should then invoke `relabel_batch` with `source_idx_host` and
    /// a placeholder CPU donor slice (will be overwritten by the GPU donor buffer).
    ///
    /// # Arguments
    ///
    /// * `episode_ids` - `[capacity]` i32 on device — episode index per transition
    /// * `source_idx_host` - PER-sampled source indices (CPU, length `her_batch_size`)
    /// * `episode_length` - Timesteps per episode (L)
    /// * `buffer_size` - Total flat buffer capacity (N * L)
    /// * `her_batch_size` - Number of experiences to relabel
    ///
    /// # Errors
    ///
    /// Returns `MLError::ConfigError` when called with `HerGpuStrategy::Random`
    /// (use `generate_random_donors` instead).
    pub fn relabel_batch_with_strategy(
        &mut self,
        episode_ids: &CudaSlice<i32>,
        source_idx_host: &[i32],
        episode_length: usize,
        buffer_size: usize,
        her_batch_size: usize,
    ) -> Result<(), MLError> {
        match self.config.strategy {
            HerGpuStrategy::Random => {
                return Err(MLError::ConfigError(
                    "relabel_batch_with_strategy: use generate_random_donors for Random strategy".to_owned(),
                ));
            }
            HerGpuStrategy::Future | HerGpuStrategy::Final => {}
        }

        if her_batch_size == 0 {
            return Err(MLError::InvalidInput(
                "HER batch size must be > 0".to_owned(),
            ));
        }
        let max_her = self.config.her_batch_size();
        if her_batch_size > max_her {
            return Err(MLError::InvalidInput(format!(
                "HER batch size {her_batch_size} exceeds pre-allocated max {max_her}"
            )));
        }

        // Upload source indices to GPU
        self.stream
            .memcpy_htod(&source_idx_host[..her_batch_size], &mut self.source_indices)
            .map_err(|e| MLError::ModelError(format!("HER source_indices upload: {e}")))?;

        let l_i32 = episode_length as i32;
        let capacity_i32 = buffer_size as i32;
        let her_batch_i32 = her_batch_size as i32;
        let blocks = ((her_batch_size + 255) / 256) as u32;
        let launch_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        match self.config.strategy {
            HerGpuStrategy::Future => {
                // Safety: all buffers are valid GPU allocations on the same context.
                // source_indices contains valid indices in [0, capacity). rng_states
                // is pre-allocated at her_batch_size. donor_indices is pre-allocated.
                unsafe {
                    self.stream
                        .launch_builder(&self.future_donors_func)
                        .arg(episode_ids)
                        .arg(&self.source_indices)
                        .arg(&mut self.donor_indices)
                        .arg(&mut self.rng_states)
                        .arg(&l_i32)
                        .arg(&capacity_i32)
                        .arg(&her_batch_i32)
                        .launch(launch_cfg)
                        .map_err(|e| MLError::ModelError(format!("her_sample_future_donors: {e}")))?;
                }
            }
            HerGpuStrategy::Final => {
                // Safety: all buffers are valid GPU allocations on the same context.
                // source_indices contains valid indices in [0, capacity).
                // end_indices == donor_indices (same pre-allocated buffer, re-used).
                unsafe {
                    self.stream
                        .launch_builder(&self.episode_end_func)
                        .arg(episode_ids)
                        .arg(&self.source_indices)
                        .arg(&mut self.donor_indices)
                        .arg(&l_i32)
                        .arg(&capacity_i32)
                        .arg(&her_batch_i32)
                        .launch(launch_cfg)
                        .map_err(|e| MLError::ModelError(format!("her_find_episode_end: {e}")))?;
                }
            }
            HerGpuStrategy::Random => unreachable!(),
        }

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Kernel compilation
// ---------------------------------------------------------------------------

/// Compile the HER episode boundary kernels.
///
/// Returns `(future_donors_func, episode_end_func)`.
fn compile_her_episode_kernels(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    let kernel_src = include_str!("her_episode_kernel.cu");
    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context)
        .map_err(|e| MLError::ModelError(format!("her_episode_kernel compilation: {e}")))?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("her_episode_kernel module load: {e}"))
    })?;
    let future_donors = module.load_function("her_sample_future_donors").map_err(|e| {
        MLError::ModelError(format!("load her_sample_future_donors: {e}"))
    })?;
    let episode_end = module.load_function("her_find_episode_end").map_err(|e| {
        MLError::ModelError(format!("load her_find_episode_end: {e}"))
    })?;
    Ok((future_donors, episode_end))
}

/// Compile the HER relabel kernel with dimension defines.
fn compile_her_kernel(
    stream: &Arc<CudaStream>,
    config: &GpuHerConfig,
) -> Result<CudaFunction, MLError> {
    // Inject compile-time constants
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define GOAL_DIM {goal_dim}\n\
         #define GOAL_THRESHOLD {threshold:.6}f\n",
        state_dim = config.state_dim,
        goal_dim = config.goal_dim,
        threshold = config.goal_threshold,
    );

    let kernel_src = include_str!("her_relabel_kernel.cu");
    let full_source = format!("{dim_overrides}\n{kernel_src}");

    info!(
        state_dim = config.state_dim,
        goal_dim = config.goal_dim,
        goal_threshold = config.goal_threshold,
        "GpuHer: compiling HER relabel kernel"
    );

    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| {
            MLError::ModelError(format!("her_relabel_kernel compilation failed: {e}"))
        })?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("her_relabel module load: {e}"))
    })?;
    module.load_function("her_relabel_kernel").map_err(|e| {
        MLError::ModelError(format!("her_relabel_kernel function load: {e}"))
    })
}

// ---------------------------------------------------------------------------
// Allocation helpers
// ---------------------------------------------------------------------------

fn alloc_f32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuHer alloc {name} ({n} f32): {e}"))
    })
}

fn alloc_i32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<i32>, MLError> {
    stream.alloc_zeros::<i32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuHer alloc {name} ({n} i32): {e}"))
    })
}

/// Extract raw CUDA device pointer from a CudaSlice<f32>.
///
/// Same pattern as `raw_device_ptr` in `gpu_dqn_trainer.rs` -- wraps the
/// `SyncOnDrop` guard in `ManuallyDrop` to skip read event recording (safe
/// on same stream, same context).
fn raw_ptr_f32(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Extract raw CUDA device pointer from a CudaSlice<i32>.
fn raw_ptr_i32(slice: &CudaSlice<i32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Clone a sub-range of a CudaSlice<f32> via async D2D memcpy.
fn slice_clone_f32(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<f32>,
    n: usize,
) -> Result<CudaSlice<f32>, MLError> {
    let dst = alloc_f32(stream, n, "her_clone_f32")?;
    let num_bytes = n * std::mem::size_of::<f32>();
    let src_ptr = raw_ptr_f32(src, stream);
    let dst_ptr = raw_ptr_f32(&dst, stream);
    // Safety: both src and dst are valid device memory on the same context.
    // n elements * 4 bytes does not exceed either allocation.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
        ).map_err(|e| MLError::ModelError(format!("HER D2D clone f32: {e}")))?;
    }
    Ok(dst)
}

/// Clone a sub-range of a CudaSlice<i32> via async D2D memcpy.
fn slice_clone_i32(
    stream: &Arc<CudaStream>,
    src: &CudaSlice<i32>,
    n: usize,
) -> Result<CudaSlice<i32>, MLError> {
    let dst = alloc_i32(stream, n, "her_clone_i32")?;
    let num_bytes = n * std::mem::size_of::<i32>();
    let src_ptr = raw_ptr_i32(src, stream);
    let dst_ptr = raw_ptr_i32(&dst, stream);
    // Safety: both src and dst are valid device memory on the same context.
    unsafe {
        cudarc::driver::result::memcpy_dtod_async(
            dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
        ).map_err(|e| MLError::ModelError(format!("HER D2D clone i32: {e}")))?;
    }
    Ok(dst)
}

// ---------------------------------------------------------------------------
// Strategy parsing
// ---------------------------------------------------------------------------

/// Parse a HER strategy string into the GPU enum variant.
pub fn parse_her_strategy(s: &str) -> HerGpuStrategy {
    match s.to_lowercase().as_str() {
        "final" => HerGpuStrategy::Final,
        "future" => HerGpuStrategy::Future,
        _ => HerGpuStrategy::Random,
    }
}
