//-----------------------------------------------------------------------
//Copyright 2013 Centrum Wiskunde & Informatica, Amsterdam
//
//Author: Daniel M. Pelt
//Contact: D.M.Pelt@cwi.nl
//-----------------------------------------------------------------------


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fftw3.h>
#include "CCspace.h"
#include "edftools.h"
#include <math.h>
#ifndef M_PI
    #define M_PI 3.14159265358979323846
#endif
 

#define myfread( a,b,c ,d ) {int res=fread(a,b,c,d); if(res<-1) printf(" gcc casse les pieds\n");}


void nnfbp_setup_fft(CCspace *self, fcomplex **fftSino, fcomplex **fftMult,fcomplex **fftSinoOrig, int dim_fft,int num_bins, float *data_orig, fftwf_plan *sinoB) {


    int fftSize[] = {dim_fft};
    int i,j;

    *fftSino = (fcomplex*) fftwf_malloc(sizeof(fftwf_complex)*dim_fft*self->params.nprojs_span );
    if(*fftSino==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }
    *fftMult = (fcomplex*) fftwf_malloc(sizeof(fftwf_complex)*dim_fft*self->params.nprojs_span );
    if(*fftMult==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }
    *fftSinoOrig = (fcomplex*) fftwf_malloc(sizeof(fftwf_complex) *dim_fft*self->params.nprojs_span );
    if(*fftSinoOrig==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }

    fftwf_plan sinoF;

    sem_wait( &(self->fftw_sem));
    fftwf_plan_with_nthreads(1 );
    // Plan to forward fft original sinogram
    sinoF = fftwf_plan_many_dft(1,fftSize,self->params.nprojs_span,(fftwf_complex*)*fftSino,fftSize,1,dim_fft,(fftwf_complex*)*fftSinoOrig,fftSize,1,dim_fft,FFTW_FORWARD,FFTW_ESTIMATE);
    // Plan to backward fft filtered sinogram
    *sinoB = fftwf_plan_many_dft(1,fftSize,self->params.nprojs_span,(fftwf_complex*)*fftMult,fftSize,1,dim_fft,(fftwf_complex*)*fftSino,fftSize,1,dim_fft,FFTW_BACKWARD,FFTW_ESTIMATE);
    sem_post( &(self->fftw_sem) );

    // Zero the arrays
    for(i=0; i<dim_fft*self->params.nprojs_span; i++) {
        (*fftSino)[i]=0;
        (*fftMult)[i]=0;
        (*fftSinoOrig)[i]=0;
    }


    // Set fftSino to original sinogram (with reduced angles when training)
    for(i=0; i<self->params.nprojs_span; i++) {
        for(j=0; j<num_bins; j++) {
            (*fftSino)[i*dim_fft+j] = data_orig[i*num_bins+j];
        }
    }

    // Calculate forward fft of sinogram
    fftwf_execute(sinoF);

    sem_wait( &(self->fftw_sem));
    fftwf_destroy_plan(sinoF);
    sem_post( &(self->fftw_sem) );
}

void nnfbp_cleanup(CCspace *self,fcomplex **fftSino, fcomplex **fftMult,fcomplex **fftSinoOrig, fftwf_plan *sinoB) {
    sem_wait( &(self->fftw_sem));
    fftwf_destroy_plan(*sinoB);
    sem_post( &(self->fftw_sem) );

    fftwf_free(*fftSino);
    fftwf_free(*fftMult);
    fftwf_free(*fftSinoOrig);
}

