use std::sync::Arc;

use cudarc::driver::{
    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceRepr, LaunchConfig,
    PushKernelArg, ValidAsZeroBits,
};
use cudarc::nvrtc::compile_ptx;

use traj_core::error::{TrajError, TrajResult};
use traj_kernels::KERNELS_SRC;

#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct Float4 {
    pub x: f32,
    pub y: f32,
    pub z: f32,
    pub w: f32,
}

unsafe impl DeviceRepr for Float4 {}
unsafe impl ValidAsZeroBits for Float4 {}

pub fn convert_coords(coords: &[[f32; 4]]) -> Vec<Float4> {
    coords
        .iter()
        .map(|c| Float4 {
            x: c[0],
            y: c[1],
            z: c[2],
            w: c[3],
        })
        .collect()
}

#[derive(Clone)]
pub struct GpuContext {
    inner: Arc<GpuContextInner>,
}

struct GpuContextInner {
    stream: Arc<CudaStream>,
    #[allow(dead_code)]
    module: Arc<CudaModule>,
    kernels: Kernels,
}

struct Kernels {
    rg_accum: Arc<CudaFunction>,
    rg_sumsq: Arc<CudaFunction>,
    rg_finalize: Arc<CudaFunction>,
    msd_accum: Arc<CudaFunction>,
    msd_finalize: Arc<CudaFunction>,
    rmsd_centroid: Arc<CudaFunction>,
    rmsd_cov: Arc<CudaFunction>,
    rmsd_raw_accum: Arc<CudaFunction>,
    rmsd_finalize: Arc<CudaFunction>,
    rdf_hist: Arc<CudaFunction>,
    polymer_end_to_end: Arc<CudaFunction>,
    polymer_contour_length: Arc<CudaFunction>,
    polymer_chain_rg: Arc<CudaFunction>,
    polymer_bond_hist: Arc<CudaFunction>,
    polymer_angle_hist: Arc<CudaFunction>,
    group_com_accum: Arc<CudaFunction>,
    group_com_finalize: Arc<CudaFunction>,
    group_dipole_accum: Arc<CudaFunction>,
    group_dipole_finalize: Arc<CudaFunction>,
    group_ke_accum: Arc<CudaFunction>,
    orientation_plane: Arc<CudaFunction>,
    orientation_vector: Arc<CudaFunction>,
    water_count: Arc<CudaFunction>,
    hbond_count: Arc<CudaFunction>,
    hbond_count_angle: Arc<CudaFunction>,
    msd_time_lag: Arc<CudaFunction>,
    rotacf_time_lag: Arc<CudaFunction>,
    conductivity_total: Arc<CudaFunction>,
    conductivity_transference: Arc<CudaFunction>,
    ion_pair_frame_cat: Arc<CudaFunction>,
    ion_pair_frame_ani: Arc<CudaFunction>,
    ion_pair_corr_cat: Arc<CudaFunction>,
    ion_pair_corr_ani: Arc<CudaFunction>,
}

impl Kernels {
    fn load(module: &Arc<CudaModule>) -> TrajResult<Self> {
        let load = |name: &str| -> TrajResult<Arc<CudaFunction>> {
            module
                .load_function(name)
                .map_err(map_driver_err)
                .map(Arc::new)
        };
        Ok(Self {
            rg_accum: load("rg_accum")?,
            rg_sumsq: load("rg_sumsq")?,
            rg_finalize: load("rg_finalize")?,
            msd_accum: load("msd_accum")?,
            msd_finalize: load("msd_finalize")?,
            rmsd_centroid: load("rmsd_centroid")?,
            rmsd_cov: load("rmsd_cov")?,
            rmsd_raw_accum: load("rmsd_raw_accum")?,
            rmsd_finalize: load("rmsd_finalize")?,
            rdf_hist: load("rdf_hist")?,
            polymer_end_to_end: load("polymer_end_to_end")?,
            polymer_contour_length: load("polymer_contour_length")?,
            polymer_chain_rg: load("polymer_chain_rg")?,
            polymer_bond_hist: load("polymer_bond_hist")?,
            polymer_angle_hist: load("polymer_angle_hist")?,
            group_com_accum: load("group_com_accum")?,
            group_com_finalize: load("group_com_finalize")?,
            group_dipole_accum: load("group_dipole_accum")?,
            group_dipole_finalize: load("group_dipole_finalize")?,
            group_ke_accum: load("group_ke_accum")?,
            orientation_plane: load("orientation_plane")?,
            orientation_vector: load("orientation_vector")?,
            water_count: load("water_count")?,
            hbond_count: load("hbond_count")?,
            hbond_count_angle: load("hbond_count_angle")?,
            msd_time_lag: load("msd_time_lag")?,
            rotacf_time_lag: load("rotacf_time_lag")?,
            conductivity_total: load("conductivity_total")?,
            conductivity_transference: load("conductivity_transference")?,
            ion_pair_frame_cat: load("ion_pair_frame_cat")?,
            ion_pair_frame_ani: load("ion_pair_frame_ani")?,
            ion_pair_corr_cat: load("ion_pair_corr_cat")?,
            ion_pair_corr_ani: load("ion_pair_corr_ani")?,
        })
    }
}

impl GpuContext {
    pub fn new(device: usize) -> TrajResult<Self> {
        let ctx = CudaContext::new(device).map_err(map_driver_err)?;
        let stream = ctx.default_stream();
        let ptx = compile_ptx(KERNELS_SRC).map_err(map_compile_err)?;
        let module = ctx.load_module(ptx).map_err(map_driver_err)?;
        let kernels = Kernels::load(&module)?;
        Ok(Self {
            inner: Arc::new(GpuContextInner {
                stream,
                module,
                kernels,
            }),
        })
    }

