///////////  CZT Chirp Z-transform Source Code  ////////////////////////////////
///
///  gcc -O3 -Wall -W -Wextra -o czt czt.c -lm
//
//
////////////////////////////////////////////////////////////////////////////////
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <math.h>

#include <complex.h>

#undef  PI
#define PI  3.1415926535897932384626433832795L  // 円周率(π)

#define Rotor(a)      (double complex)( cos(a) + I*sin(a) )
#define COMPLEX(a,b)  (double complex)( (a) + I*(b) )

enum { DFT, IDFT };
enum { FFT=DFT, IFFT=IDFT };
enum { CZT=DFT, ICZT=IDFT };

////////////////////////////////////////////////////////////////////////////////
////
//// CZT Chirp Z-transform, Bluestein's algorithm
////
//// Structure to store data for CZT Chirp Z-transform
////  samples    : number of actual data
////  samples_out: Actual data count (for output, same as input one)
////  samples_fft: Number of data used in CZT conversion (power of 2)
////  flag_fft   : flag indicating if "samples" is a power of 2
////  w          : The number of weight data elements is (samples)
////  v          : The number of impulse response data elements is (samples_fft)
////  
typedef struct CZT_CONTROL {
    int  samples;      // 標本点の数
    int  samples_out;  // 出力する標本点の数
    int  samples_fft;  // 2の整数乗値：flag_fft?samples:2*get_czt_size(samples)
    int  flag_fft;     // 直接FFTを実行できるサイズかの判定
    double complex *w; // 重みデータ             要素数は(samples)
    double complex *v; // インパルス応答データ   要素数は(samples_fft)
} CZT_CONTROL;

////////////////////////////////////////////////////////////////////////////////
//// checks if the argument is a power of 2 ( Returns 1 if it is a power of 2 )
////
int ft_size_check( int n ){ if( n<4 ) return 0;
    for(int e=n;e>1;e>>=1 ) if( e&1 ) return 0;
    return 1;
}

////////////////////////////////////////////////////////////////////////////////
//// A function that returns the smallest power of 2 that is greater than or 
//// equal to twice its argument
////
int get_czt_size( int n ){
    if( ft_size_check(n) ) return n;
    int a=1, b= n<<1;
    for( a=1; a<b; a<<=1 );
    return a;
}

////////////////////////////////////////////////////////////////////////////////
//// FFT routine for simple one-dimensional data
////     double complex *c: Pointer to complex array for input/output
////     elements  : number of data in complex array
////     flag      : Flag indicating forward transform or inverse transform
////                 { FFT | IFFT }
////
void fft( double complex *c, int elements, int flag ){
    if( !ft_size_check(elements) ) return; //2の冪(べき)のみ、power of two, ONLY
    double complex *rotor = (double complex *)calloc( elements, sizeof(double complex) );
    if(!rotor) return; // Processing is aborted because the twiddle factor table
                       // used for calculation could not be secured.
    double  unit = (flag==IFFT) ? 2*M_PI/elements : -2*M_PI/elements;;
    for(int e=0; e<elements; rotor[e]=Rotor(unit*e), e++ );
    
    for(int i=0, j=1; j<elements-1; j++ ) {// scrambler, for initial placement
        for(int k=elements>>1; k>(i^=k); k>>=1 );
        if( j<i ){double complex temp=c[j]; c[j]=c[i], c[i]=temp; }  // exchange
    }
    for(int m, mh=1; (m=mh<<1)<=elements; mh=m ) {// Operations
        for(int ir=0, i=0; i<elements; i+=m ) {   // ir=0: Rotation angle is 0rad.
            register double complex temp, W=rotor[ir];   // W: Rotation factor
            for(int k=elements>>2; k>(ir^=k); k>>=1 );
            for(int j=i; j<mh+i; j++ ) { int a = j,  b = j + mh;
                temp = c[a] - c[b];  // For butterfly operation
                c[a] = c[a] + c[b];  // a ->- (a + b) ->- a
                c[b] = temp * W;     // b ->- (a-b)*W ->- b
    }   }   }   double reciprocal = 1./elements;
    if(flag==IFFT)for(int i=0; i<elements; c[i++]*=reciprocal );
    free(rotor);
}

