pub const KERNELS_SRC: &str = r#"
extern "C" {
struct Float4 { float x; float y; float z; float w; };

__global__ void rg_accum(const Float4* coords,
                         const unsigned int* sel,
                         const float* masses,
                         int n_sel,
                         int n_atoms,
                         int n_frames,
                         int mass_weighted,
                         float* sum_x,
                         float* sum_y,
                         float* sum_z,
                         float* mass_sum) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    float m = mass_weighted ? masses[idx] : 1.0f;
    atomicAdd(&sum_x[frame], p.x * m);
    atomicAdd(&sum_y[frame], p.y * m);
    atomicAdd(&sum_z[frame], p.z * m);
    atomicAdd(&mass_sum[frame], m);
}

__global__ void rg_sumsq(const Float4* coords,
                         const unsigned int* sel,
                         const float* masses,
                         int n_sel,
                         int n_atoms,
                         int n_frames,
                         int mass_weighted,
                         const float* sum_x,
                         const float* sum_y,
                         const float* sum_z,
                         const float* mass_sum,
                         float* sum_sq) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    float msum = mass_sum[frame];
    if (msum == 0.0f) return;
    float cx = sum_x[frame] / msum;
    float cy = sum_y[frame] / msum;
    float cz = sum_z[frame] / msum;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    float dx = p.x - cx;
    float dy = p.y - cy;
    float dz = p.z - cz;
    float m = mass_weighted ? masses[idx] : 1.0f;
    atomicAdd(&sum_sq[frame], m * (dx * dx + dy * dy + dz * dz));
}

__global__ void rg_finalize(const float* mass_sum,
                            const float* sum_sq,
                            int n_frames,
                            float* out) {
    int frame = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames) return;
    float m = mass_sum[frame];
    if (m == 0.0f) {
        out[frame] = 0.0f;
    } else {
        out[frame] = sqrtf(sum_sq[frame] / m);
    }
}

__global__ void msd_accum(const Float4* coords,
                          const Float4* origin,
                          const unsigned int* sel,
                          int n_sel,
                          int n_atoms,
                          int n_frames,
                          float* sum_sq) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    Float4 r0 = origin[idx];
    float dx = p.x - r0.x;
    float dy = p.y - r0.y;
    float dz = p.z - r0.z;
    atomicAdd(&sum_sq[frame], dx * dx + dy * dy + dz * dz);
}

__global__ void msd_finalize(const float* sum_sq,
                             int n_frames,
                             int n_sel,
                             float* out) {
    int frame = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames) return;
    if (n_sel == 0) {
        out[frame] = 0.0f;
    } else {
        out[frame] = sum_sq[frame] / (float)n_sel;
    }
}

__global__ void rmsd_centroid(const Float4* coords,
                              const Float4* ref,
                              const unsigned int* sel,
                              int n_sel,
                              int n_atoms,
                              int n_frames,
                              float* sum_x,
                              float* sum_y) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    Float4 r = ref[idx];
    atomicAdd(&sum_x[frame * 3 + 0], p.x);
    atomicAdd(&sum_x[frame * 3 + 1], p.y);
    atomicAdd(&sum_x[frame * 3 + 2], p.z);
    atomicAdd(&sum_y[frame * 3 + 0], r.x);
    atomicAdd(&sum_y[frame * 3 + 1], r.y);
    atomicAdd(&sum_y[frame * 3 + 2], r.z);
}

__global__ void rmsd_cov(const Float4* coords,
                         const Float4* ref,
                         const unsigned int* sel,
                         int n_sel,
                         int n_atoms,
                         int n_frames,
                         const float* sum_x,
                         const float* sum_y,
                         float* cov,
                         float* sum_x2,
                         float* sum_y2) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    float cx = sum_x[frame * 3 + 0] / (float)n_sel;
    float cy = sum_x[frame * 3 + 1] / (float)n_sel;
    float cz = sum_x[frame * 3 + 2] / (float)n_sel;
    float rx = sum_y[frame * 3 + 0] / (float)n_sel;
    float ry = sum_y[frame * 3 + 1] / (float)n_sel;
    float rz = sum_y[frame * 3 + 2] / (float)n_sel;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    Float4 r = ref[idx];
    float x0 = p.x - cx;
    float x1 = p.y - cy;
    float x2 = p.z - cz;
    float y0 = r.x - rx;
    float y1 = r.y - ry;
    float y2 = r.z - rz;
    atomicAdd(&cov[frame * 9 + 0], x0 * y0);
    atomicAdd(&cov[frame * 9 + 1], x0 * y1);
    atomicAdd(&cov[frame * 9 + 2], x0 * y2);
    atomicAdd(&cov[frame * 9 + 3], x1 * y0);
    atomicAdd(&cov[frame * 9 + 4], x1 * y1);
    atomicAdd(&cov[frame * 9 + 5], x1 * y2);
    atomicAdd(&cov[frame * 9 + 6], x2 * y0);
    atomicAdd(&cov[frame * 9 + 7], x2 * y1);
    atomicAdd(&cov[frame * 9 + 8], x2 * y2);
    atomicAdd(&sum_x2[frame], x0 * x0 + x1 * x1 + x2 * x2);
    atomicAdd(&sum_y2[frame], y0 * y0 + y1 * y1 + y2 * y2);
}