    pub fn selection(&self, indices: &[u32], masses: Option<&[f32]>) -> TrajResult<GpuSelection> {
        let sel = self.inner.stream.clone_htod(indices).map_err(map_driver_err)?;
        let masses_host: Vec<f32> = match masses {
            Some(values) => values.to_vec(),
            None => vec![1.0; indices.len()],
        };
        let masses_dev = self
            .inner
            .stream
            .clone_htod(&masses_host)
            .map_err(map_driver_err)?;
        Ok(GpuSelection {
            sel,
            masses: masses_dev,
            n_sel: indices.len(),
        })
    }

    pub fn upload_f32(&self, data: &[f32]) -> TrajResult<GpuBufferF32> {
        let buf = self.inner.stream.clone_htod(data).map_err(map_driver_err)?;
        Ok(GpuBufferF32 { inner: buf })
    }

    pub fn upload_u32(&self, data: &[u32]) -> TrajResult<GpuBufferU32> {
        let buf = self.inner.stream.clone_htod(data).map_err(map_driver_err)?;
        Ok(GpuBufferU32 { inner: buf })
    }

    pub fn upload_coords(&self, coords: &[Float4]) -> TrajResult<GpuCoords> {
        let buf = self.inner.stream.clone_htod(coords).map_err(map_driver_err)?;
        Ok(GpuCoords { inner: buf })
    }

    pub fn groups(
        &self,
        offsets: &[u32],
        indices: &[u32],
        max_len: usize,
    ) -> TrajResult<GpuGroups> {
        if offsets.len() < 2 {
            return Err(TrajError::Mismatch(
                "group offsets must include at least one group".into(),
            ));
        }
        let n_groups = offsets.len() - 1;
        let stream = &self.inner.stream;
        let offsets_dev = stream.clone_htod(offsets).map_err(map_driver_err)?;
        let indices_dev = stream.clone_htod(indices).map_err(map_driver_err)?;
        Ok(GpuGroups {
            offsets: offsets_dev,
            indices: indices_dev,
            n_groups,
            max_len,
        })
    }

    pub fn anchors(&self, anchors: &[u32]) -> TrajResult<GpuAnchors> {
        if anchors.len() % 3 != 0 {
            return Err(TrajError::Mismatch(
                "anchors length must be a multiple of 3".into(),
            ));
        }
        let n_groups = anchors.len() / 3;
        let stream = &self.inner.stream;
        let anchors_dev = stream.clone_htod(anchors).map_err(map_driver_err)?;
        Ok(GpuAnchors {
            anchors: anchors_dev,
            n_groups,
        })
    }

    pub fn reference(&self, coords: &[Float4]) -> TrajResult<GpuReference> {
        let coords_dev = self.inner.stream.clone_htod(coords).map_err(map_driver_err)?;
        Ok(GpuReference {
            coords: coords_dev,
            n_sel: coords.len(),
        })
    }

    pub fn polymer_data(
        &self,
        chain_offsets: &[u32],
        chain_indices: &[u32],
        bond_pairs: Option<&[u32]>,
        angle_triplets: Option<&[u32]>,
    ) -> TrajResult<GpuPolymer> {
        if chain_offsets.len() < 2 {
            return Err(TrajError::Mismatch(
                "polymer chain offsets must include at least one chain".into(),
            ));
        }
        let n_chains = chain_offsets.len() - 1;
        let last = *chain_offsets.last().unwrap() as usize;
        if last > chain_indices.len() {
            return Err(TrajError::Mismatch(
                "polymer chain offsets exceed index buffer".into(),
            ));
        }
        let stream = &self.inner.stream;
        let offsets_dev = stream
            .clone_htod(chain_offsets)
            .map_err(map_driver_err)?;
        let indices_dev = stream
            .clone_htod(chain_indices)
            .map_err(map_driver_err)?;
        let (bond_dev, n_bonds) = match bond_pairs {
            Some(pairs) => {
                if pairs.len() % 2 != 0 {
                    return Err(TrajError::Mismatch(
                        "bond_pairs length must be even".into(),
                    ));
                }
                let n_bonds = pairs.len() / 2;
                let dev = stream.clone_htod(pairs).map_err(map_driver_err)?;
                (Some(dev), n_bonds)
            }
            None => (None, 0),
        };
        let (angle_dev, n_angles) = match angle_triplets {
            Some(trips) => {
                if trips.len() % 3 != 0 {
                    return Err(TrajError::Mismatch(
                        "angle_triplets length must be a multiple of 3".into(),
                    ));
                }
                let n_angles = trips.len() / 3;
                let dev = stream.clone_htod(trips).map_err(map_driver_err)?;
                (Some(dev), n_angles)
            }
            None => (None, 0),
        };
        Ok(GpuPolymer {
            chain_offsets: offsets_dev,
            chain_indices: indices_dev,
            n_chains,
            bond_pairs: bond_dev,
            n_bonds,
            angle_triplets: angle_dev,
            n_angles,
        })
    }

