#ifndef VIENNACL_LINALG_QR_METHOD_HPP_
#define VIENNACL_LINALG_QR_METHOD_HPP_

/* =========================================================================
   Copyright (c) 2010-2012, Institute for Microelectronics,
                            Institute for Analysis and Scientific Computing,
                            TU Wien.
   Portions of this software are copyright by UChicago Argonne, LLC.

                            -----------------
                  ViennaCL - The Vienna Computing Library
                            -----------------

   Project Head:    Karl Rupp                   rupp@iue.tuwien.ac.at

   (A list of authors and contributors can be found in the PDF manual)

   License:         MIT (X11), see file LICENSE in the base directory
============================================================================= */

#include "viennacl/vector.hpp"
#include "viennacl/matrix.hpp"

#include <examples/benchmarks/benchmark-utils.hpp>

#include "viennacl/linalg/qr-method-common.hpp"

#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>

/** @file viennacl/linalg/qr-method.hpp
    @brief Implementation of the QR method for eigenvalue computations. Experimental.
*/

namespace viennacl
{
  namespace linalg
  {
    namespace detail
    {
        template<typename MatrixType, typename VectorType>
        void givens_next(MatrixType& matrix,
                        VectorType& tmp1,
                        VectorType& tmp2,
                        int l,
                        int m
                      )
        {
          typedef typename MatrixType::value_type                                   ScalarType;
          typedef typename viennacl::result_of::cpu_value_type<ScalarType>::type    CPU_ScalarType;

          viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<CPU_ScalarType, 1>::program_name(), SVD_GIVENS_NEXT_KERNEL);

          kernel.global_work_size(0, viennacl::tools::roundUpToNextMultiple<cl_uint>(viennacl::traits::size1(matrix), 256));
          kernel.local_work_size(0, 256);

          viennacl::ocl::enqueue(kernel(
                                        matrix,
                                        tmp1,
                                        tmp2,
                                        static_cast<cl_uint>(matrix.size1()),
                                        static_cast<cl_uint>(matrix.internal_size1()),
                                        static_cast<cl_uint>(l),
                                        static_cast<cl_uint>(m - 1)
                                ));
        }

        // Symmetric tridiagonal QL algorithm.
        // This is derived from the Algol procedures tql2, by Bowdler, Martin, Reinsch, and Wilkinson,
        // Handbook for Auto. Comp., Vol.ii-Linear Algebra, and the corresponding Fortran subroutine in EISPACK.
        template <typename SCALARTYPE, unsigned int ALIGNMENT>
        void tql2(viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT> & Q,
                  boost::numeric::ublas::vector<SCALARTYPE> & d,
                  boost::numeric::ublas::vector<SCALARTYPE> & e)
        {
            int n = Q.size1();

            boost::numeric::ublas::vector<SCALARTYPE> cs(n), ss(n);
            viennacl::vector<SCALARTYPE> tmp1(n), tmp2(n);

            for (int i = 1; i < n; i++)
                e(i - 1) = e(i);

            e(n - 1) = 0;

            SCALARTYPE f = 0;
            SCALARTYPE tst1 = 0;
            SCALARTYPE eps = 2 * EPS;

            for (int l = 0; l < n; l++)
            {
                // Find small subdiagonal element.
                tst1 = std::max<SCALARTYPE>(tst1, std::fabs(d(l)) + std::fabs(e(l)));
                int m = l;
                while (m < n)
                {
                    if (std::fabs(e(m)) <= eps * tst1)
                        break;
                    m++;
                }

                // If m == l, d(l) is an eigenvalue, otherwise, iterate.
                if (m > l)
                {
                    int iter = 0;
                    do
                    {
                        iter = iter + 1;  // (Could check iteration count here.)

                        // Compute implicit shift
                        SCALARTYPE g = d(l);
                        SCALARTYPE p = (d(l + 1) - g) / (2 * e(l));
                        SCALARTYPE r = pythag<SCALARTYPE>(p, 1);
                        if (p < 0)
                        {
                            r = -r;
                        }

                        d(l) = e(l) / (p + r);
                        d(l + 1) = e(l) * (p + r);
                        SCALARTYPE dl1 = d(l + 1);
                        SCALARTYPE h = g - d(l);
                        for (int i = l + 2; i < n; i++)
                        {
                            d(i) -= h;
                        }

                        f = f + h;

                        // Implicit QL transformation.
                        p = d(m);
                        SCALARTYPE c = 1;
                        SCALARTYPE c2 = c;
                        SCALARTYPE c3 = c;
                        SCALARTYPE el1 = e(l + 1);
                        SCALARTYPE s = 0;
                        SCALARTYPE s2 = 0;
                        for (int i = m - 1; i >= l; i--)
                        {
                            c3 = c2;
                            c2 = c;
                            s2 = s;
                            g = c * e(i);
                            h = c * p;
                            r = pythag(p, e(i));
                            e(i + 1) = s * r;
                            s = e(i) / r;
                            c = p / r;
                            p = c * d(i) - s * g;
                            d(i + 1) = h + s * (c * g + s * d(i));

                            cs[i] = c;
                            ss[i] = s;
                        }

                        p = -s * s2 * c3 * el1 * e(l) / dl1;
                        e(l) = s * p;
                        d(l) = c * p;

                        {
                            viennacl::copy(cs, tmp1);
                            viennacl::copy(ss, tmp2);

                            givens_next(Q, tmp1, tmp2, l, m);
                        }

                        // Check for convergence.
                    }
                    while (std::fabs(e(l)) > eps * tst1);
                }
                d(l) = d(l) + f;
                e(l) = 0;
            }
        }

        template <typename SCALARTYPE, typename MatrixT>
        void final_iter_update_gpu(MatrixT& A,
                                int n,
                                int last_n,
                                SCALARTYPE q,
                                SCALARTYPE p
                                )
        {
            viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<SCALARTYPE, 1>::program_name(), SVD_FINAL_ITER_UPDATE_KERNEL);

            viennacl::ocl::enqueue(kernel(
                                          A,
                                          static_cast<cl_uint>(A.internal_size1()),
                                          static_cast<cl_uint>(n),
                                          static_cast<cl_uint>(last_n),
                                          q,
                                          p
                                  ));
        }

        template <typename SCALARTYPE, typename MatrixT>
        void update_float_QR_column_gpu(MatrixT& A,
                                const std::vector<SCALARTYPE>& buf,
                                viennacl::vector<SCALARTYPE>& buf_vcl,
                                int m,
                                int n,
                                int last_n,
                                bool //is_triangular
                                )
        {
            viennacl::fast_copy(buf, buf_vcl);

            viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<SCALARTYPE, 1>::program_name(), SVD_UPDATE_QR_COLUMN_KERNEL);