__global__ void rmsd_raw_accum(const Float4* coords,
                               const Float4* ref,
                               const unsigned int* sel,
                               int n_sel,
                               int n_atoms,
                               int n_frames,
                               float* sum_sq) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_sel) return;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    Float4 r = ref[idx];
    float dx = p.x - r.x;
    float dy = p.y - r.y;
    float dz = p.z - r.z;
    atomicAdd(&sum_sq[frame], dx * dx + dy * dy + dz * dz);
}

__global__ void rmsd_finalize(const float* sum_sq,
                              int n_frames,
                              int n_sel,
                              float* out) {
    int frame = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames) return;
    if (n_sel == 0) {
        out[frame] = 0.0f;
    } else {
        out[frame] = sqrtf(sum_sq[frame] / (float)n_sel);
    }
}

__global__ void rdf_hist(const Float4* coords,
                         const unsigned int* sel_a,
                         const unsigned int* sel_b,
                         int n_a,
                         int n_b,
                         int n_atoms,
                         int n_frames,
                         float r_max,
                         float bin_width,
                         int pbc,
                         const float* box_l,
                         unsigned long long* counts,
                         int same_sel) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames) return;
    long total = (long)n_a * (long)n_b;
    if (idx >= total) return;
    int ia = idx / n_b;
    int ib = idx - ia * n_b;
    if (same_sel && ia == ib) return;
    unsigned int atom_a = sel_a[ia];
    unsigned int atom_b = sel_b[ib];
    Float4 pa = coords[frame * n_atoms + atom_a];
    Float4 pb = coords[frame * n_atoms + atom_b];
    float dx = pb.x - pa.x;
    float dy = pb.y - pa.y;
    float dz = pb.z - pa.z;
    if (pbc) {
        float lx = box_l[frame * 3 + 0];
        float ly = box_l[frame * 3 + 1];
        float lz = box_l[frame * 3 + 2];
        dx -= roundf(dx / lx) * lx;
        dy -= roundf(dy / ly) * ly;
        dz -= roundf(dz / lz) * lz;
    }
    float r = sqrtf(dx * dx + dy * dy + dz * dz);
    if (r < r_max) {
        int bin = (int)(r / bin_width);
        atomicAdd(&counts[bin], 1ULL);
    }
}

__global__ void polymer_end_to_end(const Float4* coords,
                                   const unsigned int* chain_offsets,
                                   const unsigned int* chain_indices,
                                   int n_chains,
                                   int n_atoms,
                                   int n_frames,
                                   float* out) {
    int frame = blockIdx.y;
    int chain = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || chain >= n_chains) return;
    unsigned int start = chain_offsets[chain];
    unsigned int end = chain_offsets[chain + 1];
    if (end <= start) {
        out[frame * n_chains + chain] = 0.0f;
        return;
    }
    unsigned int first = chain_indices[start];
    unsigned int last = chain_indices[end - 1];
    Float4 p = coords[frame * n_atoms + first];
    Float4 q = coords[frame * n_atoms + last];
    float dx = q.x - p.x;
    float dy = q.y - p.y;
    float dz = q.z - p.z;
    out[frame * n_chains + chain] = sqrtf(dx * dx + dy * dy + dz * dz);
}

__global__ void polymer_contour_length(const Float4* coords,
                                       const unsigned int* chain_offsets,
                                       const unsigned int* chain_indices,
                                       int n_chains,
                                       int n_atoms,
                                       int n_frames,
                                       float* out) {
    int frame = blockIdx.y;
    int chain = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || chain >= n_chains) return;
    unsigned int start = chain_offsets[chain];
    unsigned int end = chain_offsets[chain + 1];
    float sum = 0.0f;
    if (end > start + 1) {
        for (unsigned int i = start; i + 1 < end; ++i) {
            unsigned int a = chain_indices[i];
            unsigned int b = chain_indices[i + 1];
            Float4 pa = coords[frame * n_atoms + a];
            Float4 pb = coords[frame * n_atoms + b];
            float dx = pb.x - pa.x;
            float dy = pb.y - pa.y;
            float dz = pb.z - pa.z;
            sum += sqrtf(dx * dx + dy * dy + dz * dz);
        }
    }
    out[frame * n_chains + chain] = sum;
}

__global__ void polymer_chain_rg(const Float4* coords,
                                 const unsigned int* chain_offsets,
                                 const unsigned int* chain_indices,
                                 int n_chains,
                                 int n_atoms,
                                 int n_frames,
                                 float* out) {
    int frame = blockIdx.y;
    int chain = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || chain >= n_chains) return;
    unsigned int start = chain_offsets[chain];
    unsigned int end = chain_offsets[chain + 1];
    unsigned int n = (end > start) ? (end - start) : 0;
    if (n == 0) {
        out[frame * n_chains + chain] = 0.0f;
        return;
    }
    float cx = 0.0f;
    float cy = 0.0f;
    float cz = 0.0f;
    for (unsigned int i = start; i < end; ++i) {
        unsigned int idx = chain_indices[i];
        Float4 p = coords[frame * n_atoms + idx];
        cx += p.x;
        cy += p.y;
        cz += p.z;
    }
    float inv = 1.0f / (float)n;
    cx *= inv;
    cy *= inv;
    cz *= inv;
    float sum = 0.0f;
    for (unsigned int i = start; i < end; ++i) {
        unsigned int idx = chain_indices[i];
        Float4 p = coords[frame * n_atoms + idx];
        float dx = p.x - cx;
        float dy = p.y - cy;
        float dz = p.z - cz;
        sum += dx * dx + dy * dy + dz * dz;
    }
    out[frame * n_chains + chain] = sqrtf(sum * inv);
}

