#include <stdio.h>
#include <stdlib.h>
#include<math.h>
#include<math_constants.h>
#include <cuda.h>
#include <cublas.h>
#include <cuComplex.h>
#include <cufft.h>
#include<time.h>
#include "pdwt/src/wt.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 fftbunch 128
#define FW_DEBUG 0
/// --------------------------------------------------------------------------------------------------


int iDivUp(int a, int b);



#if FW_DEBUG == 1
#include <string.h>
#define LowByteFirst 1
#define HighByteFirst  0
#define UnsignedShort  0
#define FloatValue     1
#define SignedInteger  2
int fw_byteorder ( void )
{ short int one = 1;
  int value;
  switch ((int) *(char *) &one) {
  case 1: value = LowByteFirst; break;
  case 0: value = HighByteFirst; break;
  default: fprintf(stderr,"Invalid byte order \n");
    exit(1);
  }
  return( value );
}

void fw_write_data_to_edf( float *SLICE, int num_y,int  num_x, char * nomeout ) {
  FILE *output = fopen(nomeout,"w") ;
  if(!output) {
    printf(" error opening output file for slice now stopping\n");
    fprintf(stderr, " error opening output file for slice now stopping\n");
    exit(1);
  }

  {
    char s[4000];
    int len,i;

    if( fw_byteorder()== LowByteFirst ) {

      sprintf(s,"{\nHeaderID       = EH:000001:000000:000000 ;\nImage          = 1 ;\nByteOrder = LowByteFirst ;\nSize = %ld ;\nDim_1 = %d ;\nDim_2 = %d ;\nDataType = Float ;\n",num_y*((long)num_x)*4,num_x,num_y);
    } else {
      sprintf(s,"{\nHeaderID        =  EH:000001:000000:000000 ;\nImage           =  1 ;\nByteOrder = HighByteFirst ;\nSize = %ld ;\nDim_1 = %d ;\nDim_2 = %d ;\nDataType = Float ;\n",num_y*((long)num_x)*4,num_x,num_y);
    }
    len=strlen(s);
    fwrite(s,1,len,output);
    for(i=len; i<1022; i++) {
      fwrite(" ",1,1,output);
    }
    fwrite("}\n",1,2,output);
  }
  fwrite(SLICE , sizeof(float), num_y*num_x , output);
  fclose(output);
}

void fw_write_device_data_to_edf( float *SLICE, int num_y,int  num_x, char * nomeout ) {
  float* h_arr = (float*) calloc(num_x*num_y, sizeof(float));
  cudaMemcpy(h_arr, SLICE, num_y*num_x*sizeof(float), cudaMemcpyDeviceToHost);
  fw_write_data_to_edf(h_arr, num_y, num_x, nomeout);
  free(h_arr);
}

__global__ void fw_kern_convert_float2_to_float(float2* inArray, float* outArray, int numels) {
  int gid = threadIdx.x + blockIdx.x*blockDim.x;
  if(gid<numels) {
    outArray[gid] = inArray[gid].x;
    outArray[numels+gid] = inArray[gid].y;
  }
}

int fw_write_complex_device_array(float2* d_array, int numels, const char* fname) {
  float* h_array = new float[2*numels];
  float* d_array_f;
  cudaMalloc(&d_array_f, 2*numels*sizeof(float));
  dim3 blk, grd;
  blk = dim3( 16 , 1 , 1 );
  grd = dim3( iDivUp(numels ,16) , 1 , 1 );
  fw_kern_convert_float2_to_float<<<grd,blk>>> (d_array,d_array_f,numels);
  cudaMemcpy(h_array,d_array_f, 2*numels*sizeof(float), cudaMemcpyDeviceToHost);

  FILE* fid = fopen(fname, "wb");
  fwrite(h_array, 2*sizeof(float), numels, fid);
  fclose(fid);

  delete[] h_array;
  cudaFree(d_array_f);
  return 0;
}

int fw_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));
  cudaMemcpy(h_array, d_array,  numels*sizeof(float), cudaMemcpyDeviceToHost);
  fwrite(h_array, sizeof(float), numels, fid);
  fclose(fid);
  free(h_array);
  return 0;
}
#endif




// Efficient array transpose from the NVIDIA SDK
#define BLOCK_DIM 16

