/**

  Chambolle-Pock algorithm implementation for TV regularized tomographic reconstruction


**/

//------------------------------------------------------------------------------

#include <stdio.h>
#include <stdlib.h>
#include<math.h>
#include<math_constants.h>
#include <cuda.h>
#include <cublas.h>
#include <cuComplex.h>
#include<time.h>
#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 WKSIZE 256

#  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);                                            \


#include <cufft.h>
#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_cp(int a, int b){
    return (a % b != 0) ? (a / b + 1) : (a / b);
}
int nextpow2_cp_padded(int v) {
  int vold=v;
  v--;
  v |= v >> 1;
  v |= v >> 2;
  v |= v >> 4;
  v |= v >> 8;
  v |= v >> 16;
  v++;
  if(v<vold*1.5) v*=2;
  return v;
}
int nextpow2_cp(int v) {
  v--;
  v |= v >> 1;
  v |= v >> 2;
  v |= v >> 4;
  v |= v >> 8;
  v |= v >> 16;
  v++;
  return v;
}
int ilog2_cp(int i) {
  int l = 0;
  while (i >>= 1) { ++l; }
  return l;
}
#define fftbunch 128
#define blsize_cufft 32



__global__ void  transition_kernel(float * d_r_sino_error,int num_bins,int  np2,int  numpjs , float axis_position    ) ;

int iDivUp(int a, int b);

//Align a to nearest higher multiple of b
int iAlignUp(int a, int b);



//------------------------------------------------------------------------------


float* global_sino_tmp;
float* global_slice_tmp;



//-------------------Various utils. ---------------------------------------------


/**
 * @brief compute_histogram : Compute the histogram of a data
 * @param data : Data which we want to compute the histogram from
 * @param size : number of elements of the data
 * @param nbins : number of bins of the resulting histogram
 * @param minv : (result, given by address if not NULL) minimum value
 * @param maxv : (result, given by address if not NULL) maximum value
 * @return hist : the histogram of size "nbins"
 */
int* cp_compute_histogram(float* data, int size, int nbins, float* minv = NULL, float* maxv = NULL) {
  float vmin = data[0], vmax = data[0];
  int i;
  for (i=0; i<size; i++) {
    if (data[i] < vmin) vmin = data[i];
    if (data[i] > vmax) vmax = data[i];
  }
  int* hist = (int*) calloc(nbins,sizeof(int));
  float binsize = (vmax-vmin)/nbins;
  for (i=0;i<size;i++) {
    hist[(int) ((data[i]-vmin)/binsize)]++;
  }
  if (minv) *minv = vmin;
  if (maxv) *maxv = vmax;
  return hist;
}


double cp_kullback_leibler(float* arr1, float* arr2, int size) {
  int nbins = 256;
  int* h1 = cp_compute_histogram(arr1, size, nbins);
  int* h2 = cp_compute_histogram(arr2, size, nbins);
  double dkl = 0;
  int Sx = 0, Sy = 0;
  for (int i=0; i < nbins; i++) {
    Sx += h1[i];
    Sy += h2[i];
    if (h2[i] == 0 || h1[i] == 0) continue;
    dkl += h1[i] * log(((double) h1[i]) / ((double) h2[i]));
  }
  dkl = dkl/Sx + (nbins/(Sx*1.0))*log(Sy/(Sx*1.0));

  free(h1);
  free(h2);
  return dkl;
}




//------------------------------------------------------------------------------






extern "C" {
int chambolle_pock_driver(Gpu_Context* self, float* data, float* SLICE, float DETECTOR_DUTY_RATIO, int DETECTOR_DUTY_OVERSAMPLING, float beta, float beta_rings, float rings_height, float alpha_rings);

}







#define CP_DEBUG 1
#define CP_VERBOSE 1
#define AHMOD 0



__global__  void cp_kern_compute_discrete_ramp(int length, cufftReal* oArray) {
    int gid = threadIdx.x + blockIdx.x*blockDim.x;
    if(gid<=length/2) {
        float val = ((gid & 1) ? (-1.0f/M_PI/M_PI/gid/gid) : (0.0f));
        if (gid == 0) oArray[gid] = 0.25f;
        else if (gid == length/2) oArray[gid] = val;
        else {
            oArray[gid] = val;
            oArray[length-gid] = val;
        }
    }
}
int print_device_array(float* d_array, int numels, char* format);

cufftComplex* cp_compute_discretized_ramp_filter(
        int length,
        cufftReal* d_r,
        cufftComplex* d_i,
        cufftHandle myplan)
{
    int hlen = length/2+1;
    dim3 blk, grd;
    blk = dim3( blsize_cufft , 1 , 1 );
    grd = dim3( iDivUp_cp(length ,blsize_cufft) , 1 , 1 );
    cp_kern_compute_discrete_ramp<<<grd,blk>>> (length, d_r);
    CUDA_SAFE_FFT(cufftExecR2C(myplan,(cufftReal *) d_r,(cufftComplex *) d_i));
    cufftComplex* filterCoefs;
    CUDA_SAFE_CALL(cudaMalloc(&filterCoefs, hlen*sizeof(cufftComplex)));
    CUDA_SAFE_CALL(cudaMemcpy(filterCoefs,d_i, hlen*sizeof(cufftComplex), cudaMemcpyDeviceToDevice));
    
//    print_device_array(( float *  ) filterCoefs ,100, "%e \n"  );
    return filterCoefs;
}

__global__ void cp_kern_fourier_filter(cufftComplex* inArray, cufftComplex* filter, int sizeX, int sizeY) {
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int gidy = threadIdx.y + blockIdx.y*blockDim.y;
    if (gidx < sizeX && gidy < sizeY) {
        inArray[gidy*sizeX+gidx].x *= filter[gidx].x /(2*(sizeX-1));
        inArray[gidy*sizeX+gidx].y *= filter[gidx].x /(2*(sizeX-1));
    }
}

/**
  In-place subtraction elementwise
  array <- array-array2
**/
__global__ void subtract_kernel(float* array, float* array2, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < sizeX && gidy < sizeY) {

    array[gidy*sizeX+gidx] -= array2[gidy*sizeX+gidx];

  }
}
/**
  Compute the SQUARED norm of a gradient
  (each element of the output is input[i].x **2 + input[i].y **2)
  float2*  -->  float*
**/
__global__ void norm_kernel(float2* input, float* output, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < sizeX && gidy < sizeY) {
    int idx = (gidy)*sizeX+gidx;
    output[idx] = (input[idx].x * input[idx].x) + (input[idx].y * input[idx].y);
  }
}
__global__ void  padda_kernel_cp(float * d_r_sino_error,int num_bins,int  np2,int  numpjs     )   {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  if (gidx < np2-num_bins && gidy < numpjs) {
    d_r_sino_error[ gidy*np2 + (num_bins+gidx) ] =  d_r_sino_error[ gidy*np2 + (num_bins-1)*( gidx<(np2-num_bins)/2 )];
  }
}

__global__ void gradient_kernel(float* slice, float2* slice_grad, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  float val_x = 0, val_y = 0;

  if (gidx < sizeX && gidy < sizeY) {
    if (gidx == sizeX-1) val_y = 0;
    else val_y = slice[(gidy)*sizeX+gidx+1] - slice[(gidy)*sizeX+gidx];
    if (gidy == sizeY-1) val_x = 0;
    else val_x = slice[(gidy+1)*sizeX+gidx] - slice[(gidy)*sizeX+gidx];

    slice_grad[(gidy)*sizeX+gidx].x = val_x;
    slice_grad[(gidy)*sizeX+gidx].y = val_y;
  }
}


