// Include for Visual studio, 'cos reasons.
#define _USE_MATH_DEFINES
#include <cmath>
#include <numeric>
#include <set>
#include <omp.h>

#include "hoMatrix.h"
#include "hoNDArray_elemwise.h"
#include "hoNDArray_math.h"
#include "hoNDFFT.h"
#include <boost/container/flat_set.hpp>

namespace Gadgetron {

    template <typename T> hoNDFFT<T>* hoNDFFT<T>::instance() {
        if (!instance_)
            instance_ = new hoNDFFT<T>();
        return instance_;
    }

    template <class T> hoNDFFT<T>* hoNDFFT<T>::instance_ = NULL;

    namespace {

        template <class T> struct fftw_types {};

        template <> struct fftw_types<float> {
            using complex                      = fftwf_complex;
            using plan                         = fftwf_plan_s;
            static constexpr auto plan_guru    = fftwf_plan_guru64_dft;
            static constexpr auto plan_dft     = fftwf_plan_dft;
            static constexpr auto execute_dft  = fftwf_execute_dft;
            static constexpr auto destroy_plan = fftwf_destroy_plan;
        };

        template <> struct fftw_types<double> {
            using complex                      = fftw_complex;
            using plan                         = fftw_plan_s;
            static constexpr auto plan_guru    = fftw_plan_guru64_dft;
            static constexpr auto plan_dft     = fftw_plan_dft;
            static constexpr auto execute_dft  = fftw_execute_dft;
            static constexpr auto destroy_plan = fftw_destroy_plan;
        };
        class FFTLock {
        protected:
            static std::mutex lock;
        };
        std::mutex FFTLock::lock;
        template <class T> class SingleFFTPlan : FFTLock {
        public:
            using FFTWComplex = typename fftw_types<T>::complex;

            SingleFFTPlan(int dimension, const hoNDArray<std::complex<T>>& input, hoNDArray<std::complex<T>>& output,
                bool forward) {
                std::lock_guard<std::mutex> guard(lock);

                const auto& dimensions = input.dimensions();
                size_t stride
                    = std::accumulate(dimensions.begin(), dimensions.begin() + dimension, 1, std::multiplies<>());

                auto fftw_dimensions = fftw_iodim64{ static_cast<ptrdiff_t>(dimensions[dimension]), static_cast<ptrdiff_t>(stride), static_cast<ptrdiff_t>(stride) };

                plan = fftw_types<T>::plan_guru(1, &fftw_dimensions, 0, nullptr, (FFTWComplex*)input.data(),
                    (FFTWComplex*)output.data(), forward ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE);
            }

            ~SingleFFTPlan() {
                std::lock_guard<std::mutex> guard(lock);
                fftw_types<T>::destroy_plan(plan);
            }

            void execute(const std::complex<T>* input, std::complex<T>* output) {
                fftw_types<T>::execute_dft(plan, (FFTWComplex*)input, (FFTWComplex*)output);
            }

        private:
            typename fftw_types<T>::plan* plan;
        };

        template <class T> class ContigousFFTPlan : FFTLock {
        public:
            using FFTWComplex = typename fftw_types<T>::complex;
            ContigousFFTPlan(
                int rank, const hoNDArray<std::complex<T>>& input, hoNDArray<std::complex<T>>& output, bool forward) {
                std::lock_guard<std::mutex> guard(lock);

                const auto& dimensions = input.dimensions();

                auto strides = std::vector<size_t>(rank + 1, 1);
                std::partial_sum(
                    dimensions.begin(), dimensions.begin() + rank, strides.begin() + 1, std::multiplies<>());

                auto fftw_dimensions = std::vector<fftw_iodim64>(rank);

                for (int i = 0; i < rank; i++) {
                    fftw_dimensions[i] = { (int64_t)dimensions[i], (int64_t)strides[i], (int64_t)strides[i] };
                }
                std::reverse(fftw_dimensions.begin(),fftw_dimensions.end());
                plan = fftw_types<T>::plan_guru(rank, fftw_dimensions.data(), 0, nullptr, (FFTWComplex*)input.data(),
                    (FFTWComplex*)output.data(), forward ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE);

                if (plan == nullptr) throw std::runtime_error("Illegal FFT plan created");
            }
            ~ContigousFFTPlan() {
                std::lock_guard<std::mutex> guard(lock);
                fftw_types<T>::destroy_plan(plan);
            }

            void execute(const std::complex<T>* input, std::complex<T>* output) {
                fftw_types<T>::execute_dft(plan, (FFTWComplex*)input, (FFTWComplex*)output);
            }

        private:
            typename fftw_types<T>::plan* plan;
        };

        int contigous_rank(const boost::container::flat_set<int>& dimensions) {
            if (!dimensions.count(0))
                return 0;

            int rank = std::distance(std::adjacent_find(dimensions.begin(), dimensions.end(),
                                         [](auto val1, auto val2) { return val1 != val2 - 1; }),
                dimensions.end());
            return rank;
        }

        template <typename T>
        static void contigous_fftn(const hoNDArray<std::complex<T>>& input, hoNDArray<std::complex<T>>& output, int rank,
            bool forward, bool normalize) {

            auto plan = ContigousFFTPlan<T>(rank, input, output, forward);
            size_t batch_size
                = std::accumulate(input.dimensions().begin(), input.dimensions().begin() + rank, 1, std::multiplies<>());
            size_t batches = input.size() / batch_size;

#pragma omp parallel for default(none) shared(plan,  input, output, batches, batch_size)
            for (long long i = 0; i < batches; i++) {

                plan.execute(input.data() + i * batch_size, output.data() + i * batch_size);
            }

            if (normalize)
                output *= T(1) / std::sqrt<T>(batch_size);
        }

        template <typename T>
        static void single_fft(int dimension, const hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r,
            bool forward, bool normalize) {
            assert(dimension >= 0);
            auto plan              = SingleFFTPlan<T>(dimension, a, r, forward);
            const auto& dimensions = a.dimensions();
            size_t inner_batches
                = std::accumulate(dimensions.begin(), dimensions.begin() + dimension, 1, std::multiplies<>());
            size_t outer_batches
                = std::accumulate(dimensions.begin() + dimension + 1, dimensions.end(), 1, std::multiplies<>());
            size_t outer_batchsize = inner_batches * dimensions[dimension];

#pragma omp parallel for default(none) shared(plan, a, r , outer_batches, inner_batches, outer_batchsize ) collapse(2)
            for (long long outer = 0; outer < outer_batches; outer++) {
                for (long long inner = 0; inner < inner_batches; inner++) {
                    plan.execute(
                        a.data() + inner + outer * outer_batchsize, r.data() + inner + outer * outer_batchsize);
                }
            }

            if (normalize)
                r *= T(1) / std::sqrt<T>(dimensions[dimension]);
        }
    }

    static inline size_t fftshiftPivot(size_t x) {
        return (x + 1) / 2;
    }

    static inline size_t ifftshiftPivot(size_t x) {
        return x - (x + 1) / 2;
    }

    namespace {
        template <typename T>
        inline void fftshift(const std::complex<T>* a, std::complex<T>* r, size_t stride, size_t x, size_t pivot) {
            std::rotate_copy(a, a + pivot, a + x, r);
        }

        template <typename T, typename... INDICES>
        inline void fftshift(const std::complex<T>* a, std::complex<T>* r, size_t stride, size_t n, size_t pivot,
            size_t n2, INDICES... indices) {
            for (size_t i = 0; i < n; i++) {
                auto line_begin  = a + i * stride;
                size_t new_y     = i < pivot ? i + pivot : i - pivot;
                auto output_line = r + new_y * stride;
                fftshift(line_begin, output_line, stride / n2, n2, indices...);
            }
        }

        template <typename T>
        inline void fftshift(std::complex<T>* a, std::complex<T>* a2, std::vector<std::complex<T>>& buffer,
            size_t stride, size_t ny, size_t pivoty, size_t nx, size_t pivotx) {
            for (size_t iy = 0; iy < (ny + 1) / 2; iy++) {
                auto line_begin = a + iy * stride;
                auto line_pivot = line_begin + pivotx;
                auto line_end   = line_begin + stride;

                auto line_begin2 = a2 + ((iy + pivoty) % ny) * stride;
                auto line_pivot2 = line_begin2 + pivotx;
                auto line_end2   = line_begin2 + stride;
                std::rotate_copy(line_begin2, line_pivot2, line_end2, buffer.begin());
                std::rotate_copy(line_begin, line_pivot, line_end, line_begin2);
                std::copy(buffer.begin(), buffer.end(), line_begin);
            }
        }

        template <typename T, typename... INDICES>
        inline void fftshift(std::complex<T>* a, std::complex<T>* a2, std::vector<std::complex<T>>& buffer,
            size_t stride, size_t n, size_t pivot, size_t n2, INDICES... indices) {

            for (size_t i = 0; i < n; i++) {
                auto line_begin  = a + i * stride;
                size_t new_y     = (i + pivot) % n;
                auto line_begin2 = a2 + new_y * stride;
                fftshift(line_begin, line_begin2, buffer, stride / n2, n2, indices...);
            }
        }

    }

    template <typename T> static void fftshiftPivot1D(std::complex<T>* a, size_t x, size_t n, size_t pivot) {

#pragma omp parallel for shared(n, x, pivot, a) if (n > 256) default(none)
        for (long long counter = 0; counter < (long long)n; counter++) {
            std::rotate(a + counter * x, a + counter * x + pivot, a + x + counter * x);
        }
    }

    template <typename T>
    static void fftshiftPivot1D(const std::complex<T>* a, std::complex<T>* r, size_t x, size_t n, size_t pivot) {

#pragma omp parallel for shared(n, x, pivot, a, r) if (n > 256) default(none)
        for (long long counter = 0; counter < (long long)n; counter++) {
            std::rotate_copy(a + counter * x, a + counter * x + pivot, a + x + counter * x, r + counter * x);
        }
    }

    template <typename T> void hoNDFFT<T>::fftshift1D(hoNDArray<std::complex<T>>& a) {
        size_t x           = a.get_size(0);
        size_t pivot       = fftshiftPivot(x);
        size_t numOfShifts = a.get_number_of_elements() / x;
        fftshiftPivot1D(a.begin(), x, numOfShifts, pivot);
    }

    template <typename T>
    void hoNDFFT<T>::fftshift1D(const hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }

        size_t x           = a.get_size(0);
        size_t pivot       = fftshiftPivot(x);
        size_t numOfShifts = a.get_number_of_elements() / x;
        fftshiftPivot1D(a.begin(), r.begin(), x, numOfShifts, pivot);
    }

    template <typename T> void hoNDFFT<T>::ifftshift1D(hoNDArray<std::complex<T>>& a) {
        size_t x           = a.get_size(0);
        size_t pivot       = ifftshiftPivot(x);
        size_t numOfShifts = a.get_number_of_elements() / x;

        fftshiftPivot1D(a.begin(), x, numOfShifts, pivot);
    }

    template <typename T>
    void hoNDFFT<T>::ifftshift1D(const hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }

        size_t x           = a.get_size(0);
        size_t pivot       = ifftshiftPivot(x);
        size_t numOfShifts = a.get_number_of_elements() / x;

        fftshiftPivot1D(a.begin(), r.begin(), x, numOfShifts, pivot);
    }

    template <typename T>
    static void fftshiftPivot2D(
        const std::complex<T>* a, std::complex<T>* r, size_t x, size_t y, size_t n, size_t pivotx, size_t pivoty) {
        if (a == NULL || r == NULL)
            throw std::runtime_error("hoNDFFT::fftshiftPivot2D: void ptr provided");
        assert(a != r);

#pragma omp parallel for shared(a, r, x, y, n, pivotx, pivoty) if (n > 16) default(none)
        for (long long tt = 0; tt < (long long)n; tt++) {
            const std::complex<T>* ac = a + tt * x * y;
            std::complex<T>* rc       = r + tt * x * y;
            fftshift(ac, rc, x, y, pivoty, x, pivotx);
        }
    }

    template <typename T>
    static void fftshiftPivot2D(std::complex<T>* a, size_t x, size_t y, size_t n, size_t pivotx, size_t pivoty) {

        if (a == NULL)
            throw std::runtime_error("hoNDFFT::fftshiftPivot2D: void ptr provided");

#pragma omp parallel if (n > 16) default(shared)
        {
            std::vector<std::complex<T>> buffer(x);

#pragma omp for
            for (long long tt = 0; tt < (long long)n; tt++) {
                std::complex<T>* ac = a + tt * x * y;
                fftshift(ac, ac, buffer, x, y, pivoty, x, pivotx);
            }
        }
    }

    template <typename T>
    static inline void fftshift2D(const std::complex<T>* a, std::complex<T>* r, size_t x, size_t y, size_t n) {

        if (a == NULL || r == NULL)
            throw std::runtime_error("hoNDFFT::fftshift2D: void ptr provided");

        size_t pivotx = fftshiftPivot(x);
        size_t pivoty = fftshiftPivot(y);

        fftshiftPivot2D(a, r, x, y, n, pivotx, pivoty);
    }

    template <typename T> inline void hoNDFFT<T>::fftshift2D(hoNDArray<ComplexType>& a) {
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1));
        fftshiftPivot2D(
            a.data(), a.get_size(0), a.get_size(1), n, fftshiftPivot(a.get_size(0)), fftshiftPivot(a.get_size(1)));
    }

    template <typename T>
    inline void hoNDFFT<T>::fftshift2D(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }

        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1));
        fftshiftPivot2D(a.begin(), r.begin(), a.get_size(0), a.get_size(1), n, fftshiftPivot(a.get_size(0)),
            fftshiftPivot(a.get_size(1)));
    }

    template <typename T> inline void hoNDFFT<T>::ifftshift2D(hoNDArray<ComplexType>& a) {
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1));
        fftshiftPivot2D(
            a.begin(), a.get_size(0), a.get_size(1), n, ifftshiftPivot(a.get_size(0)), ifftshiftPivot(a.get_size(1)));
    }

    template <typename T>
    inline void hoNDFFT<T>::ifftshift2D(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1));
        fftshiftPivot2D(a.begin(), r.begin(), a.get_size(0), a.get_size(1), n, ifftshiftPivot(a.get_size(0)),
            ifftshiftPivot(a.get_size(1)));
    }

    template <typename T>
    void fftshiftPivot3D(const std::complex<T>* a, std::complex<T>* r, size_t x, size_t y, size_t z, size_t n,
        size_t pivotx, size_t pivoty, size_t pivotz) {

        if (a == NULL || r == NULL)
            throw std::runtime_error("hoNDFFT::fftshift2D: void ptr provided");

        long long tt;

#pragma omp parallel for private(tt) shared(a, r, x, y, z, n, pivotx, pivoty, pivotz) if (n > 16) default(none)
        for (tt = 0; tt < (long long)n; tt++) {
            const std::complex<T>* ac = a + tt * x * y * z;
            std::complex<T>* rc       = r + tt * x * y * z;
            fftshift(ac, rc, x * y, z, pivotz, y, pivoty, x, pivotx);
        }
    }

    template <typename T>
    void fftshiftPivot3D(
        std::complex<T>* a, size_t x, size_t y, size_t z, size_t n, size_t pivotx, size_t pivoty, size_t pivotz) {

        if (a == NULL)
            throw std::runtime_error("hoNDFFT::fftshiftPivot3D: void ptr provided");

        long long tt;

#pragma omp parallel private(tt)  if (n > 16) default(shared)
        {
            std::vector<std::complex<T>> buffer(x);

#pragma omp for
            for (tt = 0; tt < (long long)n; tt++) {
                std::complex<T>* ac = a + tt * x * y * z;
                fftshift(ac, ac, buffer, x * y, z, pivotz, y, pivoty, x, pivotx);
            }
        }
    }

    template <typename T>
    inline void fftshift3D(const std::complex<T>* a, std::complex<T>* r, size_t x, size_t y, size_t z, size_t n) {

        if (a == NULL || r == NULL)
            throw std::runtime_error("hoNDFFT::fftshift3D: void ptr provided");

        size_t pivotx = fftshiftPivot(x);
        size_t pivoty = fftshiftPivot(y);
        size_t pivotz = fftshiftPivot(z);

        fftshiftPivot3D(a, r, x, y, z, n, pivotx, pivoty, pivotz);
    }

    template <typename T>
    inline void ifftshift3D(const std::complex<T>* a, std::complex<T>* r, size_t x, size_t y, size_t z, size_t n) {

        if (a == NULL || r == NULL)
            throw std::runtime_error("hoNDFFT::ifftshift3D: void ptr provided");

        size_t pivotx = ifftshiftPivot(x);
        size_t pivoty = ifftshiftPivot(y);
        size_t pivotz = ifftshiftPivot(z);

        fftshiftPivot3D(a, r, x, y, z, n, pivotx, pivoty, pivotz);
    }

    template <typename T> inline void fftshift3D(std::complex<T>* a, size_t x, size_t y, size_t z, size_t n) {
        if (a == NULL)
            throw std::runtime_error("hoNDFFT::fftshift3D: void ptr provided");
        size_t pivotx = fftshiftPivot(x);
        size_t pivoty = fftshiftPivot(y);
        size_t pivotz = fftshiftPivot(z);
        fftshiftPivot3D(a, x, y, z, n, pivotx, pivoty, pivotz);
    }

    template <typename T> inline void ifftshift3D(std::complex<T>* a, size_t x, size_t y, size_t z, size_t n) {
        if (a == NULL)
            throw std::runtime_error("hoNDFFT::ifftshift3D: void ptr provided");

        size_t pivotx = ifftshiftPivot(x);
        size_t pivoty = ifftshiftPivot(y);
        size_t pivotz = ifftshiftPivot(z);

        fftshiftPivot3D(a, x, y, z, n, pivotx, pivoty, pivotz);
    }

    template <typename T> inline void hoNDFFT<T>::fftshift3D(hoNDArray<ComplexType>& a) {
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1) * a.get_size(2));

        fftshiftPivot3D(a.begin(), a.get_size(0), a.get_size(1), a.get_size(2), n,
                        fftshiftPivot(a.get_size(0)),
                        fftshiftPivot(a.get_size(1)),
                        fftshiftPivot(a.get_size(2)));
    }

    template <typename T>
    inline void hoNDFFT<T>::fftshift3D(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }

        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1) * a.get_size(2));
        fftshiftPivot3D(a.begin(), r.begin(), a.get_size(0), a.get_size(1), a.get_size(2), n,
            fftshiftPivot(a.get_size(0)), fftshiftPivot(a.get_size(1)), fftshiftPivot(a.get_size(2)));
    }

    template <typename T> inline void hoNDFFT<T>::ifftshift3D(hoNDArray<ComplexType>& a) {
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1) * a.get_size(2));
        fftshiftPivot3D(a.begin(), a.get_size(0), a.get_size(1), a.get_size(2), n, ifftshiftPivot(a.get_size(0)),
            ifftshiftPivot(a.get_size(1)), ifftshiftPivot(a.get_size(2)));
    }

    template <typename T>
    inline void hoNDFFT<T>::ifftshift3D(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r = a;
        }
        size_t n = a.get_number_of_elements() / (a.get_size(0) * a.get_size(1) * a.get_size(2));
        fftshiftPivot3D(a.begin(), r.begin(), a.get_size(0), a.get_size(1), a.get_size(2), n,
            ifftshiftPivot(a.get_size(0)), ifftshiftPivot(a.get_size(1)), ifftshiftPivot(a.get_size(2)));
    }

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

    template <typename T> inline void hoNDFFT<T>::fft1(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 1, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft1(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 1, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft1(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }

        contigous_fftn(a, r, 1, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft1(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }
        contigous_fftn(a, r, 1, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft1c(hoNDArray<ComplexType>& a) {
        ifftshift1D(a);
        fft1(a);
        fftshift1D(a);
    }

    template <typename T> inline void hoNDFFT<T>::ifft1c(hoNDArray<ComplexType>& a) {
        ifftshift1D(a);
        ifft1(a);
        fftshift1D(a);
    }

    template <typename T> inline void hoNDFFT<T>::fft1c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift1D(a, r);
        fft1(r);
        fftshift1D(r);
    }

    template <typename T> inline void hoNDFFT<T>::ifft1c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift1D(a, r);
        ifft1(r);
        fftshift1D(r);
    }

    template <typename T>
    inline void hoNDFFT<T>::fft1c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift1D(a, r);
        fft1(r, buf);
        fftshift1D(buf, r);
    }

    template <typename T>
    inline void hoNDFFT<T>::ifft1c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift1D(a, r);
        ifft1(r, buf);
        fftshift1D(buf, r);
    }

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

    template <typename T> inline void hoNDFFT<T>::fft2(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 2, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft2(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 2, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft2(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }

        contigous_fftn(a, r, 2, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft2(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }

        contigous_fftn(a, r, 2, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft2c(hoNDArray<ComplexType>& a) {
        ifftshift2D(a);
        fft2(a);
        fftshift2D(a);
    }

    template <typename T> inline void hoNDFFT<T>::ifft2c(hoNDArray<ComplexType>& a) {
        ifftshift2D(a);
        ifft2(a);
        fftshift2D(a);
    }

    template <typename T> inline void hoNDFFT<T>::fft2c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift2D(a, r);
        fft2(r);
        fftshift2D(r);
    }

    template <typename T> inline void hoNDFFT<T>::ifft2c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift2D(a, r);
        ifft2(r);
        fftshift2D(r);
    }

    template <typename T>
    inline void hoNDFFT<T>::fft2c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift2D(a, r);
        fft2(r, buf);
        fftshift2D(buf, r);
    }

    template <typename T>
    inline void hoNDFFT<T>::ifft2c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift2D(a, r);
        ifft2(r, buf);
        fftshift2D(buf, r);
    }

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

    template <typename T> inline void hoNDFFT<T>::fft3(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 3, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft3(hoNDArray<ComplexType>& a) {
        contigous_fftn(a, a, 3, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft3(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }

        contigous_fftn(a, r, 3, true, true);
    }

    template <typename T> inline void hoNDFFT<T>::ifft3(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        if (!r.dimensions_equal(&a)) {
            r.create(a.dimensions());
        }

        contigous_fftn(a, r, 3, false, true);
    }

    template <typename T> inline void hoNDFFT<T>::fft3c(hoNDArray<ComplexType>& a) {
        ifftshift3D(a);
        fft3(a);
        fftshift3D(a);
    }

    template <typename T> inline void hoNDFFT<T>::ifft3c(hoNDArray<ComplexType>& a) {
        ifftshift3D(a);
        ifft3(a);
        fftshift3D(a);
    }

    template <typename T> inline void hoNDFFT<T>::fft3c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift3D(a, r);
        fft3(r);
        fftshift3D(r);
    }

    template <typename T> inline void hoNDFFT<T>::ifft3c(const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r) {
        ifftshift3D(a, r);
        ifft3(r);
        fftshift3D(r);
    }

    template <typename T>
    inline void hoNDFFT<T>::fft3c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift3D(a, r);
        fft3(r, buf);
        fftshift3D(buf, r);
    }

    template <typename T>
    inline void hoNDFFT<T>::ifft3c(
        const hoNDArray<ComplexType>& a, hoNDArray<ComplexType>& r, hoNDArray<ComplexType>& buf) {
        ifftshift3D(a, r);
        ifft3(r, buf);
        fftshift3D(buf, r);
    }

    template <typename T> void fft1(hoNDArray<std::complex<T>>& a, bool forward) {
        hoNDArray<std::complex<T>> res(a);
        fft1(res, a, forward);
    }

    template <typename T> void fft2(hoNDArray<std::complex<T>>& a, bool forward) {
        hoNDArray<std::complex<T>> res(a);
        fft2(res, a, forward);
    }

    template <typename T> void fft3(hoNDArray<std::complex<T>>& a, bool forward) {
        hoNDArray<std::complex<T>> res(a);
        fft3(res, a, forward);
    }

    template <typename T> void fft1(hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r, bool forward) {
        contigous_fftn(a, r, 1, forward, true);
    }

    template <typename T> void fft2(hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r, bool forward) {
        contigous_fftn(a, r, 2, forward, true);
    }

    template <typename T> void fft3(hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r, bool forward) {
        contigous_fftn(a, r, 3, forward, true);
    }

    template <typename T>
    void fftn(const hoNDArray<std::complex<T>>& a, hoNDArray<std::complex<T>>& r, const std::vector<size_t>& dimensions,
        bool forward) {

        auto dimensions_set = boost::container::flat_set<int>(dimensions.begin(), dimensions.end());
        int c_rank          = contigous_rank(dimensions_set);
        if (c_rank)
            contigous_fftn(a, r, c_rank, forward, false);

        for (auto it = dimensions_set.lower_bound(c_rank); it != dimensions_set.end(); ++it) {
            single_fft(*it,a, r, forward, false);
        }

        size_t batch_size = 1;
        for (auto d : dimensions)
            batch_size *= a.get_size(d);

        T fftRatio = T(1.0 / std::sqrt(T(batch_size)));
        r *= fftRatio;
    }

    // TODO: implement more optimized threading strategy
    inline int get_num_threads_fft1(size_t n0, size_t num) {
        if (omp_get_max_threads() == 1)
            return 1;

        if (n0 * num > 1024 * 128) {
            return omp_get_max_threads();
        } else if (n0 * num > 512 * 128) {
            return ((omp_get_max_threads() > 8) ? 8 : omp_get_max_threads());
        } else if (n0 * num > 256 * 128) {
            return ((omp_get_max_threads() > 4) ? 4 : omp_get_max_threads());
        } else if (n0 * num > 128 * 128) {
            return 2;
        }

        return 1;
    }

    inline int get_num_threads_fft2(size_t n0, size_t n1, size_t num) {
        if (omp_get_max_threads() == 1)
            return 1;

        if (n0 * n1 * num > 128 * 128 * 64) {
            return omp_get_max_threads();
        } else if (n0 * n1 * num > 128 * 128 * 32) {
            return ((omp_get_max_threads() > 8) ? 8 : omp_get_max_threads());
        } else if (n0 * n1 * num > 128 * 128 * 16) {
            return ((omp_get_max_threads() > 4) ? 4 : omp_get_max_threads());
        } else if (n0 * n1 * num > 128 * 128 * 8) {
            return 2;
        }

        return 1;
    }

    inline int get_num_threads_fft3(size_t n0, size_t n1, size_t n2, size_t num) {
        if (omp_get_max_threads() == 1)
            return 1;

        if (num >= omp_get_max_threads()) {
            return omp_get_max_threads();
        }

        return 1;
    }

    int get_num_threads_fftn(const std::vector<int>& dimensions, size_t num) {

        switch (dimensions.size()) {
        case 1:
            return get_num_threads_fft1(dimensions[0], num);
        case 2:
            return get_num_threads_fft2(dimensions[0], dimensions[1], num);
        case 3:
            return get_num_threads_fft3(dimensions[0], dimensions[1], dimensions[2], num);
        default:
            return std::min<long long>(omp_get_max_threads(), num);
        }
    }

    template <typename T> void hoNDFFT<T>::fft(hoNDArray<ComplexType>* input, unsigned int dim_to_transform) {
        single_fft(dim_to_transform, *input, *input, true, true);
    }
    template <typename T> void hoNDFFT<T>::ifft(hoNDArray<ComplexType>* input, unsigned int dim_to_transform) {
        single_fft(dim_to_transform, *input, *input, false, true);
    }
    template <typename T> void hoNDFFT<T>::fft(hoNDArray<ComplexType>* input) {
        contigous_fftn(*input, *input, input->get_number_of_dimensions(), true, true);
    }
    template <typename T> void hoNDFFT<T>::ifft(hoNDArray<ComplexType>* input) {
        contigous_fftn(*input, *input, input->get_number_of_dimensions(), true, true);
    }
    template <typename T> void hoNDFFT<T>::fft(hoNDArray<complext<T>>* input, unsigned int dim_to_transform) {
        fft(reinterpret_cast<hoNDArray<ComplexType>*>(input), dim_to_transform);
    }
    template <typename T> void hoNDFFT<T>::ifft(hoNDArray<complext<T>>* input, unsigned int dim_to_transform) {
        ifft(reinterpret_cast<hoNDArray<ComplexType>*>(input), dim_to_transform);
    }
    template <typename T> void hoNDFFT<T>::fft(hoNDArray<complext<T>>* input) {
        fft(reinterpret_cast<hoNDArray<ComplexType>*>(input));
    }
    template <typename T> void hoNDFFT<T>::ifft(hoNDArray<complext<T>>* input) {
        ifft(reinterpret_cast<hoNDArray<ComplexType>*>(input));
    }

    template <class ComplexType, class ENABLER>
    void FFT::fft(hoNDArray<ComplexType>& data, std::vector<size_t> dimensions) {
        std::sort(dimensions.begin(), dimensions.end());
        if (std::adjacent_find(dimensions.begin(), dimensions.end()) != dimensions.end()) {
            throw std::runtime_error("Duplicate dimensions in list to be transformed");
        }
        fftn(data,data,dimensions,true);
    }

    template <class ComplexType, class ENABLER> void FFT::fft(hoNDArray<ComplexType>& data, size_t dimension) {
        single_fft(dimension,data,data,true,true);
    }
    template <class ComplexType, class ENABLER>
    void FFT::ifft(hoNDArray<ComplexType>& data, std::vector<size_t> dimensions) {
        std::sort(dimensions.begin(), dimensions.end());
        if (std::adjacent_find(dimensions.begin(), dimensions.end()) != dimensions.end()) {
            throw std::runtime_error("Duplicate dimensions in list to be transformed");
        }
        fftn(data,data,dimensions,false);
    }

    template <class ComplexType, class ENABLER> void FFT::ifft(hoNDArray<ComplexType>& data, size_t dimension) {
        single_fft(dimension,data,data,false,true);
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::fft1c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->fft1c(data,output);
      return output;
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::fft2c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->fft2c(data,output);
      return output;
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::fft3c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->fft3c(data,output);
      return output;
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::ifft1c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->ifft1c(data, output);
      return output;
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::ifft2c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->ifft2c(data, output);
      return output;
    }
    template <class ComplexType, class ENABLER>
    hoNDArray<ComplexType> FFT::ifft3c(const hoNDArray<ComplexType> &data) {
      hoNDArray<ComplexType> output(data.dimensions());
      hoNDFFT<realType_t<ComplexType>>::instance()->ifft3c(data, output);
      return output;
    }

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

    //
    // Instantiation
    //

    template class EXPORTCPUFFT hoNDFFT<float>;
    template class EXPORTCPUFFT hoNDFFT<double>;

    template void FFT::fft<std::complex<float>>(hoNDArray<std::complex<float>>& data, std::vector<size_t> dimensions);
    template void FFT::fft<std::complex<double>>(hoNDArray<std::complex<double>>& data, std::vector<size_t> dimensions);
    template void FFT::fft<std::complex<float>>(hoNDArray<std::complex<float>>& data, size_t dimensions);
    template void FFT::fft<std::complex<double>>(hoNDArray<std::complex<double>>& data, size_t dimensions);

    template void FFT::ifft<std::complex<float>>(hoNDArray<std::complex<float>>& data, std::vector<size_t> dimensions);
    template void FFT::ifft<std::complex<double>>(hoNDArray<std::complex<double>>& data, std::vector<size_t> dimensions);
    template void FFT::ifft<std::complex<float>>(hoNDArray<std::complex<float>>& data, size_t dimensions);
    template void FFT::ifft<std::complex<double>>(hoNDArray<std::complex<double>>& data, size_t dimensions);

    template hoNDArray<std::complex<float>> FFT::fft1c(const hoNDArray<std::complex<float>> &data);
    template hoNDArray<std::complex<float>> FFT::fft2c(const hoNDArray<std::complex<float>> &data);
    template hoNDArray<std::complex<float>> FFT::fft3c(const hoNDArray<std::complex<float>> &data);
    template hoNDArray<std::complex<float>> FFT::ifft1c(const hoNDArray<std::complex<float>> &data);
    template hoNDArray<std::complex<float>> FFT::ifft2c(const hoNDArray<std::complex<float>> &data);
    template hoNDArray<std::complex<float>> FFT::ifft3c(const hoNDArray<std::complex<float>> &data);

    template hoNDArray<std::complex<double>> FFT::fft1c(const hoNDArray<std::complex<double>> &data);
    template hoNDArray<std::complex<double>> FFT::fft2c(const hoNDArray<std::complex<double>> &data);
    template hoNDArray<std::complex<double>> FFT::fft3c(const hoNDArray<std::complex<double>> &data);
    template hoNDArray<std::complex<double>> FFT::ifft1c(const hoNDArray<std::complex<double>> &data);
    template hoNDArray<std::complex<double>> FFT::ifft2c(const hoNDArray<std::complex<double>> &data);
    template hoNDArray<std::complex<double>> FFT::ifft3c(const hoNDArray<std::complex<double>> &data);

}