// This kernel is optimized to ensure all global reads and writes are coalesced,
// and to avoid bank conflicts in shared memory.  This kernel is up to 11x faster
// than the naive kernel below.  Note that the shared memory array is sized to
// (BLOCK_DIM+1)*BLOCK_DIM.  This pads each row of the 2D block in shared memory
// so that bank conflicts do not occur when threads address the array column-wise.
__global__ void transpose_kernel(float *odata, float *idata, int width, int height)
{
    __shared__ float block[BLOCK_DIM][BLOCK_DIM+1];

    // read the matrix tile into shared memory
    unsigned int xIndex = blockIdx.x * BLOCK_DIM + threadIdx.x;
    unsigned int yIndex = blockIdx.y * BLOCK_DIM + threadIdx.y;
    if((xIndex < width) && (yIndex < height))
    {
        unsigned int index_in = yIndex * width + xIndex;
        block[threadIdx.y][threadIdx.x] = idata[index_in];
    }

    __syncthreads();

    // write the transposed matrix tile to global memory
    xIndex = blockIdx.y * BLOCK_DIM + threadIdx.x;
    yIndex = blockIdx.x * BLOCK_DIM + threadIdx.y;
    if((xIndex < height) && (yIndex < width))
    {
        unsigned int index_out = yIndex * height + xIndex;
        odata[index_out] = block[threadIdx.x][threadIdx.y];
    }
}

int call_transpose(float* out, float* in, int Nx, int Ny) {
  dim3 n_threads_per_block = dim3(BLOCK_DIM, BLOCK_DIM, 1);
  dim3 n_blocks = dim3(iDivUp(Nx, BLOCK_DIM), iDivUp(Ny, BLOCK_DIM), 1);
  transpose_kernel<<<n_blocks, n_threads_per_block>>>(out, in, Nx, Ny);
  return 0;
}
#undef BLOCK_DIM








__global__ void kern_fourierwavelets(float2* sinoF, int Nx, int Ny, float wsigma) {
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int gidy = threadIdx.y + blockIdx.y*blockDim.y;
    int Nfft = Ny/2+1;
    if (gidx < Nx && gidy < Nfft) {
        // gidy_shift = fftshift(gidy)
        int Ny2 = Ny/2 + (Ny & 1);
        int gidy_shift = (gidy + Ny2)%Ny;

        int y = -(Ny/2) + gidy_shift;
        float factor = 1.0f - expf(-((y*y) / (wsigma*wsigma)));

        int tid = gidy*Nx + gidx;
        // do not forget the scale factor (here Ny)
        sinoF[tid].x *= factor/Ny;
        sinoF[tid].y *= factor/Ny;
    }
}

int fw_call_damping_kernel(cufftComplex* complexData, int Nx, int Ny, float wsigma) {
    int blksize = 16; // TODO: optimize
    dim3 n_blocks = dim3(iDivUp(Nx, blksize), iDivUp(Ny, blksize), 1);
    dim3 n_threads_per_block = dim3(blksize, blksize, 1);
    kern_fourierwavelets<<<n_blocks, n_threads_per_block>>>(complexData, Nx, Ny, wsigma);
    return 0;
}












/**
    Pre-compute the FFT plans for sinogram filtering.
    plans1d_r2c_batched: array of (cufftHandle), each being a R2C plan for batched 1D vertical FT of a sinogram
    plans1d_c2r_batched: array of (cufftHandle), each being a C2R plan for batched 1D vertical IFT of a sinogram
    plan1d_r2c: pointer to a (cufftHandle), R2C plan for a 1D FT of a sinogram line
    plan1d_c2r: pointer of a (cufftHandle), C2R plan for a 1D IFT of a sinogram line
**/
int fw_precompute_fft_plans(cufftHandle* plans1d_r2c_batched, cufftHandle* plans1d_c2r_batched, cufftHandle* plan1d_r2c, cufftHandle* plan1d_c2r, int nbins, int nprojs, int levels, int do_swt) {
    cufftResult res;

    // Compute plan for 1D (I)FFT of a sinogram line
    // ----------------------------------------------
    cufftPlan1d(plan1d_r2c, nbins, CUFFT_R2C, 1); // TODO: check result
    cufftPlan1d(plan1d_c2r, nbins, CUFFT_C2R, 1);

    // Compute plans for batched 1D (I)FFT of sinogram columns
    // ---------------------------------------------------------
    // See "Advanced data layout": http://docs.nvidia.com/cuda/cufft/#advanced-data-layout
    // See also: http://stackoverflow.com/a/26962357
    // For 1D:
    //     input[b * idist + x * istride]
    //     output[b * odist + x * ostride]
    //     b = signal number
    //     x = element of the b-th signal
    int rank = 1;
    int* n = &nprojs;
    int istride = nbins;
    int idist = 1;
    int ostride = nbins;
    int odist = 1;
    // Ignored for 1D
    int inembed[] = { 0 };
    int onembed[] = { 0 };

    // if SWT: size is not halved at each level, so there is only one plan
    if (do_swt) {
        res = cufftPlanMany(plans1d_r2c_batched, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_R2C, nbins);
        if (res != CUFFT_SUCCESS) {
            puts("ERROR: fw_precompute_fft_plans(): failed to create R2C plan");
            return 2;
        }
        res = cufftPlanMany(plans1d_c2r_batched, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_C2R, nbins);
        if (res != CUFFT_SUCCESS) {
            puts("ERROR: fw_precompute_fft_plans(): failed to create C2R plan");
            return 3;
        }
    }
    else {
        for (int i = 0; i < levels; i++) {
            // TODO: for new versions of pdwt, it is (n + (n & 1))/2
            nprojs /= 2;
            nbins /= 2;
            n = &nprojs;
            istride = nbins;
            ostride = nbins;
            res = cufftPlanMany(plans1d_r2c_batched + i, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_R2C, nbins);
            if (res != CUFFT_SUCCESS) {
                puts("ERROR: fw_precompute_fft_plans(): failed to create R2C plan");
                return 2;
            }
            res = cufftPlanMany(plans1d_c2r_batched + i, rank, n, inembed, istride, idist, onembed, ostride, odist, CUFFT_C2R, nbins);
            if (res != CUFFT_SUCCESS) {
                puts("ERROR: fw_precompute_fft_plans(): failed to create C2R plan");
                return 3;
            }
        }
    }
    return 0;
}