__global__ void divergence_kernel(float2* slice_grad, float* slice, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  float val_x = 0, val_y = 0;

  if (gidx < sizeX && gidy < sizeY) {
    if (gidx == 0) val_y = slice_grad[(gidy)*sizeX+gidx].y;
    else val_y = slice_grad[(gidy)*sizeX+gidx].y - slice_grad[(gidy)*sizeX+gidx-1].y;
    if (gidy == 0) val_x = slice_grad[(gidy)*sizeX+gidx].x;
    else val_x = slice_grad[(gidy)*sizeX+gidx].x - slice_grad[(gidy-1)*sizeX+gidx].x;
    slice[(gidy)*sizeX+gidx] = val_x + val_y;
  }
}


//p = (p + sigma*(x_tilde_proj - data))/(1+sigma)
__global__ void dual_shrink_kernel(float* dual_p, float* sino, float* data, float sigma, int num_bins, int nprojs_span) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < num_bins && gidy < nprojs_span) {
    int idx = (gidy)*num_bins+gidx;
    dual_p[idx] =  (dual_p[idx] + sigma*(sino[idx] - data[idx]))/(1.0f+sigma);
  }
}

//Projection onto the L-infinity unit ball
//q = proj_linf(q + sigma*gx, lambda_)
__global__ void dual_proj_linf_kernel(float2* dual_q, float2* slice_grad, float sigma, float lambda, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < sizeX && gidy < sizeY) {
    int idx = (gidy)*sizeX+gidx;
    float val_x = dual_q[idx].x + sigma*slice_grad[idx].x;
    float val_y = dual_q[idx].y + sigma*slice_grad[idx].y;
    dual_q[idx].x = copysignf(min(fabsf(val_x), lambda), val_x);
    dual_q[idx].y = copysignf(min(fabsf(val_y), lambda), val_y);
  }
}

//Convert a gradient-like array to a slice-like array, taking the sum of absolute values of each components
__global__ void reduce_gradient_kernel(float* output, float2* input, int sizeX, int sizeY) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < sizeX && gidy < sizeY) {
    int idx = gidy*sizeX+gidx;
    output[idx] = fabsf(input[idx].x)+fabsf(input[idx].y);
  }
}


//add a ring-like array to a sinogram-like array.
__global__  void add_rings_to_sinogram_kernel(float *sino, float* rings, float alpha_rings, int num_bins, int nprojs_span) {

  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if( gidx < num_bins && gidy < nprojs_span) {
    sino[gidy*num_bins +gidx] += alpha_rings*rings[gidx]; //gidy ~ idoppio*nprojs_span+iproj
  }
}

//Projection onto the L-infinity unit ball (ring-like arrays)
//v = proj_linf(v + sigma*r, beta_r)
__global__ void sino_proj_linf_kernel(float* dual_v, float* rings, float sigma, float beta_r, int num_bins) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if (gidx < num_bins && gidy == 0) {
    int idx = gidy*num_bins+gidx;
    float val = dual_v[idx] + sigma*rings[idx];
    dual_v[gidy] = copysignf(min(fabsf(val), beta_r), val);
  }
}


//r = r - tau*(p+v) where p is sinogram-like and v is rings-like
__global__ void update_rings_kernel(float* rings, float* dual_p, float* dual_v,float tau, float rings_height, int num_bins, int nprojs_span) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if( gidx< num_bins && gidy==0) { // FIXME: call with 1D grid and block !
    float sum = 0.0f;
    for(int ipro=0; ipro < nprojs_span; ipro++) {
      sum += dual_p[(gidy*nprojs_span+ipro)*num_bins + gidx];
    }
    rings[gidx + gidy*num_bins] -= tau*(sum + dual_v[gidx + gidy*num_bins]);
    rings[gidx + gidy*num_bins] = min(rings[gidx + gidy*num_bins], rings_height); //prevent rings to take too big values
  }
}


/// performs arr1 /= max(arr2, vmin)
__global__ void division_kernel(float* arr1, float* arr2, int Nx, int Ny, float vmin) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  int tid = gidy*Nx + gidx;

  if (gidx < Nx && gidy < Ny) {
    float val = arr2[tid];
    if (fabsf(val) < vmin) val = vmin;
    arr1[tid] /= val;
  }
}

int cp_call_division(float* d_arr1, float* d_arr2, int Nx, int Ny, float vmin) {
  dim3 blk = dim3(blsize_cufft, blsize_cufft, 1);
  dim3 grd = dim3(iDivUp_cp(Nx, blsize_cufft), iDivUp_cp(Ny, blsize_cufft), 1);
  division_kernel<<<grd, blk>>>(d_arr1, d_arr2, Nx, Ny, vmin);
  return 0;
}



int cp_normalize_mean(float* d_arr, int Nx, int Ny) {

  float mean;
  float* d_ones, one=1.0f;
  cudaMalloc(&d_ones, sizeof(float));
  cudaMemcpy(d_ones, &one, sizeof(float), cudaMemcpyHostToDevice);

  mean = cublasSdot(Nx*Ny, d_arr, 1, d_ones, 0);
  mean /= (Nx*Ny);
  cublasSscal(Nx*Ny, 1.0f/mean, d_arr, 1);

  cudaFree(d_ones);

  return 0;
}



/**
 * @brief calculate_l1_norm : calculate the L1 norm of a gradient-like array.
 * TODO : use a parallel reduction
 */
float calculate_l1_norm(float2* slice_grad, int dimslice_0, int dimslice_1) {
  float* slice_tmp;
//  CUDA_SAFE_CALL(cudaMalloc(&slice_tmp, dimslice_0*dimslice_1*sizeof(float)));
  slice_tmp = global_slice_tmp;
  dim3 blk, grd;
  blk = dim3( blsize_cufft , blsize_cufft , 1 );
  grd = dim3( iDivUp_cp(dimslice_0, blsize_cufft) , iDivUp_cp(dimslice_1, blsize_cufft) , 1 );
  reduce_gradient_kernel<<<grd,blk>>>(slice_tmp, slice_grad, dimslice_1, dimslice_0);
  float l1_norm = cublasSasum(dimslice_0*dimslice_1, slice_tmp, 1);
//  CUDA_SAFE_CALL(cudaFree(slice_tmp));
  return l1_norm;
}



// q2 = (q2 + sigma*U(x_tilde))/(1.0 + sigma/Lambda2)
__global__ void shrink_gradient_kernel(float2* dual_q2, float2* slice_grad, int Nx, int Ny, float sigma, float beta) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  if( gidx< Nx && gidy < Ny) {
    float2 val = dual_q2[gidy*Nx+gidx];
    val.x = (val.x + sigma*slice_grad[gidy*Nx+gidx].x)/(1.0f + sigma/beta);
    val.y = (val.y + sigma*slice_grad[gidy*Nx+gidx].y)/(1.0f + sigma/beta);
    dual_q2[gidy*Nx+gidx] = val;
  }
}


int call_shrink_gradient(float2* dual_q2, float2* slice_grad, int Nx, int Ny, float sigma, float beta_L2) {
  dim3 grd, blk;
  blk = dim3(blsize_cufft, blsize_cufft, 1);
  grd = dim3(iDivUp_cp(Nx, blsize_cufft), iDivUp_cp(Ny, blsize_cufft), 1);
  shrink_gradient_kernel<<<grd, blk>>>(dual_q2, slice_grad, Nx, Ny, sigma, beta_L2);
  return 0;
}


