/**

  Conjugate gradient solver for (approximate) TV regularized tomographic reconstruction.

  See: Weiss, Blanc-Feraud, Aubert,
    "Efficient schemes for total variation minimization under constraints in image processing"

**/


#include <stdio.h>
#include <stdlib.h>
#include<math.h>
#include <cuda.h>
#include <cublas.h>
#include <cuComplex.h>


/// ****************************************************************************
/// ******************** Code from other files  ********************************
/// ****************************************************************************


#define FROMCU
extern "C" {
#include<CCspace.h>
}
#  define CUDACHECK \
{ cudaThreadSynchronize(); \
  cudaError_t last = cudaGetLastError();\
  if(last!=cudaSuccess) {\
  printf("ERRORX: %s  %s  %i \n", cudaGetErrorString( last),    __FILE__, __LINE__    );    \
  exit(1);\
  }\
}

#  define CUDA_SAFE_CALL_NO_SYNC( call) {                                    \
  cudaError err = call;                                                    \
  if( cudaSuccess != err) {                                                \
  fprintf(stderr, "Cuda error in file '%s' in line %i : %s.\n",        \
  __FILE__, __LINE__, cudaGetErrorString( err) );              \
  exit(EXIT_FAILURE);                                                  \
  } }

#  define CUDA_SAFE_CALL( call)     CUDA_SAFE_CALL_NO_SYNC(call);                                            \


#define CUDA_SAFE_FFT(call){                                                   \
  cufftResult err = call;                                                    \
  if( CUFFT_SUCCESS != err) {                                                \
  fprintf(stderr, "Cuda error in file '%s' in line %i : %d.\n",          \
  __FILE__, __LINE__, err );                                     \
  exit(EXIT_FAILURE);                                                    \
  } }


typedef struct ParamsForTomo {
  Gpu_Context *ctxstruct;
  float DETECTOR_DUTY_RATIO;
  int DETECTOR_DUTY_OVERSAMPLING;
} ParamsForTomo ;


int iDivUp(int a, int b);

int nextpow2_cp_padded(int v);

int nextpow2_cp(int v);

void proj_wrapper(ParamsForTomo p4t, float* d_sino, float* d_image, int dimslice);

void backproj_wrapper(ParamsForTomo p4t, float* d_sino, float* d_image);

__global__  void cp_kern_compute_discrete_ramp(int length, cufftReal* oArray);

cufftComplex* cp_compute_discretized_ramp_filter(int length, cufftReal* d_r, cufftComplex* d_i, cufftHandle myplan);

__global__ void cp_kern_fourier_filter(cufftComplex* inArray, cufftComplex* filter, int sizeX, int sizeY);

#define fftbunch 128

// For proj/backproj if precondition
extern float* global_sino_tmp;
extern float* global_slice_tmp;




/// ****************************************************************************
/// ******************** Kernels, calls ****************************************
/// ****************************************************************************

__global__ void kern_anscombe(float* arr, int Nr, int Nc) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  if (gidx < Nc && gidy < Nr) {
    float val = arr[gidy*Nc + gidx];
    arr[gidy*Nc + gidx] = 2.0f * sqrtf(val + 0.375f);
  }
}

#define SQRT32 1.2247448713915889
#define IANSCOMBE(y) 0.25*y*y + 0.25*SQRT32/y - 1.375/y/y + 0.625*SQRT32/y/y/y - 0.125

__global__ void kern_ianscombe(float* arr, int Nr, int Nc) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  if (gidx < Nc && gidy < Nr) {
    float val = arr[gidy*Nc + gidx];
    arr[gidy*Nc + gidx] = IANSCOMBE(val);
  }
}

int call_anscombe(float* d_arr, int Nr, int Nc, int direction=1) {
  int tpb = 16; // tune for max perfs.
  dim3 n_blocks = dim3(iDivUp(Nc, tpb), iDivUp(Nr, tpb), 1);
  dim3 n_threads_per_block = dim3(tpb, tpb, 1);
  if (direction >= 0) kern_anscombe<<<n_blocks, n_threads_per_block>>>(d_arr, Nr, Nc);
  else kern_ianscombe<<<n_blocks, n_threads_per_block>>>(d_arr, Nr, Nc);
  return 0;
}