void nnfbp_train(CCspace *self,int dim_fft,int num_bins, float *data_orig, int dimslice,float *SLICEcalc, float * SINOGRAMMA, float **WORK, float *WORKbis, float *dumf, fcomplex *dumfC, float *WORK_perproje, float *OVERSAMPLE, int oversampling, int ncpus,float cpu_offset_x,float  cpu_offset_y, int Nfirstslice, int islice, char *nomeout) {
    fcomplex *fftSino,*fftMult,*fftSinoOrig;
    fftwf_plan sinoB;
    int curMult=0;
    int endMult;
    int binWidth=1;
    int i,j;
    int nlinear = self->params.NEURAL_TRAINING_NLINEAR;

    // Setup FFT
    nnfbp_setup_fft(self, &fftSino, &fftMult,&fftSinoOrig, dim_fft, num_bins, data_orig, &sinoB);

    //Dry run to calculate number of weights
    int nW = 0;
    while(curMult<=num_bins) {
        if(nlinear>0) {
            curMult++;
            nlinear--;
        } else {
            curMult+=binWidth;
            binWidth*=2;
        }
        nW++;
    }
    curMult=0;
    binWidth=1;
    nlinear = self->params.NEURAL_TRAINING_NLINEAR;

    // Training data array (one column for each weight, final column is 'correct' pixel value
    float *trainingData = (float *) malloc((nW+1)*self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE*sizeof(float));
    if(trainingData==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }

    // Random pixel arrays
    int *xPicked = (int *)malloc(self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE*sizeof(int));
    if(xPicked==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }
    int *yPicked = (int *)malloc(self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE*sizeof(int));
    if(yPicked==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }
    int xP,yP;
    float cP = (dimslice-1)/2.;
    float bnd = dimslice*dimslice/4.;

    // Optional mask
    unsigned char* mask = (unsigned char*)malloc(dimslice*dimslice*sizeof(unsigned char));
    if(mask==NULL) {
        fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
        exit(EXIT_FAILURE);
    }
    for(i=0; i<dimslice*dimslice; i++) mask[i]=1;
    if(self->params.NEURAL_TRAINING_USEMASK) {
        FILE *fpM = fopen(self->params.NEURAL_TRAINING_MASK_FILE,"rb");
        myfread(mask,sizeof(unsigned char),dimslice*dimslice,fpM);
        fclose(fpM);
    }

    for(i=0; i<self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE; i++) {
        // Pick random pixel in unit circle (and mask[x,y]==1)
        xP = dimslice*drand48();
        yP = dimslice*drand48();
        while(((xP-cP)*(xP-cP)+(yP-cP)*(yP-cP))>=bnd || mask[yP*dimslice+xP]==0) {
            xP = dimslice*drand48();
            yP = dimslice*drand48();
        }
        xPicked[i]=xP;
        yPicked[i]=yP;
    }

    free(mask);

    int cW=0;
    sem_wait(& (self->fbp_sem) );
    if(self->fbp_precalculated.NEURALFILTER==0) {
        // Not freed
        self->fbp_precalculated.NEURALFILTER= (float*)malloc(nW*dim_fft*sizeof(float));
        if(self->fbp_precalculated.NEURALFILTER==NULL) {
            fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
            exit(EXIT_FAILURE);
        }
        memset(self->fbp_precalculated.NEURALFILTER,0,nW*dim_fft);
        while(curMult<=num_bins) {
            if(nlinear>0) {
                endMult=curMult+1;
                nlinear--;
            } else {
                endMult = curMult+binWidth;
                binWidth*=2;
            }

            for(; curMult<endMult; curMult++) {
                if(curMult>num_bins) break;
                for(i=0; i<dim_fft; i++) self->fbp_precalculated.NEURALFILTER[i+cW*dim_fft]+=cos(curMult*2*M_PI*i/dim_fft);
            }
            cW++;
        }
    }
    sem_post(& (self->fbp_sem) );
    for(cW=0; cW<nW; cW++) {
        // Multiply original sinogram with filter in Fourier space
        for(i=0; i<self->params.nprojs_span; i++) {
            for(j=0; j<dim_fft; j++) {
                fftMult[i*dim_fft+j] = fftSinoOrig[i*dim_fft+j]*self->fbp_precalculated.NEURALFILTER[j+cW*dim_fft];
            }
        }
        // Backward fft the result
        fftwf_execute(sinoB);

        // Set sinogram to real part of backward fft
        for(i=0; i<self->params.nprojs_span; i++) {
            for(j=0; j<num_bins; j++) {
                SINOGRAMMA[i*num_bins+j] = creal(fftSino[i*dim_fft+j])/dim_fft;
            }
        }

        // Backproject the sinogram
        rec_driver(self,  WORK,WORKbis,SLICEcalc,  num_bins,  dim_fft,
                   dumf, dumfC , WORK_perproje,OVERSAMPLE, oversampling,
                   SINOGRAMMA , ncpus   , cpu_offset_x, cpu_offset_y  ,
                   0 ,1 , NULL,0,0);
        //   sprintf(nomeout,"%s_nnfbpTEST_%04d.edf",self->params.OUTPUT_FILE_ORIG, cW);
        // write_data_to_edf(SLICEcalc , dimslice, dimslice  , nomeout);

        // Set training data column to the pixel values
        for(i=0; i<self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE; i++) {
            trainingData[i*(nW+1)+cW]=SLICEcalc[yPicked[i]*dimslice+xPicked[i]];
        }
    }

    // Set 'correct' image
    FILE *fp = fopen(self->params.NEURAL_TRAINING_RECONSTRUCTION_FILE,"rb");
    fseek(fp,((long int)Nfirstslice+islice)*sizeof(float)*self->params.num_y*self->params.num_x,SEEK_SET);
    myfread(SLICEcalc,sizeof(float),self->params.num_y*self->params.num_x,fp);
    fclose(fp);

    // Set the pixels of the correct image as last column
    for(i=0; i<self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE; i++) {
        trainingData[i*(nW+1)+nW]=SLICEcalc[yPicked[i]*dimslice+xPicked[i]];
    }

    // Save training data to file
    sprintf(nomeout,"%s_nnfbp_%04d.edf",self->params.OUTPUT_FILE_ORIG, Nfirstslice+islice+self->reading_infos.pos_s[0]);
    write_data_to_edf(trainingData , self->params.NEURAL_TRAINING_PIXELS_PER_IMAGE, nW+1  , nomeout);

    free(trainingData);
    free(xPicked);
    free(yPicked);

    // Cleanup FFT
    nnfbp_cleanup(self, &fftSino, &fftMult,&fftSinoOrig, &sinoB);

}