////////////////////////////////////////////////////////////////////////////////
//// A function that generates a structure for Chirp Z-transform for a given 
//// integer value and returns its pointer.
//// You MUST call "free()" with this pointer when it is no longer needed.
////
CZT_CONTROL *GenerateCZT_CONTROL( int n ){
    int samples     = n;
    int samples_out = n;
    int samples_fft = ft_size_check(n) ? n : get_czt_size(n);
    int flag_fft    = ft_size_check(n);
    CZT_CONTROL  *a = (CZT_CONTROL*)malloc( sizeof(CZT_CONTROL) + sizeof(double complex) *
                                          (samples + (samples_fft +1)) );  // +1 For dummy
    if( !a ) return NULL;
    
    *a = (CZT_CONTROL){
            .samples     = samples,     // 標本点の数
            .samples_out = samples_out, // 出力する標本点の数
            .samples_fft = samples_fft, // ２の冪：２の整数乗値
            .flag_fft    = flag_fft,    // samples が直接FFTを実行できるサイズかの判定
            .w           = (double complex *)&a[1],       // 重みデータ           要素数は(samples)
            .v           = &((double complex *)&a[1])[n], // インパルス応答データ 要素数は(samples_fft) +dummy
           };
    double rad = PI/n;   double complex *w=a->w, *v=a->v;  // as alias
    for(int i=n;i<samples_fft;i++) v[i]=COMPLEX(0,0);
    for(int tail=samples_fft, n2=n<<1, j=0, i=0;i<n;i++, j=(j+(2*i-1))%n2 ){
        double c=cos(rad*j), s=sin(rad*j);
        w[i]      = COMPLEX( c, -s );// complex conjugate:  0<= { w[i] <-> v[i] } < n
        v[i]      = COMPLEX( c,  s );
        v[tail-i] = COMPLEX( c,  s );// v[tail] was keeded as dummy area
    }   fft( v, samples_fft, FFT );
    return a;
}

////////////////////////////////////////////////////////////////////////////////
//// Chirp Z-transform routine for fast Fourier transform of arbitrary size data
//// Argument structure is the same as the FFT routine
////
void czt( double complex *c, int n, int flag ){
    if( ft_size_check(n) ) return fft( c, n, flag );
    CZT_CONTROL    *czt = NULL;
    double complex *tmp = NULL;   // Temporary work area used by CZT

    if( !(czt = GenerateCZT_CONTROL( n ) ) )                       goto GOOD_BY;
    int    nx = czt->samples_fft, dc_size = sizeof(double complex);
    double complex *w = czt->w;
    double complex *v = czt->v;
    
    if( !(tmp = (double complex*)calloc( nx, dc_size )) )          goto GOOD_BY;

    // Division is slow, so we replace it with multiplication using the reciprocal.
    double   m = 1.0/n;  // a*m, same as a/n. if m is m = 1.0/n;

    //for(int i=n;i<nx;tmp[i++]=0+I*0);    // from n to nx are already zero.
    if( flag==ICZT ) { // 逆変換
      // ICZT part: Inverse transform processing part, inverse transform can be 
      // easily realized by using complex conjugate numbers.
      for(int i=0;i< n;i++) tmp[i] =   c[i] * conj(w[i]);
      fft( tmp, nx,  FFT );
      for(int i=0;i<nx;i++) tmp[i] = tmp[i] * conj(v[i]);
      fft( tmp, nx, IFFT );
      for(int i=0;i< n;i++)   c[i] = tmp[i] * conj(w[i]) * m;
      // Note: For the inverse transform, divide the result after the transform 
      // process by the number of elements.
    } else {// flag==CZT  順変換
      // CZT part: Forward transform processing unit
      for(int i=0;i< n;i++) tmp[i] =   c[i] * w[i]; // weight data multiplication
      fft( tmp, nx,  FFT );                         // Fourier transform
      for(int i=0;i<nx;i++) tmp[i] = tmp[i] * v[i]; // convolution
      fft( tmp, nx, IFFT );                         // inverse Fourier transform                                
      for(int i=0;i< n;i++)   c[i] = tmp[i] * w[i]; // weight data multiplication
    }

  GOOD_BY:;
    free( tmp );
    free( czt );
}