#pragma once

#include <Eigen/Core>
#include <vector>
#include "PointNormal.hh"
#include "../adjacent_pairs.hh"

namespace kt84 {

template <int N>
struct PolylineT {
    typedef Eigen::Matrix<double, N, 1> Point;
    typedef typename std::vector<Point>::iterator               iterator;
    typedef typename std::vector<Point>::reverse_iterator       reverse_iterator;
    typedef typename std::vector<Point>::const_iterator         const_iterator;
    typedef typename std::vector<Point>::const_reverse_iterator const_reverse_iterator;

    std::vector<Point> points;
    bool is_loop;

    PolylineT()
        : is_loop()
    {}
    PolylineT(const std::vector<Point>& points_, bool is_loop_)
        : points (points_ )
        , is_loop(is_loop_)
    {}
    PolylineT(int size_, bool is_loop_)
        : points (size_   )
        , is_loop(is_loop_)
    {}

    int size() const { return static_cast<int>(points.size()); }
    bool empty() const { return points.empty(); }
    void resize(int size_)                     { points.resize(size_); }
    void resize(int size_, const Point& point) { points.resize(size_, point); }
    void clear() { points.clear(); is_loop = false; }
    void push_back(const Point& point) { points.push_back(point); }
    void pop_back ()                   { points.pop_back (); }
    iterator erase(iterator where_) { return points.erase(where_); }
    iterator erase(iterator first, iterator last) { return points.erase(first, last); }
    iterator insert(iterator where_, const Point& val) { return points.insert(where_, val); }
    void     insert(iterator where_, int count, const Point& val) { points.insert(where_, count, val); }
    template <class InputIterator>
    void     insert(iterator where_, InputIterator first, InputIterator last) { points.insert(where_, first, last); }
    iterator                begin()       { return points. begin(); }
    iterator                end  ()       { return points. end  (); }
    reverse_iterator       rbegin()       { return points.rbegin(); }
    reverse_iterator       rend  ()       { return points.rend  (); }
    const_iterator          begin() const { return points. begin(); }
    const_iterator          end  () const { return points. end  (); }
    const_reverse_iterator rbegin() const { return points.rbegin(); }
    const_reverse_iterator rend  () const { return points.rend  (); }
          Point& front()       { return points.front(); }
    const Point& front() const { return points.front(); }
          Point& back ()       { return points.back (); }
    const Point& back () const { return points.back (); }
          Point& operator[](int index)       { return points[index]; }
    const Point& operator[](int index) const { return points[index]; }