/**
  Gradient of the Moreau-Yosida approximation of Total Variation :
      grad_J = -div(psi)  where
                   | grad(x)/norm(grad(x))  if norm(grad(x)) >= mu
        psi(x) = < |
                   | grad(x)/mu             otherwise

  See Weiss, Blanc-Feraud, Aubert,
  "EFFICIENT SCHEMES FOR TOTAL VARIATION MINIMIZATION UNDER CONSTRAINTS IN IMAGE PROCESSING",
   p. 8

*/
// Computes the (opposite) divergence of the normalized gradient.
// The gradient is divided by its magnitude |g| if |g| > mu,  and divided by mu otherwise.
// Note that if there is no scaling, this computes -div(grad) (i.e the opposite of the Laplacian)
// The gradient and divergence steps are done in the same kernel.
// Grad and Div are computed using the following conventions:
//     grad([a1, a2, ..., aN]) = [a2-a1, ..., 0]
//     div([g1, ..., gN]) = [-g1, -g2+g1, ... g(N-1)]
//
__global__ void kern_grad_tv_smoothed(float* arr, float* res, int Nr, int Nc, float mu) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx >= Nc || gidy >= Nr) return;

  float gradx, grady;
  float grad_prevline_x, grad_prevline_y;
  float grad_prevcol_x, grad_prevcol_y;
  float n = 1.0f;

  if (gidx < Nc-1) gradx = arr[gidy*Nc + gidx+1] - arr[gidy*Nc +gidx];
  else gradx = 0;
  if (gidy < Nr-1) grady = arr[(gidy+1)*Nc + gidx] - arr[gidy*Nc + gidx];
  else grady = 0;

  if (gidy > 0) {
    if (gidx < Nc-1) grad_prevline_x = arr[(gidy-1)*Nc + gidx+1] - arr[(gidy-1)*Nc + gidx];
    else grad_prevline_x = 0;
    grad_prevline_y = arr[gidy*Nc + gidx] - arr[(gidy-1)*Nc + gidx];
  }
  else {
    grad_prevline_x = 0;
    grad_prevline_y = 0;
  }
  if (gidx > 0) {
    grad_prevcol_x = arr[gidy*Nc + gidx] - arr[gidy*Nc + gidx-1];
    if (gidy < Nr-1) grad_prevcol_y = arr[(gidy+1)*Nc + gidx-1] - arr[gidy*Nc + gidx-1];
    else grad_prevcol_y = 0;
  }
  else {
    grad_prevcol_x = 0;
    grad_prevcol_y = 0;
  }

  n = max(sqrtf(gradx*gradx + grady*grady), mu);
  gradx /= n;
  grady /= n;
  n = max(sqrtf(grad_prevcol_x*grad_prevcol_x + grad_prevcol_y*grad_prevcol_y), mu);
  grad_prevcol_x /= n;
  n = max(sqrtf(grad_prevline_x*grad_prevline_x + grad_prevline_y*grad_prevline_y), mu);
  grad_prevline_y /= n;

  res[gidy*Nc + gidx] = (-gradx + grad_prevcol_x) + (-grady + grad_prevline_y);
}


int call_grad_tv_smoothed(float* d_in, float* d_out, int Nr, int Nc, float mu) {
  int tpb = 16; // tune for max perfs.
  dim3 n_blocks = dim3(iDivUp(Nc, tpb), iDivUp(Nr, tpb), 1);
  dim3 n_threads_per_block = dim3(tpb, tpb, 1);
  kern_grad_tv_smoothed<<<n_blocks, n_threads_per_block>>>(d_in, d_out, Nr, Nc, mu);

  return 0;
}




/// Moreau-Yosida approximation of the TV. Used only to compute the objective function
// arr is modified !
__global__ void kern_tv_smoothed(float* arr, int Nr, int Nc, float mu) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  if (gidx < Nc && gidy < Nr) {
    // Compute the two components (vertic and horiz) of the gradient
    // [a1, ..., aN] -->  [a2-a1, ..., 0]
    float gradx, grady; // .x: horizontal component (axis=1) ; .y: vertical component (axis=0)
    if (gidx < Nc-1) gradx = arr[gidy*Nc+gidx+1] - arr[gidy*Nc+gidx];
    else gradx = 0.0f;
    if (gidy < Nr-1) grady = arr[(gidy+1)*Nc+gidx] - arr[gidy*Nc+gidx];
    else grady = 0.0f;
    float n = sqrtf(gradx*gradx + grady*grady);
    // Huber function of the gradient magnitude
    if (n < mu) n = n*n/(2.0f*mu) + mu/2.0f;
    arr[gidy*Nc + gidx] = n;
  }
}

