// Types.h
#pragma once
#include <stdlib.h>

typedef unsigned char uchar;
typedef unsigned      uint;
typedef double        flt;

typedef unsigned char  uint8;
typedef unsigned short uint16;
typedef unsigned       uint32;
typedef signed char    int8;
typedef short          int16;
typedef int            int32;
typedef float          flt32;
typedef double         flt64;

template<class T> void kill(T* &t) { if (t) { delete t; t=NULL; } }
template<class T> void kill_array(T* &t) { if (t) { delete[] t; t=NULL; } }

// fft.h
#pragma once
#define _USE_MATH_DEFINES
#include <math.h>
#include "Types.h"

inline flt sqr(flt x) { return x*x; }

struct cplx { // minimal comple numbers support
    flt x,y; 
    cplx() {}
    cplx(flt x) : x(x), y(0) {}
    cplx(flt x,flt y) : x(x), y(y) {}
    cplx(const cplx &v) : x(v.x), y(v.y) {}

    cplx operator - (const cplx &rv) const { return cplx( x-rv.x, y-rv.y ); }
    cplx operator + (const cplx &rv) const { return cplx( x+rv.x, y+rv.y ); }
    cplx operator * (const cplx &rv) const { return cplx( x*rv.x-y*rv.y, x*rv.y+y*rv.x ); }
    cplx operator / (const flt  &rv) const { return cplx( x/rv, y/rv);     }
    cplx& operator = (const flt& rv)       { x=rv; y=0; return *this;      }
    cplx& operator = (const cplx& rv)      { x=rv.x; y=rv.y; return *this; }
    cplx& operator *= (const flt& rv)      { x*=rv; y*=rv; return *this;   }
    cplx& operator *= (const cplx &rv)     { flt t=x*rv.x-y*rv.y; y=x*rv.y+y*rv.x; x=t; return *this;   }
    cplx& operator += (const cplx &rv)     { x+=rv.x; y+=rv.y; return *this; }
    void conj()                            { y=-y; }
    flt  norm() const                      { return x*x+y*y; }
    static cplx expi(flt x)                { return cplx(cos(x),sin(x)); }
    friend cplx operator * (const flt &lv,const cplx &rv) { return cplx(lv*rv.x,lv*rv.y); }
};

struct FFT {
    FFT();
    FFT(int power) { w=x=z=NULL; reset(power); }
    ~FFT();
    void reset(int P);                                // init set power
    void fft(flt *f,flt *g,bool forward=true);        // single transform
    void cor(const flt *f, flt *r);                   // calculate correlation r[k]=sum(sqr(f[i+k]-f[i]),i)  
    void corr(const flt *f1, const flt *f2, flt *r);  // calculate correlation r[k]=sum(f1[i+k]-f2[i]),i)  
    void scor(const flt *f, flt *sr);                 // fft(cor)
    void scorr(const flt *f1, const flt *f2, flt *sr);// fft(corr)
    void conv(const flt* f1,const flt* f2,flt *r);    // convolution of f1,f2  r[k]=sum(f1[i+k]*f2[i],i)
protected: //internals
    int M;   // power
    int N;   // number of points N = 2^M
    cplx *w; // [N/2]
    cplx *x; // [N]
    cplx *z; // [N] - could be NULL
    void putF(const flt *f);
    void getF(flt *f);
    void packG(const flt *g);
    void unpackG(flt *g);
    void reorder();
    void scale(flt mf);
    void do_fft(bool inv); // it changes norm by N times |x| => N*|x'|
public: // for validation only
    void tst_ft(flt *f,flt *g,bool forward=true);
    void tst_cor(const flt *f, flt *r);
    void tst_corr(const flt *f1, const flt *f2, flt *r);
    void tst_conv(const flt* f1,const flt* f2,flt *r);
};

void tst_fft();

// fft.cpp
#include "fft.h"