////////////////////////////////////////////////////////////////////////////////
////
//// Sample using CZT
//// 
////////////////////////////////////////////////////////////////////////////////
int NumberOfData=19, quiet=0, verbose=0;

#include <string.h>
void argumentAnalysis(int argc, char *argv[]){ //, char **iname, char **oname){
  int  i, st=1;
    for( i=1; i<argc; i++ ) {
        if( argv[i][0]=='-' ) {
            switch(argv[i][1]) {
                case 'v': verbose=1;     break;
                case 'q': quiet=1;       break;
                case 'N': NumberOfData = atoi(&argv[i][2]);
                          if( NumberOfData<1 ) goto GOOD_BY;
                          break;
                case 'h': st=0;
                default : goto GOOD_BY;
            }
        }
        else goto GOOD_BY;
    }
    return;
  GOOD_BY:;
    char *p=argv[0], *op="[-N{Integer number of data}] [-q] [-h] ";
    for( i=strlen(p); i>0 && p[i]!='\\' && p[i]!='/'; i-- );; if(i) i++;
    fprintf( stderr, "\nusage: %s %s\n\n", &p[i], op );
    fprintf( stderr, "\t-N: The number of data.\n" );
    fprintf( stderr, "\t-v: verbose mode.\n" );
    fprintf( stderr, "\t-q: quiet mode.\n" );
    fprintf( stderr, "\t-h: This option shows only this message.\n" );
    exit(st);
}

int main(int argc, char *argv[]){
    argumentAnalysis( argc,  argv );
    int data = NumberOfData;

    double complex *base= (double complex*)malloc( (3*data)*sizeof(double complex) );
    if( !base ) return 1;
    double complex *c1 = &base[0];
    double complex *c2 = &base[data];
    double complex *c3 = &base[data+data];

    for(int i=0;i<data;i++) { // Generating dummy data
        double real = 6 * cos(3 * 2 * PI * i / data) + 4 * sin(7 * 2 * PI * i / data);
        c1[i] = c2[i] = real + I * 0;
    }
    ////////////////////////////////////////////////////////////////////////////

    czt( c2, data,  CZT);              // 順変換: CZT
    
    for(int i=0;i<data;i++) c3[i]=c2[i];
    
    czt( c3, data, ICZT);              // 逆変換: ICZT
    
    ////////////////////////////////////////////////////////////////////////////
    
    printf( "\n          %s              %s          %s\n", 
            "Original data", "Fourier transform", "Inverse transformation");
    
    int is_same = 1;
    for(int i = 0; i < data; i++) {
        if( !quiet )printf("%5d | %11.7f %11.7f | %12.7f %12.7f | %11.7f %11.7f\n",
                           i, creal(c1[i]),cimag(c1[i]), creal(c2[i]),cimag(c2[i]),
                           creal(c3[i]),cimag(c3[i]) );
        if( fabs(creal(c1[i]) - creal(c3[i])) > 0.0000001 ) is_same = 0;
        if( fabs(cimag(c1[i]) - cimag(c3[i])) > 0.0000001 ) is_same = 0;
    }
    char *msg = "The result of comparing the original data and";
    printf( "\n\n %s the inverse transformed data\n is %s.",  msg, 
                              is_same ? "Matched" : "there was a difference" );
    if(verbose) printf("\t\t\t ( % 6d  % 6d )", data, get_czt_size(data) );
    printf("\n\n");

    free(base);
    return EXIT_SUCCESS;
}