#pragma once

///// TRUE GAUSSIAN BLUR ///// 

int sigma_to_kernel_radius(float sigma)
{
    // capture almost all the gaussian extent
    return std::ceil(2.57f*sigma);
}

float gaussian( float x, float mu, float sigma ) 
{
    const float a = ( x - mu ) / sigma;
    return std::exp( -0.5f * a * a );
}

void make_gaussian_kernel(float sigma, float *& kernel, int& k) 
{
    // kernel radius from sigma
    k = sigma_to_kernel_radius(sigma);
    // kernel alloc
    kernel = new float[(2*k+1)*(2*k+1)];
    // unnormalized kernel weights 
    for (int row = 0, index = 0; row < 2*k+1; row++)
        for (int col = 0; col < 2*k+1; col++, index++) 
            kernel[index] = gaussian(row, k, sigma) * gaussian(col, k, sigma);
    // we do not need to normalize since it will be done during the convolution
}

template<typename T, int C>
void convolve(const T * in, T * out, const int w, const int h, const float * kernel, const int k)
{
    #pragma omp parallel for collapse(2)
    for (int row = 0; row < h; row++)
    for (int col = 0; col < w; col++) 
    {
        // temporary data
        float val[C];
        for(int ch = 0; ch < C; ++ch)
            val[ch] = 0;

        // compute filtered value
        float sum = 0;
        for(int irow = row-k, kid = 0; irow <= row+k; ++irow)
        for(int icol = col-k; icol <= col+k; ++icol, ++kid)
        {

        #if 0
            // skip policy
            if(irow < 0 || icol < 0 || irow > h-1 || icol > w-1) 
                continue;
            int iindex = (irow*w+icol)*C;
        #else
            // closest policy
            int crow = std::min(h-1, std::max(0, irow));
            int ccol = std::min(w-1, std::max(0, icol));
            int iindex = (crow*w+ccol)*C;
        #endif

            // accumulate
            sum += kernel[kid];
            for(int ch = 0; ch < C; ++ch)
                val[ch] += in[iindex+ch] * kernel[kid];
        }

        // normalize and copy to output buffer
        for(int ch = 0; ch < C; ++ch)
            out[(row*w+col)*C+ch] = val[ch]/sum;
    }
}

///// FAST GAUSSIAN BLUR ///// 

// edit1: avoid std::round by adding 0.5f and cast to integer type (2.5 times faster)
// edit2: it seems using the round is wrong the result has slightly more intensity than the true Gaussian blur result so I removed it. 
// and it make the code lighter at the same time since we have a unique version independant of the data type.
// edit3: integer versions acculumate way to much errors, every variable should be float 
template<typename T, int C>
void horizontal_blur(const T * in, T * out, const int w, const int h, const int r)
{
    float iarr = 1.f / (r+r+1);
    #pragma omp parallel for
    for(int i=0; i<h; i++) 
    {
        int ti = i*w; 
        int li = ti;  
        int ri = ti+r;

        float fv[C], lv[C], val[C];

        for(int ch = 0; ch < C; ++ch)
        {
            fv[ch] = in[ti*C+ch];
            lv[ch] = in[(ti+w-1)*C+ch];
            val[ch] = (r+1)*fv[ch];
        }

        for(int j=0; j<r; j++) 
        for(int ch = 0; ch < C; ++ch)
        {
            val[ch] += in[(ti+j)*C+ch]; 
        }

        for(int j=0; j<=r; j++, ri++, ti++) 
        for(int ch = 0; ch < C; ++ch)
        { 
            val[ch] += in[ri*C+ch] - fv[ch]; 
            out[ti*C+ch] = val[ch]*iarr;
        }

        for(int j=r+1; j<w-r; j++, ri++, ti++, li++) 
        for(int ch = 0; ch < C; ++ch)
        { 
            val[ch] += in[ri*C+ch] - in[li*C+ch]; 
            out[ti*C+ch] = val[ch]*iarr;
        }

        for(int j=w-r; j<w; j++, ti++, li++) 
        for(int ch = 0; ch < C; ++ch)
        { 
            val[ch] += lv[ch] - in[li*C+ch]; 
            out[ti*C+ch] = val[ch]*iarr;
        }
    }
}