    void reverse() {
        if (is_loop) {
            points.push_back(points.front());
            std::reverse(points.begin(), points.end());
            points.pop_back();
        } else {
            std::reverse(points.begin(), points.end());
        }
    }

    void offset_front(const Point& target) {
        Point offset = target - front();
        front() = target;
        for (int i = 1; i < size() - 1; ++i) {
            points[i] += offset * (size() - 1 - i) / (size() - 1.0);
        }
    }

    void offset_back (const Point& target) {
        Point offset = target - back();
        back() = target;
        for (int i = 1; i < size() - 1; ++i) {
            points[i] += offset * i / (size() - 1.0);
        }
    }

    double length() const {
        double result = 0;
        for (auto p : adjacent_pairs(*this, is_loop))
            result += (p.first - p.second).norm();

        return result;
    }

    void resample() { resample(size()); }
    void resample_by_length(double segment_length) { resample(std::max<int>(static_cast<int>(length() / segment_length), 3)); }
    void resample(int target_num_points) {
        int NN = size();

        if (NN < 2 || target_num_points < 2)
            return;

        if (target_num_points == 2) {
            points[1] = back();
            resize(2);
            return;
        }

        PolylineT result(target_num_points, is_loop);

        if (is_loop) {
            push_back(front());
            ++NN;
        }

        std::vector<double> src_len_acc(NN, 0);
        for (int i = 1; i < NN; ++i)
            src_len_acc[i] = src_len_acc[i - 1] + (points[i] - points[i - 1]).norm();

        double tgt_len_segment = src_len_acc.back() / (target_num_points - 1.);
        int src_i = 0;
        double tgt_len_acc = 0;
        int index = 0;
        result[index++] = front();
        while (true) {
            while (tgt_len_acc + tgt_len_segment <= src_len_acc[src_i]) {
                tgt_len_acc += tgt_len_segment;
                double w1 = (tgt_len_acc - src_len_acc[src_i - 1]) / (src_len_acc[src_i] - src_len_acc[src_i - 1]);
                double w0 = 1 - w1;
                Point& p0 = points[src_i - 1];
                Point& p1 = points[src_i];
                Point p = p0 * w0 + p1 * w1;
                result[index++] = p;
                if (index == target_num_points - 1) break;
            }

            if (index == target_num_points - 1) break;

            while (src_len_acc[src_i] <= tgt_len_acc + tgt_len_segment) {
                ++src_i;
            }
        }
        result[index++] = back();

        if (is_loop)
            result.pop_back();

        *this = result;
        return;
    }
    void insert_points_per_segment(int num_inserted_points_per_segment = 3) {
        auto points_temp = points;
        auto insert_pos = points_temp.begin();
        for (auto i : adjacent_pairs(*this, is_loop)) {
            auto d = (i.second - i.first) / (num_inserted_points_per_segment + 1.0);
            Point p = i.first + d;
            ++insert_pos;
            for (int j = 0; j < num_inserted_points_per_segment; ++j, ++insert_pos, p += d)
                insert_pos = points_temp.insert(insert_pos, p);
        }
        points = points_temp;
    }
    void smooth(int num_iter = 1, double weight_first_order = 1.0, double weight_second_order = 0.0, double damping = 0.5) {
        for (int k = 0; k < num_iter; ++k) {
            auto points_old = points;
            for (int i = 0; i < size(); ++i) {
                if (!is_loop && i == 0 || i == size() - 1)
                    continue;

                int i_next  = (i + 1) % size();
                int i_prev  = (i + size() - 1) % size();
                int i_next2 = (i + 2) % size();
                int i_prev2 = (i + size() - 2) % size();

                auto p_prev  = points_old[i_prev ];
                auto p_next  = points_old[i_next ];
                auto p_prev2 = points_old[i_prev2];
                auto p_next2 = points_old[i_next2];

                Point p_first_order  = (p_prev + p_next) / 2.0;
                Point p_second_order = (-p_prev2 + 4.0 * p_prev + 4.0 * p_next - p_next2) / 6.0;

                if (!is_loop && i == 1 || i == size() - 2)
                    p_second_order = p_first_order;         // undefined

                Point p_target = (weight_first_order * p_first_order + weight_second_order * p_second_order) / (weight_first_order + weight_second_order);

                points[i] = damping * points_old[i] + (1 - damping) * p_target;
            }
        }
    }
    Point point_at(double arc_length_parameter) const {
        if (arc_length_parameter < 0 || 1 < arc_length_parameter)
            // arc length parameter should be between 0 and 1
            return Point::Zero();

        if (arc_length_parameter == 0) return front();
        if (arc_length_parameter == 1) return is_loop ? front() : back();

        const int NN = size();
        const double target_length = length() * arc_length_parameter;

        double length_acc = 0;
        for (int i = 0; i < NN; ++i) {
            if (i == NN - 1 && !is_loop)
                break;

            auto& p0 = points[i];
            auto& p1 = points[(i + 1) % NN];

            double segment_length = (p1 - p0).norm();
            if (length_acc <= target_length && target_length < length_acc + segment_length) {
                double t = (target_length - length_acc) / segment_length;
                return (1 - t) * p0 + t * p1;
            }

            length_acc += segment_length;
        }

        // something wrong happened
        return Point::Zero();
    }
};

typedef PolylineT<2> Polyline2d;
typedef PolylineT<3> Polyline3d;

struct Polyline_PointNormal {
    typedef std::vector<PointNormal>::iterator               iterator;
    typedef std::vector<PointNormal>::reverse_iterator       reverse_iterator;
    typedef std::vector<PointNormal>::const_iterator         const_iterator;
    typedef std::vector<PointNormal>::const_reverse_iterator const_reverse_iterator;

    std::vector<PointNormal> points;
    bool is_loop;

    Polyline_PointNormal()
        : is_loop()
    {}
    Polyline_PointNormal(const std::vector<PointNormal>& points_, bool is_loop_)
        : points (points_ )
        , is_loop(is_loop_)
    {}
    Polyline_PointNormal(int size_, bool is_loop_)
        : points (size_   )
        , is_loop(is_loop_)
    {}