__global__ void polymer_bond_hist(const Float4* coords,
                                  const unsigned int* bond_pairs,
                                  int n_bonds,
                                  int n_atoms,
                                  int n_frames,
                                  float r_max,
                                  float bin_width,
                                  int bins,
                                  unsigned long long* counts) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_bonds) return;
    unsigned int a = bond_pairs[idx * 2];
    unsigned int b = bond_pairs[idx * 2 + 1];
    Float4 pa = coords[frame * n_atoms + a];
    Float4 pb = coords[frame * n_atoms + b];
    float dx = pb.x - pa.x;
    float dy = pb.y - pa.y;
    float dz = pb.z - pa.z;
    float r = sqrtf(dx * dx + dy * dy + dz * dz);
    if (r < r_max) {
        int bin = (int)(r / bin_width);
        if (bin < bins) {
            atomicAdd(&counts[bin], 1ULL);
        }
    }
}

__global__ void polymer_angle_hist(const Float4* coords,
                                   const unsigned int* angle_triplets,
                                   int n_angles,
                                   int n_atoms,
                                   int n_frames,
                                   float max_angle,
                                   float bin_width,
                                   int bins,
                                   int degrees,
                                   unsigned long long* counts) {
    int frame = blockIdx.y;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (frame >= n_frames || idx >= n_angles) return;
    unsigned int a = angle_triplets[idx * 3];
    unsigned int b = angle_triplets[idx * 3 + 1];
    unsigned int c = angle_triplets[idx * 3 + 2];
    Float4 pa = coords[frame * n_atoms + a];
    Float4 pb = coords[frame * n_atoms + b];
    Float4 pc = coords[frame * n_atoms + c];
    float v1x = pa.x - pb.x;
    float v1y = pa.y - pb.y;
    float v1z = pa.z - pb.z;
    float v2x = pc.x - pb.x;
    float v2y = pc.y - pb.y;
    float v2z = pc.z - pb.z;
    float n1 = sqrtf(v1x * v1x + v1y * v1y + v1z * v1z);
    float n2 = sqrtf(v2x * v2x + v2y * v2y + v2z * v2z);
    if (n1 == 0.0f || n2 == 0.0f) return;
    float dot = (v1x * v2x + v1y * v2y + v1z * v2z) / (n1 * n2);
    dot = fminf(1.0f, fmaxf(-1.0f, dot));
    float angle = acosf(dot);
    if (degrees) {
        angle *= 57.29577951308232f;
    }
    if (angle < max_angle) {
        int bin = (int)(angle / bin_width);
        if (bin < bins) {
            atomicAdd(&counts[bin], 1ULL);
        }
    }
}
}

__global__ void group_com_accum(const Float4* coords,
                                const unsigned int* group_offsets,
                                const unsigned int* group_indices,
                                const float* masses,
                                int n_groups,
                                int n_atoms,
                                int n_frames,
                                int max_len,
                                float* sum_x,
                                float* sum_y,
                                float* sum_z,
                                float* mass_sum) {
    int atom_in_group = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    int group = blockIdx.z;
    if (frame >= n_frames || group >= n_groups) return;
    unsigned int start = group_offsets[group];
    unsigned int end = group_offsets[group + 1];
    unsigned int len = end - start;
    if (atom_in_group >= len) return;
    unsigned int atom = group_indices[start + atom_in_group];
    Float4 p = coords[frame * n_atoms + atom];
    float m = masses[atom];
    int idx = frame * n_groups + group;
    atomicAdd(&sum_x[idx], p.x * m);
    atomicAdd(&sum_y[idx], p.y * m);
    atomicAdd(&sum_z[idx], p.z * m);
    atomicAdd(&mass_sum[idx], m);
}

__global__ void group_com_finalize(const float* sum_x,
                                   const float* sum_y,
                                   const float* sum_z,
                                   const float* mass_sum,
                                   int n_groups,
                                   int n_frames,
                                   float scale,
                                   Float4* out) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = n_groups * n_frames;
    if (idx >= total) return;
    float m = mass_sum[idx];
    if (m == 0.0f) {
        out[idx].x = 0.0f;
        out[idx].y = 0.0f;
        out[idx].z = 0.0f;
        out[idx].w = 1.0f;
        return;
    }
    out[idx].x = sum_x[idx] / m * scale;
    out[idx].y = sum_y[idx] / m * scale;
    out[idx].z = sum_z[idx] / m * scale;
    out[idx].w = 1.0f;
}

