// Copyright (C) 2013  Davis E. King (davis@dlib.net)
// License: Boost Software License   See LICENSE.txt for the full license.
#ifndef DLIB_FFt_Hh_
#define DLIB_FFt_Hh_

#include "matrix_fft_abstract.h"
#include "matrix_utilities.h"
#include "../hash.h"
#include "../algs.h"
#include "../math.h"

#ifdef DLIB_USE_MKL_FFT
#include "mkl_fft.h"
#else
#include "kiss_fft.h"
#endif

namespace dlib
{     
// ----------------------------------------------------------------------------------------
    
    constexpr bool is_power_of_two (const unsigned long n)
    {
        return n == 0 ? true : (n & (n - 1)) == 0;
    }
    
// ----------------------------------------------------------------------------------------
    
    constexpr long fftr_nc_size(long nc)
    {
        return nc == 0 ? 0 : nc/2+1;
    }
    
// ----------------------------------------------------------------------------------------   
    
    constexpr long ifftr_nc_size(long nc)
    {
        return nc == 0 ? 0 : 2*(nc-1);
    }
  
// ----------------------------------------------------------------------------------------
    
    template < typename T, typename Alloc >
    matrix<std::complex<T>,0,1> fft (const std::vector<std::complex<T>, Alloc>& in)
    {
        //complex FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        matrix<std::complex<T>,0,1> out(in.size());
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({(long)in.size()}, &in[0], &out(0,0), false);
#else
            kiss_fft({(long)in.size()}, &in[0], &out(0,0), false);
#endif
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template < typename T, long NR, long NC, typename MM, typename L >
    matrix<std::complex<T>,NR,NC,MM,L> fft (const matrix<std::complex<T>,NR,NC,MM,L>& in)
    {
        //complex FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        matrix<std::complex<T>,NR,NC,MM,L> out(in.nr(), in.nc());
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({in.nr(),in.nc()}, &in(0,0), &out(0,0), false);
#else
            kiss_fft({in.nr(),in.nc()}, &in(0,0), &out(0,0), false);
#endif
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template <typename EXP>
    typename EXP::matrix_type fft (const matrix_exp<EXP>& data)
    {
        //complex FFT for expression template
        static_assert(is_complex<typename EXP::type>::value, "input should be complex");
        typename EXP::matrix_type in(data);
        return fft(in);
    }
    
// ----------------------------------------------------------------------------------------
    
    template < typename T, typename Alloc >
    matrix<std::complex<T>,0,1> ifft (const std::vector<std::complex<T>, Alloc>& in)
    {
        //complex FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        matrix<std::complex<T>,0,1> out(in.size());
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({(long)in.size()}, &in[0], &out(0,0), true);
#else
            kiss_fft({(long)in.size()}, &in[0], &out(0,0), true);
#endif
            out /= out.size();
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template < typename T, long NR, long NC, typename MM, typename L >
    matrix<std::complex<T>,NR,NC,MM,L> ifft (const matrix<std::complex<T>,NR,NC,MM,L>& in)
    {
        //inverse complex FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        matrix<std::complex<T>,NR,NC,MM,L> out(in.nr(), in.nc());
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({in.nr(),in.nc()}, &in(0,0), &out(0,0), true);
#else
            kiss_fft({in.nr(),in.nc()}, &in(0,0), &out(0,0), true);
#endif
            out /= out.size();
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template <typename EXP>
    typename EXP::matrix_type ifft (const matrix_exp<EXP>& data)
    {
        //inverse complex FFT for expression template
        static_assert(is_complex<typename EXP::type>::value, "input should be complex");
        typename EXP::matrix_type in(data);
        return ifft(in);
    }

// ----------------------------------------------------------------------------------------
    
    template<typename T, long NR, long NC, typename MM, typename L>
    matrix<std::complex<T>,NR,fftr_nc_size(NC),MM,L> fftr (const matrix<T,NR,NC,MM,L>& in)
    {
        //real FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        DLIB_ASSERT(in.nc() % 2 == 0, "last dimension " << in.nc() << " needs to be even otherwise ifftr(fftr(data)) won't have matching dimensions");
        matrix<std::complex<T>,NR,fftr_nc_size(NC),MM,L> out(in.nr(), fftr_nc_size(in.nc()));
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fftr({in.nr(),in.nc()}, &in(0,0), &out(0,0));
#else
            kiss_fftr({in.nr(),in.nc()}, &in(0,0), &out(0,0));
#endif
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template <typename EXP>
    matrix<add_complex_t<typename EXP::type>> fftr (const matrix_exp<EXP>& data)
    {
        //real FFT for expression template
        static_assert(std::is_floating_point<typename EXP::type>::value, "input should be real");
        matrix<typename EXP::type> in(data);
        return fftr(in);
    }
    
// ----------------------------------------------------------------------------------------
    
    template<typename T, long NR, long NC, typename MM, typename L>
    matrix<T,NR,ifftr_nc_size(NC),MM,L> ifftr (const matrix<std::complex<T>,NR,NC,MM,L>& in)
    {
        //inverse real FFT
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        matrix<T,NR,ifftr_nc_size(NC),MM,L> out(in.nr(), ifftr_nc_size(in.nc()));
        if (in.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_ifftr({out.nr(),out.nc()}, &in(0,0), &out(0,0));
#else
            kiss_ifftr({out.nr(),out.nc()}, &in(0,0), &out(0,0));
#endif
            out /= out.size();
        }
        return out;
    }
    
// ----------------------------------------------------------------------------------------
    
    template <typename EXP>
    matrix<remove_complex_t<typename EXP::type>> ifftr (const matrix_exp<EXP>& data)
    {
        //inverse real FFT for expression template
        static_assert(is_complex<typename EXP::type>::value, "input should be complex");        
        matrix<typename EXP::type> in(data);
        return ifftr(in);
    }
    
// ----------------------------------------------------------------------------------------
    
    template < typename T, typename Alloc >
    void fft_inplace (std::vector<std::complex<T>, Alloc>& data)
    {
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        if (data.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({(long)data.size()}, &data[0], &data[0], false);
#else
            kiss_fft({(long)data.size()}, &data[0], &data[0], false);
#endif
        }
    }
    
// ----------------------------------------------------------------------------------------
    
    template < typename T, long NR, long NC, typename MM, typename L >
    void fft_inplace (matrix<std::complex<T>,NR,NC,MM,L>& data)
    {
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        if (data.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({data.nr(),data.nc()}, &data(0,0), &data(0,0), false);
#else
            kiss_fft({data.nr(),data.nc()}, &data(0,0), &data(0,0), false);
#endif
        }
    }

// ----------------------------------------------------------------------------------------

    template < typename T, typename Alloc >
    void ifft_inplace (std::vector<std::complex<T>, Alloc>& data)
    {
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        if (data.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({(long)data.size()}, &data[0], &data[0], true);
#else
            kiss_fft({(long)data.size()}, &data[0], &data[0], true);
#endif
        }
    }
    
// ----------------------------------------------------------------------------------------

    template < typename T, long NR, long NC, typename MM, typename L >
    void ifft_inplace (matrix<std::complex<T>,NR,NC,MM,L>& data)
    {
        static_assert(std::is_floating_point<T>::value, "only support floating point types");
        if (data.size() != 0)
        {
#ifdef DLIB_USE_MKL_FFT
            mkl_fft({data.nr(),data.nc()}, &data(0,0), &data(0,0), true);
#else
            kiss_fft({data.nr(),data.nc()}, &data(0,0), &data(0,0), true);
#endif
        }
    }

// ----------------------------------------------------------------------------------------

    namespace details
    {
        struct fft_func
        {
            template<typename MAT, typename T = typename MAT::type, typename std::enable_if<is_complex<T>::value, bool>::type = true>
            auto operator()(const MAT& mat) const { return dlib::fft(mat); }

            template<typename MAT, typename T = typename MAT::type, typename std::enable_if<!is_complex<T>::value, bool>::type = true>
            auto operator()(const MAT& mat) const { return dlib::fft(dlib::complex_matrix(mat)); }

            static constexpr std::size_t freqsize(std::size_t fftsize) { return fftsize; }
        };

        struct fftr_func
        {
            template<typename MAT>
            auto operator()(const MAT& mat) const { return dlib::fftr(mat); }

            static constexpr std::size_t freqsize(std::size_t fftsize) { return dlib::fftr_nc_size(fftsize); }
        };

        struct ifft_func
        {
            template<typename MAT>
            auto operator()(const MAT& mat) const { return dlib::ifft(mat); }
        };

        struct ifftr_func
        {
            template<typename MAT>
            auto operator()(const MAT& mat) const { return dlib::ifftr(mat); }
        };

        template <
            typename EXP,
            typename WINDOW,
            typename FFT_FUNC
        >
        auto stft_impl (
            const matrix_exp<EXP>& signal,
            const WINDOW& w,
            std::size_t fftsize,
            std::size_t wlen,
            std::size_t hoplen,
            const FFT_FUNC& fft_obj
        )
        {
            using T = typename EXP::type;
            using R = remove_complex_t<T>;
            using C = add_complex_t<T>;

            static_assert(std::is_floating_point<R>::value, "underlying type must be real or complex floating point type");
            DLIB_ASSERT(is_vector(signal), "input must be a vector type");
            DLIB_ASSERT(signal.size() >= (long)wlen, "signal.size() >= wlen not satisfied");
            DLIB_ASSERT(fftsize >= wlen, "fftsize >= wlen not satisfied");
            DLIB_ASSERT(wlen >= hoplen, "wlen >= hoplen not satisfied");

            // Input is left-padded by wlen/2 and right-padded wlen/2
            const std::size_t total_padding = wlen;
            const std::size_t overlap       = wlen - hoplen;
            const std::size_t nframes       = (signal.size() + total_padding - overlap) / hoplen;
            matrix<C> stft = zeros_matrix<C>(nframes,
                                             FFT_FUNC::freqsize(fftsize));
            matrix<R> win(1,wlen);
            for (std::size_t i = 0 ; i < wlen ; ++i)
                win(0, i) = w(i, wlen);

            // TODO: reduce extra buffers, e.g. padded
            matrix<T> padded;

            if (is_row_vector(signal))
                padded = join_rows(join_rows(zeros_matrix<T>(1, wlen/2), signal), zeros_matrix<T>(1, wlen/2));
            else
                padded = join_rows(join_rows(zeros_matrix<T>(1, wlen/2), trans(signal)), zeros_matrix<T>(1, wlen/2));

            for (long i = 0 ; i < stft.nr() ; ++i)
            {
                set_rowm(stft, i) = fft_obj(join_rows(pointwise_multiply(win, subm(padded, 0, i*hoplen, 1, wlen)),
                                                      zeros_matrix<T>(1, fftsize - wlen)));
            }

            return stft;
        }

        template <
            typename ReturnType,
            typename EXP,
            typename WINDOW,
            typename IFFT_FUNC
        >
        auto istft_impl (
            const matrix_exp<EXP>& stft,
            const WINDOW& w,
            std::size_t wlen,
            std::size_t hoplen,
            const IFFT_FUNC& ifft_obj
        )
        {
            using T = typename EXP::type;
            using R = remove_complex_t<T>;

            static_assert(is_complex<T>::value, "matrix type must be complex");
            static_assert(std::is_floating_point<R>::value, "underlying type must be complex floating point type");
            DLIB_ASSERT(stft.nc() > 0 && stft.nr() > 0, "stft must be non-empty");
            DLIB_ASSERT(stft.nc() >= (long)wlen, "fftsize >= wlen not satisfied");
            DLIB_ASSERT(wlen >= hoplen, "wlen >= hoplen not satisfied");

            const size_t ntime = (stft.nr() - 1) * hoplen + wlen;
            matrix<ReturnType> signal = zeros_matrix<ReturnType>(1, ntime);
            matrix<R> norm = zeros_matrix<R>(1, ntime);
            matrix<R> win(1, wlen);
            for (std::size_t i = 0 ; i < wlen ; ++i)
                win(0, i) = w(i, wlen);
            matrix<R> win2 = squared(win);

            for (long t = 0 ; t < stft.nr() ; ++t)
            {
                set_subm(signal, 0, t*hoplen, 1, wlen) += pointwise_multiply(win, subm(ifft_obj(rowm(stft, t)), 0, 0, 1, wlen));
                set_subm(norm,   0, t*hoplen, 1, wlen) += win2;
            }

            // Remove padding of wlen/2 and wlen/2 on either end
            DLIB_ASSERT(sum(subm(norm, 0, wlen/2, 1, ntime - wlen) < 1e-13) == 0, "NOLA constraint not satisfied");
            signal = pointwise_divide(subm(signal, 0, wlen/2, 1, ntime - wlen),
                                      subm(norm,   0, wlen/2, 1, ntime - wlen));

            return signal;
        }
    }

// ----------------------------------------------------------------------------------------

    inline auto make_hann()
    {
        return [](std::size_t i, std::size_t N) {return hann(i, N, PERIODIC); };
    }

    inline auto make_blackman()
    {
        return [](std::size_t i, std::size_t N) {return blackman(i, N, PERIODIC);};
    }

    inline auto make_blackman_nuttall()
    {
        return [](std::size_t i, std::size_t N) {return blackman_nuttall(i, N, PERIODIC);};
    }

    inline auto make_blackman_harris()
    {
        return [](std::size_t i, std::size_t N) { return blackman_harris(i, N, PERIODIC); };
    }

    inline auto make_blackman_harris7()
    {
        return [](std::size_t i, std::size_t N) { return blackman_harris7(i, N, PERIODIC); };
    }

    inline auto make_kaiser(beta_t beta)
    {
        return [=](std::size_t i, std::size_t N){return kaiser(i, N, beta, PERIODIC);};
    }

// ----------------------------------------------------------------------------------------

    template <typename EXP, typename WINDOW>
    auto stft (
        const matrix_exp<EXP>& signal,
        const WINDOW& w,
        std::size_t fftsize,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        return details::stft_impl(signal, w, fftsize, wlen, hoplen, details::fft_func{});
    }

// ----------------------------------------------------------------------------------------

    template <typename T, typename Alloc, typename WINDOW>
    auto stft (
        const std::vector<T, Alloc>& signal,
        const WINDOW& w,
        std::size_t fftsize,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        return stft(dlib::mat(signal), w, fftsize, wlen, hoplen);
    }

// ----------------------------------------------------------------------------------------

    template <typename EXP,typename WINDOW>
    auto istft (
        const matrix_exp<EXP>& stft,
        const WINDOW& w,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        using T = typename EXP::type;
        return details::istft_impl<T>(stft, w, wlen, hoplen, details::ifft_func{});
    }

// ----------------------------------------------------------------------------------------

    template <typename EXP, typename WINDOW>
    auto stftr (
        const matrix_exp<EXP>& signal,
        const WINDOW& w,
        std::size_t fftsize,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        return details::stft_impl(signal, w, fftsize, wlen, hoplen, details::fftr_func{});
    }

// ----------------------------------------------------------------------------------------

    template <typename T, typename Alloc, typename WINDOW>
    auto stftr (
        const std::vector<T, Alloc>& signal,
        const WINDOW& w,
        std::size_t fftsize,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        return stftr(dlib::mat(signal), w, fftsize, wlen, hoplen);
    }

// ----------------------------------------------------------------------------------------

    template <typename EXP, typename WINDOW>
    auto istftr (
        const matrix_exp<EXP>& stft,
        const WINDOW& w,
        std::size_t wlen,
        std::size_t hoplen
    )
    {
        using R = remove_complex_t<typename EXP::type>;
        return details::istft_impl<R>(stft, w, wlen, hoplen, details::ifftr_func{});
    }

// ----------------------------------------------------------------------------------------
}

#endif // DLIB_FFt_Hh_