int fw_destroy_fft_plans(cufftHandle* plans1d_r2c_batched, cufftHandle* plans1d_c2r_batched, cufftHandle plan1d_r2c, cufftHandle plan1d_c2r, int levels, int do_swt) {
    if (do_swt) {
        cufftDestroy(plans1d_r2c_batched[0]);
        cufftDestroy(plans1d_c2r_batched[0]);
    }
    else {
        for (int i = 0; i < levels; i++) {
            cufftDestroy(plans1d_r2c_batched[i]);
            cufftDestroy(plans1d_c2r_batched[i]);
        }
    }
    cufftDestroy(plan1d_r2c);
    cufftDestroy(plan1d_c2r);
    free(plans1d_r2c_batched);
    free(plans1d_c2r_batched);
    return 0;
}








// batched FFT along axis 0 of sinogram
// plans cannot be re-used in this case
int fw_perform_fft(cufftReal* realData, cufftComplex* complexData, int level, int direction, cufftHandle* plans) {
    cufftResult res;
    if (direction > 0) { // Forward
        res = cufftExecR2C(plans[level], realData, complexData);
        if (res != CUFFT_SUCCESS) {
            puts("ERROR: fw_perform_fft: failed to perform R2C transform");
            return 4;
        }
    }
    else if (direction < 0) { // Inverse
        res = cufftExecC2R(plans[level], complexData, realData);
        if (res != CUFFT_SUCCESS) {
            puts("ERROR: fw_perform_fft: failed to perform C2R transform");
            return 5;
        }
    }
    else {
        puts("ERROR: fw_perform_fft: unknown direction (please specify positive (forward) or negative (inverse))");
        return 1;
    }
    return 0;
}


int fourierwavelet_filter(float* d_sino, int nbins, int nprojs, int wlevels, float wsigma, char* wname, int do_swt, cufftHandle* plans1d_r2c_batched, cufftHandle* plans1d_c2r_batched) {
    int nprojs2 = nprojs/2;
    int nbins2 = nbins/2;

    if (do_swt) {
        nprojs2 = nprojs;
        nbins2 = nbins;
    }

    // Pre-allocate memory
    cufftReal* realData;
    cudaMalloc(&realData, (nbins2)*(nprojs2)*sizeof(cufftReal));
    cufftComplex* complexData;
    cudaMalloc(&complexData, (nbins2)*(nprojs2/2+1)*sizeof(cufftComplex));

    if (FW_DEBUG) puts("Entered fourierwavelet_filter");

    // Forward DWT of the sinogram
    int memIsOnHost = 0;
    Wavelets W(d_sino, nprojs, nbins, wname, wlevels, memIsOnHost, 1, 0, do_swt);
    wlevels = W.winfos.nlevels; // can be < wlevels
    if (FW_DEBUG) W.print_informations();
    W.forward();

    // FFT of each horizontal frequency band
    for (int i = 0; i < wlevels; i++) {
        if (FW_DEBUG) printf("Filtering at scale %d\n", i+1);
        cudaMemcpy(realData, W.d_coeffs[3*i+2], nprojs2*nbins2*sizeof(float), cudaMemcpyDeviceToDevice);

        // batched FFT along axis 0
        fw_perform_fft(realData, complexData, i*(1-do_swt), 1, plans1d_r2c_batched);
        if (FW_DEBUG) puts("FFT OK");

//        fw_write_complex_device_array(complexData, (nprojs2/2+1)*nbins2, "sino_f.dat"); // DEBUG

        // Damping of vertical stripes
        // fw_call_damping_kernel(complexData, nbins2, nprojs2, wsigma/(2**i) );
        fw_call_damping_kernel(complexData, nbins2, nprojs2, wsigma );
        if (FW_DEBUG) puts("Filtering OK");

        // Inverse FFT
        fw_perform_fft(realData, complexData, i*(1-do_swt), -1, plans1d_c2r_batched);
        if (FW_DEBUG) puts("IFFT OK");

//        fw_write_device_data_to_edf(realData, nprojs2, nbins2, "sino_filtered.dat"); // DEBUG

        // Update wavelet coeffs
        cudaMemcpy(W.d_coeffs[3*i+2], realData, nprojs2*nbins2*sizeof(float), cudaMemcpyDeviceToDevice);
        if (FW_DEBUG) puts("Update OK");

        if (!do_swt) {
            nprojs2 /= 2;
            nbins2 /= 2;
        }
    }
    // Inverse DWT
    W.inverse();

    cudaMemcpy(d_sino, W.d_image, nbins*nprojs*sizeof(float), cudaMemcpyDeviceToDevice);

    // Free memory
    cudaFree(realData);
    cudaFree(complexData);

    return 0;
}