FFT::FFT() { 
    M=N=0; w=x=z=NULL; 
}
FFT::~FFT() { 
    kill_array(z); kill_array(x); kill_array(w); 
}
void FFT::do_fft(bool inv) {// this transform changes norm by N times |x| => N*|x'| 
    int a=N;
    for(int c=0;c<M;c++) { int d=a; a>>=1;
        for(int f=0;f<a;f++) { cplx wf=w[f<<c]; if (inv) wf.conj();
            int g=0; while (g<N) { int h=g+f, j=h+a; g+=d; 
            cplx t=x[h]-x[j]; x[h]=x[h]+x[j]; x[j]=t*wf; 
            }
        }
    }
}
void FFT::reorder() {
    int m=0,p=N>>1,q=N-1;
    for(int r=0;r<q;r++) {
        if (r<m) { cplx t=x[r]; x[r]=x[m]; x[m]=t; }
        int s=p; while (s<=m) { m-=s; s>>=1; }
        m+=s;
    }
}
void FFT::reset(int P) {
    M=P; N=1<<M; int n2=N>>1; 
    kill_array(z); kill_array(x); kill_array(w); 
    if (n2<1) return;
    w=new cplx[n2];
    x=new cplx[N];
    flt k=-2*M_PI/N; for(int i=0;i<n2;i++) w[i]=cplx::expi(k*i);
}
void FFT::scale(flt mf) {
    for(int i=0;i<N;i++) x[i]*=mf; 
}
void FFT::putF(const flt *f) {
    for(int i=0;i<N;i++) x[i]=f[i]; 
}
void FFT::getF(flt *f) {
    for(int i=0;i<N;i++) f[i]=x[i].x; 
}
void FFT::packG(const flt *g) {
    int n2=N>>1;
    x[0]=g[0];    // background component
    for(int i=1;i<n2;i++) { // intermediate freq pairs
        x[i].x=g[2*i-1]; x[N-i].x= x[i].x;
        x[i].y=g[2*i  ]; x[N-i].y=-x[i].y;
    }
    x[n2]=g[N-1]; // highest freq component
}
void FFT::unpackG(flt *g) {
    int n2=N>>1;
    g[0]=x[0].x;     // background component
    for(int i=1;i<n2;i++) {
        g[2*i-1]=x[i].x; // intermediate freq pairs
        g[2*i  ]=x[i].y;
    }
    g[N-1]=x[n2].x; // highest freq component
}
void FFT::fft(flt *f,flt *g,bool forward) {
    if (forward) {
        putF(f);
        do_fft(false);
        reorder();
        scale(1.0/N);
        unpackG(g);
    } else {
        packG(g);
        do_fft(true);
        reorder();
        getF(f);
    }
}
void FFT::cor(const flt *f, flt *r) { // same to tst_Cor
    putF(f); do_fft(false); reorder();
    flt s=0; int n2=N>>1;
    for(int i=1;i<=n2;i++) { flt t=x[i].norm(); s+=t; x[i]=t; x[N-i]=t; }
    x[0]=x[n2].x-2*s;
    do_fft(true); reorder();
    { flt t=-2.0/N; for(int i=0;i<N;i++) { r[i]=t*x[i].x; } }
}
void FFT::corr(const flt *f1, const flt *f2, flt *r) { // same to tst_Corr
    int n2=N>>1;
    putF(f2); do_fft(false); reorder();
    if (!z) z=new cplx[N];
    {for(int i=0;i<N;i++) z[i]=x[i];}
    putF(f1); do_fft(false); reorder();
    flt s=0;
    for(int i=1;i<=n2;i++) {
        s+=x[i].norm()+z[i].norm();
        x[i]*=z[N-i];
        x[N-i].x= x[i].x;
        x[N-i].y=-x[i].y;
    }
    x[0]=x[n2].x-s;
    do_fft(true); reorder();
    { flt t=-2.0/N; for(int i=0;i<N;i++) r[i]=t*x[i].x; }
}
void FFT::scor(const flt *f, flt *sr) {
    putF(f); do_fft(false); reorder();
    flt s=0, k=1.0/N; int n2=N>>1;
    {for(int i=1;i<=n2;i++) { flt t=x[i].norm()*k; s+=t; x[i]=t; }}
    x[0]=x[n2].x-2*s;
    sr[0]=-2*x[0].x;
    for(int i=1;i<n2;i++) {
        sr[2*i-1]=-2*x[i].x;
        sr[2*i  ]=-2*x[i].y;
    }
    sr[N-1]=-2*x[n2].x;
}
void FFT::scorr(const flt *f1, const flt *f2, flt *sr) {
    int n2=N>>1;
    putF(f2); do_fft(false); reorder();
    if (!z) z=new cplx[N];
    for(int i=0;i<N;i++) z[i]=x[i];
    putF(f1); do_fft(false); reorder();
    flt s=0, k=1.0/N;
    {for(int i=1;i<=n2;i++) { s+=x[i].norm()+z[i].norm(); x[i]*=z[N-i]*k; }}
    x[0]=x[n2].x-s*k;
    sr[0]=-2*x[0].x;
    {for(int i=1;i<n2;i++)  { sr[2*i-1]=-2*x[i].x; sr[2*i  ]=-2*x[i].y; }}
    sr[N-1]=-2*x[n2].x;
}
void FFT::conv(const flt* f1,const flt* f2,flt *r) {
    int n2=N>>1;
    putF(f2); do_fft(false); reorder();
    if (!z) z=new cplx[N];
    flt k=1.0/N;
    {for(int i=0;i<N;i++) z[i]=x[i]*k; }
    putF(f1); do_fft(false); reorder();
    for(int i=1;i<=n2;i++) {
        x[i]*=z[N-i];
        x[N-i].x= x[i].x;
        x[N-i].y=-x[i].y;
    }
    x[0].x*=z[0].x;
    do_fft(true); reorder();
    getF(r);
}
void FFT::tst_conv(const flt* f1,const flt* f2,flt *r) {
    int m=N-1; 
    for(int i=0;i<N;i++) {
        flt s=0; for(int j=0;j<N;j++) s+=f1[(i+j)&m]*f2[j];
        r[i]=s;
    }
}
void ft(int n,cplx* f,cplx* g,bool forward) {
    flt df=2*M_PI/n; if (forward) df=-df;
    for(int i=0;i<n;i++) {
        cplx s=0; for(int k=0;k<n;k++) { s+=f[k]*cplx::expi(i*k*df); }
        g[i]=forward?s/n:s;
    }
}
void FFT::tst_ft(flt *f,flt *g,bool forward) {
    cplx *fc=new cplx[N];
    cplx *gc=new cplx[N];
    int n2=N/2;
    if (forward) {
        for(int k=0;k<N;k++) fc[k]=f[k];
        ft(N,fc,gc,true);
        for(int i=1;i<n2;i++) {
            g[2*i-1]=gc[i].x;
            g[2*i  ]=gc[i].y;
        }
        g[0  ]=gc[0 ].x;
        g[N-1]=gc[n2].x;
    } else {
        for(int i=1;i<n2;i++) {
            gc[N-i].x=  gc[i].x=g[2*i-1];
            gc[N-i].y=-(gc[i].y=g[2*i  ]);
        }
        gc[0 ]=g[0  ]; 
        gc[n2]=g[N-1];
        ft(N,gc,fc,false);
        for(int k=0;k<N;k++) f[k]=fc[k].x;
    }
    delete[] gc;
    delete[] fc;
}
void FFT::tst_cor(const flt *f, flt *r) {
    for(int i=0,m=N-1;i<N;i++) {
        flt s=0; for(int j=0;j<N;j++) s+=sqr(f[(i+j)&m]-f[j]);
        r[i]=s;
    }
}
void FFT::tst_corr(const flt *f1, const flt *f2, flt *r) {
    int m=N-1;
    for(int i=0;i<N;i++) {
        flt s=0; for(int j=0;j<N;j++) s+=sqr(f1[(i+j)&m]-f2[j]);
        r[i]=s;
    }
}