// positivity constraint : projection on the positive subspace
__global__ void positivity_kernel(float* slice, int Nx, int Ny) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;

  if( gidx< Nx && gidy < Ny) {
    if (slice[gidy*Nx+gidx] < 0) slice[gidy*Nx+gidx] = 0;
  }

}













int write_device_array(float* d_array, int numels, const char* fname) {
  FILE* fid = fopen(fname, "wb");
  if (fid == NULL) {
    printf("ERROR : could not open %s\n",fname);
    return -1;
  }
  float* h_array = (float*) calloc(numels,sizeof(float));
  CUDA_SAFE_CALL(cudaMemcpy(h_array, d_array,  numels*sizeof(float), cudaMemcpyDeviceToHost));
  fwrite(h_array, numels, sizeof(float), fid);
  fclose(fid);
  free(h_array);
  return 0;
}

//-----------------------------------------------------------------------------


void filter_projections(ParamsForTomo p4t, float* d_sino_tmp, int num_bins, int num_projs) {
  cufftHandle planRamp_forward = (cufftHandle) p4t.ctxstruct->precond_params_dl.planRamp_forward;
  cufftHandle planRamp_backward = (cufftHandle) p4t.ctxstruct->precond_params_dl.planRamp_backward;
  cufftReal* d_r_sino_error = (cufftReal*) p4t.ctxstruct->precond_params_dl.d_r_sino_error;
  cufftComplex* d_i_sino_error = (cufftComplex*) p4t.ctxstruct->precond_params_dl.d_i_sino_error;


  int sino_size = num_projs*num_bins;
  int dim_fft_ramp = nextpow2_cp_padded(num_bins)/2+1;
  dim3 blk, grd;
  blk = dim3( blsize_cufft , blsize_cufft , 1 );
  grd = dim3( iDivUp_cp(dim_fft_ramp ,blsize_cufft) , iDivUp_cp(fftbunch ,blsize_cufft) , 1 );

  //~ printf("before filtering : %e\n", cublasSasum(sino_size, d_sino_tmp, 1));
  for (int offset = 0 ; offset < sino_size;  offset += num_bins*fftbunch)  {
    int numels = min(sino_size-offset,fftbunch*num_bins);
    //The following memset is important for "weird" sizes...
    CUDA_SAFE_CALL(cudaMemset(d_r_sino_error, 0, fftbunch*nextpow2_cp_padded(num_bins)*sizeof(cufftReal)));
    //~ printf("memset : %e\n", cublasSasum(fftbunch*nextpow2_cp_padded(num_bins), d_r_sino_error, 1));

    // CUDA_SAFE_CALL(cudaMemcpy(d_r_sino_error, d_sino_tmp+offset, numels*sizeof(cufftReal), cudaMemcpyDeviceToDevice));
    CUDA_SAFE_CALL(cudaMemcpy2D(d_r_sino_error, // dst
                                nextpow2_cp_padded(num_bins)*sizeof(cufftReal), // width (bytes) of dst
                                d_sino_tmp+offset,  // src
                                num_bins*sizeof(cufftReal), // width (bytes) of src
                                num_bins*sizeof(cufftReal),  // width (bytes) of matrix transfer
                                numels/num_bins,  //height (pixels !)
                                cudaMemcpyDeviceToDevice));

    // {
    //   dim3 blk, grd;
    //   blk = dim3( blsize_cufft , blsize_cufft , 1 );
    //   grd = dim3( iDivUp_cp( nextpow2_cp_padded(num_bins) - num_bins,blsize_cufft) , iDivUp_cp( numels/num_bins,blsize_cufft) , 1 );
    //   padda_kernel_cp<<<grd,blk>>>( d_r_sino_error, num_bins,  nextpow2_cp_padded(num_bins),    numels/num_bins     )  ;
    // }

    //~ printf("() : %e\n", cublasSasum(numels, d_r_sino_error, 1));
    //FFT, ramp, IFFT
    CUDA_SAFE_FFT(cufftExecR2C(planRamp_forward,(cufftReal *) d_r_sino_error,(cufftComplex *) d_i_sino_error));
    //~ printf("[] : %e\n", cublasScnrm2(fftbunch*nextpow2_cp_padded(num_bins), d_i_sino_error, 1));
    cp_kern_fourier_filter<<<grd,blk>>>(d_i_sino_error, p4t.ctxstruct->precond_params_dl.filter_coeffs, dim_fft_ramp, fftbunch);

    CUDA_SAFE_FFT(cufftExecC2R(planRamp_backward,(cufftComplex *) d_i_sino_error,(cufftReal *) d_r_sino_error));

    // CUDA_SAFE_CALL(cudaMemcpy(d_sino_tmp+offset,d_r_sino_error, numels*sizeof(cufftReal), cudaMemcpyDeviceToDevice));
    CUDA_SAFE_CALL(cudaMemcpy2D(
                     d_sino_tmp+offset,
                     num_bins*sizeof(cufftReal),
                     d_r_sino_error	,
                     nextpow2_cp_padded(num_bins)*sizeof(cufftReal),
                     num_bins*sizeof(cufftReal),
                     numels/num_bins,
                     cudaMemcpyDeviceToDevice));

  }

  //~ printf("After filtering : %e\n", cublasSasum(num_bins*num_projs, d_sino_tmp, 1));
}


void memset_ignored_projections(ParamsForTomo p4t, float* d_sino) {
  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int curr_ignored_angle = 0;
  int i;
  for (i = 0; i < nprojs_span; i++) {
    if (i == p4t.ctxstruct->ignore_angles[curr_ignored_angle]) {
      cudaMemset(d_sino + i*num_bins, 0, num_bins*sizeof(float));
      curr_ignored_angle++;
    }
  }
}









void proj_wrapper(ParamsForTomo p4t, float* d_sino, float* d_image, int dimslice) {
    int memisonhost=0;
    p4t.ctxstruct->gpu_project(p4t.ctxstruct->gpuctx,
                               p4t.ctxstruct->num_bins,
                               p4t.ctxstruct->nprojs_span,
                               p4t.ctxstruct->angles_per_proj,
                               p4t.ctxstruct->axis_position  ,
                               d_sino    ,
                               d_image   ,
                               dimslice,
                               p4t.ctxstruct->axis_corrections,
                               p4t.ctxstruct->gpu_offset_x ,
                               p4t.ctxstruct->gpu_offset_y ,
                               p4t.ctxstruct->JOSEPHNOCLIP,
                               p4t.DETECTOR_DUTY_RATIO,
                               p4t.DETECTOR_DUTY_OVERSAMPLING,
                               memisonhost, p4t.ctxstruct->FAN_FACTOR,
                               p4t.ctxstruct->SOURCE_X
                               );
    if (p4t.ctxstruct->do_ignore_projections) memset_ignored_projections(p4t, d_sino);
    if (p4t.ctxstruct->DATA_IS_FILTERED) filter_projections(p4t, d_sino, p4t.ctxstruct->num_bins, p4t.ctxstruct->nprojs_span);
}