    int size() const { return static_cast<int>(points.size()); }
    bool empty() const { return points.empty(); }
    void resize(int size_)                     { points.resize(size_); }
    void resize(int size_, const PointNormal& point) { points.resize(size_, point); }
    void clear() { points.clear(); }
    void push_back(const PointNormal& point) { points.push_back(point); }
    void pop_back ()                   { points.pop_back (); }
    // Changed const_iterator to iterator to meet "defect" in standard
    iterator  erase(iterator where_) { return points.erase(where_); }
    iterator  erase(iterator first, iterator last) { return points.erase(first, last); }
    iterator insert(iterator where_, const PointNormal& val) { return points.insert(where_, val); }
    void     insert(iterator where_, int count, const PointNormal& val) { points.insert(where_, count, val); }
    template <class InputIterator>
    void     insert(iterator where_, InputIterator first, InputIterator last) { points.insert(where_, first, last); }
    iterator                begin()       { return points. begin(); }
    iterator                end  ()       { return points. end  (); }
    reverse_iterator       rbegin()       { return points.rbegin(); }
    reverse_iterator       rend  ()       { return points.rend  (); }
    const_iterator          begin() const { return points. begin(); }
    const_iterator          end  () const { return points. end  (); }
    const_reverse_iterator rbegin() const { return points.rbegin(); }
    const_reverse_iterator rend  () const { return points.rend  (); }
          PointNormal& front()       { return points.front(); }
    const PointNormal& front() const { return points.front(); }
          PointNormal& back ()       { return points.back (); }
    const PointNormal& back () const { return points.back (); }
          PointNormal& operator[](int index)       { return points[index]; }
    const PointNormal& operator[](int index) const { return points[index]; }

    void reverse() {
        if (is_loop) {
            std::reverse(points.begin(), points.end());
        } else {
            points.push_back(points.front());
            std::reverse(points.begin(), points.end());
            points.pop_back();
        }
    }

    void offset_front(const PointNormal& target) {
        PointNormal offset = target - front();
        front() = target;
        for (int i = 1; i < size() - 1; ++i) {
            points[i] += offset * (size() - 1 - i) / (size() - 1.0);
            pn_normalize(points[i]);
        }
    }

    void offset_back (const PointNormal& target) {
        PointNormal offset = target - back();
        back() = target;
        for (int i = 1; i < size() - 1; ++i) {
            points[i] += offset * i / (size() - 1.0);
            pn_normalize(points[i]);
        }
    }

    double length() const {
        double result = 0;

        for (auto p : adjacent_pairs(*this, is_loop))
            result += pn_norm(p.first - p.second);

        return result;
    }

    void resample() { resample(size()); }
    void resample_by_length(double segment_length) { resample(std::max<int>(static_cast<int>(length() / segment_length), 3)); }
    void resample(int target_num_points) {
        int N = size();

        if (N < 2 || target_num_points < 2)
            return;

        if (target_num_points == 2) {
            points[1] = back();
            resize(2);
            return;
        }

        Polyline_PointNormal result(target_num_points, is_loop);

        if (is_loop) {
            push_back(front());
            ++N;
        }

        std::vector<double> src_len_acc(N, 0);
        for (int i = 1; i < N; ++i)
            src_len_acc[i] = src_len_acc[i - 1] + pn_norm(points[i] - points[i - 1]);

        double tgt_len_segment = src_len_acc.back() / (target_num_points - 1.);
        int src_i = 0;
        double tgt_len_acc = 0;
        int index = 0;
        result[index++] = front();
        while (true) {
            while (tgt_len_acc + tgt_len_segment <= src_len_acc[src_i]) {
                tgt_len_acc += tgt_len_segment;
                double w1 = (tgt_len_acc - src_len_acc[src_i - 1]) / (src_len_acc[src_i] - src_len_acc[src_i - 1]);
                double w0 = 1 - w1;
                PointNormal& pn0 = points[src_i - 1];
                PointNormal& pn1 = points[src_i];
                PointNormal pn = pn0 * w0 + pn1 * w1;
                pn_normalize(pn);
                result[index++] = pn;
                if (index == target_num_points - 1) break;
            }

            if (index == target_num_points - 1) break;

            while (src_len_acc[src_i] <= tgt_len_acc + tgt_len_segment) {
                ++src_i;
            }
        }
        result[index++] = back();

        if (is_loop)
            result.pop_back();

        *this = result;
        return;
    }
    void insert_points_per_segment(int num_inserted_points_per_segment = 3) {
        auto points_temp = points;
        auto insert_pos = points_temp.begin();
        for (auto i : adjacent_pairs(*this, is_loop)) {
            auto d = (i.second - i.first) / (num_inserted_points_per_segment + 1.0);
            PointNormal pn = i.first + d;
            ++insert_pos;
            for (int j = 0; j < num_inserted_points_per_segment; ++j, ++insert_pos, pn += d)
                insert_pos = points_temp.insert(insert_pos, pn_normalized(pn));
        }
        points = points_temp;
    }
    void smooth(int num_iter = 1, double weight_first_order = 1.0, double weight_second_order = 0.0, double damping = 0.5) {
        for (int k = 0; k < num_iter; ++k) {
            auto points_old = points;
            for (int i = 0; i < size(); ++i) {
                if (!is_loop && i == 0 || i == size() - 1)
                    continue;

                int i_next  = (i + 1) % size();
                int i_prev  = (i + size() - 1) % size();
                int i_next2 = (i + 2) % size();
                int i_prev2 = (i + size() - 2) % size();

                auto p_prev  = points_old[i_prev ];
                auto p_next  = points_old[i_next ];
                auto p_prev2 = points_old[i_prev2];
                auto p_next2 = points_old[i_next2];

                PointNormal p_first_order  = (p_prev + p_next) / 2.0;
                PointNormal p_second_order = (-p_prev2 + 4.0 * p_prev + 4.0 * p_next - p_next2) / 6.0;

                if (!is_loop && i == 1 || i == size() - 2)
                    p_second_order = p_first_order;         // undefined

                PointNormal p_target = (weight_first_order * p_first_order + weight_second_order * p_second_order) / (weight_first_order + weight_second_order);

                points[i] = damping * points_old[i] + (1 - damping) * p_target;
                pn_normalize(points[i]);
            }
        }
    }
    PointNormal point_at(double arc_length_parameter) const {
        if (arc_length_parameter < 0 || 1 < arc_length_parameter)
            // arc length parameter should be between 0 and 1
            return PointNormal::Zero();

        if (arc_length_parameter == 0) return front();
        if (arc_length_parameter == 1) return is_loop ? front() : back();

        const int N = size();
        const double target_length = length() * arc_length_parameter;

        double length_acc = 0;
        for (int i = 0; i < N; ++i) {
            if (i == N - 1 && !is_loop)
                break;

            auto& pn0 = points[i];
            auto& pn1 = points[(i + 1) % N];

            double segment_length = pn_norm(pn1 - pn0);
            if (length_acc <= target_length && target_length < length_acc + segment_length) {
                double t = (target_length - length_acc) / segment_length;
                PointNormal pn = (1 - t) * pn0 + t * pn1;
                pn_normalize(pn);
                return pn;
            }

            length_acc += segment_length;
        }

        // something wrong happened
        return PointNormal::Zero();
    }
};

}


#ifndef KINT_H_
#define KINT_H_

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/sum.hpp>
#include <boost/accumulators/framework/accumulator_base.hpp>
#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/framework/parameters/sample.hpp>
#include <boost/accumulators/framework/parameters/weight.hpp>
#include <boost/accumulators/framework/accumulators/external_accumulator.hpp>
#include <boost/accumulators/statistics/weighted_sum.hpp>
#include <boost/accumulators/statistics/sum_kahan.hpp>

#include <iostream>

/*!	This class provides methods for integration over momenta or energies
 *
 *	Problem:
 *		Parameters are:
 *		- a: lattice spacing
 *		- \f$ D(\eps) \f$: density of states (DOS)
 *		- Momentum, energy dependent function. In most cases this will be the impurity Green's \f$ G(k,i\omega_n) \f$
 *		- Some constant offset. In most cases \f$ i \omega_n + \mu - \Simga(i \omega_n)\f$
 *		- function and limits for the integrand e.g.:
 *		\f[
 *			\int\limits_{-\frac{\pi}{a}}^\frac{\pi}{a} d^D k \frac{1}{i \omega_n - \epsilon_k + mu - \Sigma (i \omega_n)}
 *		\f]
 *		or
 *		\f[
 *			\int\limits_{-\infty}^\infty d^D \epsilon \frac{d(\epsilon)}{i \omega_n - \epsilon_k + mu - \Sigma (i \omega_n)}
 *		\f]
 *
 *	Methods:
 *
 *	- summation (this constructs some general weighted D dimensional sum)
 *		- riemann sum
 *		- gaussian smearing
 *		.
 *
 *	- reformulation to ODE and solution via boost ODE int
 *
 *	- Linear tetrahedral method for 1,2,3 D
 * 		- issues: can break symmetry, gamma point must be included
 *		- use MFEM?
 *		.
 *
 *	- TODO: let user decide wether to use kahan summation or not
 */
namespace utility{
using KahanTag = boost::accumulators::tag::weighted_sum_kahan;
using KahanSumAccT = boost::accumulators::stats<KahanTag>;
template <typename T>
using AccT = boost::accumulators::accumulator_set<T, KahanSumAccT, T >;

namespace detail
{
	/*!	this recursively constructs the depth of the nested for loops
	 *	for the innermost loop the partially specialized struct Internal<D,0> is called
	 */
	template<unsigned int D, unsigned int ND, typename T, typename RetT >
	struct Internal
	{
		static void sumKPoints(RetT (*integrand)(std::array<T, D> x),const std::array<T, D> &min,const std::array<T, D> &incs,\
			 const std::array<unsigned long, D> &N,std::array<T, D> &xVec, AccT<RetT> &acc)
		{
			RetT xi=min[ND];
			for(unsigned int n=0; n<N[ND];n++)
			{
				xVec[ND] = xi;
				Internal<D,ND-1, T, RetT>::sumKPoints(integrand, min, incs, N, xVec, acc);
				acc(boost::accumulators::extract_result<KahanTag>(acc), boost::accumulators::weight = incs[ND]);

				xi+=incs[ND];
			}
		}
	};

