#include "analysis/imfilter.h"
#include "ProfilerApp.h"
#include <math.h>
#include <string.h>


#define IMFILTER_INSIST INSIST
#define IMFILTER_ASSERT ASSERT
#define IMFILTER_ERROR  ERROR


// Function to convert an index
static inline int imfilter_index( int index, const int N, const imfilter::BC bc )
{
    if ( index < 0 || index >= N ) {
        if ( bc == imfilter::BC::symmetric ) {
            index = ( 2 * N - index ) % N;
        } else if ( bc == imfilter::BC::replicate ) {
            index = index < 0 ? 0 : N - 1;
        } else if ( bc == imfilter::BC::circular ) {
            index = ( index + N ) % N;
        } else if ( bc == imfilter::BC::fixed ) {
            index = -1;
        }
    }
    return index;
}


// Function to copy a 1D array and pad with the appropriate BC
template<class TYPE>
static inline void copy_array( const int N, const int Ns, const int Nh,
    const TYPE *A, const imfilter::BC BC, const TYPE X, TYPE *B )
{
    // Fill the center with a memcpy
    for (int i=0; i<N; i++ )
        B[i+Nh] = A[i*Ns];
    // Fill the boundaries
    for (int i=0; i<Nh; i++ ) {
        int j1 = imfilter_index( -(i+1), N, BC );
        int j2 = imfilter_index(   N+i, N, BC );
        B[Nh-i-1] = j1==-1 ? X : B[Nh+j1];
        B[N+Nh+i] = j2==-1 ? X : B[Nh+j2];
    }
}


/********************************************************
* Perform a 1D filter in a single direction             *
********************************************************/
template<class TYPE>
static void filter_direction( int Ns, int N, int Ne, int Nh, const TYPE *H,
    imfilter::BC boundary, TYPE X, TYPE *A )
{
    if ( Nh < 0 )
        IMFILTER_ERROR("Invalid filter size");
    if ( Nh == 0 ) {
        for (int i=0; i<Ns*N*Ne; i++)
            A[i] *= H[0];
        return;
    }
    TYPE *tmp = new TYPE[N+2*Nh];
    for (int j=0; j<Ne; j++) {
        for (int i=0; i<Ns; i++) {
            copy_array( N, Ns, Nh, &A[i+j*Ns*N], boundary, X, tmp );
            for (int k=0; k<N; k++) {
                TYPE tmp2 = 0;
                for (int m=0; m<=2*Nh; m++)
                    tmp2 += H[m] * tmp[k+m];
                A[i+k*Ns+j*Ns*N] = tmp2;
            }
        }
    }
    delete[] tmp;
}
template<class TYPE>
static void filter_direction( int Ns, int N, int Ne, int Nh,
    std::function<TYPE(const Array<TYPE>&)> H, imfilter::BC boundary, TYPE X, TYPE *A )
{
    if ( Nh < 0 )
        IMFILTER_ERROR("Invalid filter size");
    TYPE *tmp = new TYPE[N+2*Nh];
    Array<TYPE> tmp2(2*Nh+1);
    for (int j=0; j<Ne; j++) {
        for (int i=0; i<Ns; i++) {
            copy_array( N, Ns, Nh, &A[i+j*Ns*N], boundary, X, tmp );
            for (int k=0; k<N; k++) {
                for (int m=0; m<=2*Nh; m++)
                    tmp2(m) = tmp[k+m];
                A[i+k*Ns+j*Ns*N] = H(tmp2);
            }
        }
    }
    delete[] tmp;
}
template<class TYPE>
static void filter_direction( int Ns, int N, int Ne, int Nh,
    std::function<TYPE(int, const TYPE*)> H, imfilter::BC boundary, TYPE X, TYPE *A )
{
    if ( Nh < 0 )
        IMFILTER_ERROR("Invalid filter size");
    TYPE *tmp = new TYPE[N+2*Nh];
    int Nh2 = 2*Nh+1;
    for (int j=0; j<Ne; j++) {
        for (int i=0; i<Ns; i++) {
            copy_array( N, Ns, Nh, &A[i+j*Ns*N], boundary, X, tmp );
            for (int k=0; k<N; k++)
                A[i+k*Ns+j*Ns*N] = H(Nh2,&tmp[k]);
        }
    }
    delete[] tmp;
}