__global__ void kern_rows_sum(float* odata, float* idata, int Nx, int Ny) {
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    if (gidx < Nx) {
        float sum = 0;
        for (int y = 0; y < Ny; y++) {
            sum += idata[y*Nx+gidx];  //TODO: Kahan summation ?
        }
        odata[gidx] = sum;
    }
}
int fw_call_rows_sum(float* odata, float* idata, int Nx, int Ny) {
    int blksize = 16; // TODO: optimize
    dim3 n_blocks = dim3(iDivUp(Nx, blksize), 1, 1);
    dim3 n_threads_per_block = dim3(blksize, 1, 1);
    kern_rows_sum<<<n_blocks, n_threads_per_block>>>(odata, idata, Nx, Ny);
    return 0;
}




__global__ void kern_filter_line(float2* data, int Nx, float param) {
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int Nfft = Nx/2+1;
    if (gidx < Nfft) {
        float nu = (0.5/(Nx-1))*gidx; // normalized frequency
        float val = 1.0 - expf(-nu/param);
        // do not forget the scale factor
        data[gidx].x *= val/Nx;
        data[gidx].y *= val/Nx;
    }
}


int fw_call_filter_sinogram_sum(float2* data, int Nx, float wfilterparam) {
    int blksize = 16; // TODO: optimize
    dim3 n_blocks = dim3(iDivUp(Nx, blksize), 1, 1);
    dim3 n_threads_per_block = dim3(blksize, 1, 1);
    kern_filter_line<<<n_blocks, n_threads_per_block>>>(data, Nx, wfilterparam);
    return 0;
}


__global__ void kern_substract_line(float* data, float* line, int Nx, int Ny, float scale) {
    int gidx = threadIdx.x + blockIdx.x*blockDim.x;
    int gidy = threadIdx.y + blockIdx.y*blockDim.y;
    if (gidx < Nx && gidy < Ny) {
        data[gidy*Nx+gidx] -= line[gidx]/scale;
    }
}



int fw_call_substract_lines(float* data, float* line, int Nx, int Ny, float scale) {
    int blksize = 16; // TODO: optimize
    dim3 n_blocks = dim3(iDivUp(Nx, blksize), iDivUp(Ny, blksize), 1);
    dim3 n_threads_per_block = dim3(blksize, blksize, 1);
    kern_substract_line<<<n_blocks, n_threads_per_block>>>(data, line, Nx, Ny, scale);
    return 0;
}



int sino_normalization(float* d_sino, int num_bins, int num_projs, float wscaling, float wfilterparam, cufftHandle planr2c, cufftHandle planc2r) {

    if (FW_DEBUG) puts("Entering sino_normalization()");

    // Compute the sum along rows
    float* d_sum;
    cudaMalloc(&d_sum, num_bins*sizeof(float));
    if (FW_DEBUG) cudaMemset(d_sum, 0, num_bins*sizeof(float));
    fw_call_rows_sum(d_sum, d_sino, num_bins, num_projs);

    // Filter this sum, if requested
    if (fabsf(wfilterparam) > 1e-8) {
        // FT
        cufftComplex* d_sum_f;
        cudaMalloc(&d_sum_f, num_bins*sizeof(cufftComplex));
        if (FW_DEBUG) cudaMemset(d_sum_f, 0, num_bins*sizeof(cufftComplex));

        cufftExecR2C(planr2c, d_sum, d_sum_f);
        // Filter...
        fw_call_filter_sinogram_sum(d_sum_f, num_bins, wfilterparam);

        // IFFT
        cufftExecC2R(planc2r, d_sum_f, d_sum);

        cudaFree(d_sum_f);
    }
    // Substract sum/wscaling from each sinogram line
    fw_call_substract_lines(d_sino, d_sum, num_bins, num_projs, wscaling);

    cudaFree(d_sum);
    return 0;
}