__global__ void group_dipole_accum(const Float4* coords,
                                   const unsigned int* group_offsets,
                                   const unsigned int* group_indices,
                                   const float* charges,
                                   int n_groups,
                                   int n_atoms,
                                   int n_frames,
                                   int max_len,
                                   float* sum_x,
                                   float* sum_y,
                                   float* sum_z) {
    int atom_in_group = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    int group = blockIdx.z;
    if (frame >= n_frames || group >= n_groups) return;
    unsigned int start = group_offsets[group];
    unsigned int end = group_offsets[group + 1];
    unsigned int len = end - start;
    if (atom_in_group >= len) return;
    unsigned int atom = group_indices[start + atom_in_group];
    Float4 p = coords[frame * n_atoms + atom];
    float q = charges[atom];
    int idx = frame * n_groups + group;
    atomicAdd(&sum_x[idx], p.x * q);
    atomicAdd(&sum_y[idx], p.y * q);
    atomicAdd(&sum_z[idx], p.z * q);
}

__global__ void group_dipole_finalize(const float* sum_x,
                                      const float* sum_y,
                                      const float* sum_z,
                                      int n_groups,
                                      int n_frames,
                                      float scale,
                                      Float4* out) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = n_groups * n_frames;
    if (idx >= total) return;
    out[idx].x = sum_x[idx] * scale;
    out[idx].y = sum_y[idx] * scale;
    out[idx].z = sum_z[idx] * scale;
    out[idx].w = 0.0f;
}

__global__ void group_ke_accum(const Float4* coords,
                               const unsigned int* group_offsets,
                               const unsigned int* group_indices,
                               const float* masses,
                               int n_groups,
                               int n_atoms,
                               int n_frames,
                               int max_len,
                               float vel_scale,
                               float* sum_ke) {
    int atom_in_group = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    int group = blockIdx.z;
    if (frame >= n_frames || group >= n_groups) return;
    unsigned int start = group_offsets[group];
    unsigned int end = group_offsets[group + 1];
    unsigned int len = end - start;
    if (atom_in_group >= len) return;
    unsigned int atom = group_indices[start + atom_in_group];
    Float4 v = coords[frame * n_atoms + atom];
    float vx = v.x * vel_scale;
    float vy = v.y * vel_scale;
    float vz = v.z * vel_scale;
    float m = masses[atom];
    float ke = 0.5f * m * (vx * vx + vy * vy + vz * vz);
    int idx = frame * n_groups + group;
    atomicAdd(&sum_ke[idx], ke);
}

__global__ void orientation_plane(const Float4* coords,
                                  const unsigned int* anchors,
                                  int n_groups,
                                  int n_atoms,
                                  int n_frames,
                                  float scale,
                                  Float4* out) {
    int group = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (frame >= n_frames || group >= n_groups) return;
    unsigned int a = anchors[group * 3 + 0];
    unsigned int b = anchors[group * 3 + 1];
    unsigned int c = anchors[group * 3 + 2];
    Float4 pa = coords[frame * n_atoms + a];
    Float4 pb = coords[frame * n_atoms + b];
    Float4 pc = coords[frame * n_atoms + c];
    float v1x = (pa.x - pb.x) * scale;
    float v1y = (pa.y - pb.y) * scale;
    float v1z = (pa.z - pb.z) * scale;
    float v2x = (pa.x - pc.x) * scale;
    float v2y = (pa.y - pc.y) * scale;
    float v2z = (pa.z - pc.z) * scale;
    float cx = v1y * v2z - v1z * v2y;
    float cy = v1z * v2x - v1x * v2z;
    float cz = v1x * v2y - v1y * v2x;
    float norm = sqrtf(cx * cx + cy * cy + cz * cz);
    int idx = frame * n_groups + group;
    if (norm == 0.0f) {
        out[idx].x = 0.0f;
        out[idx].y = 0.0f;
        out[idx].z = 0.0f;
        out[idx].w = 0.0f;
    } else {
        out[idx].x = cx / norm;
        out[idx].y = cy / norm;
        out[idx].z = cz / norm;
        out[idx].w = 0.0f;
    }
}

__global__ void orientation_vector(const Float4* coords,
                                   const unsigned int* anchors,
                                   int n_groups,
                                   int n_atoms,
                                   int n_frames,
                                   float scale,
                                   Float4* out) {
    int group = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (frame >= n_frames || group >= n_groups) return;
    unsigned int a = anchors[group * 3 + 0];
    unsigned int b = anchors[group * 3 + 1];
    Float4 pa = coords[frame * n_atoms + a];
    Float4 pb = coords[frame * n_atoms + b];
    float vx = (pb.x - pa.x) * scale;
    float vy = (pb.y - pa.y) * scale;
    float vz = (pb.z - pa.z) * scale;
    float norm = sqrtf(vx * vx + vy * vy + vz * vz);
    int idx = frame * n_groups + group;
    if (norm == 0.0f) {
        out[idx].x = 0.0f;
        out[idx].y = 0.0f;
        out[idx].z = 0.0f;
        out[idx].w = 0.0f;
    } else {
        out[idx].x = vx / norm;
        out[idx].y = vy / norm;
        out[idx].z = vz / norm;
        out[idx].w = 0.0f;
    }
}