#include <stdio.h>
void tst_fft() {
    enum { N=8 };
    flt z[N]={7,8,1,2,3,4,5,6};
    flt f[N]={1,2,3,4,5,6,7,8};
    flt g[N],q[N],r[N],r1[N],sr[N],sr1[N];
    FFT fft(3);

    printf("FT\n");
    fft.tst_ft(f,g,true);
    fft.tst_ft(q,g,false);
    printf(" N    f[k]   g[k]   q[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f\n",i,f[i],g[i],q[i]);

    printf("FFT\n");
    fft.fft(f,g,true);
    fft.fft(q,g,false);
    printf(" N    f[k]   g[k]   q[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f\n",i,f[i],g[i],q[i]);

    printf("Cor\n");
    fft.tst_cor(f,r1);
    fft.cor(f,r);
    printf(" N    f[k]   r[k]   r1[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f\n",i,f[i],r[i],r1[i]);

    printf("Corr\n");
    fft.tst_corr(f,z,r1);
    fft.corr(f,z,r);
    printf(" N    f[k]   z[k]   r[k]   r1[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f %6.2f\n",i,f[i],z[i],r[i],r1[i]);

    printf("SCor\n");
    fft.tst_cor(f,r1); fft.tst_ft(r1,sr1);
    fft.scor(f,sr);
    printf(" N    f[k]   sr[k]  sr1[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f\n",i,f[i],sr[i],sr1[i]);

    printf("SCorr\n");
    fft.tst_corr(f,z,r1); fft.tst_ft(r1,sr1);
    fft.scorr(f,z,sr);
    printf(" N    f[k]   z[k]   sr[k]  sr1[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f %6.2f\n",i,f[i],z[i],sr[i],sr1[i]);

    printf("Conv\n");
    fft.tst_conv(f,z,r1);
    fft.conv(f,z,r);
    printf(" N    f[k]   z[k]   r[k]   r1[k]\n");
    for(int i=0;i<N;i++) printf("%2d. %6.2f %6.2f %6.2f %6.2f\n",i,f[i],z[i],r[i],r1[i]);
}