    pub fn rg(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        sel: &GpuSelection,
        mass_weighted: bool,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut sum_x = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut sum_y = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut sum_z = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut mass_sum = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut sum_sq = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut out = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;

        let n_sel = sel.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;
        let mass_flag = if mass_weighted { 1i32 } else { 0i32 };
        let masses_dev = &sel.masses;

        let block = 256u32;
        let grid_x = ceil_div(sel.n_sel, block);
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (grid_x, n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rg_accum);
            builder.arg(&coords_dev);
            builder.arg(&sel.sel);
            builder.arg(masses_dev);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&mass_flag);
            builder.arg(&mut sum_x);
            builder.arg(&mut sum_y);
            builder.arg(&mut sum_z);
            builder.arg(&mut mass_sum);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rg_sumsq);
            builder.arg(&coords_dev);
            builder.arg(&sel.sel);
            builder.arg(masses_dev);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&mass_flag);
            builder.arg(&sum_x);
            builder.arg(&sum_y);
            builder.arg(&sum_z);
            builder.arg(&mass_sum);
            builder.arg(&mut sum_sq);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let cfg_out = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames as usize, block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rg_finalize);
            builder.arg(&mass_sum);
            builder.arg(&sum_sq);
            builder.arg(&n_frames_i32);
            builder.arg(&mut out);
            builder.launch(cfg_out).map_err(map_driver_err)?;
        }

        let mut host = vec![0.0f32; n_frames as usize];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn msd(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        sel: &GpuSelection,
        origin: &GpuReference,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut sum_sq = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut out = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;

        let n_sel = sel.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;

        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(sel.n_sel, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.msd_accum);
            builder.arg(&coords_dev);
            builder.arg(&origin.coords);
            builder.arg(&sel.sel);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&mut sum_sq);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let cfg_out = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames as usize, block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.msd_finalize);
            builder.arg(&sum_sq);
            builder.arg(&n_frames_i32);
            builder.arg(&n_sel);
            builder.arg(&mut out);
            builder.launch(cfg_out).map_err(map_driver_err)?;
        }

        let mut host = vec![0.0f32; n_frames as usize];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn rmsd_covariance(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        sel: &GpuSelection,
        reference: &GpuReference,
    ) -> TrajResult<RmsdCovariance> {
        if n_frames == 0 {
            return Ok(RmsdCovariance {
                cov: Vec::new(),
                sum_x2: Vec::new(),
                sum_y2: Vec::new(),
            });
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut sum_x = stream
            .alloc_zeros::<f32>(n_frames * 3)
            .map_err(map_driver_err)?;
        let mut sum_y = stream
            .alloc_zeros::<f32>(n_frames * 3)
            .map_err(map_driver_err)?;
        let mut cov = stream
            .alloc_zeros::<f32>(n_frames * 9)
            .map_err(map_driver_err)?;
        let mut sum_x2 = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut sum_y2 = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;

        let n_sel = sel.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;

        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(sel.n_sel, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rmsd_centroid);
            builder.arg(&coords_dev);
            builder.arg(&reference.coords);
            builder.arg(&sel.sel);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&mut sum_x);
            builder.arg(&mut sum_y);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rmsd_cov);
            builder.arg(&coords_dev);
            builder.arg(&reference.coords);
            builder.arg(&sel.sel);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&sum_x);
            builder.arg(&sum_y);
            builder.arg(&mut cov);
            builder.arg(&mut sum_x2);
            builder.arg(&mut sum_y2);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let mut cov_host = vec![0.0f32; n_frames as usize * 9];
        let mut sum_x2_host = vec![0.0f32; n_frames as usize];
        let mut sum_y2_host = vec![0.0f32; n_frames as usize];
        stream.memcpy_dtoh(&cov, &mut cov_host).map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&sum_x2, &mut sum_x2_host)
            .map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&sum_y2, &mut sum_y2_host)
            .map_err(map_driver_err)?;

        let mut cov_frames = Vec::with_capacity(n_frames as usize);
        for frame in 0..(n_frames as usize) {
            let mut mat = [0.0f32; 9];
            let offset = frame * 9;
            mat.copy_from_slice(&cov_host[offset..offset + 9]);
            cov_frames.push(mat);
        }

        Ok(RmsdCovariance {
            cov: cov_frames,
            sum_x2: sum_x2_host,
            sum_y2: sum_y2_host,
        })
    }

    pub fn rmsd_raw(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        sel: &GpuSelection,
        reference: &GpuReference,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut sum_sq = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;
        let mut out = stream.alloc_zeros::<f32>(n_frames).map_err(map_driver_err)?;

        let n_sel = sel.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;

        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(sel.n_sel, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rmsd_raw_accum);
            builder.arg(&coords_dev);
            builder.arg(&reference.coords);
            builder.arg(&sel.sel);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&mut sum_sq);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let cfg_out = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames as usize, block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rmsd_finalize);
            builder.arg(&sum_sq);
            builder.arg(&n_frames_i32);
            builder.arg(&n_sel);
            builder.arg(&mut out);
            builder.launch(cfg_out).map_err(map_driver_err)?;
        }

        let mut host = vec![0.0f32; n_frames as usize];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn rdf_accum(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        sel_a: &GpuSelection,
        sel_b: &GpuSelection,
        r_max: f32,
        bins: usize,
        pbc: bool,
        box_l: Option<&[f32]>,
        same_sel: bool,
        counts: &mut GpuCounts,
    ) -> TrajResult<()> {
        if n_frames == 0 || bins == 0 {
            return Ok(());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let bin_width = r_max / bins as f32;
        let pbc_flag = if pbc { 1i32 } else { 0i32 };
        let same_flag = if same_sel { 1i32 } else { 0i32 };
        let n_a = sel_a.n_sel as i32;
        let n_b = sel_b.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;
        let box_dev = if pbc {
            let box_len =
                box_l.ok_or_else(|| TrajError::Mismatch("missing box lengths".into()))?;
            Some(stream.clone_htod(box_len).map_err(map_driver_err)?)
        } else {
            None
        };

        let total_pairs = sel_a.n_sel * sel_b.n_sel;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(total_pairs, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        let dummy_box = if !pbc {
            Some(stream.alloc_zeros::<f32>(1).map_err(map_driver_err)?)
        } else {
            None
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rdf_hist);
            builder.arg(&coords_dev);
            builder.arg(&sel_a.sel);
            builder.arg(&sel_b.sel);
            builder.arg(&n_a);
            builder.arg(&n_b);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&r_max);
            builder.arg(&bin_width);
            builder.arg(&pbc_flag);
            if let Some(box_dev) = &box_dev {
                builder.arg(box_dev);
            } else if let Some(dummy) = &dummy_box {
                builder.arg(dummy);
            }
            builder.arg(&mut counts.inner);
            builder.arg(&same_flag);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        Ok(())
    }

    pub fn polymer_end_to_end(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        polymer: &GpuPolymer,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 || polymer.n_chains == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut out = stream
            .alloc_zeros::<f32>(n_frames * polymer.n_chains)
            .map_err(map_driver_err)?;
        let n_chains = polymer.n_chains as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(polymer.n_chains, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.polymer_end_to_end);
            builder.arg(&coords_dev);
            builder.arg(&polymer.chain_offsets);
            builder.arg(&polymer.chain_indices);
            builder.arg(&n_chains);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&mut out);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![0.0f32; n_frames as usize * polymer.n_chains];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn polymer_contour_length(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        polymer: &GpuPolymer,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 || polymer.n_chains == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut out = stream
            .alloc_zeros::<f32>(n_frames * polymer.n_chains)
            .map_err(map_driver_err)?;
        let n_chains = polymer.n_chains as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(polymer.n_chains, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.polymer_contour_length);
            builder.arg(&coords_dev);
            builder.arg(&polymer.chain_offsets);
            builder.arg(&polymer.chain_indices);
            builder.arg(&n_chains);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&mut out);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![0.0f32; n_frames as usize * polymer.n_chains];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn polymer_chain_rg(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        polymer: &GpuPolymer,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 || polymer.n_chains == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut out = stream
            .alloc_zeros::<f32>(n_frames * polymer.n_chains)
            .map_err(map_driver_err)?;
        let n_chains = polymer.n_chains as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(polymer.n_chains, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.polymer_chain_rg);
            builder.arg(&coords_dev);
            builder.arg(&polymer.chain_offsets);
            builder.arg(&polymer.chain_indices);
            builder.arg(&n_chains);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&mut out);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![0.0f32; n_frames as usize * polymer.n_chains];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn polymer_bond_hist(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        polymer: &GpuPolymer,
        r_max: f32,
        bins: usize,
        counts: &mut GpuCounts,
    ) -> TrajResult<()> {
        if n_frames == 0 || bins == 0 || polymer.n_bonds == 0 {
            return Ok(());
        }
        let bond_pairs = polymer
            .bond_pairs
            .as_ref()
            .ok_or_else(|| TrajError::Mismatch("missing bond pairs on GPU".into()))?;
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let bin_width = r_max / bins as f32;
        let n_bonds = polymer.n_bonds as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let bins_i = bins as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(polymer.n_bonds, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.polymer_bond_hist);
            builder.arg(&coords_dev);
            builder.arg(bond_pairs);
            builder.arg(&n_bonds);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&r_max);
            builder.arg(&bin_width);
            builder.arg(&bins_i);
            builder.arg(&mut counts.inner);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        Ok(())
    }

    pub fn polymer_angle_hist(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        polymer: &GpuPolymer,
        max_angle: f32,
        bins: usize,
        degrees: bool,
        counts: &mut GpuCounts,
    ) -> TrajResult<()> {
        if n_frames == 0 || bins == 0 || polymer.n_angles == 0 {
            return Ok(());
        }
        let angle_triplets = polymer
            .angle_triplets
            .as_ref()
            .ok_or_else(|| TrajError::Mismatch("missing angle triplets on GPU".into()))?;
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let bin_width = max_angle / bins as f32;
        let n_angles = polymer.n_angles as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let bins_i = bins as i32;
        let deg_flag = if degrees { 1i32 } else { 0i32 };
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(polymer.n_angles, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.polymer_angle_hist);
            builder.arg(&coords_dev);
            builder.arg(angle_triplets);
            builder.arg(&n_angles);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&max_angle);
            builder.arg(&bin_width);
            builder.arg(&bins_i);
            builder.arg(&deg_flag);
            builder.arg(&mut counts.inner);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        Ok(())
    }

    pub fn group_com(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        groups: &GpuGroups,
        masses: &GpuBufferF32,
        scale: f32,
    ) -> TrajResult<Vec<Float4>> {
        if n_frames == 0 || groups.n_groups == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let total = n_frames * groups.n_groups;
        let mut sum_x = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut sum_y = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut sum_z = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut mass_sum = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut out = stream
            .alloc_zeros::<Float4>(total)
            .map_err(map_driver_err)?;

        let n_groups = groups.n_groups as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let max_len = groups.max_len as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (
                ceil_div(groups.max_len.max(1), block),
                n_frames as u32,
                n_groups as u32,
            ),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.group_com_accum);
            builder.arg(&coords_dev);
            builder.arg(&groups.offsets);
            builder.arg(&groups.indices);
            builder.arg(&masses.inner);
            builder.arg(&n_groups);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&max_len);
            builder.arg(&mut sum_x);
            builder.arg(&mut sum_y);
            builder.arg(&mut sum_z);
            builder.arg(&mut mass_sum);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let cfg_out = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(total, block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.group_com_finalize);
            builder.arg(&sum_x);
            builder.arg(&sum_y);
            builder.arg(&sum_z);
            builder.arg(&mass_sum);
            builder.arg(&n_groups);
            builder.arg(&n_frames);
            builder.arg(&scale);
            builder.arg(&mut out);
            builder.launch(cfg_out).map_err(map_driver_err)?;
        }

        let mut host = vec![Float4::default(); total];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn group_dipole(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        groups: &GpuGroups,
        charges: &GpuBufferF32,
        scale: f32,
    ) -> TrajResult<Vec<Float4>> {
        if n_frames == 0 || groups.n_groups == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let total = n_frames * groups.n_groups;
        let mut sum_x = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut sum_y = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut sum_z = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;
        let mut out = stream
            .alloc_zeros::<Float4>(total)
            .map_err(map_driver_err)?;

        let n_groups = groups.n_groups as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let max_len = groups.max_len as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (
                ceil_div(groups.max_len.max(1), block),
                n_frames as u32,
                n_groups as u32,
            ),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.group_dipole_accum);
            builder.arg(&coords_dev);
            builder.arg(&groups.offsets);
            builder.arg(&groups.indices);
            builder.arg(&charges.inner);
            builder.arg(&n_groups);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&max_len);
            builder.arg(&mut sum_x);
            builder.arg(&mut sum_y);
            builder.arg(&mut sum_z);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let cfg_out = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(total, block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.group_dipole_finalize);
            builder.arg(&sum_x);
            builder.arg(&sum_y);
            builder.arg(&sum_z);
            builder.arg(&n_groups);
            builder.arg(&n_frames);
            builder.arg(&scale);
            builder.arg(&mut out);
            builder.launch(cfg_out).map_err(map_driver_err)?;
        }

        let mut host = vec![Float4::default(); total];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn group_ke(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        groups: &GpuGroups,
        masses: &GpuBufferF32,
        vel_scale: f32,
    ) -> TrajResult<Vec<f32>> {
        if n_frames == 0 || groups.n_groups == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let total = n_frames * groups.n_groups;
        let mut sum_ke = stream.alloc_zeros::<f32>(total).map_err(map_driver_err)?;

        let n_groups = groups.n_groups as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let max_len = groups.max_len as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (
                ceil_div(groups.max_len.max(1), block),
                n_frames as u32,
                n_groups as u32,
            ),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.group_ke_accum);
            builder.arg(&coords_dev);
            builder.arg(&groups.offsets);
            builder.arg(&groups.indices);
            builder.arg(&masses.inner);
            builder.arg(&n_groups);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&max_len);
            builder.arg(&vel_scale);
            builder.arg(&mut sum_ke);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let mut host = vec![0.0f32; total];
        stream.memcpy_dtoh(&sum_ke, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn orientation_plane(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        anchors: &GpuAnchors,
        scale: f32,
    ) -> TrajResult<Vec<Float4>> {
        if n_frames == 0 || anchors.n_groups == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let total = n_frames * anchors.n_groups;
        let mut out = stream
            .alloc_zeros::<Float4>(total)
            .map_err(map_driver_err)?;
        let n_groups = anchors.n_groups as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(anchors.n_groups, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.orientation_plane);
            builder.arg(&coords_dev);
            builder.arg(&anchors.anchors);
            builder.arg(&n_groups);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&scale);
            builder.arg(&mut out);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![Float4::default(); total];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn orientation_vector(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        anchors: &GpuAnchors,
        scale: f32,
    ) -> TrajResult<Vec<Float4>> {
        if n_frames == 0 || anchors.n_groups == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let total = n_frames * anchors.n_groups;
        let mut out = stream
            .alloc_zeros::<Float4>(total)
            .map_err(map_driver_err)?;
        let n_groups = anchors.n_groups as i32;
        let n_atoms = n_atoms as i32;
        let n_frames = n_frames as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(anchors.n_groups, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.orientation_vector);
            builder.arg(&coords_dev);
            builder.arg(&anchors.anchors);
            builder.arg(&n_groups);
            builder.arg(&n_atoms);
            builder.arg(&n_frames);
            builder.arg(&scale);
            builder.arg(&mut out);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![Float4::default(); total];
        stream.memcpy_dtoh(&out, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn water_count_frame(
        &self,
        coords: &GpuCoords,
        n_atoms: usize,
        frame: usize,
        selection: &GpuSelection,
        center: [f32; 3],
        box_unit: [f32; 3],
        region: [f32; 3],
        dims: [i32; 3],
        scale: f32,
        counts: &mut GpuCountsU32,
    ) -> TrajResult<()> {
        let stream = &self.inner.stream;
        let n_sel = selection.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let frame = frame as i32;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(selection.n_sel.max(1), block), 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.water_count);
            builder.arg(&coords.inner);
            builder.arg(&selection.sel);
            builder.arg(&n_sel);
            builder.arg(&n_atoms);
            builder.arg(&frame);
            builder.arg(&center[0]);
            builder.arg(&center[1]);
            builder.arg(&center[2]);
            builder.arg(&box_unit[0]);
            builder.arg(&box_unit[1]);
            builder.arg(&box_unit[2]);
            builder.arg(&region[0]);
            builder.arg(&region[1]);
            builder.arg(&region[2]);
            builder.arg(&dims[0]);
            builder.arg(&dims[1]);
            builder.arg(&dims[2]);
            builder.arg(&scale);
            builder.arg(&mut counts.inner);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        Ok(())
    }

    pub fn hbond_counts(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        donors: &GpuSelection,
        acceptors: &GpuSelection,
        dist2: f32,
    ) -> TrajResult<Vec<u32>> {
        if n_frames == 0 || donors.n_sel == 0 || acceptors.n_sel == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut counts = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;
        let n_donors = donors.n_sel as i32;
        let n_acceptors = acceptors.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;
        let total_pairs = donors.n_sel * acceptors.n_sel;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(total_pairs.max(1), block), n_frames_i32 as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.hbond_count);
            builder.arg(&coords_dev);
            builder.arg(&donors.sel);
            builder.arg(&acceptors.sel);
            builder.arg(&n_donors);
            builder.arg(&n_acceptors);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&dist2);
            builder.arg(&mut counts);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![0u32; n_frames];
        stream.memcpy_dtoh(&counts, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn hbond_counts_angle(
        &self,
        coords: &[Float4],
        n_atoms: usize,
        n_frames: usize,
        donors: &GpuSelection,
        hydrogens: &GpuSelection,
        acceptors: &GpuSelection,
        dist2: f32,
        cos_cutoff: f32,
    ) -> TrajResult<Vec<u32>> {
        if n_frames == 0 || donors.n_sel == 0 || acceptors.n_sel == 0 {
            return Ok(Vec::new());
        }
        let stream = &self.inner.stream;
        let coords_dev = stream.clone_htod(coords).map_err(map_driver_err)?;
        let mut counts = stream.alloc_zeros::<u32>(n_frames).map_err(map_driver_err)?;
        let n_donors = donors.n_sel as i32;
        let n_acceptors = acceptors.n_sel as i32;
        let n_atoms = n_atoms as i32;
        let n_frames_i32 = n_frames as i32;
        let total_pairs = donors.n_sel * acceptors.n_sel;
        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(total_pairs.max(1), block), n_frames_i32 as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.hbond_count_angle);
            builder.arg(&coords_dev);
            builder.arg(&donors.sel);
            builder.arg(&hydrogens.sel);
            builder.arg(&acceptors.sel);
            builder.arg(&n_donors);
            builder.arg(&n_acceptors);
            builder.arg(&n_atoms);
            builder.arg(&n_frames_i32);
            builder.arg(&dist2);
            builder.arg(&cos_cutoff);
            builder.arg(&mut counts);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut host = vec![0u32; n_frames];
        stream.memcpy_dtoh(&counts, &mut host).map_err(map_driver_err)?;
        Ok(host)
    }

    pub fn msd_time_lag(
        &self,
        com: &[f32],
        times: &[f32],
        type_ids: &[u32],
        type_counts: &[u32],
        n_groups: usize,
        n_types: usize,
        ndframe: usize,
        axis: Option<[f32; 3]>,
        frame_decimation: Option<(usize, usize)>,
        dt_decimation: Option<(usize, usize, usize, usize)>,
        time_binning: (f32, f32),
    ) -> TrajResult<(Vec<f32>, Vec<u32>)> {
        let n_frames = times.len();
        if n_frames < 2 || n_groups == 0 || ndframe == 0 {
            return Ok((Vec::new(), Vec::new()));
        }
        let components = if axis.is_some() { 5 } else { 4 };
        let cols = components * (n_types + 1);
        let stream = &self.inner.stream;
        let com_dev = stream.clone_htod(com).map_err(map_driver_err)?;
        let times_dev = stream.clone_htod(times).map_err(map_driver_err)?;
        let type_ids_dev = stream.clone_htod(type_ids).map_err(map_driver_err)?;
        let type_counts_dev = stream.clone_htod(type_counts).map_err(map_driver_err)?;
        let mut out_dev = stream
            .alloc_zeros::<f32>((ndframe + 1) * cols)
            .map_err(map_driver_err)?;
        let mut n_diff_dev = stream
            .alloc_zeros::<u32>(ndframe + 1)
            .map_err(map_driver_err)?;

        let [ax, ay, az] = axis.unwrap_or([0.0, 0.0, 0.0]);
        let axis_enabled = if axis.is_some() { 1i32 } else { 0i32 };
        let (frame_dec_start, frame_dec_stride, frame_dec_enabled) = match frame_decimation {
            Some((start, stride)) => (start as i32, stride as i32, 1i32),
            None => (0i32, 1i32, 0i32),
        };
        let (dt_cut1, dt_stride1, dt_cut2, dt_stride2, dt_dec_enabled) = match dt_decimation {
            Some((cut1, stride1, cut2, stride2)) => {
                (cut1 as i32, stride1 as i32, cut2 as i32, stride2 as i32, 1i32)
            }
            None => (0i32, 1i32, 0i32, 1i32, 0i32),
        };
        let dt0 = times[1] - times[0];
        let n_groups_i32 = n_groups as i32;
        let n_types_i32 = n_types as i32;
        let n_frames_i32 = n_frames as i32;
        let ndframe_i32 = ndframe as i32;

        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames, block), n_frames as u32, n_groups as u32),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.msd_time_lag);
            builder.arg(&com_dev);
            builder.arg(&times_dev);
            builder.arg(&type_ids_dev);
            builder.arg(&type_counts_dev);
            builder.arg(&n_groups_i32);
            builder.arg(&n_types_i32);
            builder.arg(&n_frames_i32);
            builder.arg(&ndframe_i32);
            builder.arg(&dt0);
            builder.arg(&time_binning.0);
            builder.arg(&time_binning.1);
            builder.arg(&frame_dec_start);
            builder.arg(&frame_dec_stride);
            builder.arg(&frame_dec_enabled);
            builder.arg(&dt_cut1);
            builder.arg(&dt_stride1);
            builder.arg(&dt_cut2);
            builder.arg(&dt_stride2);
            builder.arg(&dt_dec_enabled);
            builder.arg(&axis_enabled);
            builder.arg(&ax);
            builder.arg(&ay);
            builder.arg(&az);
            builder.arg(&mut out_dev);
            builder.arg(&mut n_diff_dev);
            builder.launch(cfg).map_err(map_driver_err)?;
        }
        let mut out = vec![0.0f32; (ndframe + 1) * cols];
        let mut n_diff = vec![0u32; ndframe + 1];
        stream.memcpy_dtoh(&out_dev, &mut out).map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&n_diff_dev, &mut n_diff)
            .map_err(map_driver_err)?;
        Ok((out, n_diff))
    }

    pub fn rotacf_time_lag(
        &self,
        orient: &[f32],
        times: &[f32],
        type_ids: &[u32],
        type_counts: &[u32],
        n_groups: usize,
        n_types: usize,
        ndframe: usize,
        frame_decimation: Option<(usize, usize)>,
        dt_decimation: Option<(usize, usize, usize, usize)>,
        time_binning: (f32, f32),
    ) -> TrajResult<(Vec<f32>, Vec<f32>, Vec<u32>)> {
        let n_frames = times.len();
        if n_frames < 2 || n_groups == 0 || ndframe == 0 {
            return Ok((Vec::new(), Vec::new(), Vec::new()));
        }
        let stride = n_types + 1;
        let stream = &self.inner.stream;
        let orient_dev = stream.clone_htod(orient).map_err(map_driver_err)?;
        let times_dev = stream.clone_htod(times).map_err(map_driver_err)?;
        let type_ids_dev = stream.clone_htod(type_ids).map_err(map_driver_err)?;
        let type_counts_dev = stream.clone_htod(type_counts).map_err(map_driver_err)?;
        let mut corr_dev = stream
            .alloc_zeros::<f32>((ndframe + 1) * stride)
            .map_err(map_driver_err)?;
        let mut corr_p2_dev = stream
            .alloc_zeros::<f32>((ndframe + 1) * stride)
            .map_err(map_driver_err)?;
        let mut n_diff_dev = stream
            .alloc_zeros::<u32>(ndframe + 1)
            .map_err(map_driver_err)?;

        let (frame_dec_start, frame_dec_stride, frame_dec_enabled) = match frame_decimation {
            Some((start, stride)) => (start as i32, stride as i32, 1i32),
            None => (0i32, 1i32, 0i32),
        };
        let (dt_cut1, dt_stride1, dt_cut2, dt_stride2, dt_dec_enabled) = match dt_decimation {
            Some((cut1, stride1, cut2, stride2)) => {
                (cut1 as i32, stride1 as i32, cut2 as i32, stride2 as i32, 1i32)
            }
            None => (0i32, 1i32, 0i32, 1i32, 0i32),
        };
        let dt0 = times[1] - times[0];
        let n_groups_i32 = n_groups as i32;
        let n_types_i32 = n_types as i32;
        let n_frames_i32 = n_frames as i32;
        let ndframe_i32 = ndframe as i32;

        let block = 256u32;
        let cfg = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames, block), n_frames as u32, n_groups as u32),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.rotacf_time_lag);
            builder.arg(&orient_dev);
            builder.arg(&times_dev);
            builder.arg(&type_ids_dev);
            builder.arg(&type_counts_dev);
            builder.arg(&n_groups_i32);
            builder.arg(&n_types_i32);
            builder.arg(&n_frames_i32);
            builder.arg(&ndframe_i32);
            builder.arg(&dt0);
            builder.arg(&time_binning.0);
            builder.arg(&time_binning.1);
            builder.arg(&frame_dec_start);
            builder.arg(&frame_dec_stride);
            builder.arg(&frame_dec_enabled);
            builder.arg(&dt_cut1);
            builder.arg(&dt_stride1);
            builder.arg(&dt_cut2);
            builder.arg(&dt_stride2);
            builder.arg(&dt_dec_enabled);
            builder.arg(&mut corr_dev);
            builder.arg(&mut corr_p2_dev);
            builder.arg(&mut n_diff_dev);
            builder.launch(cfg).map_err(map_driver_err)?;
        }

        let mut corr = vec![0.0f32; (ndframe + 1) * stride];
        let mut corr_p2 = vec![0.0f32; (ndframe + 1) * stride];
        let mut n_diff = vec![0u32; ndframe + 1];
        stream.memcpy_dtoh(&corr_dev, &mut corr).map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&corr_p2_dev, &mut corr_p2)
            .map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&n_diff_dev, &mut n_diff)
            .map_err(map_driver_err)?;
        Ok((corr, corr_p2, n_diff))
    }

    pub fn conductivity_time_lag(
        &self,
        com: &[f32],
        times: &[f32],
        charges: &[f32],
        type_ids: &[u32],
        type_charge: &[f32],
        n_groups: usize,
        n_types: usize,
        ndframe: usize,
        transference: bool,
        frame_decimation: Option<(usize, usize)>,
        dt_decimation: Option<(usize, usize, usize, usize)>,
        time_binning: (f32, f32),
    ) -> TrajResult<(Vec<f32>, Vec<u32>, usize)> {
        let n_frames = times.len();
        if n_frames < 2 || n_groups == 0 || ndframe == 0 {
            return Ok((Vec::new(), Vec::new(), 0));
        }
        let cols = if transference {
            n_types * n_types + 1
        } else {
            1
        };
        let stream = &self.inner.stream;
        let com_dev = stream.clone_htod(com).map_err(map_driver_err)?;
        let times_dev = stream.clone_htod(times).map_err(map_driver_err)?;
        let charges_dev = stream.clone_htod(charges).map_err(map_driver_err)?;
        let type_ids_dev = stream.clone_htod(type_ids).map_err(map_driver_err)?;
        let type_charge_dev = stream.clone_htod(type_charge).map_err(map_driver_err)?;
        let mut out_dev = stream
            .alloc_zeros::<f32>((ndframe + 1) * cols)
            .map_err(map_driver_err)?;
        let mut n_diff_dev = stream
            .alloc_zeros::<u32>(ndframe + 1)
            .map_err(map_driver_err)?;

        let (frame_dec_start, frame_dec_stride, frame_dec_enabled) = match frame_decimation {
            Some((start, stride)) => (start as i32, stride as i32, 1i32),
            None => (0i32, 1i32, 0i32),
        };
        let (dt_cut1, dt_stride1, dt_cut2, dt_stride2, dt_dec_enabled) = match dt_decimation {
            Some((cut1, stride1, cut2, stride2)) => {
                (cut1 as i32, stride1 as i32, cut2 as i32, stride2 as i32, 1i32)
            }
            None => (0i32, 1i32, 0i32, 1i32, 0i32),
        };
        let dt0 = times[1] - times[0];
        let n_groups_i32 = n_groups as i32;
        let n_types_i32 = n_types as i32;
        let n_frames_i32 = n_frames as i32;
        let ndframe_i32 = ndframe as i32;
        let cols_i32 = cols as i32;

        let block = 256u32;
        let cfg_total = LaunchConfig {
            block_dim: (block, 1, 1),
            grid_dim: (ceil_div(n_frames, block), n_frames as u32, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            let mut builder = stream.launch_builder(&self.inner.kernels.conductivity_total);
            builder.arg(&com_dev);
            builder.arg(&times_dev);
            builder.arg(&charges_dev);
            builder.arg(&n_groups_i32);
            builder.arg(&n_frames_i32);
            builder.arg(&ndframe_i32);
            builder.arg(&dt0);
            builder.arg(&time_binning.0);
            builder.arg(&time_binning.1);
            builder.arg(&frame_dec_start);
            builder.arg(&frame_dec_stride);
            builder.arg(&frame_dec_enabled);
            builder.arg(&dt_cut1);
            builder.arg(&dt_stride1);
            builder.arg(&dt_cut2);
            builder.arg(&dt_stride2);
            builder.arg(&dt_dec_enabled);
            builder.arg(&cols_i32);
            builder.arg(&mut out_dev);
            builder.arg(&mut n_diff_dev);
            builder.launch(cfg_total).map_err(map_driver_err)?;
        }

        if transference {
            let total_pairs = n_groups * n_groups;
            let cfg = LaunchConfig {
                block_dim: (block, 1, 1),
                grid_dim: (
                    ceil_div(total_pairs.max(1), block),
                    n_frames as u32,
                    n_frames as u32,
                ),
                shared_mem_bytes: 0,
            };
            unsafe {
                let mut builder = stream.launch_builder(&self.inner.kernels.conductivity_transference);
                builder.arg(&com_dev);
                builder.arg(&times_dev);
                builder.arg(&charges_dev);
                builder.arg(&type_ids_dev);
                builder.arg(&type_charge_dev);
                builder.arg(&n_groups_i32);
                builder.arg(&n_types_i32);
                builder.arg(&n_frames_i32);
                builder.arg(&ndframe_i32);
                builder.arg(&dt0);
                builder.arg(&time_binning.0);
                builder.arg(&time_binning.1);
                builder.arg(&frame_dec_start);
                builder.arg(&frame_dec_stride);
                builder.arg(&frame_dec_enabled);
                builder.arg(&dt_cut1);
                builder.arg(&dt_stride1);
                builder.arg(&dt_cut2);
                builder.arg(&dt_stride2);
                builder.arg(&dt_dec_enabled);
                builder.arg(&cols_i32);
                builder.arg(&mut out_dev);
                builder.launch(cfg).map_err(map_driver_err)?;
            }
        }

        let mut out = vec![0.0f32; (ndframe + 1) * cols];
        let mut n_diff = vec![0u32; ndframe + 1];
        stream.memcpy_dtoh(&out_dev, &mut out).map_err(map_driver_err)?;
        stream
            .memcpy_dtoh(&n_diff_dev, &mut n_diff)
            .map_err(map_driver_err)?;
        Ok((out, n_diff, cols))
    }

    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)
    }
}