template<typename T>
void horizontal_blur(const T * in, T * out, const int w, const int h, const int c, const int r)
{
    switch(c)
    {
        case 1: horizontal_blur<T,1>(in, out, w, h, r); break;
        case 2: horizontal_blur<T,2>(in, out, w, h, r); break;
        case 3: horizontal_blur<T,3>(in, out, w, h, r); break;
        case 4: horizontal_blur<T,4>(in, out, w, h, r); break;
        default: printf("%d channels is not supported yet. Add a specific case if possible or fall back to the generic version.", c); break;
        // default: horizontal_blur<T>(in, out, w, h, c, r); break;
    }
}

template<typename T, int C>
void flip_block(const T * in, T * out, const int w, const int h)
{
    constexpr int block = 256/C;
    #pragma omp parallel for collapse(2)
    for(int x= 0; x < w; x+= block)
    for(int y= 0; y < h; y+= block)
    {
        const T * p = in + y*w*C + x*C;
        T * q = out + y*C + x*h*C;
        
        const int blockx= std::min(w, x+block) - x;
        const int blocky= std::min(h, y+block) - y;
        for(int xx= 0; xx < blockx; xx++)
        {
            for(int yy= 0; yy < blocky; yy++)
            {
                for(int k= 0; k < C; k++)
                    q[k]= p[k];
                p+= w*C;
                q+= C;                    
            }
            p+= -blocky*w*C + C;
            q+= -blocky*C + h*C;
        }
    }
}

template<typename T>
void flip_block(const T * in, T * out, const int w, const int h, const int c)
{
    switch(c)
    {
        case 1: flip_block<T,1>(in, out, w, h); break;
        case 2: flip_block<T,2>(in, out, w, h); break;
        case 3: flip_block<T,3>(in, out, w, h); break;
        case 4: flip_block<T,4>(in, out, w, h); break;
        default: printf("%d channels is not supported yet. Add a specific case if possible or fall back to the generic version.", c); break;
        // default: flip_block<T>(in, out, w, h, c); break;
    }
}

void sigma_to_box_radius(int boxes[], const float sigma, const int n)  
{
    // ideal filter width
    float wi = std::sqrt((12*sigma*sigma/n)+1); 
    int wl = wi; // no need std::floor  
    if(wl%2==0) wl--;
    int wu = wl+2;
                
    float mi = (12*sigma*sigma - n*wl*wl - 4*n*wl - 3*n)/(-4*wl - 4);
    int m = mi+0.5f; // avoid std::round by adding 0.5f and cast to integer type
                
    for(int i=0; i<n; i++) 
        boxes[i] = ((i < m ? wl : wu) - 1) / 2;
}

// modified version (fastest one)
template<typename T>
void fast_gaussian_blur_Nh_Nv(T *& in, T *& out, const int w, const int h, const int c, const float sigma, const int n) 
{
    // compute box kernel sizes
    std::vector<int> boxes(n);
    sigma_to_box_radius(boxes.data(), sigma, n);

    // perform N horizontal blur passes
    for(int i = 0; i < n; ++i)
    {
        horizontal_blur(in, out, w, h, c, boxes[i]);
        std::swap(in, out);
    }   

    // flip buffer
    flip_block(in, out, w, h, c);
    std::swap(in, out);
    
    // perform N horizontal blur passes
    for(int i = 0; i < n; ++i)
    {
        horizontal_blur(in, out, h, w, c, boxes[i]);
        std::swap(in, out);
    }   
    
    // flip buffer
    flip_block(in, out, h, w, c);
}

// original version (suffer from cache incoherencies)
template<typename T>
void fast_gaussian_blur_N_hv(T *& in, T *& out, const int w, const int h, const int c, const float sigma, const int n) 
{
    // compute box kernel sizes
    std::vector<int> boxes(n);
    sigma_to_box_radius(boxes.data(), sigma, n);

    // perform N box blur passes (1 horizontal + 1 vertical)
    for(int i = 0; i < n; ++i)
    {
        horizontal_blur(in, out, w, h, c, boxes[i]);
        flip_block(out, in, w, h, c);
        horizontal_blur(in, out, h, w, c, boxes[i]);
        flip_block(out, in, h, w, c);
    }   

    // flip buffer
    std::swap(in, out);
}