            viennacl::ocl::enqueue(kernel(
                                          A,
                                          static_cast<cl_uint>(A.internal_size1()),
                                          buf_vcl,
                                          static_cast<cl_uint>(m),
                                          static_cast<cl_uint>(n),
                                          static_cast<cl_uint>(last_n)
                                  ));
        }

        template <typename SCALARTYPE, typename MatrixT>
        void final_iter_update(MatrixT& A,
                                int n,
                                int last_n,
                                SCALARTYPE q,
                                SCALARTYPE p
                                )
        {
            for (int i = 0; i < last_n; i++)
            {
                SCALARTYPE v_in = A(i, n);
                SCALARTYPE z = A(i, n - 1);
                A(i, n - 1) = q * z + p * v_in;
                A(i, n) = q * v_in - p * z;
            }
        }

        template <typename SCALARTYPE, typename MatrixT>
        void update_float_QR_column(MatrixT& A,
                                const std::vector<SCALARTYPE>& buf,
                                int m,
                                int n,
                                int last_i,
                                bool is_triangular
                                )
        {
            for (int i = 0; i < last_i; i++)
            {
                int start_k = is_triangular?std::max(i + 1, m):m;

                SCALARTYPE* a_row = A.row(i);

                SCALARTYPE a_ik   = a_row[start_k];
                SCALARTYPE a_ik_1 = 0;
                SCALARTYPE a_ik_2 = 0;

                if(start_k < n)
                    a_ik_1 = a_row[start_k + 1];

                for(int k = start_k; k < n; k++)
                {
                    bool notlast = (k != n - 1);

                    SCALARTYPE p = buf[5 * k] * a_ik + buf[5 * k + 1] * a_ik_1;

                    if (notlast)
                    {
                        a_ik_2 = a_row[k + 2];
                        p = p + buf[5 * k + 2] * a_ik_2;
                        a_ik_2 = a_ik_2 - p * buf[5 * k + 4];
                    }

                    a_row[k] = a_ik - p;
                    a_ik_1 = a_ik_1 - p * buf[5 * k + 3];

                    a_ik = a_ik_1;
                    a_ik_1 = a_ik_2;
                }

                if(start_k < n)
                    a_row[n] = a_ik;
            }
        }

        template <typename SCALARTYPE>
        class FastMatrix
        {
        public:
            FastMatrix()
            {
                size_ = 0;
            }

            FastMatrix(std::size_t sz)
            {
                size_ = sz;
                data.resize(sz * sz);
            }

            SCALARTYPE& operator()(int i, int j)
            {
                return data[i * size_ + j];
            }

            SCALARTYPE* row(int i)
            {
                return &data[i * size_];
            }

            SCALARTYPE* begin()
            {
                return &data[0];
            }

            SCALARTYPE* end()
            {
                return &data[0] + data.size();
            }

            std::vector<SCALARTYPE> data;
        private:
            std::size_t size_;
        };

        // Nonsymmetric reduction from Hessenberg to real Schur form.
        // This is derived from the Algol procedure hqr2, by Martin and Wilkinson, Handbook for Auto. Comp.,
        // Vol.ii-Linear Algebra, and the corresponding  Fortran subroutine in EISPACK.
        template <typename SCALARTYPE, unsigned int ALIGNMENT>
        void hqr2(viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT>& vcl_H,
                    viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT>& V,
                    boost::numeric::ublas::vector<SCALARTYPE>& d,
                    boost::numeric::ublas::vector<SCALARTYPE>& e)
        {
            transpose(V);

            int nn = vcl_H.size1();

            FastMatrix<SCALARTYPE> H(nn);//, V(nn);

            std::vector<SCALARTYPE> buf(5 * nn);
            viennacl::vector<SCALARTYPE> buf_vcl(5 * nn);

            viennacl::fast_copy(vcl_H, H.begin());

            int n = nn - 1;

            SCALARTYPE eps = 2 * EPS;
            SCALARTYPE exshift = 0;
            SCALARTYPE p = 0;
            SCALARTYPE q = 0;
            SCALARTYPE r = 0;
            SCALARTYPE s = 0;
            SCALARTYPE z = 0;
            SCALARTYPE t;
            SCALARTYPE w;
            SCALARTYPE x;
            SCALARTYPE y;

            SCALARTYPE out1, out2;

            // compute matrix norm
            SCALARTYPE norm = 0;
            for (int i = 0; i < nn; i++)
            {
                for (int j = std::max(i - 1, 0); j < nn; j++)
                    norm = norm + std::fabs(H(i, j));
            }

            // Outer loop over eigenvalue index
            int iter = 0;
            while (n >= 0)
            {
                // Look for single small sub-diagonal element
                int l = n;
                while (l > 0)
                {
                    s = std::fabs(H(l - 1, l - 1)) + std::fabs(H(l, l));
                    if (s == 0) s = norm;
                    if (std::fabs(H(l, l - 1)) < eps * s)
                        break;

                    l--;
                }

                // Check for convergence
                if (l == n)
                {
                    // One root found
                    H(n, n) = H(n, n) + exshift;
                    d(n) = H(n, n);
                    e(n) = 0;
                    n--;
                    iter = 0;
                }
                else if (l == n - 1)
                {
                    // Two roots found
                    w = H(n, n - 1) * H(n - 1, n);
                    p = (H(n - 1, n - 1) - H(n, n)) / 2;
                    q = p * p + w;
                    z = static_cast<SCALARTYPE>(std::sqrt(std::fabs(q)));
                    H(n, n) = H(n, n) + exshift;
                    H(n - 1, n - 1) = H(n - 1, n - 1) + exshift;
                    x = H(n, n);

                    if (q >= 0)
                    {
                        // Real pair
                        z = (p >= 0) ? (p + z) : (p - z);
                        d(n - 1) = x + z;
                        d(n) = d(n - 1);
                        if (z != 0)
                            d(n) = x - w / z;
                        e(n - 1) = 0;
                        e(n) = 0;
                        x = H(n, n - 1);
                        s = std::fabs(x) + std::fabs(z);
                        p = x / s;
                        q = z / s;
                        r = static_cast<SCALARTYPE>(std::sqrt(p * p + q * q));
                        p = p / r;
                        q = q / r;

                        // Row modification
                        for (int j = n - 1; j < nn; j++)
                        {
                            SCALARTYPE h_nj = H(n, j);
                            z = H(n - 1, j);
                            H(n - 1, j) = q * z + p * h_nj;
                            H(n, j) = q * h_nj - p * z;
                        }

                        final_iter_update(H, n, n + 1, q, p);
                        final_iter_update_gpu(V, n, nn, q, p);
                    }
                    else
                    {
                        // Complex pair
                        d(n - 1) = x + p;
                        d(n) = x + p;
                        e(n - 1) = z;
                        e(n) = -z;
                    }

                    n = n - 2;
                    iter = 0;
                }
                else
                {
                    // No convergence yet

                    // Form shift
                    x = H(n, n);
                    y = 0;
                    w = 0;
                    if (l < n)
                    {
                        y = H(n - 1, n - 1);
                        w = H(n, n - 1) * H(n - 1, n);
                    }

                    // Wilkinson's original ad hoc shift
                    if (iter == 10)
                    {
                        exshift += x;
                        for (int i = 0; i <= n; i++)
                            H(i, i) -= x;

                        s = std::fabs(H(n, n - 1)) + std::fabs(H(n - 1, n - 2));
                        x = y = 0.75 * s;
                        w = (-0.4375) * s * s;
                    }

                    // MATLAB's new ad hoc shift
                    if (iter == 30)
                    {
                        s = (y - x) / 2;
                        s = s * s + w;
                        if (s > 0)
                        {
                            s = static_cast<SCALARTYPE>(std::sqrt(s));
                            if (y < x) s = -s;
                            s = x - w / ((y - x) / 2 + s);
                            for (int i = 0; i <= n; i++)
                                H(i, i) -= s;
                            exshift += s;
                            x = y = w = (SCALARTYPE)0.964;
                        }
                    }

                    iter = iter + 1;

                    // Look for two consecutive small sub-diagonal elements
                    int m = n - 2;
                    while (m >= l)
                    {
                        SCALARTYPE h_m1_m1 = H(m + 1, m + 1);
                        z = H(m, m);
                        r = x - z;
                        s = y - z;
                        p = (r * s - w) / H(m + 1, m) + H(m, m + 1);
                        q = h_m1_m1 - z - r - s;
                        r = H(m + 2, m + 1);
                        s = std::fabs(p) + std::fabs(q) + std::fabs(r);
                        p = p / s;
                        q = q / s;
                        r = r / s;
                        if (m == l)
                            break;
                        if (std::fabs(H(m, m - 1)) * (std::fabs(q) + std::fabs(r)) < eps * (std::fabs(p) * (std::fabs(H(m - 1, m - 1)) + std::fabs(z) + std::fabs(h_m1_m1))))
                            break;
                        m--;
                    }

                    for (int i = m + 2; i <= n; i++)
                    {
                        H(i, i - 2) = 0;
                        if (i > m + 2)
                            H(i, i - 3) = 0;
                    }

                    // float QR step involving rows l:n and columns m:n
                    for (int k = m; k < n; k++)
                    {
                        bool notlast = (k != n - 1);
                        if (k != m)
                        {
                            p = H(k, k - 1);
                            q = H(k + 1, k - 1);
                            r = (notlast ? H(k + 2, k - 1) : 0);
                            x = std::fabs(p) + std::fabs(q) + std::fabs(r);
                            if (x != 0)
                            {
                                p = p / x;
                                q = q / x;
                                r = r / x;
                            }
                        }

                        if (x == 0) break;

                        s = static_cast<SCALARTYPE>(std::sqrt(p * p + q * q + r * r));
                        if (p < 0) s = -s;

                        if (s != 0)
                        {
                            if (k != m)
                                H(k, k - 1) = -s * x;
                            else
                                if (l != m)
                                    H(k, k - 1) = -H(k, k - 1);

                            p = p + s;
                            y = q / s;
                            z = r / s;
                            x = p / s;
                            q = q / p;
                            r = r / p;

                            buf[5 * k] = x;
                            buf[5 * k + 1] = y;
                            buf[5 * k + 2] = z;
                            buf[5 * k + 3] = q;
                            buf[5 * k + 4] = r;

                            SCALARTYPE* a_row_k = H.row(k);
                            SCALARTYPE* a_row_k_1 = H.row(k + 1);
                            SCALARTYPE* a_row_k_2 = H.row(k + 2);
                            // Row modification
                            for (int j = k; j < nn; j++)
                            {
                                SCALARTYPE h_kj = a_row_k[j];
                                SCALARTYPE h_k1_j = a_row_k_1[j];

                                p = h_kj + q * h_k1_j;
                                if (notlast)
                                {
                                    SCALARTYPE h_k2_j = a_row_k_2[j];
                                    p = p + r * h_k2_j;
                                    a_row_k_2[j] = h_k2_j - p * z;
                                }

                                a_row_k[j] = h_kj - p * x;
                                a_row_k_1[j] = h_k1_j - p * y;
                            }

                            //H(k + 1, nn - 1) = h_kj;

                            // Column modification
                            for (int i = k; i < std::min(nn, k + 4); i++)
                            {
                                p = x * H(i, k) + y * H(i, k + 1);
                                if (notlast)
                                {
                                    p = p + z * H(i, k + 2);
                                    H(i, k + 2) = H(i, k + 2) - p * r;
                                }

                                H(i, k) = H(i, k) - p;
                                H(i, k + 1) = H(i, k + 1) - p * q;
                            }
                        }
                        else
                        {
                            buf[5 * k] = 0;
                            buf[5 * k + 1] = 0;
                            buf[5 * k + 2] = 0;
                            buf[5 * k + 3] = 0;
                            buf[5 * k + 4] = 0;
                        }
                    }

                    // Timer timer;
                    // timer.start();

                    update_float_QR_column(H, buf, m, n, n, true);
                    update_float_QR_column_gpu(V, buf, buf_vcl, m, n, nn, false);

                    // std::cout << timer.get() << "\n";
                }
            }

            // Backsubstitute to find vectors of upper triangular form
            if (norm == 0)
            {
                return;
            }

            for (n = nn - 1; n >= 0; n--)
            {
                p = d(n);
                q = e(n);

                // Real vector
                if (q == 0)
                {
                    int l = n;
                    H(n, n) = 1;
                    for (int i = n - 1; i >= 0; i--)
                    {
                        w = H(i, i) - p;
                        r = 0;
                        for (int j = l; j <= n; j++)
                            r = r + H(i, j) * H(j, n);

                        if (e(i) < 0)
                        {
                            z = w;
                            s = r;
                        }
                        else
                        {
                            l = i;
                            if (e(i) == 0)
                            {
                                H(i, n) = (w != 0) ? (-r / w) : (-r / (eps * norm));
                            }
                            else
                            {
                                // Solve real equations
                                x = H(i, i + 1);
                                y = H(i + 1, i);
                                q = (d(i) - p) * (d(i) - p) + e(i) * e(i);
                                t = (x * s - z * r) / q;
                                H(i, n) = t;
                                H(i + 1, n) = (std::fabs(x) > std::fabs(z)) ? ((-r - w * t) / x) : ((-s - y * t) / z);
                            }

                            // Overflow control
                            t = std::fabs(H(i, n));
                            if ((eps * t) * t > 1)
                                for (int j = i; j <= n; j++)
                                    H(j, n) /= t;
                        }
                    }
                }
                else if (q < 0)
                {
                    // Complex vector
                    int l = n - 1;

                    // Last vector component imaginary so matrix is triangular
                    if (std::fabs(H(n, n - 1)) > std::fabs(H(n - 1, n)))
                    {
                        H(n - 1, n - 1) = q / H(n, n - 1);
                        H(n - 1, n) = -(H(n, n) - p) / H(n, n - 1);
                    }
                    else
                    {
                        cdiv<SCALARTYPE>(0, -H(n - 1, n), H(n - 1, n - 1) - p, q, out1, out2);

                        H(n - 1, n - 1) = out1;
                        H(n - 1, n) = out2;
                    }

                    H(n, n - 1) = 0;
                    H(n, n) = 1;
                    for (int i = n - 2; i >= 0; i--)
                    {
                        SCALARTYPE ra, sa, vr, vi;
                        ra = 0;
                        sa = 0;
                        for (int j = l; j <= n; j++)
                        {
                            SCALARTYPE h_ij = H(i, j);
                            ra = ra + h_ij * H(j, n - 1);
                            sa = sa + h_ij * H(j, n);
                        }

                        w = H(i, i) - p;

                        if (e(i) < 0)
                        {
                            z = w;
                            r = ra;
                            s = sa;
                        }
                        else
                        {
                            l = i;
                            if (e(i) == 0)
                            {
                                cdiv<SCALARTYPE>(-ra, -sa, w, q, out1, out2);
                                H(i, n - 1) = out1;
                                H(i, n) = out2;
                            }
                            else
                            {
                                // Solve complex equations
                                x = H(i, i + 1);
                                y = H(i + 1, i);
                                vr = (d(i) - p) * (d(i) - p) + e(i) * e(i) - q * q;
                                vi = (d(i) - p) * 2 * q;
                                if ( (vr == 0) && (vi == 0) )
                                    vr = eps * norm * (std::fabs(w) + std::fabs(q) + std::fabs(x) + std::fabs(y) + std::fabs(z));

                                cdiv<SCALARTYPE>(x * r - z * ra + q * sa, x * s - z * sa - q * ra, vr, vi, out1, out2);

                                H(i, n - 1) = out1;
                                H(i, n) = out2;

                                if (std::fabs(x) > (std::fabs(z) + std::fabs(q)))
                                {
                                    H(i + 1, n - 1) = (-ra - w * H(i, n - 1) + q * H(i, n)) / x;
                                    H(i + 1, n) = (-sa - w * H(i, n) - q * H(i, n - 1)) / x;
                                }
                                else
                                {
                                    cdiv<SCALARTYPE>(-r - y * H(i, n - 1), -s - y * H(i, n), z, q, out1, out2);

                                    H(i + 1, n - 1) = out1;
                                    H(i + 1, n) = out2;
                                }
                            }

                            // Overflow control
                            t = std::max(std::fabs(H(i, n - 1)), std::fabs(H(i, n)));
                            if ((eps * t) * t > 1)
                            {
                                for (int j = i; j <= n; j++)
                                {
                                    H(j, n - 1) /= t;
                                    H(j, n) /= t;
                                }
                            }
                        }
                    }
                }
            }

            viennacl::fast_copy(H.begin(), H.end(),  vcl_H);
            // viennacl::fast_copy(V.begin(), V.end(),  vcl_V);

            viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT> tmp = V;

            V = viennacl::linalg::prod(trans(tmp), vcl_H);
        }

        template <typename SCALARTYPE, unsigned int ALIGNMENT>
        bool householder_twoside(
                            viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT>& A,
                            viennacl::matrix<SCALARTYPE, row_major, ALIGNMENT>& Q,
                            viennacl::vector<SCALARTYPE, ALIGNMENT>& D,
                            std::size_t start)
        {
            if(start + 2 >= A.size1())
                return false;

            prepare_householder_vector(A, D, A.size1(), start + 1, start, start + 1, true);

            {
                viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<SCALARTYPE, 1>::program_name(), SVD_HOUSEHOLDER_UPDATE_A_LEFT_KERNEL);

                viennacl::ocl::enqueue(kernel(
                                              A,
                                              D,
                                              static_cast<cl_uint>(start + 1),
                                              static_cast<cl_uint>(start),
                                              static_cast<cl_uint>(A.size1()),
                                              static_cast<cl_uint>(A.size2()),
                                              static_cast<cl_uint>(A.internal_size2()),
                                              viennacl::ocl::local_mem(static_cast<cl_uint>(128 * 4))
                                      ));
            }

            {
                viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<SCALARTYPE, 1>::program_name(), SVD_HOUSEHOLDER_UPDATE_A_RIGHT_KERNEL);

                viennacl::ocl::enqueue(kernel(
                                              A,
                                              D,
                                              static_cast<cl_uint>(0),
                                              static_cast<cl_uint>(0),
                                              static_cast<cl_uint>(A.size1()),
                                              static_cast<cl_uint>(A.size2()),
                                              static_cast<cl_uint>(A.internal_size2()),
                                              viennacl::ocl::local_mem(static_cast<cl_uint>(128 * sizeof(SCALARTYPE)))
                                      ));
            }

            {
                viennacl::ocl::kernel& kernel = viennacl::ocl::get_kernel(viennacl::linalg::kernels::svd<SCALARTYPE, 1>::program_name(), SVD_HOUSEHOLDER_UPDATE_QL_KERNEL);

                viennacl::ocl::enqueue(kernel(
                                                Q,
                                                D,
                                                static_cast<cl_uint>(A.size1()),
                                                static_cast<cl_uint>(A.size2()),
                                                static_cast<cl_uint>(Q.internal_size2()),
                                                viennacl::ocl::local_mem(static_cast<cl_uint>(128 * sizeof(SCALARTYPE)))
                                            ));
            }

            return true;
        }

        template <typename SCALARTYPE, typename F, unsigned int ALIGNMENT>
        void tridiagonal_reduction(viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& A,
                                    viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& Q)
        {
            std::size_t sz = A.size1();

            viennacl::vector<SCALARTYPE> hh_vector(sz);

            for(std::size_t i = 0; i < sz; i++)
            {
                householder_twoside(A, Q, hh_vector, i);
            }

        }

        template <typename SCALARTYPE, typename F, unsigned int ALIGNMENT>
        void qr_method(viennacl::matrix<SCALARTYPE, F, ALIGNMENT> & A,
                       viennacl::matrix<SCALARTYPE, F, ALIGNMENT> & Q,
                       boost::numeric::ublas::vector<SCALARTYPE> & D,
                       boost::numeric::ublas::vector<SCALARTYPE> & E,
                       bool is_symmetric = true)
        {
            assert(A.size1() == A.size2());

            D.resize(A.size1());
            E.resize(A.size1());

            viennacl::linalg::kernels::svd<SCALARTYPE, 1>::init();

            detail::eye(Q);

            // reduce to tridiagonal form
            detail::tridiagonal_reduction(A, Q);

            // pack diagonal and super-diagonal
            // ublas::vector<SCALARTYPE> D(A.size1()), E(A.size1());

            bidiag_pack(A, D, E);

            // find eigenvalues
            if(is_symmetric)
            {

                detail::tql2(Q, D, E);
                transpose(Q);
            }
            else
            {
                detail::hqr2(A, Q, D, E);
            }

            // std::cout << A << "\n";

            boost::numeric::ublas::matrix<float> eigen_values(A.size1(), A.size1());
            eigen_values.clear();

            for (std::size_t i = 0; i < A.size1(); i++)
            {
                if(std::fabs(E(i)) < EPS)
                {
                    eigen_values(i, i) = D(i);
                }
                else
                {
                    eigen_values(i, i) = D(i);
                    eigen_values(i, i + 1) = E(i);
                    eigen_values(i + 1, i) = -E(i);
                    eigen_values(i + 1, i + 1) = D(i);
                    i++;
                }
            }

            copy(eigen_values, A);
        }
    }

    template <typename SCALARTYPE, typename F, unsigned int ALIGNMENT>
    void qr_method_nsm(viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& A,
                       viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& Q,
                       boost::numeric::ublas::vector<SCALARTYPE>& D,
                       boost::numeric::ublas::vector<SCALARTYPE>& E
                      )
    {
        detail::qr_method(A, Q, D, E, false);
    }

    template <typename SCALARTYPE, typename F, unsigned int ALIGNMENT>
    void qr_method_sym(viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& A,
                       viennacl::matrix<SCALARTYPE, F, ALIGNMENT>& Q,
                       boost::numeric::ublas::vector<SCALARTYPE>& D
                      )
    {
        boost::numeric::ublas::vector<SCALARTYPE> E(A.size1());

        detail::qr_method(A, Q, D, E, true);
    }

  }
}

