/*
 * Based on the original CMCMC implementation by Louis Aslett, April 2018.
 * Revised for this package by Ahmad ALQabandi and Louis Aslett, 2026.
 */

#include "kernels/cmcmc_kernel.h"
#include <cuda_runtime.h>





// Setup / Cleanup
__global__ void metrop_1gpu_sm30_setup_d(float *theta_cur, float *logdens_cur,
                                         const float *Xf, const int *Xi) {
  // Initial offset
  const int n = K / 2;
  const int off = threadIdx.x + blockIdx.x*blockDim.x;
  if (off >= n) return;

  // Shift pointers
  theta_cur   += off;
  logdens_cur += off;

  // Gather the strided theta vector into a contiguous local array for the target kernel.
  float theta_local[DIM];
  for (int i = 0; i < DIM; i++) {
    theta_local[i] = theta_cur[i * n];
  }

  // Compute log density for initial population
  *logdens_cur = logdens(theta_local, Xf, Xi);
}

extern "C" {
  void metrop_1gpu_sm30_setup_D_stream(int threads, int blocks,
                                       float *theta_cur, float *logdens_cur,
                                       const float *Xf, const int *Xi,
                                       cudaStream_t stream) {
    metrop_1gpu_sm30_setup_d<<<blocks, threads, 0, stream>>>(theta_cur, logdens_cur, Xf, Xi);
  }

  void metrop_1gpu_sm30_setup_D(int threads, int blocks,
                                float *theta_cur, float *logdens_cur,
                                const float *Xf, const int *Xi) {
    metrop_1gpu_sm30_setup_D_stream(threads, blocks, theta_cur, logdens_cur, Xf, Xi, 0);
  }
}





// Iterate
__global__ void metrop_1gpu_sm30_d(float *theta_cur, float *logdens_cur,
                                   float *theta_prop, float *logdens_nxt,
                                   const float *Xf, const int *Xi,
                                   const float *unif) {
  // Shift pointers to our particle in current/proposal streams, as well as RNG stream.
  const int n = K / 2;
  const int off = threadIdx.x + blockIdx.x*blockDim.x;
  if (off >= n) return;
  theta_cur   += off;
  logdens_cur += off;
  theta_prop  += off;
  unif        += off;
  (void)logdens_nxt;

  float theta_local[DIM];
  float theta_prop_local[DIM];

  for (int i = 0; i < DIM; i++) {
    theta_local[i] = theta_cur[i * n];
    theta_prop_local[i] = theta_local[i] + theta_prop[i * n];
  }

  const float logdens_prop = logdens(theta_prop_local, Xf, Xi);
  const float log_ratio = logdens_prop - *logdens_cur;
  if(log_ratio >= 0.0f || logf(*unif) < log_ratio) {
    for(int i = 0; i < DIM; i++) {
      theta_cur[i * n] = theta_prop_local[i];
    }
    *logdens_cur = logdens_prop;
  }
}

extern "C" {
  void metrop_1gpu_sm30_D_stream(int threads, int blocks,
                                 float *theta_cur, float *logdens_cur,
                                 float *theta_prop, float *logdens_nxt,
                                 const float *Xf, const int *Xi,
                                 const float *unif_d,
                                 cudaStream_t stream) {
    metrop_1gpu_sm30_d<<<blocks, threads, 0, stream>>>(theta_cur, logdens_cur,
                                                      theta_prop, logdens_nxt,
                                                      Xf, Xi,
                                                      unif_d);
  }

  void metrop_1gpu_sm30_D(int threads, int blocks,
                          float *theta_cur, float *logdens_cur,
                          float *theta_prop, float *logdens_nxt,
                          const float *Xf, const int *Xi,
                          const float *unif_d, int *acc) {
    (void)acc;
    metrop_1gpu_sm30_D_stream(threads, blocks,
                              theta_cur, logdens_cur,
                              theta_prop, logdens_nxt,
                              Xf, Xi,
                              unif_d,
                              0);
  }
}


// Iterate several Metropolis steps using a pre-generated proposal and uniform
// chunk. Proposals are stored as a column-major ((chunk_it * K/2) x DIM) matrix.
__global__ void metrop_1gpu_sm30_chunk_d(float *theta_cur, float *logdens_cur,
                                         const float *theta_prop_chunk,
                                         const float *Xf, const int *Xi,
                                         const float *unif_chunk,
                                         int chunk_it) {
  const int n = K / 2;
  const int off = threadIdx.x + blockIdx.x*blockDim.x;
  if (off >= n) return;

  theta_cur   += off;
  logdens_cur += off;

  float theta_local[DIM];
  float theta_prop_local[DIM];
  for(int i = 0; i < DIM; i++) {
    theta_local[i] = theta_cur[i*n];
  }

  float logdens_local = *logdens_cur;

  for(int itn = 0; itn < chunk_it; itn++) {
    const int row = itn*n + off;
    for(int i = 0; i < DIM; i++) {
      theta_prop_local[i] = theta_local[i] + theta_prop_chunk[row + i*(chunk_it*n)];
    }

    const float logdens_prop = logdens(theta_prop_local, Xf, Xi);
    const float log_ratio = logdens_prop - logdens_local;
    if(log_ratio >= 0.0f || logf(unif_chunk[row]) < log_ratio) {
      for(int i = 0; i < DIM; i++) {
        theta_local[i] = theta_prop_local[i];
      }
      logdens_local = logdens_prop;
    }
  }

  for(int i = 0; i < DIM; i++) {
    theta_cur[i*n] = theta_local[i];
  }
  *logdens_cur = logdens_local;
}

extern "C" {
  void metrop_1gpu_sm30_chunk_D_stream(int threads, int blocks,
                                       float *theta_cur, float *logdens_cur,
                                       const float *theta_prop_chunk,
                                       const float *Xf, const int *Xi,
                                       const float *unif_chunk,
                                       int chunk_it,
                                       cudaStream_t stream) {
    metrop_1gpu_sm30_chunk_d<<<blocks, threads, 0, stream>>>(theta_cur,
                                                             logdens_cur,
                                                             theta_prop_chunk,
                                                             Xf, Xi,
                                                             unif_chunk,
                                                             chunk_it);
  }

  int metrop_1gpu_sm30_occupancy_D(int *threads, int *blocks, int n,
                                   int use_chunk) {
    int min_grid_size = 0;
    int block_size = 0;
    cudaError_t err;

    if (use_chunk) {
      err = cudaOccupancyMaxPotentialBlockSize(&min_grid_size, &block_size,
                                               metrop_1gpu_sm30_chunk_d, 0, 0);
    } else {
      err = cudaOccupancyMaxPotentialBlockSize(&min_grid_size, &block_size,
                                               metrop_1gpu_sm30_d, 0, 0);
    }

    if (err != cudaSuccess) return (int)err;
    if (block_size < 1) return (int)cudaErrorInvalidConfiguration;

    *threads = block_size;
    *blocks = (n + block_size - 1) / block_size;
    return 0;
  }
}