	template<unsigned int D, typename T, typename RetT>
	struct Internal<D,0,T,RetT>
	{
		static void sumKPoints(RetT (*integrand)(std::array<T, D> x),const std::array<T, D> &min, \
			const std::array<T, D> &incs, const std::array<unsigned long, D> &N, std::array<T, D> &xVec, AccT<RetT> &acc)
		{
			RetT xi=min[0];
			for(unsigned int n=0; n<N[0];n++)
			{
				xVec[0] = xi;
				acc(integrand(xVec), boost::accumulators::weight = incs[0]);
				xi+=incs[0];
			}
		}
	};
}

//TODO: 2 test types: 3 or 4 standard intrgrals, mean over >1000 VERY large OR mean over >MAXINT number of samples numbers (check for overflow)
//TODO: let user specify weights, TODO: let user choose between kahan and normal sum

template<unsigned int D>
class KInt
{
	public:
		/*!	@brief	instantiates an accumulator for D nested for loops \f$ \sum\limits_{min_1}^{max_1} ... \sum\limits_{min_D}^{max_D} f \f$
		 *			with \f$ f(T x_1, ... ,T x_D) -> T \f$
		 *
		 *	@param	integrand	function pointer over which to sum
		 *	@param	min			array of start values (one element for each dimension)
		 *	@param	max			array of values for the upper limit (one element for each dimension)
		 *	@param	N			array with number of steps (one element for each dimension)
		 *
		 *	@return	accumulated value
		 */
		template<typename T, typename RetT>
		RetT sumKPoints(RetT (*integrand)(std::array<T, D> x), std::array<T, D> &min, \
			const std::array<T, D> &max, const std::array<unsigned long, D> &N) const
		{
			//acc<summand type, method, weight type>
			AccT<RetT> acc;
			std::array<T, D> incs; 						// increments in each dimension
			std::array<T, D> xVec 	= min;
			for(unsigned int d=0;d<D;d++) incs[d] = (max[d]-min[d])/N[d];

			detail::Internal<D,D-1,T,RetT>::sumKPoints(integrand, min, incs, N, xVec, acc);
			return boost::accumulators::extract_result<KahanTag>(acc);
		}

	private:

};

}
#endif


// Copyright (c) 2021 CNES
//
// All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
#pragma once
#include <pybind11/eigen.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

#include <Eigen/Core>
#include <atomic>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/stats.hpp>

#include "pyinterp/detail/math.hpp"
#include "pyinterp/detail/thread.hpp"
#include "pyinterp/eigen.hpp"
#include "pyinterp/grid.hpp"