#endif


// ---------------------------------------------------------------------
//
// Copyright (c) 2020 - 2020 by the IBAMR developers
// All rights reserved.
//
// This file is part of IBAMR.
//
// IBAMR is free software and is distributed under the 3-clause BSD
// license. The full text of the license can be found in the file
// COPYRIGHT at the top level directory of IBAMR.
//
// ---------------------------------------------------------------------

// Verify that we can correctly tag cells based on input information to
// FEDataManager so that a single FE mesh ends up on multiple patch levels.

// Headers for basic SAMRAI objects
#include <BergerRigoutsos.h>
#include <CartesianGridGeometry.h>
#include <HierarchyCellDataOpsInteger.h>
#include <LoadBalancer.h>
#include <StandardTagAndInitialize.h>

// Headers for basic libMesh objects
#include <libmesh/boundary_info.h>
#include <libmesh/equation_systems.h>
#include <libmesh/gmv_io.h>
#include <libmesh/mesh.h>
#include <libmesh/mesh_generation.h>
#include <libmesh/mesh_refinement.h>

// Headers for application-specific algorithm/data structure objects
#include <ibamr/IBExplicitHierarchyIntegrator.h>
#include <ibamr/IBFECentroidPostProcessor.h>
#include <ibamr/IBFEMethod.h>
#include <ibamr/INSCollocatedHierarchyIntegrator.h>
#include <ibamr/INSStaggeredHierarchyIntegrator.h>