pub struct GpuSelection {
    sel: CudaSlice<u32>,
    masses: CudaSlice<f32>,
    n_sel: usize,
}

impl GpuSelection {
    pub fn n_sel(&self) -> usize {
        self.n_sel
    }
}

pub struct GpuReference {
    coords: CudaSlice<Float4>,
    n_sel: usize,
}

impl GpuReference {
    pub fn n_sel(&self) -> usize {
        self.n_sel
    }
}

pub struct GpuPolymer {
    chain_offsets: CudaSlice<u32>,
    chain_indices: CudaSlice<u32>,
    n_chains: usize,
    bond_pairs: Option<CudaSlice<u32>>,
    n_bonds: usize,
    angle_triplets: Option<CudaSlice<u32>>,
    n_angles: usize,
}

pub struct GpuBufferF32 {
    inner: CudaSlice<f32>,
}

pub struct GpuBufferU32 {
    #[allow(dead_code)]
    inner: CudaSlice<u32>,
}

pub struct GpuGroups {
    offsets: CudaSlice<u32>,
    indices: CudaSlice<u32>,
    n_groups: usize,
    max_len: usize,
}

pub struct GpuAnchors {
    anchors: CudaSlice<u32>,
    n_groups: usize,
}

pub struct GpuCountsU32 {
    inner: CudaSlice<u32>,
}

pub struct GpuCoords {
    inner: CudaSlice<Float4>,
}

pub struct RmsdCovariance {
    pub cov: Vec<[f32; 9]>,
    pub sum_x2: Vec<f32>,
    pub sum_y2: Vec<f32>,
}

pub struct GpuCounts {
    inner: CudaSlice<u64>,
}

fn ceil_div(value: usize, block: u32) -> u32 {
    if value == 0 {
        1
    } else {
        ((value as u32) + block - 1) / block
    }
}

fn map_driver_err(err: cudarc::driver::DriverError) -> TrajError {
    TrajError::Unsupported(format!("cuda driver error: {err}"))
}

fn map_compile_err(err: cudarc::nvrtc::CompileError) -> TrajError {
    TrajError::Unsupported(format!("cuda compile error: {err}"))
}