void nnfbp_reconstruct(CCspace *self,int dim_fft,int num_bins, float *data_orig, int dimslice,float *SLICEcalc,float *SLICE, float * SINOGRAMMA, float **WORK, float *WORKbis, float *dumf, fcomplex *dumfC, float *WORK_perproje, float *OVERSAMPLE, int oversampling, int ncpus,float cpu_offset_x,float  cpu_offset_y) {
    fcomplex *fftSino,*fftMult,*fftSinoOrig;
    fftwf_plan sinoB;
    int curMult=0;
    int endMult;
    int binWidth=1;
    int i,j,ivolta;
    int nlinear = self->params.NEURAL_TRAINING_NLINEAR;

    // Setup FFT
    nnfbp_setup_fft(self, &fftSino, &fftMult,&fftSinoOrig, dim_fft, num_bins, data_orig, &sinoB);

    // Probably not needed, check with Alessandro
    for( i=0; i< dimslice*dimslice ; i++) {
        SLICE[i]=0;
    }
    sem_wait(& (self->fbp_sem) );
    if(self->fbp_precalculated.NEURALFILTER==0) {
        // Set filters to trained filters

        // Is not freed atm, should check with Alessandro
        self->fbp_precalculated.NEURALFILTER = (float *)malloc(self->params.NEURAL_NFILTERS*dim_fft*sizeof(float));
        if(self->fbp_precalculated.NEURALFILTER==NULL) {
            fprintf(stderr,"ERROR allocating memory (nnfbp)\n");
            exit(EXIT_FAILURE);
        }
        memset(self->fbp_precalculated.NEURALFILTER,0,self->params.NEURAL_NFILTERS*dim_fft);
        for(ivolta=0; ivolta<self->params.NEURAL_NFILTERS; ivolta++) {
            int cW=0;
            float cM;
            curMult=0;
            binWidth=1;
            nlinear = self->params.NEURAL_TRAINING_NLINEAR;
            while(curMult<=num_bins) {
                cM = self->params.NEURAL_FILTERS[ivolta*self->params.NEURAL_FILTERSIZE+cW];
                if(nlinear>0) {
                    endMult=curMult+1;
                    nlinear--;
                } else {
                    endMult = curMult+binWidth;
                    binWidth*=2;
                }
                for(; curMult<endMult; curMult++) {
                    if(curMult>num_bins) break;
                    for(i=0; i<dim_fft; i++) self->fbp_precalculated.NEURALFILTER[i+ivolta*dim_fft]+=cM*cos(curMult*2*M_PI*i/dim_fft);
                }
                cW++;
            }
        }
    }
    sem_post(& (self->fbp_sem) );

    // Loop over all filters
    for(ivolta=0; ivolta<self->params.NEURAL_NFILTERS; ivolta++) {
        // Multiply original sinogram with filter in Fourier space
        for(i=0; i<self->params.nprojs_span; i++) {
            for(j=0; j<dim_fft; j++) {
                fftMult[i*dim_fft+j] = fftSinoOrig[i*dim_fft+j]*self->fbp_precalculated.NEURALFILTER[j+ivolta*dim_fft];
            }
        }
        // Backward fft the result
        fftwf_execute(sinoB);

        // Set sinogram to real part of backward fft
        for(i=0; i<self->params.nprojs_span; i++) {
            for(j=0; j<num_bins; j++) {
                SINOGRAMMA[i*num_bins+j] = creal(fftSino[i*dim_fft+j])/dim_fft;
            }
        }

        // Backproject the sinogram
        rec_driver(self,  WORK,WORKbis,SLICEcalc,  num_bins,  dim_fft,
                   dumf, dumfC , WORK_perproje,OVERSAMPLE, oversampling,
                   SINOGRAMMA , ncpus   , cpu_offset_x, cpu_offset_y  ,
                   0 ,1 , NULL,0,0);
        // Add weighted sigmoid of backprojection with offset to reconstruction
        for( i=0; i< dimslice*dimslice ; i++) {
            SLICE[i]+=self->params.NEURAL_WEIGHTS[ivolta]/(1+exp(self->params.NEURAL_OFFSETS[ivolta]-SLICEcalc[i]));
        }
    }
    // Sigmoid the reconstruction, and rescale
    for( i=0; i< dimslice*dimslice ; i++) {
        SLICE[i]=1./(1+exp(self->params.NEURAL_WEIGHTS[self->params.NEURAL_NFILTERS]-SLICE[i]));
        SLICE[i]=(SLICE[i]-0.25)*2*(self->params.NEURAL_MAXIN-self->params.NEURAL_MININ) + self->params.NEURAL_MININ;
    }
    // Cleanup FFT
    nnfbp_cleanup(self, &fftSino, &fftMult,&fftSinoOrig, &sinoB);
}