// TODO: automatic determination of "factor" by TV minimization
int scale_sino_rows(float* sino, int nbins, int nprojs, float factor) {
  float* sum0 = (float*) calloc(nbins, sizeof(float));
  for (int x = 0; x < nbins; x++) {
    for (int y = 0; y < nprojs; y++) {
      sum0[x] += sino[y*nbins+x];
    }
  }

  for (int y = 0; y < nprojs; y++) {
    for (int x = 0; x < nbins; x++) {
      sino[y*nbins+x] -= sum0[x]/factor;
    }
  }
  return 0;
}









///
/// File entry point
///
extern "C" {
  int sinofilters_driver(Gpu_Context* gpuctx, float* data);
}

int sinofilters_driver(Gpu_Context* gpuctx, float* data) {

  cuCtxSetCurrent(*((CUcontext *) gpuctx->gpuctx));
  if (gpuctx->verbosity > 10) {
    puts("------------------------------------------------------------------------------");
    puts("-------------------- Entering sino filters driver ----------------------------");
    puts("------------------------------------------------------------------------------");
  }

  //Import parameters from self
  int num_bins = gpuctx->num_bins;
  int num_projs = gpuctx->nprojs_span;

  //Allocate memory
  float* d_sino;
  cudaMalloc(&d_sino, num_bins*num_projs*sizeof(float));
  cudaMemcpy(d_sino, data, num_bins*num_projs*sizeof(float), cudaMemcpyHostToDevice);

  // Retrieve relevant parameters
  int wlevels = gpuctx->FW_LEVELS;
  float wsigma = gpuctx->FW_SIGMA;
  char* wname = gpuctx->FW_WNAME;
  float wscaling = gpuctx->FW_SCALING;
  float wfilterparam = gpuctx->FW_FILTERPARAM;
  int do_swt = (gpuctx->W_SWT > 0 ? 1 : 0);

  // Pre-compute FFT plans
//   if (!gpuctx->fw_plans_ok) {
//     gpuctx->fw_plans_ok = 1;
  static int plans_are_computed = 0;
  if(!plans_are_computed) {
    plans_are_computed = 1;
    int nplans = wlevels;
    if (do_swt) nplans = 1;
    cufftHandle* plans1d_r2c_batched = (cufftHandle*) calloc(nplans, sizeof(cufftHandle));
    cufftHandle* plans1d_c2r_batched = (cufftHandle*) calloc(nplans, sizeof(cufftHandle));
    gpuctx->fw_plans1d_r2c = plans1d_r2c_batched;
    gpuctx->fw_plans1d_c2r = plans1d_c2r_batched;
    fw_precompute_fft_plans(gpuctx->fw_plans1d_r2c, gpuctx->fw_plans1d_c2r, &(gpuctx->fw_plan1d_r2c), &(gpuctx->fw_plan1d_c2r), num_bins, num_projs, wlevels, do_swt);
  }
CUDACHECK;
  // Perform sino "normalization" if requested
  if (fabsf(wscaling) > 1e-8) {
    sino_normalization(d_sino, num_bins, num_projs, wscaling, wfilterparam, gpuctx->fw_plan1d_r2c, gpuctx->fw_plan1d_c2r);
  }
  CUDACHECK;
  // Perform Munch filter if requested
  if (wlevels > 0) {
    fourierwavelet_filter(d_sino, num_bins, num_projs, wlevels, wsigma, wname, gpuctx->W_SWT, gpuctx->fw_plans1d_r2c, gpuctx->fw_plans1d_c2r);
  }

  // Copy result
  cudaMemcpy(data, d_sino, num_bins*num_projs*sizeof(float), cudaMemcpyDeviceToHost);

  // Free memory
//  fw_destroy_fft_plans(plans1d_r2c_batched, plans1d_c2r_batched, plan1d_r2c, plan1d_c2r, wlevels, do_swt); // should be done outside !
  cudaFree(d_sino);
  return 0;
}