// d_arr is modified !
float call_tv_smoothed(float* d_arr, int Nr, int Nc, float mu) {
  int tpb = 16; // tune for max perfs.
  dim3 n_blocks = dim3(iDivUp(Nc, tpb), iDivUp(Nr, tpb), 1);
  dim3 n_threads_per_block = dim3(tpb, tpb, 1);
  kern_tv_smoothed<<<n_blocks, n_threads_per_block>>>(d_arr, Nr, Nc, mu);
  return cublasSasum(Nr*Nc, d_arr, 1);
}





/// ****************************************************************************
/// ******************** Algorithms ********************************************
/// ****************************************************************************


#define DEBUG 0

int conjugate_gradient(ParamsForTomo p4t, float* d_data, float* d_result, int n_it, float Lambda, float mu) {

  // Geometry
  // --------
  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs = p4t.ctxstruct->nprojs_span;
  int Nr = p4t.ctxstruct->num_y, Nc = p4t.ctxstruct->num_x;
  size_t slicesize = Nr*Nc*sizeof(float);

  // Conjugate gradient variables (7+sino)
  // -------------------------------------
  // x : the unknown slice
  float* d_x;
  d_x = d_result;
  cudaMemset(d_x, 0, slicesize);
  // gradf : gradient of quadratic part
  float* d_gradf;
  cudaMalloc(&d_gradf, slicesize);
  // gradF : gradient of the whole function
  float* d_gradF;
  cudaMalloc(&d_gradF, slicesize);
  // d : conjugate direction
  float* d_d;
  cudaMalloc(&d_d, slicesize);
  // ATAd : computation of P^T P d at each iteration
  float* d_ATAd;
  cudaMalloc(&d_ATAd, slicesize);
  // gradf_old, gradF_old
  float* d_gradf_old, *d_gradF_old;
  cudaMalloc(&d_gradf_old, slicesize);
  cudaMalloc(&d_gradF_old, slicesize);
  // a sinogram
  float* d_sino;
  cudaMalloc(&d_sino, num_bins*nprojs*sizeof(float));
  // a temp. slice
  float* d_tmp;
  cudaMalloc(&d_tmp, slicesize);
  CUDACHECK;

  // Energy computation
  float* energies;
  if (1) {
    energies = (float*) calloc(n_it, sizeof(float));
  }

  // Conjugate gradient initialization
  // ----------------------------------
  //  grad_f = -Kadj(data)
  backproj_wrapper(p4t, d_data, d_gradf);
  cublasSscal(Nr*Nc, -1.0f, d_gradf, 1);
  // x = 0*Kadj(data) # start from 0
  // grad_F = grad_f + Lambda*grad_tv_smoothed(x, mu)
  call_grad_tv_smoothed(d_x, d_gradF, Nr, Nc, mu);
  cublasSscal(Nr*Nc, Lambda, d_gradF, 1);
  cublasSaxpy(Nr*Nc, 1.0, d_gradf, 1, d_gradF, 1);
  // d = - copy(grad_F)
  cudaMemcpy(d_d, d_gradF, slicesize, cudaMemcpyDeviceToDevice);
  cublasSscal(Nr*Nc, -1.0, d_d, 1);
  float alpha = 0.0f, beta = 0.0f;

  // Conjugate gradient main loop
  // -----------------------------
  for (int k = 0; k < n_it; k++) {
    // grad_f_old = grad_f
    cudaMemcpy(d_gradf_old, d_gradf, slicesize, cudaMemcpyDeviceToDevice);
    // grad_F_old = grad_F
    cudaMemcpy(d_gradF_old, d_gradF, slicesize, cudaMemcpyDeviceToDevice);
    // ATAd = Kadj(K(d))
    proj_wrapper(p4t, d_sino, d_d, Nc);
    backproj_wrapper(p4t, d_sino, d_ATAd);
    // step size:  alpha = dot(d, -grad_F_old)/dot(d, ATAd)
    float denom = cublasSdot(Nr*Nc, d_d, 1, d_ATAd, 1);
    float num = -cublasSdot(Nr*Nc, d_d, 1, d_gradF_old, 1);
    alpha = num/denom;
    if (DEBUG) printf("alpha = %f\n", alpha);

    // Update variables
    // x = x + alpha*d
    cublasSaxpy(Nr*Nc, alpha, d_d, 1, d_x, 1);
    // grad_f = grad_f_old + alpha*ATAd
    cublasSaxpy(Nr*Nc, alpha, d_ATAd, 1, d_gradf, 1);
    if ((k % 50) == 0) {// re-compute periodically the gradient every K iterations to avoid error accumulation
      proj_wrapper(p4t, d_sino, d_x, Nc);
      cublasSaxpy(num_bins*nprojs, -1.0f, d_data, 1, d_sino, 1);
      backproj_wrapper(p4t, d_sino, d_gradf);
    }

    // grad_F = grad_f + Lambda*grad_tv_smoothed(x,mu)
    cudaMemcpy(d_gradF, d_gradf, slicesize, cudaMemcpyDeviceToDevice);
    call_grad_tv_smoothed(d_x, d_tmp, Nr, Nc, mu);
    cublasSaxpy(Nr*Nc, Lambda, d_tmp, 1, d_gradF, 1);
    // beta = dot(grad_F, grad_F - grad_F_old)/norm2sq(grad_F_old) # Polak-Ribiere
    cudaMemcpy(d_tmp, d_gradF, slicesize, cudaMemcpyDeviceToDevice);
    cublasSaxpy(Nr*Nc, -1.0f, d_gradF_old, 1, d_tmp, 1);
    num = cublasSdot(Nr*Nc, d_gradF, 1, d_tmp, 1);
    denom = cublasSnrm2(Nr*Nc, d_gradF_old, 1);
    beta = num/denom/denom;
    if (DEBUG) printf("beta = %f/%f = %f\n", num, denom*denom, beta);
    if (beta < 0) beta = 0.0f;

    // d = -grad_F + beta*d
    cublasSscal(Nr*Nc, beta, d_d, 1);
    cublasSaxpy(Nr*Nc, -1.0f, d_gradF, 1, d_d, 1);

    // Stoping criterion
    if (fabsf(alpha) < 1e-15) { // TODO : try other bounds
      printf("Warning : minimum step reached, interrupting at iteration %d\n", k);
      break;
    }

    // Compute energy
    if (p4t.ctxstruct->verbosity > 1) {
      float fid, tv;
      // fid = norm2sq(K(x)-data)
      proj_wrapper(p4t, d_sino, d_x, Nc);
      cublasSaxpy(num_bins*nprojs, -1.0, d_data, 1, d_sino, 1);
      fid = cublasSnrm2(num_bins*nprojs, d_sino, 1);
      fid = fid*fid;
      // tv = tv_smoothed(x,mu)
      cudaMemcpy(d_tmp, d_x, slicesize, cudaMemcpyDeviceToDevice);
      tv = call_tv_smoothed(d_tmp, Nr, Nc, mu);
      // energy = fid+Lambda*tv
      energies[k] = fid + Lambda*tv;
      if ((k % 10) == 0) printf("[%d] Energy = %e \t Fidelity = %e \t TV = %e\n", k, energies[k], fid, tv);
    }
    if (DEBUG) CUDACHECK;
  }

  // Free memory
  cudaFree(d_d);
  cudaFree(d_gradf);
  cudaFree(d_gradf_old);
  cudaFree(d_gradF);
  cudaFree(d_gradF_old);
  cudaFree(d_ATAd);
  cudaFree(d_tmp);
  cudaFree(d_sino);
  CUDACHECK;
  if (1) {
    FILE* flid = fopen("energy_conjgrad.dat", "wb");
    fwrite(energies, sizeof(float), n_it, flid);
    fclose(flid);
    free(energies);
  }

  return 0;
}