/********************************************************
*  Create a filter                                      *
********************************************************/
template<class TYPE>
Array<TYPE> imfilter::create_filter( const std::vector<int>& N0, const std::string &type, const void *args )
{
    std::vector<size_t> N2(N0.size());
    for (size_t i=0; i<N2.size(); i++)
        N2[i] = 2*N0[i]+1;
    Array<TYPE> h(N2);
    h.fill(0);
    if ( type == "average" ) {
        // average
        h.fill( 1.0 / static_cast<TYPE>( h.length() ) );
    } else if ( type == "gaussian" ) {
        // gaussian
        if ( N0.size() > 3 )
            IMFILTER_ERROR( "Not implimented for dimensions > 3" );
        TYPE std[3] = { 0.5, 0.5, 0.5 };
        if ( args != NULL ) {
            const TYPE *args2 = reinterpret_cast<const TYPE*>( args );
            for ( size_t d = 0; d < N0.size(); d++ )
                std[d]  = args2[d];
        }
        auto N = N0;
        N.resize(3,0);
        for ( int k = -N[2]; k <= N[2]; k++ ) {
            for ( int j = -N[1]; j <= N[1]; j++ ) {
                for ( int i = -N[0]; i <= N[0]; i++ ) {
                    h(i+N[0],j+N[1],k+N[2]) = 
                        exp( -i * i / ( 2 * std[0] * std[0] ) ) *
                        exp( -j * j / ( 2 * std[1] * std[1] ) ) *
                        exp( -k * k / ( 2 * std[2] * std[2] ) );
                }
            }
        }
        h.scale( 1.0/h.sum() );
    } else {
        IMFILTER_ERROR( "Unknown filter" );
    }
    return h;
}


// Perform 2-D filtering
template<class TYPE>
void imfilter_2D( int Nx, int Ny, const TYPE *A, int Nhx, int Nhy, const TYPE *H,
    imfilter::BC BCx, imfilter::BC BCy, const TYPE X, TYPE *B )
{
    IMFILTER_ASSERT( A != B );
    PROFILE_START( "imfilter_2D" );
    memset( B, 0, Nx * Ny * sizeof( TYPE ) );
    for ( int j1 = 0; j1 < Ny; j1++ ) {
        for ( int i1 = 0; i1 < Nx; i1++ ) {
            TYPE tmp = 0;
            if ( i1 >= Nhx && i1 < Nx - Nhx && j1 >= Nhy && j1 < Ny - Nhy ) {
                int ijkh = 0;
                for ( int j2 = j1 - Nhy; j2 <= j1 + Nhy; j2++ ) {
                    for ( int i2 = i1 - Nhx; i2 <= i1 + Nhx; i2++, ijkh++ )
                        tmp += H[ijkh] * A[i2 + j2 * Nx];
                }
            } else {
                int ijkh = 0;
                for ( int jh = -Nhy; jh <= Nhy; jh++ ) {
                    int j2 = imfilter_index( j1+jh, Ny, BCy );
                    for ( int ih = -Nhx; ih <= Nhx; ih++ ) {
                        int i2     = imfilter_index( i1+ih, Nx, BCx );
                        bool fixed = i2 == -1 || j2 == -1;
                        TYPE A2  = fixed ? X : A[i2 + j2 * Nx];
                        tmp += H[ijkh] * A2;
                        ijkh++;
                    }
                }
            }
            B[i1 + j1 * Nx] = tmp;
        }
    }
    PROFILE_STOP( "imfilter_2D" );
}