#include <ibtk/AppInitializer.h>
#include <ibtk/IBTKInit.h>
#include <ibtk/IBTK_MPI.h>
#include <ibtk/StableCentroidPartitioner.h>
#include <ibtk/libmesh_utilities.h>
#include <ibtk/muParserCartGridFunction.h>
#include <ibtk/muParserRobinBcCoefs.h>

#include <boost/multi_array.hpp>

// Set up application namespace declarations
#include <ibamr/app_namespaces.h>

#include "../tests.h"

System&
setup_deformation_system(ReplicatedMesh& mesh, EquationSystems& equation_systems, const libMesh::Order order)
{
    // Set up the system
    auto& X_system = equation_systems.add_system<ExplicitSystem>("X");
    const auto X_sys_num = X_system.number();
    const unsigned int dim = mesh.mesh_dimension();
    for (unsigned int d = 0; d < dim; ++d)
    {
        X_system.add_variable("X_" + std::to_string(d), order, LAGRANGE);
    }
    equation_systems.init();

    // Set up the system solution to match the current mesh coordinates
    auto& X_solution = *X_system.solution;
    const auto el_begin = mesh.active_local_elements_begin();
    const auto el_end = mesh.active_local_elements_end();
    for (auto el_it = el_begin; el_it != el_end; ++el_it)
    {
        const libMesh::Elem* const elem = *el_it;
        const unsigned int n_nodes = elem->n_nodes();
        for (unsigned int k = 0; k < n_nodes; ++k)
        {
            const libMesh::Node* const node = elem->node_ptr(k);
            for (unsigned int d = 0; d < dim; ++d)
            {
                TBOX_ASSERT(node->n_dofs(X_sys_num, d) == 1);
                const auto dof = node->dof_number(X_sys_num, d, 0);
                X_solution.set(dof, (*node)(d));
            }
        }
    }
    X_solution.close();

    return X_system;
}

class TestTag : public SAMRAI::mesh::StandardTagAndInitStrategy<NDIM>
{
public:
    TestTag(const Parallel::Communicator& comm, Pointer<Database>& input_db) : d_mesh(comm), d_es(d_mesh)
    {
        MeshTools::Generation::build_sphere(d_mesh, 0.5, 4, NDIM == 2 ? TRI3 : HEX8);
        MeshBase::element_iterator el_end = d_mesh.elements_end();
        for (MeshBase::element_iterator el = d_mesh.elements_begin(); el != el_end; ++el)
        {
            Elem* elem = *el;
#if LIBMESH_VERSION_LESS_THAN(1, 7, 0)
            const libMesh::Point centroid = elem->centroid();
#else
            const libMesh::Point centroid = elem->vertex_average();
#endif
#if NDIM == 2
            if (centroid(1) > 0.0)
                elem->subdomain_id() = 1;
            else
                elem->subdomain_id() = 2;
#else

            if (centroid(1) > 0.0 && centroid(2) > 0.0)
                elem->subdomain_id() = 1;
            else
                elem->subdomain_id() = 2;
#endif
        }

        GMVIO gmv_io(d_mesh);
        gmv_io.write("mesh.gmv");

        auto& X_system = setup_deformation_system(d_mesh, d_es, FIRST);

        auto fe_data = std::make_shared<FEData>("fe_data", d_es, true);

        FEDataManager::InterpSpec interp_spec("IB_4", QGAUSS, THIRD, true, 0.9, true, false);
        FEDataManager::SpreadSpec spread_spec("IB_4", QGAUSS, THIRD, true, 0.9, false);
        FEDataManager::WorkloadSpec workload_spec;

        d_fe_data_manager = FEDataManager::getManager(fe_data,
                                                      "fe_data_manager",
                                                      input_db->getDatabase("FEDataManager"),
                                                      input_db->getInteger("MAX_LEVELS"),
                                                      interp_spec,
                                                      spread_spec,
                                                      workload_spec);
        d_fe_data_manager->COORDINATES_SYSTEM_NAME = X_system.name();
    }

    void initializeLevelData(const Pointer<BasePatchHierarchy<NDIM> > /*hierarchy*/,
                             const int /*level_number*/,
                             const double /*init_data_time*/,
                             const bool /*can_be_refined*/,
                             const bool /*initial_time*/,
                             const tbox::Pointer<hier::BasePatchLevel<NDIM> > /*old_level*/ = nullptr,
                             const bool /*allocate_data*/ = true) override
    {
    }

    void resetHierarchyConfiguration(const tbox::Pointer<hier::BasePatchHierarchy<NDIM> > /*hierarchy*/,
                                     const int /*coarsest_level*/,
                                     const int /*finest_level*/) override
    {
    }

    void applyGradientDetector(const tbox::Pointer<hier::BasePatchHierarchy<NDIM> > hierarchy,
                               const int level_number,
                               const double error_data_time,
                               const int tag_index,
                               const bool initial_time,
                               const bool uses_richardson_extrapolation_too) override
    {
        d_fe_data_manager->setPatchHierarchy(hierarchy);
        d_fe_data_manager->applyGradientDetector(
            hierarchy, level_number, error_data_time, tag_index, initial_time, uses_richardson_extrapolation_too);
    }

protected:
    FEDataManager* d_fe_data_manager;

    ReplicatedMesh d_mesh;

    EquationSystems d_es;
};

int
main(int argc, char** argv)
{
    IBTKInit ibtk_init(argc, argv, MPI_COMM_WORLD);
    const LibMeshInit& init = ibtk_init.getLibMeshInit();

    // Parse command line options, set some standard options from the input
    // file, and enable file logging.
    Pointer<AppInitializer> app_initializer = new AppInitializer(argc, argv, "multilevel_fe_01.log");
    Pointer<Database> input_db = app_initializer->getInputDatabase();
    TestTag tag(init.comm(), input_db);

    Pointer<CartesianGridGeometry<NDIM> > grid_geometry = new CartesianGridGeometry<NDIM>(
        "CartesianGeometry", app_initializer->getComponentDatabase("CartesianGeometry"));
    Pointer<PatchHierarchy<NDIM> > patch_hierarchy = new PatchHierarchy<NDIM>("PatchHierarchy", grid_geometry);
    Pointer<StandardTagAndInitialize<NDIM> > error_detector = new StandardTagAndInitialize<NDIM>(
        "StandardTagAndInitialize", &tag, app_initializer->getComponentDatabase("StandardTagAndInitialize"));
    Pointer<BergerRigoutsos<NDIM> > box_generator = new BergerRigoutsos<NDIM>();
    Pointer<LoadBalancer<NDIM> > load_balancer =
        new LoadBalancer<NDIM>("LoadBalancer", app_initializer->getComponentDatabase("LoadBalancer"));
    Pointer<GriddingAlgorithm<NDIM> > gridding_algorithm =
        new GriddingAlgorithm<NDIM>("GriddingAlgorithm",
                                    app_initializer->getComponentDatabase("GriddingAlgorithm"),
                                    error_detector,
                                    box_generator,
                                    load_balancer);

    gridding_algorithm->makeCoarsestLevel(patch_hierarchy, 0.0);
    int level_number = 0;
    while (gridding_algorithm->levelCanBeRefined(level_number))
    {
        gridding_algorithm->makeFinerLevel(patch_hierarchy, 0.0, true, 0);
        ++level_number;
    }
#if 0
    tbox::Array<int> tag_buffer;
    const int finest_hier_ln = patch_hierarchy->getFinestLevelNumber();
    tag_buffer.resizeArray(finest_hier_ln + 1);
    for (int i = 0; i < tag_buffer.getSize(); ++i)
        tag_buffer[i] = 0;
    gridding_algorithm->regridAllFinerLevels(patch_hierarchy, 0, 0.0, tag_buffer);
#endif

    // plot stuff
    VariableDatabase<NDIM>* var_db = VariableDatabase<NDIM>::getDatabase();
    Pointer<VariableContext> ctx = var_db->getContext("context");
    Pointer<CellVariable<NDIM, double> > u_cc_var = new CellVariable<NDIM, double>("u_cc");
    const int u_cc_idx = var_db->registerVariableAndContext(u_cc_var, ctx, IntVector<NDIM>(1));
    for (int ln = 0; ln <= patch_hierarchy->getFinestLevelNumber(); ++ln)
    {
        Pointer<PatchLevel<NDIM> > level = patch_hierarchy->getPatchLevel(ln);
        TBOX_ASSERT(level);
        level->allocatePatchData(u_cc_idx, 0.0);
    }

    // Register variables for plotting.
    Pointer<VisItDataWriter<NDIM> > visit_data_writer = app_initializer->getVisItDataWriter();
    TBOX_ASSERT(visit_data_writer);
    visit_data_writer->registerPlotQuantity(u_cc_var->getName(), "SCALAR", u_cc_idx);
    visit_data_writer->writePlotData(patch_hierarchy, 0, 0.0);

    print_partitioning_on_plog_0(patch_hierarchy, 0, patch_hierarchy->getFinestLevelNumber());
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_IEEE_FUNCTIONS_SCALAR_FAST_LDEXP_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_IEEE_FUNCTIONS_SCALAR_FAST_LDEXP_HPP_INCLUDED
#include <boost/simd/toolbox/ieee/functions/fast_ldexp.hpp>
#include <boost/simd/include/constants/ldexpmask.hpp>
#include <boost/simd/include/constants/nbmantissabits.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/simd/include/constants/real.hpp>
#include <boost/simd/include/functions/scalar/bitwise_andnot.hpp>
#include <boost/simd/include/functions/scalar/bitwise_or.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::fast_ldexp_, tag::cpu_
                            , (A0)(A1)
                            , (scalar_< integer_<A0> >)(scalar_< integer_<A1> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL(2)
    {
       return (a1>=0)?(a0<<a1):(a0>>a1);
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::fast_ldexp_, tag::cpu_
                            , (A0)(A1)
                            , (scalar_< floating_<A0> >)(scalar_< integer_<A1> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL(2)
    {
      // No denormal provision
      typedef typename dispatch::meta::as_integer<A0, unsigned>::type int_type;
      // clear exponent in x
      A0 const x(b_andnot(a0, Ldexpmask<A0>()));
      // extract exponent and compute the new one
      int_type e    = b_and(Ldexpmask<A0>(), a0);
      e += int_type(a1) << Nbmantissabits<A0>();
      return b_or(x, e);
    }
  };
} } }