__global__ void water_count(const Float4* coords,
                            const unsigned int* sel,
                            int n_sel,
                            int n_atoms,
                            int frame,
                            float center_x,
                            float center_y,
                            float center_z,
                            float box_x,
                            float box_y,
                            float box_z,
                            float region_x,
                            float region_y,
                            float region_z,
                            int dim_x,
                            int dim_y,
                            int dim_z,
                            float scale,
                            unsigned int* counts) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_sel) return;
    unsigned int atom = sel[idx];
    Float4 p = coords[frame * n_atoms + atom];
    float x = p.x * scale - center_x;
    float y = p.y * scale - center_y;
    float z = p.z * scale - center_z;
    if (x < 0.0f || y < 0.0f || z < 0.0f) return;
    if (x > region_x || y > region_y || z > region_z) return;
    int ix = (int)floorf(x / box_x);
    int iy = (int)floorf(y / box_y);
    int iz = (int)floorf(z / box_z);
    if (ix < 0 || iy < 0 || iz < 0) return;
    if (ix >= dim_x || iy >= dim_y || iz >= dim_z) return;
    int flat = ix + dim_x * (iy + dim_y * iz);
    atomicAdd(&counts[flat], 1);
}

__global__ void hbond_count(const Float4* coords,
                            const unsigned int* donors,
                            const unsigned int* acceptors,
                            int n_donors,
                            int n_acceptors,
                            int n_atoms,
                            int n_frames,
                            float dist2,
                            unsigned int* counts) {
    int pair = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (frame >= n_frames) return;
    int total = n_donors * n_acceptors;
    if (pair >= total) return;
    int d = pair / n_acceptors;
    int a = pair - d * n_acceptors;
    unsigned int donor = donors[d];
    unsigned int acceptor = acceptors[a];
    if (donor == acceptor) return;
    Float4 pd = coords[frame * n_atoms + donor];
    Float4 pa = coords[frame * n_atoms + acceptor];
    float dx = pa.x - pd.x;
    float dy = pa.y - pd.y;
    float dz = pa.z - pd.z;
    float r2 = dx * dx + dy * dy + dz * dz;
    if (r2 <= dist2) {
        atomicAdd(&counts[frame], 1);
    }
}

__global__ void hbond_count_angle(const Float4* coords,
                                  const unsigned int* donors,
                                  const unsigned int* hydrogens,
                                  const unsigned int* acceptors,
                                  int n_donors,
                                  int n_acceptors,
                                  int n_atoms,
                                  int n_frames,
                                  float dist2,
                                  float cos_cutoff,
                                  unsigned int* counts) {
    int pair = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (frame >= n_frames) return;
    int total = n_donors * n_acceptors;
    if (pair >= total) return;
    int d = pair / n_acceptors;
    int a = pair - d * n_acceptors;
    unsigned int donor = donors[d];
    unsigned int hydrogen = hydrogens[d];
    unsigned int acceptor = acceptors[a];
    if (donor == acceptor) return;
    Float4 pd = coords[frame * n_atoms + donor];
    Float4 ph = coords[frame * n_atoms + hydrogen];
    Float4 pa = coords[frame * n_atoms + acceptor];
    float dx = pa.x - pd.x;
    float dy = pa.y - pd.y;
    float dz = pa.z - pd.z;
    float r2 = dx * dx + dy * dy + dz * dz;
    if (r2 > dist2) return;
    float v1x = pd.x - ph.x;
    float v1y = pd.y - ph.y;
    float v1z = pd.z - ph.z;
    float v2x = pa.x - ph.x;
    float v2y = pa.y - ph.y;
    float v2z = pa.z - ph.z;
    float n1 = v1x * v1x + v1y * v1y + v1z * v1z;
    float n2 = v2x * v2x + v2y * v2y + v2z * v2z;
    if (n1 == 0.0f || n2 == 0.0f) return;
    float cos_val = (v1x * v2x + v1y * v2y + v1z * v2z) / sqrtf(n1 * n2);
    if (cos_val >= cos_cutoff) {
        atomicAdd(&counts[frame], 1);
    }
}