// Perform 3-D filtering
template<class TYPE>
void imfilter_3D( int Nx, int Ny, int Nz, const TYPE *A, int Nhx, int Nhy, int Nhz,
    const TYPE *H, imfilter::BC BCx, imfilter::BC BCy, imfilter::BC BCz,
    const TYPE X, TYPE *B )
{
    IMFILTER_ASSERT( A != B );
    PROFILE_START( "imfilter_3D" );
    memset( B, 0, Nx * Ny * Nz * sizeof( TYPE ) );
    for ( int k1 = 0; k1 < Nz; k1++ ) {
        for ( int j1 = 0; j1 < Ny; j1++ ) {
            for ( int i1 = 0; i1 < Nx; i1++ ) {
                TYPE tmp = 0;
                int ijkh   = 0;
                for ( int kh = -Nhz; kh <= Nhz; kh++ ) {
                    int k2 = imfilter_index( k1+kh, Nz, BCz );
                    for ( int jh = -Nhy; jh <= Nhy; jh++ ) {
                        int j2 = imfilter_index( j1+jh, Ny, BCy );
                        for ( int ih = -Nhx; ih <= Nhx; ih++ ) {
                            int i2     = imfilter_index( i1+ih, Nx, BCx );
                            bool fixed = i2 == -1 || j2 == -1 || k2 == -1;
                            TYPE A2  = fixed ? X : A[i2 + j2 * Nx + k2 * Nx * Ny];
                            tmp += H[ijkh] * A2;
                            ijkh++;
                        }
                    }
                }
                B[i1 + j1 * Nx + k1 * Nx * Ny] = tmp;
            }
        }
    }
    PROFILE_STOP( "imfilter_3D" );
}


/********************************************************
*  Perform N-D filtering                                *
********************************************************/
template<class TYPE>
Array<TYPE> imfilter::imfilter( const Array<TYPE>& A, 
    const Array<TYPE>& H, const std::vector<imfilter::BC>& BC, const TYPE X )
{
    IMFILTER_ASSERT( A.ndim() == H.ndim() );
    IMFILTER_ASSERT( A.ndim() == BC.size() );
    std::vector<size_t> Nh = H.size();
    for (int d=0; d<A.ndim(); d++) {
        Nh[d] = (H.size(d)-1)/2;
        IMFILTER_INSIST(2*Nh[d]+1==H.size(d),"Filter must be of size 2*N+1");
    }
    auto B = A;
    if ( A.ndim() == 1 ) {
        PROFILE_START( "imfilter_1D" );
        filter_direction( 1, A.size(0), 1, Nh[0], H.data(), BC[0], X, B.data() );
        PROFILE_STOP( "imfilter_1D" );
    } else if ( A.ndim() == 2 ) {
        imfilter_2D( A.size(0), A.size(1), A.data(), Nh[0], Nh[1], H.data(), BC[0], BC[1], X, B.data() );
    } else if ( A.ndim() == 3 ) {
        imfilter_3D( A.size(0), A.size(1), A.size(2), A.data(),
            Nh[0], Nh[1], Nh[2], H.data(), BC[0], BC[1], BC[2], X, B.data() );
    } else {
        IMFILTER_ERROR( "Arbitrary dimension not yet supported" );
    }
    return B;
}
template<class TYPE>
Array<TYPE> imfilter::imfilter( const Array<TYPE>& A, const std::vector<int>& Nh0,
    std::function<TYPE(const Array<TYPE>&)> H,
    const std::vector<imfilter::BC>& BC0, const TYPE X )
{
    PROFILE_START( "imfilter (lambda)" );
    IMFILTER_ASSERT( A.ndim() == Nh0.size() );
    IMFILTER_ASSERT( A.ndim() == BC0.size() );
    std::vector<size_t> Nh2( A.size() );
    for (int d=0; d<A.ndim(); d++)
        Nh2[d] = 2*Nh0[d]+1;
    auto B = A;
    Array<TYPE> data(Nh2);
    IMFILTER_INSIST(A.ndim()<=3,"Not programmed for more than 3 dimensions yet");
    auto N = A.size();
    auto Nh = Nh0;
    auto BC = BC0;
    N.resize(3,1);
    Nh.resize(3,0);
    BC.resize(3,imfilter::BC::fixed);
    for ( int k1 = 0; k1 < N[2]; k1++ ) {
        for ( int j1 = 0; j1 < N[1]; j1++ ) {
            for ( int i1 = 0; i1 < N[0]; i1++ ) {
                for ( int kh = -Nh[2]; kh <= Nh[2]; kh++ ) {
                    int k2 = imfilter_index( k1+kh, N[2], BC[2] );
                    for ( int jh = -Nh[1]; jh <= Nh[1]; jh++ ) {
                        int j2 = imfilter_index( j1+jh, N[1], BC[1] );
                        for ( int ih = -Nh[0]; ih <= Nh[0]; ih++ ) {
                            int i2 = imfilter_index( i1+ih, N[0], BC[0] );
                            bool fixed = i2 == -1 || j2 == -1 || k2 == -1;
                            data(ih+Nh[0],jh+Nh[1],kh+Nh[2]) = fixed ? X : A(i2,j2,k2);
                        }
                    }
                }
                B(i1,j1,k1) = H( data );
            }
        }
    }
    PROFILE_STOP( "imfilter (lambda)" );
    return B;
}


