use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream};
use crate::heads::HIDDEN_DIM;

pub const N_STEP_MAX: usize = 16;
pub const SCALARS_PER_TRANSITION: usize = 7;

pub struct GpuReplayBuffer {
    pub capacity: usize,
    pub b_size: usize,

    pub h_t_d: CudaSlice<f32>,
    pub h_tp1_d: CudaSlice<f32>,
    pub scalars_d: CudaSlice<f32>,
    pub write_head_d: CudaSlice<u32>,
    pub replay_len_d: CudaSlice<u32>,
    pub max_priority_d: CudaSlice<f32>,
    pub priority_tree_d: CudaSlice<f32>,

    pub nstep_scalars_d: CudaSlice<f32>,
    pub nstep_h_t_d: CudaSlice<f32>,
    pub nstep_write_idx_d: CudaSlice<u32>,
    pub nstep_count_d: CudaSlice<u32>,

    pub flush_flags_d: CudaSlice<i32>,
    pub write_offsets_d: CudaSlice<i32>,

    pub sample_indices_d: CudaSlice<u32>,
    pub sample_prng_d: CudaSlice<u32>,
}

impl GpuReplayBuffer {
    pub fn new(stream: &Arc<CudaStream>, capacity: usize, b_size: usize) -> Result<Self> {
        Ok(Self {
            capacity,
            b_size,
            h_t_d: stream.alloc_zeros(capacity * HIDDEN_DIM).context("gpu_replay h_t")?,
            h_tp1_d: stream.alloc_zeros(capacity * HIDDEN_DIM).context("gpu_replay h_tp1")?,
            scalars_d: stream.alloc_zeros(capacity * SCALARS_PER_TRANSITION).context("gpu_replay scalars")?,
            write_head_d: stream.alloc_zeros(1).context("gpu_replay write_head")?,
            replay_len_d: stream.alloc_zeros(1).context("gpu_replay len")?,
            max_priority_d: stream.alloc_zeros(1).context("gpu_replay max_priority")?,
            priority_tree_d: stream.alloc_zeros(2 * capacity).context("gpu_replay tree")?,
            nstep_scalars_d: stream.alloc_zeros(b_size * N_STEP_MAX * 3).context("gpu_replay nstep_scalars")?,
            nstep_h_t_d: stream.alloc_zeros(b_size * N_STEP_MAX * HIDDEN_DIM).context("gpu_replay nstep_h_t")?,
            nstep_write_idx_d: stream.alloc_zeros(b_size).context("gpu_replay nstep_write_idx")?,
            nstep_count_d: stream.alloc_zeros(b_size).context("gpu_replay nstep_count")?,
            flush_flags_d: stream.alloc_zeros(b_size).context("gpu_replay flush_flags")?,
            write_offsets_d: stream.alloc_zeros(b_size + 2).context("gpu_replay write_offsets")?,
            sample_indices_d: stream.alloc_zeros(b_size).context("gpu_replay sample_indices")?,
            sample_prng_d: stream.alloc_zeros(b_size).context("gpu_replay sample_prng")?,
        })
    }
}