#endif


/* dos_estimate_mixture.cc

Copyright 2018 Grant M. Rotskoff

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

 */
#include <armadillo>
#include "particles.h"
#include "gaussian_mixture.h"

int main(int argc, char **argv)
{

  // usage
  if (argc != 5)
  {
    fprintf(stderr,"usage ./run_mixture dim gamma n_traj n_well\n");
    exit(-1);
  }
  int dim = atoi(argv[1]);
  double gamma = atof(argv[2]);
  int n_traj = atoi(argv[3]);
  int n_well = atoi(argv[4]);

  arma_rng::set_seed(983948524);

  gaussian_mixture g;

  g.initialize(dim, n_well);

  // load in the parameters for the mog
  char buffer[50];
  sprintf(buffer, "../data/sigmas_%03d_%03d.dat", dim, n_well);
  g.sigmas.load(buffer);
  sprintf(buffer, "../data/depths_%03d_%03d.dat", dim, n_well);
  g.depths.load(buffer);
  sprintf(buffer, "../data/mus_%03d_%03d.dat", dim, n_well);
  g.mus.load(buffer);
  g.mus = g.mus.t();
  // set the integration parameters
  double dt = 1e-5;
  double tol = 1e-5;
  double hmin = -g.depths.max();
  double hmax = 450.;
  double hres = 0.2;
  int max_iter = 0;
  fprintf(stderr,"Setting hmin=%4.3f\n",hmin);

  // reinjection parameters
  g.qmin = -1.;
  g.qmax = 1.;
  g.pmin = 0; // (NB: using mc reinjection)
  g.pmax = 0;
  g.mc_stepsize = 1e-2;
  g.beta = 5;
  bool use_mc = true;

  g.initialize_integration_variables(dt, gamma);
  g.initialize_estimator(hmax, hres, hmin, n_traj);
  fprintf(stderr,"Initialized estimator\n");
  fprintf(stderr, "Running Mixture of Gaussians with");
  fprintf(stderr, "\n\t n_well=%d",g.n_well);
  fprintf(stderr, "\n\t dim=%d\n",g.dim);

  // save everything
  for (int i=0; i<n_traj; i++) {
    g.run_estimation_trajectory(tol, hmax, max_iter, i, langevin, use_mc);
    fprintf(stderr, "Ran traj %d\n", i);
  }

  char dos_filename[80];
  sprintf(dos_filename, "mog_dim=%04d_gamma=%05.3f_ntraj=%04d.dat", dim, gamma, n_traj);
  FILE *dosf = fopen(dos_filename, "w");
  g.dump_dos_estimate(n_traj, dosf);
}


#if !defined(BALSA_EIGEN_TYPES_HPP)
#define BALSA_EIGEN_TYPES_HPP

// Verbose because clang/llvm pretend to have gnuc
#if defined(__GNUC__) && !defined(__llvm__) && !defined(__INTEL_COMPILER)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wint-in-bool-context"
#endif
#include <Eigen/Dense>
#if defined(__GNUC__) && !defined(__llvm__) && !defined(__INTEL_COMPILER)
#pragma GCC diagnostic pop
#endif

namespace balsa::eigen {
template<typename T, int A, int B>
using Matrix = Eigen::Matrix<T, A, B>;

template<typename T, int A>
using SquareMatrix = Matrix<T, A, A>;
template<typename T>
using MatrixX = SquareMatrix<T, Eigen::Dynamic>;

template<typename T, int D>
using Vector = Matrix<T, D, 1>;
template<typename T, int D>
using RowVector = Matrix<T, 1, D>;
template<typename T>
using VectorX = Vector<T, Eigen::Dynamic>;
template<typename T>
using RowVectorX = RowVector<T, Eigen::Dynamic>;

template<typename T>
using Vector2 = Vector<T, 2>;
template<typename T>
using RowVector2 = RowVector<T, 2>;

template<typename T>
using Vector3 = Vector<T, 3>;
template<typename T>
using RowVector3 = RowVector<T, 3>;

template<typename T>
using Vector4 = Vector<T, 4>;
template<typename T>
using RowVector4 = RowVector<T, 4>;

template<typename T, int D>
using ColVectors = Matrix<T, D, Eigen::Dynamic>;
template<typename T, int D>
using RowVectors = Matrix<T, Eigen::Dynamic, D>;

using Mat3i = SquareMatrix<int, 3>;
using Mat2i = SquareMatrix<int, 2>;
using ColVecs4i = ColVectors<int, 4>;
using ColVecs3i = ColVectors<int, 3>;
using ColVecs2i = ColVectors<int, 2>;
using RowVecs3i = RowVectors<int, 3>;
using RowVecs2i = RowVectors<int, 2>;
using Vec3i = Vector<int, 3>;
using Vec2i = Vector<int, 2>;
using MatXi = MatrixX<int>;
using VecXi = VectorX<int>;
using RowVecXi = RowVectorX<int>;

using Mat4f = SquareMatrix<float, 4>;
using Mat3f = SquareMatrix<float, 3>;
using Mat2f = SquareMatrix<float, 2>;
using ColVecs4f = ColVectors<float, 4>;
using ColVecs3f = ColVectors<float, 3>;
using ColVecs2f = ColVectors<float, 2>;
using RowVecs3f = RowVectors<float, 3>;
using RowVecs2f = RowVectors<float, 2>;
using Vec4f = Vector<float, 4>;
using Vec3f = Vector<float, 3>;
using Vec2f = Vector<float, 2>;
using MatXf = MatrixX<float>;
using VecXf = VectorX<float>;
using RowVecXf = RowVectorX<float>;

using Mat4d = SquareMatrix<double, 4>;
using Mat3d = SquareMatrix<double, 3>;
using Mat2d = SquareMatrix<double, 2>;
using ColVecs4d = ColVectors<double, 4>;
using ColVecs3d = ColVectors<double, 3>;
using ColVecs2d = ColVectors<double, 2>;
using RowVecs3d = RowVectors<double, 3>;
using RowVecs2d = RowVectors<double, 2>;
using Vec4d = Vector<double, 4>;
using Vec3d = Vector<double, 3>;
using Vec2d = Vector<double, 2>;
using MatXd = MatrixX<double>;
using VecXd = VectorX<double>;
using RowVecXd = RowVectorX<double>;

using Vec2i = Vector2<int>;
using Vec3i = Vector3<int>;
using Vec4i = Vector4<int>;

}// namespace balsa::eigen

#endif


#ifndef TMC_UNIT_TESTS
#define BOOST_TEST_MODULE lib_test_border_system
#endif
#include <boost/test/included/unit_test.hpp>
#include <boost/property_tree/json_parser.hpp>
#include "sample_creator.h"
#include "border_system_creator.h"

BOOST_AUTO_TEST_SUITE( border_system )