/********************************************************
* imfilter with separable filter functions              *
********************************************************/
template<class TYPE>
Array<TYPE> imfilter::imfilter_separable( const Array<TYPE>& A,
    const std::vector<Array<TYPE>>& H,
    const std::vector<imfilter::BC>& boundary, const TYPE X )
{
    PROFILE_START( "imfilter_separable" );
    IMFILTER_ASSERT( A.ndim() == (int) H.size() );
    IMFILTER_ASSERT( A.ndim() == (int) boundary.size() );
    std::vector<size_t> Nh( H.size() );
    for (int d=0; d<A.ndim(); d++) {
        IMFILTER_ASSERT(H[d].ndim()==1);
        Nh[d] = (H[d].length()-1)/2;
        IMFILTER_INSIST(2*Nh[d]+1==H[d].length(),"Filter must be of size 2*N+1");
    }
    auto B = A;
    for ( int d = 0; d < A.ndim(); d++ ) {
        int N = A.size(d);
        int Ns = 1;
        int Ne = 1;
        for ( int d2 = 0; d2 < d; d2++ )
            Ns *= A.size(d2);
        for ( int d2 = d+1; d2 < A.ndim(); d2++ )
            Ne *= A.size(d2);
        filter_direction( Ns, N, Ne, Nh[d], H[d].data(), boundary[d], X, B.data() );
    }
    PROFILE_STOP( "imfilter_separable" );
    return B;
}
template<class TYPE>
Array<TYPE> imfilter::imfilter_separable( const Array<TYPE>& A, const std::vector<int>& Nh,
    std::vector<std::function<TYPE(const Array<TYPE>&)>> H,
    const std::vector<imfilter::BC>& boundary, const TYPE X )
{
    PROFILE_START( "imfilter_separable (lambda)" );
    IMFILTER_ASSERT( A.ndim() == (int) boundary.size() );
    auto B = A;
    for ( int d = 0; d < A.ndim(); d++ ) {
        int N = A.size(d);
        int Ns = 1;
        int Ne = 1;
        for ( int d2 = 0; d2 < d; d2++ )
            Ns *= A.size(d2);
        for ( int d2 = d+1; d2 < A.ndim(); d2++ )
            Ne *= A.size(d2);
        filter_direction( Ns, N, Ne, Nh[d], H[d], boundary[d], X, B.data() );
    }
    PROFILE_STOP( "imfilter_separable (lambda)" );
    return B;
}
template<class TYPE>
Array<TYPE> imfilter::imfilter_separable( const Array<TYPE>& A, const std::vector<int>& Nh,
    std::vector<std::function<TYPE(int, const TYPE*)>> H,
    const std::vector<imfilter::BC>& boundary, const TYPE X )
{
    PROFILE_START( "imfilter_separable (function)" );
    IMFILTER_ASSERT( A.ndim() == (int) boundary.size() );
    auto B = A;
    for ( int d = 0; d < A.ndim(); d++ ) {
        int N = A.size(d);
        int Ns = 1;
        int Ne = 1;
        for ( int d2 = 0; d2 < d; d2++ )
            Ns *= A.size(d2);
        for ( int d2 = d+1; d2 < A.ndim(); d2++ )
            Ne *= A.size(d2);
        filter_direction( Ns, N, Ne, Nh[d], H[d], boundary[d], X, B.data() );
    }
    PROFILE_STOP( "imfilter_separable (function)" );
    return B;
}