__global__ void msd_time_lag(const float* com,
                             const float* times,
                             const unsigned int* type_ids,
                             const unsigned int* type_counts,
                             int n_groups,
                             int n_types,
                             int n_frames,
                             int ndframe,
                             float dt0,
                             float eps_num,
                             float eps_add,
                             int frame_dec_start,
                             int frame_dec_stride,
                             int frame_dec_enabled,
                             int dt_cut1,
                             int dt_stride1,
                             int dt_cut2,
                             int dt_stride2,
                             int dt_dec_enabled,
                             int axis_enabled,
                             float ax,
                             float ay,
                             float az,
                             float* out,
                             unsigned int* n_diff) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    int g = blockIdx.z;
    if (i >= n_frames || j >= n_frames || g >= n_groups) return;
    if (j <= i) return;
    if (frame_dec_enabled) {
        int idx1 = i + 1;
        if (idx1 > frame_dec_start && (idx1 % frame_dec_stride) != 0) return;
    }
    float dt = times[j] - times[i];
    int idt = (int)floorf((dt + eps_num) / dt0 + eps_add);
    if (idt <= 0 || idt > ndframe) return;
    if (dt_dec_enabled) {
        if (idt > dt_cut2 && (idt % dt_stride2) != 0) return;
        if (idt > dt_cut1 && (idt % dt_stride1) != 0) return;
    }
    if (g == 0) {
        atomicAdd(&n_diff[idt], 1);
    }
    int idx_i = (i * n_groups + g) * 3;
    int idx_j = (j * n_groups + g) * 3;
    float dx = com[idx_i] - com[idx_j];
    float dy = com[idx_i + 1] - com[idx_j + 1];
    float dz = com[idx_i + 2] - com[idx_j + 2];
    float msd_x = dx * dx;
    float msd_y = dy * dy;
    float msd_z = dz * dz;
    float msd_tot = msd_x + msd_y + msd_z;
    unsigned int t = type_ids[g];
    float inv_type = 0.0f;
    if (type_counts[t] > 0) {
        inv_type = 1.0f / (float)type_counts[t];
    }
    float inv_groups = 1.0f / (float)n_groups;
    int components = axis_enabled ? 5 : 4;
    int stride = n_types + 1;
    int base = idt * (components * stride);
    int offset_type = base + t;
    int offset_total = base + n_types;

    atomicAdd(&out[offset_type], msd_x * inv_type);
    atomicAdd(&out[offset_total], msd_x * inv_groups);
    atomicAdd(&out[offset_type + stride], msd_y * inv_type);
    atomicAdd(&out[offset_total + stride], msd_y * inv_groups);
    atomicAdd(&out[offset_type + 2 * stride], msd_z * inv_type);
    atomicAdd(&out[offset_total + 2 * stride], msd_z * inv_groups);
    int comp = 3;
    if (axis_enabled) {
        float proj = dx * ax + dy * ay + dz * az;
        float msd_axis = proj * proj;
        atomicAdd(&out[offset_type + 3 * stride], msd_axis * inv_type);
        atomicAdd(&out[offset_total + 3 * stride], msd_axis * inv_groups);
        comp = 4;
    }
    atomicAdd(&out[offset_type + comp * stride], msd_tot * inv_type);
    atomicAdd(&out[offset_total + comp * stride], msd_tot * inv_groups);
}

__global__ void rotacf_time_lag(const float* orient,
                                const float* times,
                                const unsigned int* type_ids,
                                const unsigned int* type_counts,
                                int n_groups,
                                int n_types,
                                int n_frames,
                                int ndframe,
                                float dt0,
                                float eps_num,
                                float eps_add,
                                int frame_dec_start,
                                int frame_dec_stride,
                                int frame_dec_enabled,
                                int dt_cut1,
                                int dt_stride1,
                                int dt_cut2,
                                int dt_stride2,
                                int dt_dec_enabled,
                                float* corr,
                                float* corr_p2,
                                unsigned int* n_diff) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    int g = blockIdx.z;
    if (i >= n_frames || j >= n_frames || g >= n_groups) return;
    if (j <= i) return;
    if (frame_dec_enabled) {
        int idx1 = i + 1;
        if (idx1 > frame_dec_start && (idx1 % frame_dec_stride) != 0) return;
    }
    float dt = times[j] - times[i];
    int idt = (int)floorf((dt + eps_num) / dt0 + eps_add);
    if (idt <= 0 || idt > ndframe) return;
    if (dt_dec_enabled) {
        if (idt > dt_cut2 && (idt % dt_stride2) != 0) return;
        if (idt > dt_cut1 && (idt % dt_stride1) != 0) return;
    }
    if (g == 0) {
        atomicAdd(&n_diff[idt], 1);
    }
    int idx_i = (i * n_groups + g) * 3;
    int idx_j = (j * n_groups + g) * 3;
    float v1x = orient[idx_i];
    float v1y = orient[idx_i + 1];
    float v1z = orient[idx_i + 2];
    float v2x = orient[idx_j];
    float v2y = orient[idx_j + 1];
    float v2z = orient[idx_j + 2];
    float dot = v1x * v2x + v1y * v2y + v1z * v2z;
    float dot2 = dot * dot;
    unsigned int t = type_ids[g];
    float inv_type = 0.0f;
    if (type_counts[t] > 0) {
        inv_type = 1.0f / (float)type_counts[t];
    }
    float inv_groups = 1.0f / (float)n_groups;
    int stride = n_types + 1;
    int base = idt * stride;
    atomicAdd(&corr[base + t], dot * inv_type);
    atomicAdd(&corr[base + n_types], dot * inv_groups);
    atomicAdd(&corr_p2[base + t], dot2 * inv_type);
    atomicAdd(&corr_p2[base + n_types], dot2 * inv_groups);
}