/// ****************************************************************************
/// ******************** Entry point *******************************************
/// ****************************************************************************





extern "C" {
  int conjgrad_driver(Gpu_Context* self,
      float* data,
      float* SLICE,
      float DETECTOR_DUTY_RATIO,
      int DETECTOR_DUTY_OVERSAMPLING,
      float Lambda);
}




int conjgrad_driver(Gpu_Context* self, float* data, float* SLICE, float DETECTOR_DUTY_RATIO, int DETECTOR_DUTY_OVERSAMPLING, float Lambda) {

  cuCtxSetCurrent ( *((CUcontext *) self->gpuctx  ))  ;

  //Import parameters from self
  int num_bins = self->num_bins;
  int num_projs = self->nprojs_span;
  int dimslice = self->num_x ;
  ParamsForTomo p4t  =  (ParamsForTomo)  { (Gpu_Context*) self, DETECTOR_DUTY_RATIO, DETECTOR_DUTY_OVERSAMPLING } ;

  if (self->verbosity > 0) {
    puts("------------------------------------------------------------------------------");
    puts("------------------ Entering Conjugate gradient driver ------------------------");
    puts("------------------------------------------------------------------------------");
    printf("Nb = %d , Np = %d, d = %d, B = %f, mu = %f\n",num_bins, num_projs, dimslice, Lambda, self->CG_MU);
    if (p4t.ctxstruct->DO_PRECONDITION) puts("Ramp-filtering is ENABLED");
    else puts("Ramp-filtering is DISABLED");
  }


  //Prepare cuFFT plan for FBP
  CUDA_SAFE_CALL(cudaMalloc(&self->precond_params_dl.d_r_sino_error, fftbunch*nextpow2_cp_padded(num_bins)*sizeof(cufftReal)));
  CUDA_SAFE_CALL(cudaMalloc(&self->precond_params_dl.d_i_sino_error, fftbunch*nextpow2_cp_padded(num_bins)*sizeof(cufftComplex)));
  static int plans_are_computed = 0;
  if(!plans_are_computed) {
      plans_are_computed = 1;
      CUDA_SAFE_FFT(cufftPlan1d((cufftHandle *) &self->precond_params_dl.planRamp_forward, nextpow2_cp_padded(num_bins),CUFFT_R2C,fftbunch));
      CUDA_SAFE_FFT(cufftPlan1d((cufftHandle *) &self->precond_params_dl.planRamp_backward,nextpow2_cp_padded(num_bins),CUFFT_C2R,fftbunch));
  }
  cufftComplex* d_i_discrete_ramp = cp_compute_discretized_ramp_filter(nextpow2_cp_padded(num_bins), self->precond_params_dl.d_r_sino_error, self->precond_params_dl.d_i_sino_error, self->precond_params_dl.planRamp_forward);
  self->precond_params_dl.filter_coeffs = d_i_discrete_ramp; //size : nextpow2(num_bins)/2+1


  // Allocate memory
  float* d_result;
  cudaMalloc(&d_result, dimslice*dimslice*sizeof(float));
  float* d_data;
  CUDA_SAFE_CALL(cudaMalloc(&d_data, num_bins*num_projs*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemcpy(d_data, data,  num_bins*num_projs*sizeof(float), cudaMemcpyHostToDevice));
  if (1 || p4t.ctxstruct->DO_PRECONDITION) {
    CUDA_SAFE_CALL(cudaMalloc(&global_sino_tmp, num_bins*num_projs*sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc(&global_slice_tmp, dimslice*dimslice*sizeof(float)));
  }

  // Run
  float mu = self->CG_MU;
  Lambda = 1.0f/Lambda; // Bizarre, le role de Lambda est inverse dans cet algo. Il faut voir pourquoi.
  conjugate_gradient(p4t, d_data, d_result, self->ITERATIVE_CORRECTIONS, Lambda, mu);
  cudaMemcpy(SLICE, d_result, dimslice*dimslice*sizeof(float), cudaMemcpyDeviceToHost);

  // Free memory
  cudaFree(d_result);
  cudaFree(d_data);
  cudaFree(self->precond_params_dl.d_r_sino_error);
  cudaFree(self->precond_params_dl.d_i_sino_error);
  if (1 || p4t.ctxstruct->DO_PRECONDITION) {
    cudaFree(global_sino_tmp);
    cudaFree(global_slice_tmp);
  }
  CUDACHECK;

  return 0;
}