BOOST_AUTO_TEST_CASE( from_36points_layer_partitioning ) {
    const std::string names_file("data/4layers_36points/4layers_36points.names");
    const std::string data_file("data/4layers_36points/4layers_36points.data");
    BOOST_TEST_MESSAGE("Creating sample from file: " << data_file);
    std::shared_ptr<FeatureNames> pFN = std::make_shared<FeatureNames>(names_file);
    SampleCreator sample_creator;
    sample_creator.set_feature_names(pFN);
    std::shared_ptr<Sample> pSample = sample_creator.from_file(data_file);
    BOOST_TEST_MESSAGE("Resulting sample: " << *pSample);

    boost::dynamic_bitset<> db4(36);
    boost::dynamic_bitset<> db3(27);
    boost::dynamic_bitset<> db2(19);
    db4[27] = db4[28] = db4[29] = db4[30] = db4[31] = db4[32] = db4[33] = db4[34] = db4[35] = true;
    db3[19] = db3[20] = db3[21] = db3[22] = db3[23] = db3[24] = db3[25] = db3[26] = true;
    db2[12] = db2[13] = db2[14] = db2[15] = db2[16] = db2[17] = db2[18] = true;

    std::shared_ptr<LayerPartitioning> pLD = std::make_shared<LayerPartitioning>();
    pLD->push_back(pSample);

    BOOST_TEST_MESSAGE("Create layer partitioning via splitting the lowest layer three times into 4 layers");
    { auto it = pLD->begin(); pLD->split_layer(it, db4); }
    { auto it = pLD->begin(); pLD->split_layer(it, db3); }
    { auto it = pLD->begin(); pLD->split_layer(it, db2); }

    BOOST_CHECK_EQUAL(pLD->consistent(),true);
    auto pLayer1 = *pLD->begin();
    auto pLayer2 = *(pLD->begin() + 1);
    auto pLayer3 = *(pLD->begin() + 2);
    auto pLayer4 = *(pLD->begin() + 3);

    BOOST_TEST_MESSAGE("First layer: " << *pLayer1);
    BOOST_TEST_MESSAGE("Second layer: " << *pLayer2);
    BOOST_TEST_MESSAGE("Third layer: " << *pLayer3);
    BOOST_TEST_MESSAGE("Fourth layer: " << *pLayer4);

    std::shared_ptr<BorderSystemCreator> pBSC = std::make_shared<BorderSystemCreator>();
    std::shared_ptr<BorderSystem> pBS = pBSC->from_layer_partitioning(pLD);
    BOOST_CHECK_EQUAL(pBS->dim(), pLD->dim());
    BOOST_CHECK_EQUAL(pBS->size(), pLD->size());
    const Border & LowerBorder0 = *pBS->get_lower(0);
    const Border & LowerBorder1 = *pBS->get_lower(1);
    const Border & LowerBorder2 = *pBS->get_lower(2);
    const Border & LowerBorder3 = *pBS->get_lower(3);

    const Border & UpperBorder0 = *pBS->get_upper(0);
    const Border & UpperBorder1 = *pBS->get_upper(1);
    const Border & UpperBorder2 = *pBS->get_upper(2);
    const Border & UpperBorder3 = *pBS->get_upper(3);

    BOOST_TEST_MESSAGE("First lower border: " << LowerBorder0);
    BOOST_TEST_MESSAGE("Second lower border: " << LowerBorder1);
    BOOST_TEST_MESSAGE("Third lower border: " << LowerBorder2);
    BOOST_TEST_MESSAGE("Fourth lower border: " << LowerBorder3);

    BOOST_TEST_MESSAGE("First upper border: " << UpperBorder0);
    BOOST_TEST_MESSAGE("Second upper border: " << UpperBorder1);
    BOOST_TEST_MESSAGE("Third upper border: " << UpperBorder2);
    BOOST_TEST_MESSAGE("Fourth upper border: " << UpperBorder3);

    BOOST_CHECK(pLayer1->get_neg_pos_counts() == LowerBorder0.get_neg_pos_counts());
    BOOST_CHECK(pLayer1->get_neg_pos_counts() == UpperBorder0.get_neg_pos_counts());
    BOOST_CHECK(pLayer2->get_neg_pos_counts() == LowerBorder1.get_neg_pos_counts());
    BOOST_CHECK(pLayer2->get_neg_pos_counts() == UpperBorder1.get_neg_pos_counts());
    BOOST_CHECK(pLayer3->get_neg_pos_counts() == LowerBorder2.get_neg_pos_counts());
    BOOST_CHECK(pLayer3->get_neg_pos_counts() == UpperBorder2.get_neg_pos_counts());
    BOOST_CHECK(pLayer4->get_neg_pos_counts() == LowerBorder3.get_neg_pos_counts());
    BOOST_CHECK(pLayer4->get_neg_pos_counts() == UpperBorder3.get_neg_pos_counts());

    //lower borders
    BOOST_CHECK_EQUAL(LowerBorder0.size(), 5);
    BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{9,1})));
    BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{8,3})));
    BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{5,4})));
    BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{2,7})));
    BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{0,11})));

    BOOST_CHECK_EQUAL(LowerBorder1.size(), 3);
    BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{15,0})));
    BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{13,2})));
    BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{12,5})));

    BOOST_CHECK_EQUAL(LowerBorder2.size(), 3);
    BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{11,8})));
    BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{5,9})));
    BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{4,10})));

    BOOST_CHECK_EQUAL(LowerBorder3.size(), 5);
    BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{16,6})));
    BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{13,7})));
    BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{9,9})));
    BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{6,13})));
    BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{3,14})));

    // upper borders
    BOOST_CHECK_EQUAL(UpperBorder3.size(), 6);
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{16,6})));
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{13,7})));
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{12,11})));
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{9,12})));
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{6,15})));
    BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{5,17})));

    BOOST_CHECK_EQUAL(UpperBorder2.size(), 3);
    BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{11,8})));
    BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{8,11})));
    BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{7,12})));

    BOOST_CHECK_EQUAL(UpperBorder1.size(), 2);
    BOOST_CHECK(UpperBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{17,5})));
    BOOST_CHECK(UpperBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{14,6})));

    BOOST_CHECK_EQUAL(UpperBorder0.size(), 5);
    BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{14,1})));
    BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{12,3})));
    BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{11,7})));
    BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{3,12})));
    BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{2,15})));
}

BOOST_AUTO_TEST_CASE( ptree ) {
    const std::string names_file("data/4layers_36points/4layers_36points.names");
    const std::string data_file("data/4layers_36points/4layers_36points.data");
    BOOST_TEST_MESSAGE("Creating sample from file: " << data_file);
    std::shared_ptr<FeatureNames> pFN = std::make_shared<FeatureNames>(names_file);
    SampleCreator sample_creator;
    sample_creator.set_feature_names(pFN);
    std::shared_ptr<Sample> pSample = sample_creator.from_file(data_file);
    BOOST_TEST_MESSAGE("Resulting sample: " << *pSample);

    boost::dynamic_bitset<> db4(36);
    boost::dynamic_bitset<> db3(27);
    boost::dynamic_bitset<> db2(19);
    db4[27] = db4[28] = db4[29] = db4[30] = db4[31] = db4[32] = db4[33] = db4[34] = db4[35] = true;
    db3[19] = db3[20] = db3[21] = db3[22] = db3[23] = db3[24] = db3[25] = db3[26] = true;
    db2[12] = db2[13] = db2[14] = db2[15] = db2[16] = db2[17] = db2[18] = true;

    std::shared_ptr<LayerPartitioning> pLD = std::make_shared<LayerPartitioning>();
    pLD->push_back(pSample);

    BOOST_TEST_MESSAGE("Create layer partitioning via splitting the lowest layer three times into 4 layers");
    { auto it = pLD->begin(); pLD->split_layer(it, db4); }
    { auto it = pLD->begin(); pLD->split_layer(it, db3); }
    { auto it = pLD->begin(); pLD->split_layer(it, db2); }

    BOOST_CHECK_EQUAL(pLD->consistent(),true);
    auto pLayer1 = *pLD->begin();
    auto pLayer2 = *(pLD->begin() + 1);
    auto pLayer3 = *(pLD->begin() + 2);
    auto pLayer4 = *(pLD->begin() + 3);

    BOOST_TEST_MESSAGE("First layer: " << *pLayer1);
    BOOST_TEST_MESSAGE("Second layer: " << *pLayer2);
    BOOST_TEST_MESSAGE("Third layer: " << *pLayer3);
    BOOST_TEST_MESSAGE("Fourth layer: " << *pLayer4);

    std::shared_ptr<BorderSystemCreator> pBSC = std::make_shared<BorderSystemCreator>();
    std::shared_ptr<BorderSystem> pBS1 = pBSC->from_layer_partitioning(pLD);

    boost::property_tree::ptree pt;
    pBS1->dump_to_ptree(pt);
    BOOST_CHECK_EQUAL(pt.size(), 4);
    std::stringstream ss;
    boost::property_tree::json_parser::write_json(ss, pt);
    BOOST_TEST_MESSAGE("Property tree as json:\n" << ss.str());

    std::shared_ptr<BorderSystem> pBS = std::make_shared<BorderSystem>(pt);
    BOOST_CHECK_EQUAL(pBS->dim(), pLD->dim());
    BOOST_CHECK_EQUAL(pBS->size(), pLD->size());
    const Border & LowerBorder0 = *pBS->get_lower(0);
    const Border & LowerBorder1 = *pBS->get_lower(1);
    const Border & LowerBorder2 = *pBS->get_lower(2);
    const Border & LowerBorder3 = *pBS->get_lower(3);

    const Border & UpperBorder0 = *pBS->get_upper(0);
    const Border & UpperBorder1 = *pBS->get_upper(1);
    const Border & UpperBorder2 = *pBS->get_upper(2);
    const Border & UpperBorder3 = *pBS->get_upper(3);

    BOOST_TEST_MESSAGE("First lower border: " << LowerBorder0);
    BOOST_TEST_MESSAGE("Second lower border: " << LowerBorder1);
    BOOST_TEST_MESSAGE("Third lower border: " << LowerBorder2);
    BOOST_TEST_MESSAGE("Fourth lower border: " << LowerBorder3);

    BOOST_TEST_MESSAGE("First upper border: " << UpperBorder0);
    BOOST_TEST_MESSAGE("Second upper border: " << UpperBorder1);
    BOOST_TEST_MESSAGE("Third upper border: " << UpperBorder2);
    BOOST_TEST_MESSAGE("Fourth upper border: " << UpperBorder3);

    BOOST_CHECK(pLayer1->get_neg_pos_counts() == LowerBorder0.get_neg_pos_counts());
    BOOST_CHECK(pLayer1->get_neg_pos_counts() == UpperBorder0.get_neg_pos_counts());
    BOOST_CHECK(pLayer2->get_neg_pos_counts() == LowerBorder1.get_neg_pos_counts());
    BOOST_CHECK(pLayer2->get_neg_pos_counts() == UpperBorder1.get_neg_pos_counts());
    BOOST_CHECK(pLayer3->get_neg_pos_counts() == LowerBorder2.get_neg_pos_counts());
    BOOST_CHECK(pLayer3->get_neg_pos_counts() == UpperBorder2.get_neg_pos_counts());
    BOOST_CHECK(pLayer4->get_neg_pos_counts() == LowerBorder3.get_neg_pos_counts());
    BOOST_CHECK(pLayer4->get_neg_pos_counts() == UpperBorder3.get_neg_pos_counts());

        //lower borders
        BOOST_CHECK_EQUAL(LowerBorder0.size(), 5);
        BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{9,1})));
        BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{8,3})));
        BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{5,4})));
        BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{2,7})));
        BOOST_CHECK(LowerBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{0,11})));

        BOOST_CHECK_EQUAL(LowerBorder1.size(), 3);
        BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{15,0})));
        BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{13,2})));
        BOOST_CHECK(LowerBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{12,5})));

        BOOST_CHECK_EQUAL(LowerBorder2.size(), 3);
        BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{11,8})));
        BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{5,9})));
        BOOST_CHECK(LowerBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{4,10})));

        BOOST_CHECK_EQUAL(LowerBorder3.size(), 5);
        BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{16,6})));
        BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{13,7})));
        BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{9,9})));
        BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{6,13})));
        BOOST_CHECK(LowerBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{3,14})));

        // upper borders
        BOOST_CHECK_EQUAL(UpperBorder3.size(), 6);
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{16,6})));
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{13,7})));
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{12,11})));
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{9,12})));
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{6,15})));
        BOOST_CHECK(UpperBorder3.contains(std::make_shared<FeatureVector>(std::vector<double>{5,17})));

        BOOST_CHECK_EQUAL(UpperBorder2.size(), 3);
        BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{11,8})));
        BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{8,11})));
        BOOST_CHECK(UpperBorder2.contains(std::make_shared<FeatureVector>(std::vector<double>{7,12})));

        BOOST_CHECK_EQUAL(UpperBorder1.size(), 2);
        BOOST_CHECK(UpperBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{17,5})));
        BOOST_CHECK(UpperBorder1.contains(std::make_shared<FeatureVector>(std::vector<double>{14,6})));

        BOOST_CHECK_EQUAL(UpperBorder0.size(), 5);
        BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{14,1})));
        BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{12,3})));
        BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{11,7})));
        BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{3,12})));
        BOOST_CHECK(UpperBorder0.contains(std::make_shared<FeatureVector>(std::vector<double>{2,15})));
}