namespace pyinterp {
namespace detail {

/// Calculate the zonal average in x direction
///
/// @param grid The grid to be processed.
/// @param mask Matrix describing the undefined pixels of the grid providedNaN
/// NumberReplaces all missing (_FillValue) values in a grid with values derived
/// from solving Poisson's equation via relaxation. of threads used for the
/// calculation
///
/// @param grid
template <typename Type>
void set_zonal_average(pybind11::EigenDRef<Matrix<Type>>& grid,
                       Matrix<bool>& mask, const size_t num_threads) {
  // Captures the detected exceptions in the calculation function
  // (only the last exception captured is kept)
  auto except = std::exception_ptr(nullptr);

  detail::dispatch(
      [&](size_t y_start, size_t y_end) {
        try {
          // Calculation of longitude band means.
          for (size_t iy = y_start; iy < y_end; ++iy) {
            auto acc = boost::accumulators::accumulator_set<
                Type,
                boost::accumulators::stats<boost::accumulators::tag::count,
                                           boost::accumulators::tag::mean>>();
            for (int64_t ix = 0; ix < grid.rows(); ++ix) {
              if (!mask(ix, iy)) {
                acc(grid(ix, iy));
              }
            }

            // The masked value is replaced by the average of the longitude band
            // if it is defined; otherwise it is replaced by zero.
            auto first_guess = boost::accumulators::count(acc)
                                   ? boost::accumulators::mean(acc)
                                   : Type(0);
            for (int64_t ix = 0; ix < grid.rows(); ++ix) {
              if (mask(ix, iy)) {
                grid(ix, iy) = first_guess;
              }
            }
          }
        } catch (...) {
          except = std::current_exception();
        }
      },
      grid.cols(), num_threads);

  if (except != nullptr) {
    std::rethrow_exception(except);
  }
}

///  Replaces all undefined values (NaN) in a grid using the Gauss-Seidel
///  method by relaxation.
///
/// @param grid The grid to be processed
/// @param is_circle True if the X axis of the grid defines a circle.
/// @param relaxation Relaxation constant
/// @return maximum residual value
template <typename Type>
auto gauss_seidel(pybind11::EigenDRef<pyinterp::Matrix<Type>>& grid,
                  Matrix<bool>& mask, const bool is_circle,
                  const Type relaxation, const size_t num_threads) -> Type {
  // Maximum residual values for each thread.
  std::vector<Type> max_residuals(num_threads);

  // Shape of the grid
  auto x_size = grid.rows();
  auto y_size = grid.cols();

  // Captures the detected exceptions in the calculation function
  // (only the last exception captured is kept)
  auto except = std::exception_ptr(nullptr);

  // Gets the index of the pixel (ix, iy) in the matrix.
  auto coordinates = [](const int64_t ix, const int64_t iy,
                        const int64_t size) -> int64_t {
    return iy * size + ix;
  };

  // Thread worker responsible for processing several strips along the y-axis
  // of the grid.
  //
  // @param y_start First index y of the band to be processed.
  // @param y_end Last index y, excluded, of the band to be processed.
  // @param max_residual Maximum residual of this strip.
  // @param pipe_out Last index to be processed on in this band.
  // @param pipe_in Last index processed in the previous band.
  auto worker = [&](int64_t y_start, int64_t y_end, Type* max_residual,
                    std::atomic<int64_t>* pipe_out,
                    std::atomic<int64_t>* pipe_in) -> void {
    // Modifies the value of a masked pixel.
    auto cell_fill = [&grid, &relaxation, &max_residual](
                         int64_t ix0, int64_t ix, int64_t ix1, int64_t iy0,
                         int64_t iy, int64_t iy1) {
      auto residual = (Type(0.25) * (grid(ix0, iy) + grid(ix1, iy) +
                                     grid(ix, iy0) + grid(ix, iy1)) -
                       grid(ix, iy)) *
                      relaxation;
      grid(ix, iy) += residual;
      *max_residual = std::max(*max_residual, std::fabs(residual));
    };

    // Initialization of the maximum value of the residuals of the treated
    // strips.
    *max_residual = Type(0);

    try {
      for (auto ix = 0; ix < x_size; ++ix) {
        //
        auto ix0 = ix == 0 ? (is_circle ? x_size - 1 : 1) : ix - 1;
        auto ix1 = ix == x_size - 1 ? (is_circle ? 0 : x_size - 2) : ix + 1;

        // If necessary, check that the last index required for this block has
        // been processed in the previous band.
        if (pipe_in) {
          auto next_coordinates = coordinates(ix, y_start, y_size);
          while (*pipe_in < next_coordinates) {
            std::this_thread::sleep_for(std::chrono::nanoseconds(5));
          }
        }

        for (auto iy = y_start; iy < y_end; ++iy) {
          auto iy0 = iy == 0 ? 1 : iy - 1;
          auto iy1 = iy == y_size - 1 ? y_size - 2 : iy + 1;
          if (mask(ix, iy)) {
            cell_fill(ix0, ix, ix1, iy0, iy, iy1);
          }
        }

        // If necessary, the other thread responsible for processing the next
        // band is notified.
        if (pipe_out) {
          *pipe_out = coordinates(ix, y_end, y_size);
        }
      }
    } catch (...) {
      except = std::current_exception();
    }
  };

  if (num_threads == 1) {
    worker(0, y_size, &max_residuals[0], nullptr, nullptr);
  } else {
    assert(num_threads >= 2);
    std::vector<std::atomic<int64_t>> pipeline(num_threads);
    std::vector<std::thread> threads;

    int64_t start = 0;
    int64_t shift = y_size / num_threads;

    for (auto& item : pipeline) {
      item = std::numeric_limits<int>::min();
    }

    for (size_t index = 0; index < num_threads - 1; ++index) {
      threads.emplace_back(std::thread(
          worker, start, start + shift, &max_residuals[index], &pipeline[index],
          index == 0 ? nullptr : &pipeline[index - 1]));
      start += shift;
    }
    threads.emplace_back(std::thread(worker, start, y_size,
                                     &max_residuals[num_threads - 1], nullptr,
                                     &pipeline[num_threads - 2]));
    for (auto&& item : threads) {
      item.join();
    }
  }
  if (except != nullptr) {
    std::rethrow_exception(except);
  }
  return *std::max_element(max_residuals.begin(), max_residuals.end());
}

}  // namespace detail

namespace fill {

/// Type of first guess grid.
enum FirstGuess {
  kZero,          //!< Use 0.0 as an initial guess
  kZonalAverage,  //!< Use zonal average in x direction
};

/// Replaces all undefined values (NaN) in a grid using the Gauss-Seidel
/// method by relaxation.
///
/// @param grid The grid to be processed
/// @param is_circle True if the X axis of the grid defines a circle.
/// @param max_iterations Maximum number of iterations to be used by relaxation.
/// @param epsilon Tolerance for ending relaxation before the maximum number of
/// iterations limit.
/// @param relaxation Relaxation constant
/// @param num_threads The number of threads to use for the computation. If 0
/// all CPUs are used. If 1 is given, no parallel computing code is used at all,
/// which is useful for debugging.
/// @return A tuple containing the number of iterations performed and the
/// maximum residual value.
template <typename Type>
auto gauss_seidel(pybind11::EigenDRef<Matrix<Type>>& grid,
                  const FirstGuess first_guess, const bool is_circle,
                  const size_t max_iterations, const Type epsilon,
                  const Type relaxation, size_t num_threads)
    -> std::tuple<size_t, Type> {
  /// If the grid doesn't have an undefined value, this routine has nothing more
  /// to do.
  if (!grid.hasNaN()) {
    return std::make_tuple(0, Type(0));
  }

  /// Calculation of the maximum number of threads if the user chooses.
  if (num_threads == 0) {
    num_threads = std::thread::hardware_concurrency();
  }

  /// Calculation of the position of the undefined values on the grid.
  auto mask = Matrix<bool>(grid.array().isNaN());

  /// Calculation of the first guess with the chosen method
  switch (first_guess) {
    case FirstGuess::kZero:
      grid = (mask.array()).select(0, grid);
      break;
    case FirstGuess::kZonalAverage:
      detail::set_zonal_average(grid, mask, num_threads);
      break;
    default:
      throw std::invalid_argument("Invalid guess type: " +
                                  std::to_string(first_guess));
  }

  // Initialization of the function results.
  size_t iteration = 0;
  Type max_residual = 0;

  for (size_t it = 0; it < max_iterations; ++it) {
    ++iteration;
    max_residual = detail::gauss_seidel<Type>(grid, mask, is_circle, relaxation,
                                              num_threads);
    if (max_residual < epsilon) {
      break;
    }
  }
  return std::make_tuple(iteration, max_residual);
}

// Get the indexes that frame a given index.
inline auto frame_index(const int64_t index, const int64_t size,
                        const bool is_angle, std::vector<int64_t>& frame)
    -> void {
  // Index in the center of the window
  auto center = static_cast<int64_t>(frame.size() / 2);

  for (int64_t ix = 0; ix < static_cast<int64_t>(frame.size()); ++ix) {
    auto idx = index - center + ix;

    // Normalizing longitude?
    if (is_angle) {
      idx = detail::math::remainder(idx, size);
    } else {
      // Otherwise, the symmetrical indexes are used if the indexes are outside
      // the domain definition.
      if (idx < 0 || idx >= size) {
        auto where = detail::math::remainder(idx, (size - 1) * 2);
        if (where >= size) {
          idx = size - 2 - detail::math::remainder(where, size);
        } else {
          idx = detail::math::remainder(where, size);
        }
      }
    }
    frame[ix] = idx;
  }
}

/// Checking the size of the filter window.
inline auto check_windows_size(const std::string& name1, const uint32_t size)
    -> void {
  if (size < 1) {
    throw std::invalid_argument(name1 + " must be >= 1");
  }
}

/// Checking the size of the filter window.
template <typename... Args>
inline auto check_windows_size(const std::string& name1, uint32_t size,
                               Args... args) -> void {
  check_windows_size(name1, size);
  check_windows_size(args...);
}

/// Type of values processed by the Loess filter.
enum ValueType {
  kUndefined,  //!< Undefined values (fill undefined values)
  kDefined,    //!< Defined values (smooth values)
  kAll         //!< Smooth and fill values
};

/// Fills undefined values using a locally weighted regression function or
/// LOESS. The weight function used for LOESS is the tri-cube weight
/// function, w(x)=(1-|d|^{3})^{3}
///
/// @param grid Grid Function on a uniform 2-dimensional grid to be filled.
/// @param nx Number of points of the half-window to be taken into account
/// along the longitude axis.
/// @param nx Number of points of the half-window to be taken into account
/// along the latitude axis.
/// @param value_type Type of values processed by the filter
/// @param num_threads The number of threads to use for the computation. If
/// 0 all CPUs are used. If 1 is given, no parallel computing code is used
/// at all, which is useful for debugging.
/// @return The grid will have all the NaN filled with extrapolated values.
template <typename Type>
auto loess(const Grid2D<Type>& grid, const uint32_t nx, const uint32_t ny,
           const ValueType value_type, const size_t num_threads)
    -> pybind11::array_t<Type> {
  check_windows_size("nx", nx, "ny", ny);
  auto result = pybind11::array_t<Type>(
      pybind11::array::ShapeContainer{grid.x()->size(), grid.y()->size()});
  auto _result = result.template mutable_unchecked<2>();

  // Captures the detected exceptions in the calculation function
  // (only the last exception captured is kept)
  auto except = std::exception_ptr(nullptr);

  auto worker = [&](const size_t start, const size_t end) {
    try {
      // Access to the shared pointer outside the loop to avoid data races
      const auto& x_axis = *grid.x();
      const auto& y_axis = *grid.y();
      auto x_frame = std::vector<int64_t>(nx * 2 + 1);
      auto y_frame = std::vector<int64_t>(ny * 2 + 1);

      for (size_t ix = start; ix < end; ++ix) {
        auto x = x_axis(ix);

        // We retrieve the indexes framing the current value.
        frame_index(ix, x_axis.size(), x_axis.is_angle(), x_frame);

        // Read the first value of the calculated window.
        const auto x0 = x_axis(x_frame[0]);

        // The current value is normalized to the first value in the
        // window.
        if (x_axis.is_angle()) {
          x = detail::math::normalize_angle(x, x0, 360.0);
        }

        for (int64_t iy = 0; iy < y_axis.size(); ++iy) {
          auto z = grid.value(ix, iy);

          // If the current value is masked.
          const auto undefined = std::isnan(z);
          if (value_type == kAll || (value_type == kDefined && !undefined) ||
              (value_type == kUndefined && undefined)) {
            auto y = y_axis(iy);

            // We retrieve the indexes framing the current value.
            frame_index(iy, y_axis.size(), false, y_frame);

            // Initialization of values to calculate the extrapolated
            // value.
            auto value = Type(0);
            auto weight = Type(0);

            // For all the coordinates of the frame.
            for (auto wx : x_frame) {
              auto xi = x_axis(wx);

              // We normalize the window's coordinates to its first value.
              if (x_axis.is_angle()) {
                xi = detail::math::normalize_angle(xi, x0, 360.0);
              }

              for (auto wy : y_frame) {
                auto zi = grid.value(wx, wy);

                // If the value is not masked, its weight is calculated from
                // the tri-cube weight function
                if (!std::isnan(zi)) {
                  const auto power = 3.0;
                  auto d =
                      std::sqrt(detail::math::sqr(((xi - x)) / nx) +
                                detail::math::sqr(((y_axis(wy) - y)) / ny));
                  auto wi = d <= 1 ? std::pow((1.0 - std::pow(d, power)), power)
                                   : 0.0;
                  value += static_cast<Type>(wi * zi);
                  weight += static_cast<Type>(wi);
                }
              }
            }

            // Finally, we calculate the extrapolated value if possible,
            // otherwise we will recopy the masked original value.
            if (weight != 0) {
              z = value / weight;
            }
          }
          _result(ix, iy) = z;
        }
      }
    } catch (...) {
      except = std::current_exception();
    }
  };

  {
    pybind11::gil_scoped_release release;
    detail::dispatch(worker, grid.x()->size(), num_threads);
  }
  return result;
}

template <typename Type, typename AxisType>
auto loess(const Grid3D<Type, AxisType>& grid, const uint32_t nx,
           const uint32_t ny, const ValueType value_type,
           const size_t num_threads) -> pybind11::array_t<Type> {
  check_windows_size("nx", nx, "ny", ny);
  auto result = pybind11::array_t<Type>(pybind11::array::ShapeContainer{
      grid.x()->size(), grid.y()->size(), grid.z()->size()});
  auto _result = result.template mutable_unchecked<3>();

  // Captures the detected exceptions in the calculation function
  // (only the last exception captured is kept)
  auto except = std::exception_ptr(nullptr);

  auto worker = [&](const size_t start, const size_t end) {
    try {
      // Access to the shared pointer outside the loop to avoid data races
      const auto& x_axis = *grid.x();
      const auto& y_axis = *grid.y();
      auto x_frame = std::vector<int64_t>(nx * 2 + 1);
      auto y_frame = std::vector<int64_t>(ny * 2 + 1);

      for (size_t iz = start; iz < end; ++iz) {
        for (int64_t ix = 0; ix < x_axis.size(); ++ix) {
          auto x = x_axis(ix);

          // We retrieve the indexes framing the current value.
          frame_index(ix, x_axis.size(), x_axis.is_angle(), x_frame);

          // Read the first value of the calculated window.
          const auto x0 = x_axis(x_frame[0]);

          // The current value is normalized to the first value in the
          // window.
          if (x_axis.is_angle()) {
            x = detail::math::normalize_angle(x, x0, 360.0);
          }

          for (int64_t iy = 0; iy < y_axis.size(); ++iy) {
            auto z = grid.value(ix, iy, iz);

            // If the current value is masked.
            const auto undefined = std::isnan(z);
            if (value_type == kAll || (value_type == kDefined && !undefined) ||
                (value_type == kUndefined && undefined)) {
              auto y = y_axis(iy);

              // We retrieve the indexes framing the current value.
              frame_index(iy, y_axis.size(), false, y_frame);

              // Initialization of values to calculate the extrapolated
              // value.
              auto value = Type(0);
              auto weight = Type(0);

              // For all the coordinates of the frame.
              for (auto wx : x_frame) {
                auto xi = x_axis(wx);

                // We normalize the window's coordinates to its first value.
                if (x_axis.is_angle()) {
                  xi = detail::math::normalize_angle(xi, x0, 360.0);
                }

                for (auto wy : y_frame) {
                  auto zi = grid.value(wx, wy, iz);

                  // If the value is not masked, its weight is calculated
                  // from the tri-cube weight function
                  if (!std::isnan(zi)) {
                    const auto power = 3.0;
                    auto d =
                        std::sqrt(detail::math::sqr(((xi - x)) / nx) +
                                  detail::math::sqr(((y_axis(wy) - y)) / ny));
                    auto wi = d <= 1
                                  ? std::pow((1.0 - std::pow(d, power)), power)
                                  : 0.0;
                    value += static_cast<Type>(wi * zi);
                    weight += static_cast<Type>(wi);
                  }
                }
              }

              // Finally, we calculate the extrapolated value if possible,
              // otherwise we will recopy the masked original value.
              if (weight != 0) {
                z = value / weight;
              }
            }
            _result(ix, iy, iz) = z;
          }
        }
      }
    } catch (...) {
      except = std::current_exception();
    }
  };

  {
    pybind11::gil_scoped_release release;
    detail::dispatch(worker, grid.z()->size(), num_threads);
  }
  return result;
}

}  // namespace fill
}  // namespace pyinterp


//
// Created by daeyun on 5/10/17.
//

#include "meshdist.h"

#include <Eigen/Dense>
#include <type_traits>
#include <chrono>

#include "lib/common.h"
#include "lib/random_utils.h"
#include "lib/benchmark.h"