__global__ void conductivity_total(const float* com,
                                   const float* times,
                                   const float* charges,
                                   int n_groups,
                                   int n_frames,
                                   int ndframe,
                                   float dt0,
                                   float eps_num,
                                   float eps_add,
                                   int frame_dec_start,
                                   int frame_dec_stride,
                                   int frame_dec_enabled,
                                   int dt_cut1,
                                   int dt_stride1,
                                   int dt_cut2,
                                   int dt_stride2,
                                   int dt_dec_enabled,
                                   int cols,
                                   float* out,
                                   unsigned int* n_diff) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    if (i >= n_frames || j >= n_frames) return;
    if (j <= i) return;
    if (frame_dec_enabled) {
        int idx1 = i + 1;
        if (idx1 > frame_dec_start && (idx1 % frame_dec_stride) != 0) return;
    }
    float dt = times[j] - times[i];
    int idt = (int)floorf((dt + eps_num) / dt0 + eps_add);
    if (idt <= 0 || idt > ndframe) return;
    if (dt_dec_enabled) {
        if (idt > dt_cut2 && (idt % dt_stride2) != 0) return;
        if (idt > dt_cut1 && (idt % dt_stride1) != 0) return;
    }
    atomicAdd(&n_diff[idt], 1);
    float sx = 0.0f;
    float sy = 0.0f;
    float sz = 0.0f;
    for (int g = 0; g < n_groups; g++) {
        int idx_i = (i * n_groups + g) * 3;
        int idx_j = (j * n_groups + g) * 3;
        float dx = (com[idx_i] - com[idx_j]) * charges[g];
        float dy = (com[idx_i + 1] - com[idx_j + 1]) * charges[g];
        float dz = (com[idx_i + 2] - com[idx_j + 2]) * charges[g];
        sx += dx;
        sy += dy;
        sz += dz;
    }
    float dot = sx * sx + sy * sy + sz * sz;
    int base = idt * cols;
    atomicAdd(&out[base + (cols - 1)], dot);
}

__global__ void conductivity_transference(const float* com,
                                          const float* times,
                                          const float* charges,
                                          const unsigned int* type_ids,
                                          const float* type_charge,
                                          int n_groups,
                                          int n_types,
                                          int n_frames,
                                          int ndframe,
                                          float dt0,
                                          float eps_num,
                                          float eps_add,
                                          int frame_dec_start,
                                          int frame_dec_stride,
                                          int frame_dec_enabled,
                                          int dt_cut1,
                                          int dt_stride1,
                                          int dt_cut2,
                                          int dt_stride2,
                                          int dt_dec_enabled,
                                          int cols,
                                          float* out) {
    int pair = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    int j = blockIdx.z;
    int total_pairs = n_groups * n_groups;
    if (pair >= total_pairs || i >= n_frames || j >= n_frames) return;
    if (j <= i) return;
    if (frame_dec_enabled) {
        int idx1 = i + 1;
        if (idx1 > frame_dec_start && (idx1 % frame_dec_stride) != 0) return;
    }
    float dt = times[j] - times[i];
    int idt = (int)floorf((dt + eps_num) / dt0 + eps_add);
    if (idt <= 0 || idt > ndframe) return;
    if (dt_dec_enabled) {
        if (idt > dt_cut2 && (idt % dt_stride2) != 0) return;
        if (idt > dt_cut1 && (idt % dt_stride1) != 0) return;
    }
    int k = pair / n_groups;
    int l = pair - k * n_groups;
    if (l < k) return;
    unsigned int type_k = type_ids[k];
    unsigned int type_l = type_ids[l];
    if (type_charge[type_k] == 0.0f || type_charge[type_l] == 0.0f) return;
    int idx_i_k = (i * n_groups + k) * 3;
    int idx_j_k = (j * n_groups + k) * 3;
    int idx_i_l = (i * n_groups + l) * 3;
    int idx_j_l = (j * n_groups + l) * 3;
    float v1x = (com[idx_i_k] - com[idx_j_k]) * charges[k];
    float v1y = (com[idx_i_k + 1] - com[idx_j_k + 1]) * charges[k];
    float v1z = (com[idx_i_k + 2] - com[idx_j_k + 2]) * charges[k];
    float v2x = (com[idx_i_l] - com[idx_j_l]) * charges[l];
    float v2y = (com[idx_i_l + 1] - com[idx_j_l + 1]) * charges[l];
    float v2z = (com[idx_i_l + 2] - com[idx_j_l + 2]) * charges[l];
    float dot = v1x * v2x + v1y * v2y + v1z * v2z;
    float factor = (l == k) ? 1.0f : 2.0f;
    int idx = type_l + type_k * n_types;
    int base = idt * cols;
    atomicAdd(&out[base + idx], factor * dot);
}

__global__ void ion_pair_frame_cat(const float* com,
                                   const float* box_l,
                                   const unsigned int* cat_indices,
                                   const unsigned int* ani_indices,
                                   int n_cat,
                                   int n_ani,
                                   int n_groups,
                                   int n_frames,
                                   int max_cluster,
                                   float rclust,
                                   unsigned int* idx_pair,
                                   unsigned int* idx_cluster) {
    int cat = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (cat >= n_cat || frame >= n_frames) return;
    int pair_offset = frame * n_cat + cat;
    idx_pair[pair_offset] = 0xFFFFFFFFu;
    int cluster_offset = (frame * n_cat + cat) * max_cluster;
    for (int k = 0; k < max_cluster; k++) {
        idx_cluster[cluster_offset + k] = 0xFFFFFFFFu;
    }
    unsigned int cat_group = cat_indices[cat];
    int idx_cat = (frame * n_groups + cat_group) * 3;
    float cx = com[idx_cat];
    float cy = com[idx_cat + 1];
    float cz = com[idx_cat + 2];
    float lx = box_l[frame * 3 + 0];
    float ly = box_l[frame * 3 + 1];
    float lz = box_l[frame * 3 + 2];
    float best = 1.0e30f;
    for (int a = 0; a < n_ani; a++) {
        unsigned int ani_group = ani_indices[a];
        int idx_ani = (frame * n_groups + ani_group) * 3;
        float dx = com[idx_ani] - cx;
        float dy = com[idx_ani + 1] - cy;
        float dz = com[idx_ani + 2] - cz;
        if (lx > 0.0f) {
            dx -= roundf(dx / lx) * lx;
            dy -= roundf(dy / ly) * ly;
            dz -= roundf(dz / lz) * lz;
        }
        float r2 = dx * dx + dy * dy + dz * dz;
        float r = sqrtf(r2);
        if (r < rclust) {
            for (int k = 0; k < max_cluster; k++) {
                unsigned int* slot = &idx_cluster[cluster_offset + k];
                if (*slot == 0xFFFFFFFFu) {
                    *slot = ani_group;
                    break;
                }
            }
            if (r < best) {
                best = r;
                idx_pair[pair_offset] = ani_group;
            }
        }
    }
}