BOOST_AUTO_TEST_CASE( containing_border ) {
    const std::string names_file("data/4layers_36points/4layers_36points.names");
    const std::string data_file("data/4layers_36points/4layers_36points.data");
    BOOST_TEST_MESSAGE("Creating sample from file: " << data_file);
    std::shared_ptr<FeatureNames> pFN = std::make_shared<FeatureNames>(names_file);
    SampleCreator sample_creator;
    sample_creator.set_feature_names(pFN);
    std::shared_ptr<Sample> pSample = sample_creator.from_file(data_file);
    BOOST_TEST_MESSAGE("Resulting sample: " << *pSample);

    boost::dynamic_bitset<> db4(36);
    boost::dynamic_bitset<> db3(27);
    boost::dynamic_bitset<> db2(19);
    db4[27] = db4[28] = db4[29] = db4[30] = db4[31] = db4[32] = db4[33] = db4[34] = db4[35] = true;
    db3[19] = db3[20] = db3[21] = db3[22] = db3[23] = db3[24] = db3[25] = db3[26] = true;
    db2[12] = db2[13] = db2[14] = db2[15] = db2[16] = db2[17] = db2[18] = true;

    std::shared_ptr<LayerPartitioning> pLD = std::make_shared<LayerPartitioning>();
    pLD->push_back(pSample);

    BOOST_TEST_MESSAGE("Create layer partitioning via splitting the lowest layer three times into 4 layers");
    { auto it = pLD->begin(); pLD->split_layer(it, db4); }
    { auto it = pLD->begin(); pLD->split_layer(it, db3); }
    { auto it = pLD->begin(); pLD->split_layer(it, db2); }

    BOOST_CHECK_EQUAL(pLD->consistent(),true);
    auto pLayer1 = *pLD->begin();
    auto pLayer2 = *(pLD->begin() + 1);
    auto pLayer3 = *(pLD->begin() + 2);
    auto pLayer4 = *(pLD->begin() + 3);

    BOOST_CHECK( pLayer1->get_neg_pos_counts() == (std::pair<Sample::WeightType,Sample::WeightType>({50,10})) );
    BOOST_CHECK( pLayer2->get_neg_pos_counts() == (std::pair<Sample::WeightType,Sample::WeightType>({16,8})) );
    BOOST_CHECK( pLayer3->get_neg_pos_counts() == (std::pair<Sample::WeightType,Sample::WeightType>({9,18})) );
    BOOST_CHECK( pLayer4->get_neg_pos_counts() == (std::pair<Sample::WeightType,Sample::WeightType>({9,45})) );

    BOOST_TEST_MESSAGE("Creating border system");
    std::shared_ptr<BorderSystemCreator> pBSC = std::make_shared<BorderSystemCreator>();
    std::shared_ptr<BorderSystem> pBS = pBSC->from_layer_partitioning(pLD);
    BOOST_CHECK_EQUAL(pBS.use_count(), 1);
    std::vector<double> p;  std::pair<int, int> borders; std::pair<double, double> conf;
    p = {10,2}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {10,2}; conf = {1.0/6.0,1.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.first);
    p = {1,12}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {16,1}; borders={1,1};BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {16,1}; conf = {1.0/3.0,1.0/3.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.first);
    p = {7,10}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {7,10}; conf = {2.0/3.0,2.0/3.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.first);
    p = {11,10}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {11,10}; conf = {5.0/6.0,5.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.first);
    p = {4,15}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {0,11}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {2,15}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {2,7}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {3,12}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {5,4}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {8,3}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {9,1}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {12,3}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {11,7}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {14,1}; borders={0,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {15,0}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {17,2}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {13,2}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {12,5}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {14,6}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {17,5}; borders={1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {11,8}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {8,9}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {8,11}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {7,12}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {7,9}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {5,9}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {4,10}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {4,12}; borders={2,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {16,6}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {13,7}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {12,9}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {12,11}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {9,9}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {9,11}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {9,12}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {6,13}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {6,15}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {3,14}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {3,17}; borders={3,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {4,6}; borders={-1,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {4,6}; conf = {0,1.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.second);
    p = {7,3}; borders={-1,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {1,10}; borders={-1,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {7.9,3.9}; borders={-1,0}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {7,13}; borders={3,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {7,13}; conf = {5.0/6.0,1}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), conf.first);
    p = {6.1,13.1}; borders={3,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {10,15}; borders={3,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {14.5,0.5}; borders={-1,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {14.5,0.5}; conf = {0,1.0/3.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8)/(double)(50+10+16+8));
    p = {-1,16}; borders={-1,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {-1,16}; conf = {0,5.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8+18+45)/(double)(50+10+16+8+9+18+9+45));
    p = {-1,18}; borders={-1,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {-1,18}; conf = {0,1}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8+18+45)/(double)(50+10+16+8+9+18+9+45));
    p = {18,-1}; borders={-1,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);

    p = {12.5,1.5}; borders={0,1}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {12.5,1.5}; conf = {1.0/6.0,1.0/3.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8)/(double)(50+10+16+8));
    p = {4,8}; borders={0,2}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {4,8}; conf = {1.0/6.0,2.0/3.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8+18)/(double)(50+10+16+8+9+18));
    p = {3,13}; borders={0,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {3,13}; conf = {1.0/6.0,5.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8+18+45)/(double)(50+10+16+8+9+18+9+45));
    p = {1,18}; borders={0,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {1,18};  conf = {1.0/6.0,1}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (10+8+18+45)/(double)(50+10+16+8+9+18+9+45));

    p = {15,5.5}; borders={1,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {15,5.5}; conf = {1.0/3.0,5.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (8+18+45)/(double)(16+8+9+18+9+45));
    p = {18,1}; borders={1,4}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {18,1}; conf = {1.0/3.0,1}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (8+18+45)/(double)(16+8+9+18+9+45));

    p = {8,12}; borders={2,3}; BOOST_CHECK(pBS->containing_borders(p) == borders);
    p = {8,12}; conf = {2.0/3.0,5.0/6.0}; BOOST_CHECK(pBS->confidence_interval(p) == conf);
    BOOST_CHECK_EQUAL(pBS->confidence(p), (18+45)/(double)(9+18+9+45));
}

BOOST_AUTO_TEST_SUITE_END()


/*
 * Copyright 2009-2020 The VOTCA Development Team (http://www.votca.org)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
#include <libint2/initialize.h>
#define BOOST_TEST_MAIN

#define BOOST_TEST_MODULE moldenreader_test

// Third party includes
#include <boost/test/unit_test.hpp>

// VOTCA includes
#include <votca/tools/eigenio_matrixmarket.h>
#include <votca/tools/filesystem.h>

// Local VOTCA includes
#include "votca/xtp/logger.h"
#include "votca/xtp/molden.h"
#include "votca/xtp/orbitals.h"

using namespace votca::xtp;
using namespace votca;
using namespace std;

BOOST_AUTO_TEST_SUITE(molden_test)

BOOST_AUTO_TEST_CASE(moldenreader_test) {
  libint2::initialize();
  Eigen::MatrixXd coeffs_ref = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/orbitalsMOs_ref.mm");

  Orbitals orbitals_ref;
  orbitals_ref.QMAtoms().LoadFromFile(std::string(XTP_TEST_DATA_FOLDER) +
                                      "/molden/benzene.xyz");

  Logger log;
  Molden molden(log);
  molden.setBasissetInfo(
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/def2-tzvp.xml",
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/aux-def2-tzvp.xml");
  Orbitals orbitals;
  molden.parseMoldenFile(
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/benzene.molden.input",
      orbitals);

  // Check if eigenvalues are equal
  orbitals_ref.MOs().eigenvalues() =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) +
          "/molden/orbitals_eigenvalues.mm");

  BOOST_CHECK(orbitals.MOs().eigenvalues().isApprox(
      orbitals_ref.MOs().eigenvalues(), 1e-4));

  // Check if MO's are read correctly
  BOOST_CHECK(orbitals.MOs().eigenvectors().isApprox(coeffs_ref, 1e-4));

  // Check if atoms are read correctly
  BOOST_CHECK(orbitals.QMAtoms().size() == orbitals_ref.QMAtoms().size());
  for (int i = 0; i < orbitals.QMAtoms().size(); i++) {
    BOOST_CHECK(orbitals.QMAtoms()[i].getPos().isApprox(
        orbitals_ref.QMAtoms()[i].getPos(), 1e-3));
  }

  libint2::finalize();
}

BOOST_AUTO_TEST_CASE(moldenwriter_test) {
  libint2::initialize();
  // Setup orbitals object
  Orbitals orbitals_ref;
  orbitals_ref.QMAtoms().LoadFromFile(std::string(XTP_TEST_DATA_FOLDER) +
                                      "/molden/benzene.xyz");
  orbitals_ref.setNumberOfOccupiedLevels(21);
  orbitals_ref.MOs().eigenvectors() =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) + "/molden/orbitalsMOs_ref.mm");
  orbitals_ref.MOs().eigenvalues() =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) +
          "/molden/orbitals_eigenvalues.mm");
  orbitals_ref.SetupDftBasis(std::string(XTP_TEST_DATA_FOLDER) +
                             "/molden/def2-tzvp.xml");
  orbitals_ref.SetupAuxBasis(std::string(XTP_TEST_DATA_FOLDER) +
                             "/molden/aux-def2-tzvp.xml");

  // write orbitals object to molden file
  Logger log;
  Molden molden(log);
  molden.WriteFile("moldenFile.molden", orbitals_ref);

  // read in written molden file
  molden.setBasissetInfo(
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/def2-tzvp.xml",
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/aux-def2-tzvp.xml");
  Orbitals orbitals;
  molden.parseMoldenFile("moldenFile.molden", orbitals);

  Eigen::MatrixXd coeffs_ref = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/molden/orbitalsMOs_ref.mm");

  // Check if MO's are equal
  bool mos_equal = coeffs_ref.isApprox(orbitals.MOs().eigenvectors(), 1e-5);
  BOOST_CHECK(mos_equal);

  if (!mos_equal) {
    std::cout << "ref" << std::endl;
    std::cout << coeffs_ref << std::endl;

    std::cout << "results" << std::endl;
    std::cout << orbitals.MOs().eigenvectors() << std::endl;
  }

  // Check if atoms are equal
  BOOST_CHECK(orbitals.QMAtoms().size() == orbitals_ref.QMAtoms().size());
  for (int i = 0; i < orbitals.QMAtoms().size(); i++) {
    BOOST_CHECK(orbitals.QMAtoms()[i].getPos().isApprox(
        orbitals_ref.QMAtoms()[i].getPos(), 1e-3));
  }

  libint2::finalize();
}
}


#include <ros/ros.h>
#include <tf2_ros/transform_listener.h>
#include <geometry_msgs/TransformStamped.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include <eigen_conversions/eigen_msg.h>
#include <tf2_ros/transform_broadcaster.h>

int main(int argc, char **argv)
{
    ros::init(argc, argv, "tf_linker");

    ros::NodeHandle nh;

    tf2_ros::Buffer tfBuffer;
    tf2_ros::TransformListener tfListener(tfBuffer);
    tf2_ros::TransformBroadcaster br;

    Eigen::Affine3d T_base2camera;
    T_base2camera = Eigen::Translation3d(0.23, 0, 0);

    ros::Rate rate(10.0);
    while (nh.ok())
    {
        geometry_msgs::TransformStamped transformStamped_gmap2base;
        geometry_msgs::TransformStamped transformStamped_rtab2camera;
        Eigen::Affine3d T_gmap2base;
        Eigen::Affine3d T_rtab2camera;
        try
        {
            transformStamped_gmap2base = tfBuffer.lookupTransform("gmapping_map", "base_link",
                                                                  ros::Time(0));
            transformStamped_rtab2camera = tfBuffer.lookupTransform("rtabmap_map", "camera_link",
                                                                    ros::Time(0));
        }
        catch (tf2::TransformException &ex)
        {
            ROS_WARN("%s", ex.what());
            ros::Duration(0.02).sleep();
            continue;
        }

        tf::transformMsgToEigen(transformStamped_gmap2base.transform, T_gmap2base);
        tf::transformMsgToEigen(transformStamped_rtab2camera.transform, T_rtab2camera);

        Eigen::Affine3d T_gmap2rtab;
        T_gmap2rtab = T_gmap2base * T_base2camera * T_rtab2camera.inverse();

        geometry_msgs::TransformStamped transformStamped_gmap2rtab;
        tf::transformEigenToMsg(T_gmap2rtab, transformStamped_gmap2rtab.transform);

        transformStamped_gmap2rtab.header.stamp = ros::Time::now();
        transformStamped_gmap2rtab.header.frame_id = "gmapping_map";
        transformStamped_gmap2rtab.child_frame_id = "rtabmap_map";

        br.sendTransform(transformStamped_gmap2rtab);

        rate.sleep();
    }
    return 0;
};


/******************************************************************************
 * Vulkan-Interpolator: interpolator/rasterizer based on Vulkan API
 *
 * Copyright (c) 2019-2020 Dmitriy Korchemkin
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 ******************************************************************************/
#include <gtest/gtest.h>

#include <Eigen/Dense>
#include <atomic>
#include <chrono>
#include <random>
#include <thread>

#include <vulkan_interpolator/Interpolator.hpp>

TEST(VulkanInterpolator, InterpolationCorrectness) {
  std::mt19937 rng;
  std::uniform_int_distribution<int> rwh(1024, 8192), rn(10, 100000);
  std::uniform_real_distribution<float> runif(-1.f, 1.f);

  vulkan_interpolator::Interpolator interpolator;
  for (int i = 0; i < 100; ++i) {
    const int N = rn(rng);

    int width = rwh(rng);
    int height = rwh(rng);
    std::uniform_int_distribution<int> rw(0, width - 1);
    std::uniform_int_distribution<int> rh(0, height - 1);

    std::vector<float> points(2 * N), values(N);
    for (int i = 0; i < N; ++i) {
      points[i * 2] = rw(rng);
      points[i * 2 + 1] = rh(rng);
      values[i] = runif(rng);
    }

    std::vector<int> tris;
    vulkan_interpolator::Interpolator::PrepareInterpolation(N, points.data(),
                                                            tris);
    std::vector<float> data(width * height),
        data_golden(width * height, std::nan(""));
    auto ts_a = std::chrono::high_resolution_clock::now();
    interpolator.interpolate(N, points.data(), values.data(), tris.size() / 3,
                             tris.data(), width, height, width * sizeof(float),
                             data.data());
    auto ts_b = std::chrono::high_resolution_clock::now();
    int total = 0, totalIn = 0;
    double misfit = 0.;
    for (size_t tri = 0; tri < tris.size(); tri += 3) {
      Eigen::AlignedBox2d box2f;
      Eigen::Matrix3d M;
      Eigen::Vector3d b;

      for (int iii = 0; iii < 3; ++iii) {
        M.col(iii) = Eigen::Vector2d(points[tris[tri + iii] * 2],
                                     points[tris[tri + iii] * 2 + 1])
                         .homogeneous();
        b[iii] = values[tris[tri + iii]];
        box2f.extend(M.col(iii).template head<2>());
      }

      Eigen::Matrix3d Mi = M.inverse();

      for (int y = std::max(0., std::floor(box2f.min().y()));
           y <= std::min(height - 1.0, std::ceil(box2f.max().y())); ++y) {
        for (int x = std::max(0., std::floor(box2f.min().x()));
             x <= std::min(width - 1.0, std::ceil(box2f.max().x())); ++x) {
          Eigen::Vector3d pt = Eigen::Vector2d(x, y).homogeneous();
          Eigen::Vector3d gamma = Mi * pt;
          if (gamma.array().maxCoeff() >= 1 || gamma.array().minCoeff() <= 0.) {
            continue;
          }
          totalIn++;
          data_golden[y * width + x] = gamma.dot(b);
          if (std::isnan(data[y * width + x]))
            continue;
          double diff = gamma.dot(b) - data[y * width + x];
          misfit += diff * diff;
          total++;
        }
      }
    }
    auto ts_c = std::chrono::high_resolution_clock::now();
    std::cout << "GPU is "
              << double((ts_c - ts_b).count()) / (ts_b - ts_a).count()
              << "x faster" << std::endl;
    misfit = std::sqrt(misfit / total);
    static double minRatio = 1., maxMisfit = 0.;
    minRatio = std::min(minRatio, double(total) / totalIn);
    maxMisfit = std::max(maxMisfit, misfit);
    EXPECT_GT(total, 0.999 * totalIn);
    EXPECT_LT(misfit, 1e-6);
  }
}

TEST(VulkanInterpolator, ParallelBurnTest) {
  vulkan_interpolator::Interpolator interpolator;

  std::mt19937 rng_global;
  std::vector<std::thread> workers;
  std::atomic<int> alive(0);
  int N = 100000;
  auto start_all = std::chrono::high_resolution_clock::now();
  for (int i = 0; i < N; ++i) {
    while (alive > 200)
      std::this_thread::yield();

    const auto seed = rng_global();
    workers.emplace_back([&interpolator, &alive,seed]() {
      ++alive;
      std::mt19937 rng(seed);
      std::uniform_int_distribution<int> rwh(2048, 4096), rn(1000, 10000);
      std::uniform_real_distribution<float> runif(-1.f, 1.f);

      const int N = rn(rng);

      int width = rwh(rng);
      int height = rwh(rng);
      std::uniform_int_distribution<int> rw(0, width - 1);
      std::uniform_int_distribution<int> rh(0, height - 1);

      std::vector<float> points(2 * N), values(N);
      for (int i = 0; i < N; ++i) {
        points[i * 2] = rw(rng);
        points[i * 2 + 1] = rh(rng);
        values[i] = runif(rng);
      }
      std::vector<float> data(width * height);
      auto start = std::chrono::high_resolution_clock::now();
      interpolator.interpolate(N, points.data(), values.data(), width, height,
                               width * sizeof(float), data.data());
      auto stop = std::chrono::high_resolution_clock::now();
      std::cout << (1e9 / (stop - start).count()) << "FPS" << std::endl;
      --alive;
    });
  }
  for (auto &worker : workers)
    worker.join();
  auto stop_all = std::chrono::high_resolution_clock::now();
  std::cout << (1e9 / (stop_all - start_all).count()) << "FPS-all" << std::endl;
}


#include <iostream>
#include <boost/mpi.hpp>

namespace mpi = boost::mpi;

bool test_poisson(mpi::communicator&);
bool test_dot(mpi::communicator&);
bool test_solution_wrapper(mpi::communicator&);
bool test_scatter_mesh(mpi::communicator&);
bool test_helmholtz(mpi::communicator&);
bool test_v_wrapper(mpi::communicator&);
bool test_schrodinger_bdf2(mpi::communicator&);
bool test_schrodinger_dirk4s3(mpi::communicator&);
bool test_schrodinger_dirk4s5(mpi::communicator&);

int main(int argc, char* argv[])
{
    mpi::environment env(argc, argv);
    mpi::communicator comm;

    bool success = true;

    success = test_solution_wrapper(comm)    && success;
    success = test_v_wrapper(comm)           && success;
    success = test_scatter_mesh(comm)        && success;
    success = test_dot(comm)                 && success;
    success = test_poisson(comm)             && success;
    success = test_helmholtz(comm)           && success;
    success = test_schrodinger_bdf2(comm)    && success;
    success = test_schrodinger_dirk4s3(comm) && success;
    success = test_schrodinger_dirk4s5(comm) && success;

    if (success)
        if (comm.rank() == 0)
            std::cout << "all tests passed :)\n";

    return 0;
}