namespace scene3d {
namespace meshdist_cgal {

typedef CGAL::Simple_cartesian<float> K;
typedef std::list<K::Triangle_3>::iterator Iterator;
typedef CGAL::AABB_triangle_primitive<K, Iterator> Primitive;
typedef CGAL::AABB_traits<K, Primitive> AABB_triangle_traits;
typedef CGAL::AABB_tree<AABB_triangle_traits> Tree;

void Triangle::ApplyRt(const Mat34 &M) {
  auto R = M.leftCols<3>();
  auto t = M.col(3);
  a = (R * a + t).eval();
  b = (R * b + t).eval();
  c = (R * c + t).eval();
  OnTransformation();
}

void Triangle::ApplyR(const Mat33 &M) {
  a = (M * a).eval();
  b = (M * b).eval();
  c = (M * c).eval();
  OnTransformation();
}

void Triangle::Translate(const Vec3 &dxyz) {
  a += dxyz;
  b += dxyz;
  c += dxyz;
  OnTransformation();
}

void Triangle::Print() const {
  std::cout << a.transpose() << ",      "
            << b.transpose() << ",      "
            << c.transpose() << std::endl;
}

Vec3 Triangle::SamplePoint() const {
  Vec2 v12{scene3d::Random::Rand(), scene3d::Random::Rand()};
  if (v12.sum() > 1) {
    v12 = 1 - v12.array();
  }
  return a + (v12(0) * ab_) + (v12(1) * ac_);
}

double Triangle::Area() const {
  return ab_.cross(ac_).norm() * 0.5f;
}

void Triangle::OnTransformation() {
  assert(a.allFinite());
  assert(b.allFinite());
  assert(c.allFinite());

  ab_ = b - a;
  ac_ = c - a;
}

void SamplePointsOnTriangles(const std::vector<Triangle> &triangles, float density, Points3d *points, std::vector<size_t> *triangle_indices) {
  std::vector<double> areas;
  areas.reserve(triangles.size());
  for (auto &&triangle : triangles) {
    areas.push_back(triangle.Area());
  }

  double surface_area = std::accumulate(areas.begin(), areas.end(), static_cast<double>(0));
  Expects(surface_area > 0.0);

  int num_samples = static_cast<int>(surface_area * density);
  if (num_samples <= 0) {
    // Surface area is non-zero
    // the density parameter is not big enough. This should be a rare case. For now, print warning and just make sure to sample at least one point.
    num_samples = 1;
    LOGGER->warn("density is too small. Surface area is {}", surface_area);
  }
  Expects(num_samples > 0);

  std::discrete_distribution<> distribution(std::begin(areas), std::end(areas));

  Eigen::Matrix<int, Dynamic, 1> counts;

  Eigen::Matrix<int, Dynamic, Dynamic> local_counts;

#pragma omp parallel if (USE_OMP && num_samples > 1e6)
  {
    const int num_threads = omp_get_num_threads();
    const int i_thread = omp_get_thread_num();

#pragma omp single
    {
      local_counts.resize(triangles.size(), num_threads);
      local_counts.fill(0);
    }

#pragma omp for
    for (int i = 0; i < num_samples; ++i) {
      ++local_counts(distribution(scene3d::Random::Engine()), i_thread);
    }
  }

  counts = local_counts.rowwise().sum();

  points->resize(3, num_samples);
  if (triangle_indices) {
    triangle_indices->resize(static_cast<size_t>(num_samples));
  }
  int k = 0;
  for (int i = 0; i < triangles.size(); ++i) {
    for (int j = 0; j < counts(i); ++j) {
      points->col(k) = triangles[i].SamplePoint();
      if (triangle_indices) {
        triangle_indices->push_back((size_t) i);
      }
      ++k;
    }
  }
  Expects(k == num_samples);

  Eigen::PermutationMatrix<Dynamic, Dynamic> perm(points->cols());
  perm.setIdentity();
  std::shuffle(perm.indices().data(), perm.indices().data() + perm.indices().size(), scene3d::Random::Engine());
  *points *= perm; // permute columns
}

// TODO(daeyun): divide by mean distance

float MeshToMeshDistanceOneDirection(const std::vector<Triangle> &from,
                                     const std::vector<Triangle> &to,
                                     float sampling_density,
                                     std::vector<float> *distances) {

  Points3d points;

  SamplePointsOnTriangles(from, sampling_density, &points);

  std::list<K::Triangle_3> triangle_list;
  for (const auto &triangle : to) {
    triangle_list.emplace_back(K::Point_3{triangle.a[0], triangle.a[1], triangle.a[2]},
                               K::Point_3{triangle.b[0], triangle.b[1], triangle.b[2]},
                               K::Point_3{triangle.c[0], triangle.c[1], triangle.c[2]});
  }
  std::vector<K::Point_3> point_list;
  for (int i = 0; i < points.cols(); ++i) {
    point_list.emplace_back(points(0, i), points(1, i), points(2, i));
  }

  int num_triangles = static_cast<int>(to.size());
  int num_points = static_cast<int>(points.cols());

  LOGGER->debug("Computing minimum distances from {} points to {} triangles.", num_points, num_triangles);

  auto start = scene3d::TimeSinceEpoch<std::milli>();
  Tree tree(triangle_list.begin(), triangle_list.end());
  tree.build();
  tree.accelerate_distance_queries();
  LOGGER->debug("Time elapsed for building tree (CGAL): {}", scene3d::TimeSinceEpoch<std::milli>() - start);

  float distance_sum = 0;

#pragma omp parallel for if(USE_OMP) reduction(+:distance_sum) schedule(static)
  for (int i = 0; i < point_list.size(); ++i) {
    float dist = tree.squared_distance(point_list[i]);
    distance_sum += dist;

// TODO: needs refactoring
    if (distances) {
#pragma omp critical
      distances->push_back(dist);
    }
  }

  LOGGER->debug("distance: {}", distance_sum);
  float rms = static_cast<float>(std::sqrt(distance_sum / static_cast<double>(point_list.size())));
  LOGGER->debug("RMS: {}", rms);
  auto elapsed = scene3d::TimeSinceEpoch<std::milli>() - start;
  LOGGER->debug("Time elapsed (CGAL): {} ms", elapsed);

  return rms;
}

float PointsToMeshDistanceOneDirection(const std::vector<std::array<float, 3>> &from,
                                       const std::vector<Triangle> &to) {
  std::list<K::Triangle_3> triangle_list;
  for (const auto &triangle : to) {
    triangle_list.emplace_back(K::Point_3{triangle.a[0], triangle.a[1], triangle.a[2]},
                               K::Point_3{triangle.b[0], triangle.b[1], triangle.b[2]},
                               K::Point_3{triangle.c[0], triangle.c[1], triangle.c[2]});
  }
  std::vector<K::Point_3> point_list;
  for (int i = 0; i < from.size(); ++i) {
    point_list.emplace_back(from[i][0], from[i][1], from[i][2]);
  }

  int num_triangles = static_cast<int>(to.size());
  int num_points = static_cast<int>(from.size());

  LOGGER->debug("Computing minimum distances from {} points to {} triangles.", num_points, num_triangles);

  auto start = scene3d::TimeSinceEpoch<std::milli>();
  Tree tree(triangle_list.begin(), triangle_list.end());
  tree.build();
  tree.accelerate_distance_queries();
  LOGGER->debug("Time elapsed for building tree (CGAL): {}", scene3d::TimeSinceEpoch<std::milli>() - start);

  float distance_sum = 0;

#pragma omp parallel for if(USE_OMP) reduction(+:distance_sum) schedule(static)
  for (int i = 0; i < point_list.size(); ++i) {
    float dist = tree.squared_distance(point_list[i]);
    distance_sum += dist;
  }

  LOGGER->debug("distance: {}", distance_sum);
  float rms = static_cast<float>(std::sqrt(distance_sum / static_cast<double>(point_list.size())));

  LOGGER->debug("RMS: {}", rms);
  auto elapsed = scene3d::TimeSinceEpoch<std::milli>() - start;
  LOGGER->debug("Time elapsed (CGAL): {} ms", elapsed);

  return rms;
}

float MeshToPointsDistanceOneDirection(const std::vector<Triangle> &from,
                                       const std::vector<std::array<float, 3>> &target_points,
                                       float sampling_density) {
  std::list<K::Triangle_3> triangle_list;
  for (const auto &triangle : from) {
    triangle_list.emplace_back(K::Point_3{triangle.a[0], triangle.a[1], triangle.a[2]},
                               K::Point_3{triangle.b[0], triangle.b[1], triangle.b[2]},
                               K::Point_3{triangle.c[0], triangle.c[1], triangle.c[2]});
  }

  Points3d points_on_mesh;
  SamplePointsOnTriangles(from, sampling_density, &points_on_mesh);

  int num_source_points = static_cast<int>(points_on_mesh.cols());
  int num_target_points = static_cast<int>(target_points.size());

  LOGGER->debug("Computing minimum distances from {} points on triangles to {} points.", num_source_points, num_target_points);

  auto start = scene3d::TimeSinceEpoch<std::milli>();

  float distance_sum = 0;

#pragma omp parallel for if(USE_OMP) reduction(+:distance_sum) schedule(static)
  for (int i = 0; i < num_source_points; ++i) {
    const auto source_point = points_on_mesh.col(i);
    double min_dist = kInfinity;
    // TODO(daeyun): this is brute force
    for (int j = 0; j < num_target_points; ++j) {
      double dx = (source_point(0) - target_points[j][0]);
      double dy = (source_point(1) - target_points[j][1]);
      double dz = (source_point(2) - target_points[j][2]);
      double dist = std::sqrt(dx * dx + dy * dy + dz * dz);
      if (dist < min_dist) {
        min_dist = dist;
      }
    }
    distance_sum += min_dist;
  }

  LOGGER->debug("distance: {}", distance_sum);
  float rms = static_cast<float>(std::sqrt(distance_sum / static_cast<double>(num_source_points)));

  LOGGER->debug("RMS: {}", rms);
  auto elapsed = scene3d::TimeSinceEpoch<std::milli>() - start;
  LOGGER->debug("Time elapsed (CGAL): {} ms", elapsed);

  return rms;
}

float MeshToMeshDistance(const std::vector<Triangle> &a, const std::vector<Triangle> &b) {
  auto start = scene3d::TimeSinceEpoch<std::milli>();

  constexpr int kSamplingDensity = 300;

  float d1 = meshdist_cgal::MeshToMeshDistanceOneDirection(a, b, kSamplingDensity);
  float d2 = meshdist_cgal::MeshToMeshDistanceOneDirection(b, a, kSamplingDensity);

  auto elapsed = scene3d::TimeSinceEpoch<std::milli>() - start;
  LOGGER->debug("Time elapsed (MeshToMeshDistance): {} ms", elapsed);
  LOGGER->debug("{}, {}", d1, d2);
  return static_cast<float>((d1 + d2) * 0.5);
}

void TrianglesFromTriMesh(const TriMesh &mesh, std::vector<Triangle> *out) {
  for (const auto &face: mesh.faces) {
    const auto &v1 = mesh.vertices[face[0]];
    const auto &v2 = mesh.vertices[face[1]];
    const auto &v3 = mesh.vertices[face[2]];
    out->push_back(Triangle(Vec3{v1[0], v1[1], v1[2]}, Vec3{v2[0], v2[1], v2[2]}, Vec3{v3[0], v3[1], v3[2]}));
  }
}

}
}