__global__ void ion_pair_frame_ani(const float* com,
                                   const float* box_l,
                                   const unsigned int* cat_indices,
                                   const unsigned int* ani_indices,
                                   int n_cat,
                                   int n_ani,
                                   int n_groups,
                                   int n_frames,
                                   int max_cluster,
                                   float rclust,
                                   unsigned int* idx_pair,
                                   unsigned int* idx_cluster) {
    int ani = blockIdx.x * blockDim.x + threadIdx.x;
    int frame = blockIdx.y;
    if (ani >= n_ani || frame >= n_frames) return;
    int pair_offset = frame * n_ani + ani;
    idx_pair[pair_offset] = 0xFFFFFFFFu;
    int cluster_offset = (frame * n_ani + ani) * max_cluster;
    for (int k = 0; k < max_cluster; k++) {
        idx_cluster[cluster_offset + k] = 0xFFFFFFFFu;
    }
    unsigned int ani_group = ani_indices[ani];
    int idx_ani = (frame * n_groups + ani_group) * 3;
    float cx = com[idx_ani];
    float cy = com[idx_ani + 1];
    float cz = com[idx_ani + 2];
    float lx = box_l[frame * 3 + 0];
    float ly = box_l[frame * 3 + 1];
    float lz = box_l[frame * 3 + 2];
    float best = 1.0e30f;
    for (int a = 0; a < n_cat; a++) {
        unsigned int cat_group = cat_indices[a];
        int idx_cat = (frame * n_groups + cat_group) * 3;
        float dx = com[idx_cat] - cx;
        float dy = com[idx_cat + 1] - cy;
        float dz = com[idx_cat + 2] - cz;
        if (lx > 0.0f) {
            dx -= roundf(dx / lx) * lx;
            dy -= roundf(dy / ly) * ly;
            dz -= roundf(dz / lz) * lz;
        }
        float r2 = dx * dx + dy * dy + dz * dz;
        float r = sqrtf(r2);
        if (r < rclust) {
            for (int k = 0; k < max_cluster; k++) {
                unsigned int* slot = &idx_cluster[cluster_offset + k];
                if (*slot == 0xFFFFFFFFu) {
                    *slot = cat_group;
                    break;
                }
            }
            if (r < best) {
                best = r;
                idx_pair[pair_offset] = cat_group;
            }
        }
    }
}

__global__ void ion_pair_corr_cat(const unsigned int* idx_pair,
                                  const unsigned int* idx_cluster,
                                  int n_frames,
                                  int n_cat,
                                  int max_cluster,
                                  unsigned int* ip_out,
                                  unsigned int* cp_out) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    int cat = blockIdx.z;
    if (i >= n_frames || j >= n_frames || cat >= n_cat) return;
    if (j < i) return;
    int dt = j - i;
    unsigned int a = idx_pair[i * n_cat + cat];
    unsigned int b = idx_pair[j * n_cat + cat];
    if (a == b) {
        atomicAdd(&ip_out[dt], 1);
    }
    int off_i = (i * n_cat + cat) * max_cluster;
    int off_j = (j * n_cat + cat) * max_cluster;
    for (int k = 0; k < max_cluster; k++) {
        if (idx_cluster[off_i + k] != idx_cluster[off_j + k]) {
            return;
        }
    }
    atomicAdd(&cp_out[dt], 1);
}

__global__ void ion_pair_corr_ani(const unsigned int* idx_pair,
                                  const unsigned int* idx_cluster,
                                  int n_frames,
                                  int n_ani,
                                  int max_cluster,
                                  unsigned int* ip_out,
                                  unsigned int* cp_out) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y;
    int ani = blockIdx.z;
    if (i >= n_frames || j >= n_frames || ani >= n_ani) return;
    if (j < i) return;
    int dt = j - i;
    unsigned int a = idx_pair[i * n_ani + ani];
    unsigned int b = idx_pair[j * n_ani + ani];
    if (a == b) {
        atomicAdd(&ip_out[dt], 1);
    }
    int off_i = (i * n_ani + ani) * max_cluster;
    int off_j = (j * n_ani + ani) * max_cluster;
    for (int k = 0; k < max_cluster; k++) {
        if (idx_cluster[off_i + k] != idx_cluster[off_j + k]) {
            return;
        }
    }
    atomicAdd(&cp_out[dt], 1);
}
"#;