void smooth_transition(float *d_sino_tmp  , int num_bins, int   nprojs_span  ,  float   axis_position ) {
  int   npitch = num_bins;
  dim3 blk, grd;
  
  
  blk = dim3( 32 , 32 , 1 );
  grd = dim3( iDivUp( (num_bins)           ,32) , iDivUp( nprojs_span,32) , 1 );
  
  transition_kernel<<<grd,blk>>>(d_sino_tmp  ,num_bins,  npitch,   nprojs_span  ,    axis_position  )  ;
}


void backproj_wrapper(ParamsForTomo p4t, float* d_sino, float* d_image, float *d_sino_tmp=NULL) {
  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int sino_size = nprojs_span*num_bins;

  if (d_sino_tmp==NULL) d_sino_tmp = global_sino_tmp;

  CUDA_SAFE_CALL(cudaMemcpy(d_sino_tmp, d_sino, sino_size*sizeof(float), cudaMemcpyDeviceToDevice));

  if (p4t.ctxstruct->do_ignore_projections) memset_ignored_projections(p4t, d_sino_tmp);
	
  if(p4t.ctxstruct->fai360){smooth_transition(d_sino_tmp  , num_bins,  nprojs_span  ,  p4t.ctxstruct->axis_position );   };

  
  if (p4t.ctxstruct->DATA_IS_FILTERED || p4t.ctxstruct->DO_PRECONDITION) {
    filter_projections(p4t, d_sino_tmp, num_bins, nprojs_span);
    if(p4t.ctxstruct->fai360) smooth_transition(d_sino_tmp  , num_bins,  nprojs_span  ,  p4t.ctxstruct->axis_position );

    p4t.ctxstruct->gpu_backproj(p4t.ctxstruct, d_sino_tmp, d_image, 0, p4t.DETECTOR_DUTY_RATIO, p4t.DETECTOR_DUTY_OVERSAMPLING, 1, 0);
  }
  else {
    if(p4t.ctxstruct->fai360) smooth_transition(d_sino_tmp  , num_bins,  nprojs_span  ,  p4t.ctxstruct->axis_position );

    if ((p4t.ctxstruct->FLUO_SINO) && (p4t.ctxstruct->FLUO_step == 2)) {
      cp_call_division(d_sino_tmp, p4t.ctxstruct->d_Sigma, num_bins, nprojs_span, 1.0/100.0); //TODO : parameter for "vmin"
    }

    cublasSscal(sino_size, (M_PI*0.5f)/nprojs_span, d_sino_tmp, 1);
    p4t.ctxstruct->gpu_backproj(p4t.ctxstruct, d_sino_tmp, d_image, 0, p4t.DETECTOR_DUTY_RATIO, p4t.DETECTOR_DUTY_OVERSAMPLING, 1, 0);
    CUDA_SAFE_CALL(cudaMemcpy(d_sino_tmp, d_sino, sino_size*sizeof(float), cudaMemcpyDeviceToDevice));
  }
}

/*
void backproj_dfi_wrapper(ParamsForTomo p4t, float* d_sino, float* d_image) {

  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int dim_fft = nextpow2(num_bins);

  float *WORK[nprojs_span];
  WORK[0] =  (float*) malloc(self->params.nprojs_span*(dim_fft)*sizeof(float));
  memcpy(WORK[0],d_sino +iv*blocksino + (projection) * num_bins, num_bins * sizeof(float));



  int projection;
  for(projection=0; projection < self->params.nprojs_span; projection++) {
    WORK[projection] =  WORK[0] + projection*dim_fft;
  }
  p4t.ctxstruct->gpu_backproj_dfi(p4t.ctxstruct, d_sino, d_image);


}
*/
















//-----------------------------------------------------------------------------------------




float calculate_lipschitz(ParamsForTomo p4t, float* sino, float* slice, int n_it) {

  int verbosity = p4t.ctxstruct->verbosity;
  if (verbosity > 4) puts("Entering calculate_lipschitz()");
  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int dimslice = p4t.ctxstruct->num_x;
  int numels_slice = dimslice*dimslice;
  if (verbosity > 4) printf("Nb = %d , Np = %d , d = %d\n",num_bins, nprojs_span, dimslice);

  float* slicetmp;
  CUDA_SAFE_CALL(cudaMalloc(&slicetmp, numels_slice*sizeof(float)));
  CUDACHECK;
  float2* slice_grad;
  CUDA_SAFE_CALL(cudaMalloc( &slice_grad, numels_slice*sizeof(float2)));
  CUDACHECK;

  //~ printf("Before backproj sino =  : %e\n", cublasSasum(num_bins*nprojs_span, sino, 1));
  backproj_wrapper(p4t, sino, slice);
  //~ printf("After backproj slice =  : %e\n", cublasSasum(dimslice*dimslice, sino, 1));

  dim3 blk, grd;
  blk = dim3( blsize_cufft , blsize_cufft , 1 );
  grd = dim3( iDivUp_cp(dimslice ,blsize_cufft) , iDivUp_cp(dimslice ,blsize_cufft) , 1 );

  float norm = 0.0f;
  float Lipschitz = 0.0f;
//  float L1 = 0, L2 = 0;
  int k;
  for (k = 0; k < n_it; k++) {
    //x = P^T*(P*x) - div(grad(x))
    proj_wrapper(p4t, sino, slice, dimslice);
    gradient_kernel<<<grd,blk>>>(slice, slice_grad, dimslice, dimslice);
    backproj_wrapper(p4t, sino, slice);
    divergence_kernel<<<grd,blk>>>(slice_grad, slicetmp, dimslice, dimslice);
    cublasSaxpy (dimslice*dimslice, -1.0f, slicetmp, 1, slice, 1);

    //renormalize variables
    norm = cublasSnrm2 (numels_slice, slice, 1);
    cublasSscal(numels_slice, 1.0f/norm, slice , 1);

/*
    //Apply operator K on (x)
    proj_wrapper(p4t, sino, slice, dimslice);
    gradient_kernel<<<grd,blk>>>(slice, slice_grad, dimslice, dimslice);
    //L = norm(K*(x),'fro')
    norm_kernel<<<grd,blk>>>(slice_grad, slicetmp, dimslice, dimslice);
    L1 = cublasSnrm2(num_bins*nprojs_span, sino, 1);
    L2 = cublasSasum(numels_slice, slicetmp, 1);
    Lipschitz = sqrtf(L1*L1 + L2);
*/

    Lipschitz = sqrt(norm);

    if ((verbosity > 3) && (k % 10 == 0)) printf("Lipschitz (%d) : %e\n",k,Lipschitz);
  }
  CUDA_SAFE_CALL(cudaFree(slicetmp));
  CUDA_SAFE_CALL(cudaFree(slice_grad));
  CUDACHECK;
  return Lipschitz;
}


