use super::*;

impl GpuContext {
    pub fn ion_pair_correlations(
        &self,
        com: &[f32],
        box_l: &[f32],
        cat_indices: &[u32],
        ani_indices: &[u32],
        n_groups: usize,
        n_frames: usize,
        max_cluster: usize,
        rclust_cat: f32,
        rclust_ani: f32,
    ) -> TrajResult<(Vec<u32>, Vec<u32>, Vec<u32>, Vec<u32>)> {
        let n_cat = cat_indices.len();
        let n_ani = ani_indices.len();
        if n_frames == 0 || n_cat == 0 || n_ani == 0 {
            return Ok((Vec::new(), Vec::new(), Vec::new(), Vec::new()));
        }
        let stream = &self.inner.stream;
        let com_dev = stream.clone_htod(com).map_err(map_driver_err)?;
        let box_dev = stream.clone_htod(box_l).map_err(map_driver_err)?;
        let cat_dev = stream.clone_htod(cat_indices).map_err(map_driver_err)?;
        let ani_dev = stream.clone_htod(ani_indices).map_err(map_driver_err)?;
        let mut idx_pair_cat = stream
            .alloc_zeros::<u32>(n_frames * n_cat)
            .map_err(map_driver_err)?;
        let mut idx_cluster_cat = stream
            .alloc_zeros::<u32>(n_frames * n_cat * max_cluster)
            .map_err(map_driver_err)?;
        let mut idx_pair_ani = stream
            .alloc_zeros::<u32>(n_frames * n_ani)
            .map_err(map_driver_err)?;
        let mut idx_cluster_ani = stream
            .alloc_zeros::<u32>(n_frames * n_ani * max_cluster)
            .map_err(map_driver_err)?;

        let n_cat_i32 = n_cat as i32;
        let n_ani_i32 = n_ani as i32;
        let n_groups_i32 = n_groups as i32;
        let n_frames_i32 = n_frames as i32;
        let max_cluster_i32 = max_cluster as i32;

        let block = 256u32;
        let cfg_cat = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_cat, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.ion_pair_frame_cat);
            builder.arg(&com_dev);
            builder.arg(&box_dev);
            builder.arg(&cat_dev);
            builder.arg(&ani_dev);
            builder.arg(&n_cat_i32);
            builder.arg(&n_ani_i32);
            builder.arg(&n_groups_i32);
            builder.arg(&n_frames_i32);
            builder.arg(&max_cluster_i32);
            builder.arg(&rclust_cat);
            builder.arg(&mut idx_pair_cat);
            builder.arg(&mut idx_cluster_cat);
            builder.launch(cfg_cat).map_err(map_driver_err)?;
        }

        let cfg_ani = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_ani, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.ion_pair_frame_ani);
            builder.arg(&com_dev);
            builder.arg(&box_dev);
            builder.arg(&cat_dev);
            builder.arg(&ani_dev);
            builder.arg(&n_cat_i32);
            builder.arg(&n_ani_i32);
            builder.arg(&n_groups_i32);
            builder.arg(&n_frames_i32);
            builder.arg(&max_cluster_i32);
            builder.arg(&rclust_ani);
            builder.arg(&mut idx_pair_ani);
            builder.arg(&mut idx_cluster_ani);
            builder.launch(cfg_ani).map_err(map_driver_err)?;
        }

        let mut ip_cat_dev = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;
        let mut cp_cat_dev = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;
        let mut ip_ani_dev = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;
        let mut cp_ani_dev = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;

        let cfg_corr_cat = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames, block), n_frames as u32, n_cat as u32),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.ion_pair_corr_cat);
            builder.arg(&idx_pair_cat);
            builder.arg(&idx_cluster_cat);
            builder.arg(&n_frames_i32);
            builder.arg(&n_cat_i32);
            builder.arg(&max_cluster_i32);
            builder.arg(&mut ip_cat_dev);
            builder.arg(&mut cp_cat_dev);
            builder.launch(cfg_corr_cat).map_err(map_driver_err)?;
        }

        let cfg_corr_ani = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames, block), n_frames as u32, n_ani as u32),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.ion_pair_corr_ani);
            builder.arg(&idx_pair_ani);
            builder.arg(&idx_cluster_ani);
            builder.arg(&n_frames_i32);
            builder.arg(&n_ani_i32);
            builder.arg(&max_cluster_i32);
            builder.arg(&mut ip_ani_dev);
            builder.arg(&mut cp_ani_dev);
            builder.launch(cfg_corr_ani).map_err(map_driver_err)?;
        }

        let mut ip_cat = vec![0u32; n_frames];
        let mut cp_cat = vec![0u32; n_frames];
        let mut ip_ani = vec![0u32; n_frames];
        let mut cp_ani = vec![0u32; n_frames];
        stream.memcpy_dtoh(&ip_cat_dev, &mut ip_cat).map_err(map_driver_err)?;
        stream.memcpy_dtoh(&cp_cat_dev, &mut cp_cat).map_err(map_driver_err)?;
        stream.memcpy_dtoh(&ip_ani_dev, &mut ip_ani).map_err(map_driver_err)?;
        stream.memcpy_dtoh(&cp_ani_dev, &mut cp_ani).map_err(map_driver_err)?;
        Ok((ip_cat, ip_ani, cp_cat, cp_ani))
    }


    pub fn read_counts(&self, counts: &GpuCounts) -> TrajResult<Vec<u64>> {
        let mut host = vec![0u64; counts.inner.len()];
        self.inner
            .stream
            .memcpy_dtoh(&counts.inner, &mut host)
            .map_err(map_driver_err)?;
        Ok(host)
    }


    pub fn alloc_counts(&self, bins: usize) -> TrajResult<GpuCounts> {
        let inner = self
            .inner
            .stream
            .alloc_zeros::<u64>(bins)
            .map_err(map_driver_err)?;
        Ok(GpuCounts { inner })
    }


    pub fn reset_counts(&self, counts: &mut GpuCounts) -> TrajResult<()> {
        self.inner
            .stream
            .memset_zeros(&mut counts.inner)
            .map_err(map_driver_err)
    }


    pub fn alloc_counts_u32(&self, len: usize) -> TrajResult<GpuCountsU32> {
        let inner = self
            .inner
            .stream
            .alloc_zeros::<u32>(len)
            .map_err(map_driver_err)?;
        Ok(GpuCountsU32 { inner })
    }


    pub fn reset_counts_u32(&self, counts: &mut GpuCountsU32) -> TrajResult<()> {
        self.inner
            .stream
            .memset_zeros(&mut counts.inner)
            .map_err(map_driver_err)
    }


    pub fn read_counts_u32(&self, counts: &GpuCountsU32) -> TrajResult<Vec<u32>> {
        let mut host = vec![0u32; counts.inner.len()];
        self.inner
            .stream
            .memcpy_dtoh(&counts.inner, &mut host)
            .map_err(map_driver_err)?;
        Ok(host)
    }

}