float calculate_lipschitz_rings(ParamsForTomo p4t, float* sino, float* slice, float rings_height, float alpha_rings, int n_it) {

  int verbosity = p4t.ctxstruct->verbosity;
  if (verbosity > 4) puts("Entering calculate_lipschitz_rings()");
  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int dimslice = p4t.ctxstruct->num_x;
  int numels_slice = dimslice*dimslice;
  if (verbosity > 4) printf("Nb = %d , Np = %d , d = %d\n",num_bins, nprojs_span, dimslice);

  float* slicetmp;
  CUDA_SAFE_CALL(cudaMalloc(&slicetmp, numels_slice*sizeof(float)));
  CUDACHECK;
  float2* slice_grad;
  CUDA_SAFE_CALL(cudaMalloc( &slice_grad, numels_slice*sizeof(float2)));
  float* rings, *null_array;
  CUDA_SAFE_CALL(cudaMalloc(&rings, num_bins*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemset(rings, 0, num_bins*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&null_array, num_bins*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemset(null_array, 0, num_bins*sizeof(float)));
  CUDACHECK;

  backproj_wrapper(p4t, sino, slice);

  dim3 blk, grd;
  blk = dim3( blsize_cufft , blsize_cufft , 1 );
  grd = dim3( iDivUp_cp(dimslice ,blsize_cufft) , iDivUp_cp(dimslice ,blsize_cufft) , 1 );
  dim3 grd_rings = dim3(iDivUp_cp(num_bins ,blsize_cufft), 1, 1);
  dim3 grd_rings2 = dim3(iDivUp_cp(num_bins ,blsize_cufft), iDivUp_cp(nprojs_span, blsize_cufft), 1);

  float norm = 0.0f, norm_r = 0.0f;
  float Lipschitz = 0.0f;
//  float L1 = 0, L2 = 0, L3 = 0;
  int k;
  for (k = 0; k < n_it; k++) {
    //x = P^T*(P*x+r) - div(grad(x))
    proj_wrapper(p4t, sino, slice, dimslice);
    gradient_kernel<<<grd,blk>>>(slice, slice_grad, dimslice, dimslice);
    add_rings_to_sinogram_kernel<<<grd_rings2,blk>>>(sino, rings, alpha_rings, num_bins, nprojs_span);
    backproj_wrapper(p4t, sino, slice);
    divergence_kernel<<<grd,blk>>>(slice_grad, slicetmp, dimslice, dimslice);
    cublasSaxpy (dimslice*dimslice, -1.0f, slicetmp, 1, slice, 1);
    //r = P*x + 2*r
    if (fabsf(alpha_rings - 1) > 0.001) cublasSscal(num_bins, alpha_rings, rings, 1);
    update_rings_kernel<<<grd_rings,blk>>>(rings, sino, null_array, -1.0f, rings_height, num_bins, nprojs_span);

    //renormalize variables
    /*
    norm = cublasSnrm2 (numels_slice, slice, 1);
    cublasSscal(numels_slice, 1.0f/norm, slice , 1);
    norm_r = cublasSnrm2(num_bins, rings, 1);
    cublasSscal(num_bins, 1.0f/norm_r, rings, 1);
    */
    // All variables should be normalized with the same norm (norm of K)
    norm = cublasSnrm2(numels_slice, slice, 1);
    norm_r = cublasSnrm2(num_bins, rings, 1);
    norm = sqrt(norm*norm + norm_r*norm_r);
    cublasSscal(numels_slice, 1.0f/norm, slice , 1);
    cublasSscal(num_bins, 1.0f/norm, rings, 1);



    /*
    //Apply operator K on (x, r)
    proj_wrapper(p4t, sino, slice, dimslice);
    add_rings_to_sinogram_kernel<<<grd_rings2,blk>>>(sino, rings, alpha_rings, num_bins, nprojs_span);
    gradient_kernel<<<grd,blk>>>(slice, slice_grad, dimslice, dimslice);

    //L = norm(K*(x,r),'fro')
    norm_kernel<<<grd,blk>>>(slice_grad, slicetmp, dimslice, dimslice);
    L1 = cublasSnrm2(num_bins*nprojs_span, sino, 1);
    L2 = cublasSasum(numels_slice, slicetmp, 1);
    L3 = cublasSnrm2(num_bins, rings, 1);
    Lipschitz = sqrtf(L1*L1 + L2 + L3*L3);
    */
    Lipschitz = sqrt(norm);


    if (verbosity > 3) if (k % 10 == 0) printf("Lipschitz (%d) : %e\n",k,Lipschitz);
  }
  CUDA_SAFE_CALL(cudaFree(slicetmp));
  CUDA_SAFE_CALL(cudaFree(slice_grad));
  CUDA_SAFE_CALL(cudaFree(null_array));
  CUDA_SAFE_CALL(cudaFree(rings));
  CUDACHECK;
  return Lipschitz;
}






int chambolle_pock_main_rings(ParamsForTomo p4t, float* sino, float* slice, float* data, int n_it, float beta, float beta_r, float rings_height, float alpha_rings, float* last_l2 = NULL, float* last_tv = NULL) {

  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs_span = p4t.ctxstruct->nprojs_span;
  int dimslice = p4t.ctxstruct->num_x;
  char DO_RING_CORR = (rings_height > 0.00001 ? 1 : 0);
  float beta_L2 = p4t.ctxstruct->BETA_L2;
  char DO_L2_REG = (beta_L2 > 1e-7 ? 1 : 0);
  int verbosity = p4t.ctxstruct->verbosity;
  char POS_CONSTRAINT = (p4t.ctxstruct->ITER_POS_CONSTRAINT > 0 ? 1 : 0);

  dim3 blk, grd;
  blk = dim3(blsize_cufft , blsize_cufft , 1 );
  grd = dim3(iDivUp_cp(dimslice, blsize_cufft), iDivUp_cp(dimslice, blsize_cufft), 1);
  dim3 grd2 = dim3(iDivUp_cp(num_bins ,blsize_cufft), iDivUp_cp(nprojs_span, blsize_cufft), 1);
  dim3 grd_rings = dim3(iDivUp_cp(num_bins ,blsize_cufft), 1, 1);
  dim3 grd_rings2 = dim3(iDivUp_cp(num_bins ,blsize_cufft), iDivUp_cp(nprojs_span, blsize_cufft), 1);

  int lip_iter = p4t.ctxstruct->LIPSCHITZ_ITERATIONS;
  float L;
  if (DO_RING_CORR) L = calculate_lipschitz_rings(p4t, sino, slice, rings_height, alpha_rings, lip_iter);
  else L = calculate_lipschitz(p4t, sino, slice, lip_iter);


  L *= p4t.ctxstruct->LIPSCHITZFACTOR;
  if (DO_L2_REG) L = sqrt(L*L + 2.0*1.4143);

  if (verbosity > 5) printf("Lipschitz = %e\n", L);

//Initial guess
  backproj_wrapper(p4t, data, slice);
  //if (CP_DEBUG) write_device_array(slice, dimslice*dimslice, "firstguess.dat");

  if (n_it == 0) {
    puts("(CP) No iterations, returning filtered back-projection result");
    return 0;
  }
  //TODO : work to reduce the memory usage. For eg. dual_p_backproj be removed using slice_tmp
  int numels_slice = dimslice*dimslice;
  int numels_sino = num_bins*nprojs_span;
  float* dual_p, *slice_tilde, *slice_old, *dual_p_backproj, *slice_tmp;
  float2* dual_q2;
  float2* slice_grad, *dual_q;
  CUDA_SAFE_CALL(cudaMalloc(&dual_p, numels_sino*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemset(dual_p, 0, numels_sino*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&dual_q, numels_slice*sizeof(float2)));
  CUDA_SAFE_CALL(cudaMemset(dual_q, 0, numels_slice*sizeof(float2)));
  CUDA_SAFE_CALL(cudaMalloc(&slice_grad, numels_slice*sizeof(float2)));
  CUDA_SAFE_CALL(cudaMemset(slice_grad, 0, numels_slice*sizeof(float2)));
  CUDA_SAFE_CALL(cudaMalloc(&slice_tilde, numels_slice*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemcpy(slice_tilde, slice,  numels_slice*sizeof(float), cudaMemcpyDeviceToDevice));
  CUDA_SAFE_CALL(cudaMalloc(&slice_old, numels_slice*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&dual_p_backproj, numels_slice*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&slice_tmp, numels_slice*sizeof(float)));
  if (DO_L2_REG) {
    cudaMalloc(&dual_q2, numels_slice*sizeof(float2));
    cudaMemset(dual_q2, 0, numels_sino*sizeof(float2));
  }

  float* dual_v, *rings, *rings_tilde, *rings_old;
  if (DO_RING_CORR) {
    CUDA_SAFE_CALL(cudaMalloc(&rings, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc(&rings_tilde, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc(&dual_v, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc(&rings_old, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMemset(rings, 0, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMemset(rings_tilde, 0, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMemset(dual_v, 0, num_bins*sizeof(float)));
    CUDA_SAFE_CALL(cudaMemset(rings_old, 0, num_bins*sizeof(float)));
  }
  CUDACHECK;

  //TODO : check allocations (cudaSuccess, cudaErrorMemoryAllocation)
//  if (dual_p == NULL || dual_q == NULL || slice_grad == NULL || slice_tilde == NULL || slice_old == NULL || dual_p_backproj == NULL || slice_tmp == NULL) {
//    puts("ERROR : out of memory, could not allocate enough memory for all device arrays");
//    return -1;
//  }

  if (verbosity > 3) {
    const char* status[2] = {"DISABLED", "ENABLED"};
    puts("Now executing Chambolle-Pock main loop");
    printf("Nb = %d , Np = %d, d = %d, B = %f, Br = %f\n",num_bins, nprojs_span, dimslice, beta, beta_r);
    printf("Rings correction is %s\n", status[DO_RING_CORR]);
    printf("Ramp filtering is %s\n", status[p4t.ctxstruct->DO_PRECONDITION]);
    printf("L2 regularization is %s\n", status[DO_L2_REG]);
  }


  /*
   * Chambolle-Pock parameters for tomography.
   * TODO : do not hardcode these
   */
  float lambda = beta;
  float tau, gamma, theta, sigma;
//  float rho;
  if (AHMOD) {
    tau = 0.02f;
    gamma = 0.7f*lambda;
    theta = sqrtf(1 + 2*gamma*tau);
    sigma = 4.0/(tau * L*L);
  }
  else {
    sigma = 1.0f/L;
    theta = 1.0f;
    tau = 1.0f/L;
//    rho = 1.9f;
  }
  //-------------

  float* energies = (float*) malloc(n_it*sizeof(float)); //DEBUG

  float fidelity = 0, l1_norm = 0, l1_norm_rings = 0;
  for (int k=0; k < n_it; k++) {
    //update dual variables (dual_p, dual_q)
      //q = proj_linf(q + sigma*grad(x))
      //p = shrink(p + sigma*(P*x + r))
      //v = proj_linf(v + sigma*r)
    gradient_kernel<<<grd,blk>>>(slice_tilde, slice_grad, dimslice, dimslice); //CUDACHECK;
    proj_wrapper(p4t, sino, slice_tilde, dimslice); //CUDACHECK;
    if (DO_RING_CORR) add_rings_to_sinogram_kernel<<<grd_rings2,blk>>>(sino, rings_tilde, alpha_rings, num_bins, nprojs_span); //CUDACHECK;
    dual_shrink_kernel<<<grd2,blk>>>(dual_p, sino, data, sigma, num_bins, nprojs_span);  //CUDACHECK;
    dual_proj_linf_kernel<<<grd,blk>>>(dual_q, slice_grad, sigma, lambda, dimslice, dimslice);  //CUDACHECK;
    if (DO_RING_CORR) sino_proj_linf_kernel<<<grd_rings,blk>>>(dual_v, rings_tilde, sigma, beta_r, num_bins); //CUDACHECK;
    // If this is working well, the following should be merged with another kernel
    if (DO_L2_REG) call_shrink_gradient(dual_q2, slice_grad, dimslice, dimslice, sigma, beta_L2);
    //update primal variables
      //x = x - tau*p_backproj + tau*div(q)
    backproj_wrapper(p4t, dual_p, dual_p_backproj); //CUDACHECK;
    CUDA_SAFE_CALL(cudaMemcpy(slice_old, slice,  numels_slice*sizeof(float), cudaMemcpyDeviceToDevice));
    if (DO_RING_CORR) CUDA_SAFE_CALL(cudaMemcpy(rings_old, rings,  num_bins*sizeof(float), cudaMemcpyDeviceToDevice));
    divergence_kernel<<<grd,blk>>>(dual_q, slice_tmp, dimslice, dimslice); //CUDACHECK;
    cublasSaxpy(numels_slice, -tau, dual_p_backproj, 1, slice, 1); //CUDACHECK;
    cublasSaxpy(numels_slice, tau, slice_tmp, 1, slice, 1); //CUDACHECK;
    if (DO_L2_REG) {
      divergence_kernel<<<grd,blk>>>(dual_q2, slice_tmp, dimslice, dimslice);
      cublasSaxpy(numels_slice, tau, slice_tmp, 1, slice, 1);
    }
    if (POS_CONSTRAINT) {
      positivity_kernel<<<grd,blk>>>(slice, dimslice, dimslice);
    }
      //r = r - tau*(p+v)
    if (DO_RING_CORR && fabsf(alpha_rings - 1) > 0.0001) cublasSscal(num_bins, alpha_rings, rings, 1);
    if (DO_RING_CORR) update_rings_kernel<<<grd_rings,blk>>>(rings, dual_p, dual_v, tau, rings_height, num_bins, nprojs_span);
      //x_tilde = x + theta*(x - x_old)  =  (1+theta)*x - theta*x_old
    CUDA_SAFE_CALL(cudaMemset(slice_tilde, 0, numels_slice*sizeof(float))); //CUDACHECK;
    cublasSaxpy(numels_slice, 1+theta, slice, 1, slice_tilde, 1); //CUDACHECK;
    cublasSaxpy(numels_slice, -theta, slice_old, 1, slice_tilde, 1); //CUDACHECK;
      //rings_tilde = (1+theta)*rings - theta*rings_old
    if (DO_RING_CORR) {
      CUDA_SAFE_CALL(cudaMemset(rings_tilde, 0, num_bins*sizeof(float))); //CUDACHECK;
      cublasSaxpy(num_bins, 1+theta, rings, 1, rings_tilde, 1); //CUDACHECK;
      cublasSaxpy(num_bins, -theta, rings_old, 1, rings_tilde, 1); //CUDACHECK;
    }




    //Norms
    if ((verbosity > 1) || (k == n_it - 1)) {
      cublasSaxpy(numels_sino, -1.0f, data, 1, sino, 1); //CUDACHECK;
      fidelity = cublasSnrm2(numels_sino, sino, 1); //CUDACHECK;
      fidelity *= fidelity/2;
      l1_norm =  calculate_l1_norm(slice_grad, dimslice, dimslice);
      if (DO_RING_CORR) l1_norm_rings = cublasSasum(num_bins, rings, 1);
      energies[k] = fidelity+lambda*l1_norm+beta_r*l1_norm_rings; //DEBUG
      if (k % 10 == 0) printf("Iteration %d : Energy = %e \t Fidelity = %e \t L1 norm = %e \t rings = %e\n",k,fidelity+lambda*l1_norm+beta_r*l1_norm_rings,fidelity,l1_norm,l1_norm_rings);
    }
  }
  if (last_l2) *last_l2 = fidelity;
  if (last_tv) *last_tv = l1_norm;

  CUDA_SAFE_CALL(cudaFree(dual_p));
  CUDA_SAFE_CALL(cudaFree(dual_q));
  CUDA_SAFE_CALL(cudaFree(slice_grad));
  CUDA_SAFE_CALL(cudaFree(slice_tilde));
  CUDA_SAFE_CALL(cudaFree(slice_old));
  CUDA_SAFE_CALL(cudaFree(dual_p_backproj));
  CUDA_SAFE_CALL(cudaFree(slice_tmp));
  CUDACHECK;

  FILE* fid = fopen("energy_CP.dat", "wb"); //DEBUG
  fwrite(energies, sizeof(float), n_it, fid);
  fclose(fid);
  free(energies);

  if (DO_RING_CORR) {


    //DEBUG
//    float* h_rings = (float*) malloc(num_bins*sizeof(float));
//    CUDA_SAFE_CALL(cudaMemcpy(h_rings, rings, num_bins*sizeof(float), cudaMemcpyDeviceToHost));
//    FILE* filedebug = fopen("rings.dat","w");
//    int wrote = fwrite(h_rings,num_bins*sizeof(float),1,filedebug);
//    fclose(filedebug);
//    free(h_rings);
    //-----

    CUDA_SAFE_CALL(cudaFree(rings));
    CUDA_SAFE_CALL(cudaFree(rings_tilde));
    CUDA_SAFE_CALL(cudaFree(dual_v));
    CUDA_SAFE_CALL(cudaFree(rings_old));
  }

  return 0;
}


// avoid NaN in sqrt(sino)
# define FLUO_DO_CLIP 1


__global__ void sqrt_kernel(float* d_out, float* d_in, int Nx, int Ny, int clip) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  int tid = gidy*Nx + gidx;

  if (gidx < Nx && gidy < Ny) {
    float val = d_in[tid];
    if (clip && val < 0) val = 0;
    d_out[tid] = sqrtf(val);
  }
}


int cp_call_sqrt(float* d_out, float* d_in, int Nx, int Ny) {
  dim3 blk = dim3(blsize_cufft, blsize_cufft, 1);
  dim3 grd = dim3(iDivUp_cp(Nx, blsize_cufft), iDivUp_cp(Ny, blsize_cufft), 1);
  sqrt_kernel<<<grd, blk>>>(d_out, d_in, Nx, Ny, FLUO_DO_CLIP);
  return 0;
}


__global__ void abs_kernel(float* d_out, float* d_in, int Nx, int Ny, int clip) {
  int gidx = threadIdx.x + blockIdx.x*blockDim.x;
  int gidy = threadIdx.y + blockIdx.y*blockDim.y;
  int tid = gidy*Nx + gidx;

  if (gidx < Nx && gidy < Ny) {
    float val = d_in[tid];
    if (clip && val < 0) val = 0;
    d_out[tid] = fabsf(val);
  }
}

int cp_call_abs(float* d_out, float* d_in, int Nx, int Ny) {
  dim3 blk = dim3(blsize_cufft, blsize_cufft, 1);
  dim3 grd = dim3(iDivUp_cp(Nx, blsize_cufft), iDivUp_cp(Ny, blsize_cufft), 1);
  abs_kernel<<<grd, blk>>>(d_out, d_in, Nx, Ny, FLUO_DO_CLIP);
  return 0;
}





int cp_fluo(ParamsForTomo p4t, float* d_sino, float* d_image, float* d_data, int n_it, float beta, float beta_rings, float rings_height, float alpha_rings) {

  int num_bins = p4t.ctxstruct->num_bins;
  int nprojs = p4t.ctxstruct->nprojs_span;
  int dimslice = p4t.ctxstruct->num_x;
  int verbosity = p4t.ctxstruct->verbosity;
//  char POS_CONSTRAINT = (p4t.ctxstruct->ITER_POS_CONSTRAINT > 0 ? 1 : 0);


  // Perform a standard TV reconstruction
  if (verbosity > 2) puts("[FLUO 1/2] Performing standard TV reconstruction");
  // put verbosity to 0 ?
  // force positivity constraint to avoid nan in sqrt ?
  p4t.ctxstruct->FLUO_step = 1;
  chambolle_pock_main_rings(p4t, d_sino, d_image, d_data, n_it, beta, beta_rings, rings_height, alpha_rings);

  // Project result and compute the estimated (Diagonal) STD matrix
  float* d_Sigma;
  cudaMalloc(&d_Sigma, num_bins*nprojs*sizeof(float));
  proj_wrapper(p4t, d_sino, d_image, dimslice);
//  cp_call_sqrt(d_Sigma, d_sino, num_bins, nprojs);
  cp_call_abs(d_Sigma, d_sino, num_bins, nprojs);
  // Renormalize Sigma so that mean(Sigma) = 1
  cp_normalize_mean(d_Sigma, num_bins, nprojs);


//  // DEBUG
//  float* h_Sigma = (float*) calloc(num_bins*nprojs, sizeof(float));
//  FILE* fid = fopen("Sigma.dat", "wb");
//  cudaMemcpy(h_Sigma, d_Sigma, num_bins*nprojs*sizeof(float), cudaMemcpyDeviceToHost);
//  fwrite(h_Sigma, sizeof(float), num_bins*nprojs, fid);
//  fclose(fid);

//  fid = fopen("data.dat", "wb");
//  cudaMemcpy(h_Sigma, d_data, num_bins*nprojs*sizeof(float), cudaMemcpyDeviceToHost);
//  fwrite(h_Sigma, sizeof(float), num_bins*nprojs, fid);
//  fclose(fid);
//  // -------





  // Run TV reconstruction with modified fidelity term incorporating the "Sigma" metric
  p4t.ctxstruct->FLUO_step = 2;
  p4t.ctxstruct->d_Sigma = d_Sigma;
  for (int k = 0; k < p4t.ctxstruct->FLUO_ITERS; k++) {
    p4t.ctxstruct->verbosity = 0;
    chambolle_pock_main_rings(p4t, d_sino, d_image, d_data, n_it, beta, beta_rings, rings_height, alpha_rings);
    p4t.ctxstruct->verbosity = verbosity;
    if (verbosity > 2) {
      float sigmanorm = cublasSasum(num_bins*nprojs, d_Sigma, 1);
//      sigmanorm *= sigmanorm;
      printf("[FLUO 2/2] Reconstruction %d : S = %e\n", k, sigmanorm);
    }

    // Update Sigma
    proj_wrapper(p4t, d_sino, d_image, dimslice);
//    cp_call_sqrt(d_Sigma, d_sino, num_bins, nprojs);
    cp_call_abs(d_Sigma, d_sino, num_bins, nprojs);
    cp_normalize_mean(d_Sigma, num_bins, nprojs);
  }

  cudaFree(d_Sigma);

  return 0;
}













///-----------------------------------------------------------------------------
///------------------------ Param/noise Estimation -----------------------------
///-----------------------------------------------------------------------------



/// pass 1: horizontal convolution with [1, -2,  1]
__global__ void convolve_laplacian_kernel_pass1(
    float * input,
    float * output,
    int IMG_W,
    int IMG_H)
{
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int gidy = threadIdx.y + blockIdx.y*blockDim.y;
    if (gidy < IMG_H && gidx < IMG_W) {
        int tid = gidy*IMG_W + gidx;
        // "Valid" convolution, ignore the edges
        if (1 <= gidx && gidx <= IMG_W-2) {
            output[tid] = input[tid-1] - 2*input[tid] + input[tid+1];
        }
        else output[tid] = 0;
     }
}


/// pass 2: vertical convolution with [1, -2,  1]
__global__ void convolve_laplacian_kernel_pass2(
    float * input,
    float * output,
    int IMG_W,
    int IMG_H)
{
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int gidy = threadIdx.y + blockIdx.y*blockDim.y;
    if (gidy < IMG_H && gidx < IMG_W) {
        // "Valid" convolution, ignore the edges
        if (1 <= gidy && gidy <= IMG_H-2) {
            output[gidy*IMG_W + gidx] = input[(gidy-1)*IMG_W + gidx] - 2*input[gidy*IMG_W + gidx] + input[(gidy+1)*IMG_W + gidx];
        }
        else output[gidy*IMG_W + gidx] = 0;
     }
}





int call_convolve_laplacian(float* d_data_out, float* d_data_in, float* d_data_tmp, int Nx, int Ny) {
    dim3 grd, blk;
    blk = dim3(blsize_cufft, blsize_cufft, 1);
    grd = dim3(iDivUp_cp(Nx, blsize_cufft), iDivUp_cp(Ny, blsize_cufft), 1);
    convolve_laplacian_kernel_pass1<<<grd, blk>>>(d_data_in, d_data_tmp, Nx, Ny);
    convolve_laplacian_kernel_pass2<<<grd, blk>>>(d_data_tmp, d_data_out, Nx, Ny);
    return 0;
}


/// Estimate the noise std for Gaussian noise data.
/// Reference
/// ----------
/// Fast Noise Variance Estimation
/// COMPUTER VISION AND IMAGE UNDERSTANDING
/// Vol. 64, No. 2, September, pp. 300–302, 1996
/// ARTICLE NO 0060
float estimate_noise_std(float* d_data, int Nx, int Ny) {

    float* d_data_out, *d_data_tmp;
    cudaMalloc(&d_data_out, Nx*Ny*sizeof(float));
    cudaMalloc(&d_data_tmp, Nx*Ny*sizeof(float));

    call_convolve_laplacian(d_data_out, d_data, d_data_tmp, Nx, Ny);

    // DEBUG
    write_device_array(d_data_out, Nx*Ny, "laplacian.edf");
    //


    float res = cublasSasum(Nx*Ny, d_data_out, 1);

    res *= sqrtf(M_PI_2)/(6.0f*(Nx-2)*(Ny-2));

    cudaFree(d_data_out);
    cudaFree(d_data_tmp);

    return res;
}


/// Experimental !
/// arcsinh(2/3 * sigma * sqrt(2*log(N)))
///     sigma = estimate of noise std
///     N = total number of samples (pixels)
float estimate_regularization_parameter(float* d_fbp_slice, int Nx, int Ny) {

    float sigma = estimate_noise_std(d_fbp_slice, Nx, Ny);
    return asinh( 2.0/3 * sigma * sqrt(2*log(Nx*Ny)) );

}


///-----------------------------------------------------------------------------
///-----------------------------------------------------------------------------
///-----------------------------------------------------------------------------




///-----------------------------------------------------------------------------
/// ---------------------- chambollepock.cu entry point ------------------------
///-----------------------------------------------------------------------------

int chambolle_pock_driver(Gpu_Context* self, float* data, float* SLICE, float DETECTOR_DUTY_RATIO, int DETECTOR_DUTY_OVERSAMPLING, float beta, float beta_rings, float rings_height, float alpha_rings) {

  if (CP_VERBOSE) {
    puts("------------------------------------------------------------------------------");
    puts("------------------ Entering Chambolle-Pock driver ----------------------------");
    puts("------------------------------------------------------------------------------");
  }
  cuCtxSetCurrent ( *((CUcontext *) self->gpuctx  ))  ;
  //Import parameters from self
  int num_bins = self->num_bins;
  int nprojs_span = self->nprojs_span;
  int num_projs = self->nprojs_span;
  int dimslice = self->num_x ;
  ParamsForTomo p4t  =  (ParamsForTomo)  { (Gpu_Context*) self, DETECTOR_DUTY_RATIO, DETECTOR_DUTY_OVERSAMPLING } ;

  //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) {
      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));
      
      plans_are_computed = 1;
  }
  
  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_sino, *d_image, *d_data;
  CUDA_SAFE_CALL(cudaMalloc(&d_sino, num_bins*nprojs_span*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&d_image, dimslice*dimslice*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&d_data, num_bins*nprojs_span*sizeof(float)));
  CUDA_SAFE_CALL(cudaMemcpy(d_data, data,  num_bins*nprojs_span*sizeof(float), cudaMemcpyHostToDevice  ));
  CUDA_SAFE_CALL(cudaMemcpy(d_sino, d_data,  num_bins*nprojs_span*sizeof(float), cudaMemcpyDeviceToDevice  ));
  CUDA_SAFE_CALL(cudaMemset(d_image, 0, dimslice*dimslice*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&global_sino_tmp, num_bins*num_projs*sizeof(float)));
  CUDA_SAFE_CALL(cudaMalloc(&global_slice_tmp, dimslice*dimslice*sizeof(float)));

  // Estimate beta ?
  if (p4t.ctxstruct->ESTIMATE_BETA) {
    puts("----------------------------------------------");
    puts("Estimating the regularization parameter...");
    // TODO: if this works, clean-up !
    float* d_image2;
    cudaMalloc(&d_image2, dimslice*dimslice*sizeof(float));
    int doprec = p4t.ctxstruct->DO_PRECONDITION;
    p4t.ctxstruct->DO_PRECONDITION = 1;
    backproj_wrapper(p4t, d_data, d_image2);
    p4t.ctxstruct->DO_PRECONDITION = doprec;

    float beta2 = estimate_regularization_parameter(d_image2, dimslice, dimslice);
    cudaFree(d_image2);
    printf("Computed beta = %f\n", beta2);
    puts("----------------------------------------------");
    beta = beta2;
  }

  // Run the algorithm
  if (!p4t.ctxstruct->FLUO_SINO) {
    chambolle_pock_main_rings(p4t, d_sino, d_image, d_data, self->ITERATIVE_CORRECTIONS, beta, beta_rings, rings_height, alpha_rings);
  }
  else {
    cp_fluo(p4t, d_sino, d_image, d_data, self->ITERATIVE_CORRECTIONS, beta, beta_rings, rings_height, alpha_rings);
  }
  cudaMemcpy( SLICE, d_image, dimslice*dimslice*sizeof(float), cudaMemcpyDeviceToHost);

  // Free memory
  CUDA_SAFE_CALL(cudaFree(d_sino));
  CUDA_SAFE_CALL(cudaFree(d_image));
  CUDA_SAFE_CALL(cudaFree(d_data));
  CUDA_SAFE_CALL(cudaFree(self->precond_params_dl.d_r_sino_error));
  CUDA_SAFE_CALL(cudaFree(self->precond_params_dl.d_i_sino_error));
  CUDA_SAFE_CALL(cudaFree(self->precond_params_dl.filter_coeffs));
  cudaFree(global_sino_tmp);
  cudaFree(global_slice_tmp);

  return 0;
}