#include <armadillo>
#include <iostream>
#include <fstream>
#include <stdio.h>
#include <string>
#include "mcpdft.h"
#include "openrdmConfig.h"

#ifdef WITH_OPENMP // _OPENMP
   #include <omp.h>
#endif

namespace mcpdft {

   MCPDFT::MCPDFT(std::string test_case)  { common_init(test_case); }
   MCPDFT::MCPDFT() {}
   MCPDFT::~MCPDFT() {}

   void MCPDFT::common_init(std::string test_case) {
       print_banner();
       read_grids_from_file(test_case);
       read_orbitals_from_file(test_case);
       read_energies_from_file(test_case);
       if(test_case == "h2_tpbe_sto3g") {
	  is_gga_ = true;
	  read_gradients_from_file(test_case);
       }else{
	  is_gga_ = false;
       }

       read_opdm_from_file(test_case);
       // read_cmat_from_file();
   }

   void MCPDFT::build_rho() {
      build_density_functions();
      if ( is_gga_ ) {
         build_density_gradients();
      }
   }

   void MCPDFT::build_density_functions() {
      int nbfs = get_nbfs();
      size_t npts = get_npts();
      arma::mat phi(get_phi());
      arma::mat D1a(get_D1a());
      arma::mat D1b(get_D1b());
      arma::vec W(get_w());
      arma::vec rhoa(npts, arma::fill::zeros);
      arma::vec rhob(npts, arma::fill::zeros);
      arma::vec rho(npts, arma::fill::zeros);
      double dum_a = 0.0;
      double dum_b = 0.0;
      double dum_tot = 0.0;
      size_t chunk_size = 0;
      int p{0}, mu{0}, nu{0};
      #ifdef WITH_OPENMP
         int nthrds{0};
         printf("\n++++++++++++++++++++++++++++++++++++++++++++++++++++++\n");
         printf("                   *** Warning ***\n");
         printf("   Calculating the density (gradients) using OpenMP");
         printf("\n++++++++++++++++++++++++++++++++++++++++++++++++++++++\n");
         nthrds = omp_get_max_threads();
         nthrds /= 2;
         omp_set_num_threads(nthrds);
      #endif

      #pragma omp parallel default(shared) \
                           private(p, mu, nu)
      {
         #pragma omp for schedule(static) \
                         reduction(+:dum_a, dum_b, dum_tot) \
	                 nowait
            for(p = 0; p < npts; p++) {
               double tempa = 0.0;
               double tempb = 0.0;
               #pragma omp parallel for schedule(static) \
                                        reduction(+:tempa,tempb) \
	                                num_threads(2) \
                                        collapse(2)
                  for(mu = 0; mu < nbfs; mu++) {
                     for(int nu = 0; nu < nbfs; nu++) {
                        tempa += D1a(mu, nu) * phi(p, mu) * phi(p, nu);
                        tempb += D1b(mu, nu) * phi(p, mu) * phi(p, nu);
                     }
                  }
                  rhoa(p) = tempa;
                  rhob(p) = tempb;
                  rho(p) = rhoa(p) + rhob(p);

                  dum_a += rhoa(p) * W(p);
                  dum_b += rhob(p) * W(p);
                  dum_tot += ( rhoa(p) + rhob(p) ) * W(p) ;
            } /* end of omp parallel for loop */
      } /* end of omp parallel region */
      set_rhoa(rhoa);
      set_rhob(rhob);
      set_rho(rho);

      printf("\n");
      printf("  Integrated total density = %20.12lf\n",dum_tot);
      printf("  Integrated alpha density = %20.12lf\n",dum_a);
      printf("  Integrated beta density  = %20.12lf\n",dum_b);
      printf("\n");
   }

   void MCPDFT::build_density_gradients() {
      int nbfs = get_nbfs();
      size_t npts = get_npts();
      arma::mat phi(get_phi());
      arma::mat D1a(get_D1a());
      arma::mat D1b(get_D1b());
      int p{0}, nu{0}, sigma{0};
      arma::mat phi_x(get_phi_x());
      arma::mat phi_y(get_phi_y());
      arma::mat phi_z(get_phi_z());
      arma::vec rho_a_x(npts, arma::fill::zeros);
      arma::vec rho_b_x(npts, arma::fill::zeros);
      arma::vec rho_a_y(npts, arma::fill::zeros);
      arma::vec rho_b_y(npts, arma::fill::zeros);
      arma::vec rho_a_z(npts, arma::fill::zeros);
      arma::vec rho_b_z(npts, arma::fill::zeros);
      arma::vec sigma_aa(npts, arma::fill::zeros);
      arma::vec sigma_ab(npts, arma::fill::zeros);
      arma::vec sigma_bb(npts, arma::fill::zeros);
      #pragma omp parallel default(shared) \
                           private(p, nu, sigma)
      {
         #pragma omp for schedule(static)
             for (int p = 0; p < npts; p++) {
                 double duma_x = 0.0;
                 double dumb_x = 0.0;
                 double duma_y = 0.0;
                 double dumb_y = 0.0;
                 double duma_z = 0.0;
                 double dumb_z = 0.0;
                 #pragma omp parallel for schedule(static) \
                   	                  reduction(+:duma_x, duma_y, duma_z,\
                   		                      dumb_x, dumb_y, dumb_z)\
                                          shared(p, nbfs, \
    				                 sigma_aa, sigma_bb, sigma_ab, \
    				                 rho_a_x, rho_a_y, rho_a_z, \
    				                 rho_b_x, rho_b_y, rho_b_z) \
    			                  num_threads(2) \
                                          collapse(2)
                    for (int sigma = 0; sigma < nbfs; sigma++) {
                        for (int nu = 0; nu < nbfs; nu++) {
                            duma_x += ( phi_x(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_x(p, nu) ) * D1a(sigma, nu);
                            dumb_x += ( phi_x(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_x(p, nu) ) * D1b(sigma, nu);
                            duma_y += ( phi_y(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_y(p, nu) ) * D1a(sigma, nu);
                            dumb_y += ( phi_y(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_y(p, nu) ) * D1b(sigma, nu);
                            duma_z += ( phi_z(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_z(p, nu) ) * D1a(sigma, nu);
                            dumb_z += ( phi_z(p, sigma) * phi(p, nu) + phi(p, sigma) * phi_z(p, nu) ) * D1b(sigma, nu);
                        }
                    }
                    rho_a_x(p) = duma_x;
                    rho_b_x(p) = dumb_x;
                    rho_a_y(p) = duma_y;
                    rho_b_y(p) = dumb_y;
                    rho_a_z(p) = duma_z;
                    rho_b_z(p) = dumb_z;
                    sigma_aa(p) = ( rho_a_x(p) * rho_a_x(p) ) +  ( rho_a_y(p) * rho_a_y(p) ) + ( rho_a_z(p) * rho_a_z(p) );
                    sigma_bb(p) = ( rho_b_x(p) * rho_b_x(p) ) +  ( rho_b_y(p) * rho_b_y(p) ) + ( rho_b_z(p) * rho_b_z(p) );
                    sigma_ab(p) = ( rho_a_x(p) * rho_b_x(p) ) +  ( rho_a_y(p) * rho_b_y(p) ) + ( rho_a_z(p) * rho_b_z(p) );
             }
      }
      set_rhoa_x(rho_a_x);
      set_rhob_x(rho_b_x);
      set_rhoa_y(rho_a_y);
      set_rhob_y(rho_b_y);
      set_rhoa_z(rho_a_z);
      set_rhob_z(rho_b_z);
      set_sigma_aa(sigma_aa);
      set_sigma_ab(sigma_ab);
      set_sigma_bb(sigma_bb);
   }

   void MCPDFT::build_pi(const arma::mat &D2ab) {
      build_ontop_pair_density(D2ab);
      if ( is_gga_ ) {
         build_ontop_pair_density_gradients(D2ab);
      }
   }

   void MCPDFT::build_ontop_pair_density(const arma::mat &D2ab) {
      int nbfs = get_nbfs();
      size_t npts = get_npts();
      arma::vec temp(npts);
      arma::mat phi(get_phi());
      int p{0},
	  mu{0}, nu{0},
	  lambda{0}, sigma{0};
      for (int p = 0; p < npts; p++) {
          double dum = 0.0;
          // pi(r,r) = D(mu,nu; lambda,sigma) * phi(r,mu) * phi(r,nu) * phi(r,lambda) * phi(r,sigma)
          #pragma omp parallel for default(shared) \
          	                       private(mu,nu,lambda,sigma) \
          	                       reduction(+:dum) \
          	                       collapse(4)
             for (int mu = 0; mu < nbfs; mu++) {
                 for (int nu = 0; nu < nbfs; nu++) {
                     for (int lambda = 0; lambda < nbfs; lambda++) {
                         for (int sigma = 0; sigma < nbfs; sigma++) {
                             dum += phi(p, mu) * phi(p, nu) * phi(p, lambda) * phi(p, sigma) * D2ab(nu*nbfs+mu, sigma*nbfs+lambda);
                         }
                     }
                 }
             } /* end of omp parallel for loop */
             temp(p) = dum;
      }
      set_pi(temp);
   }

   void MCPDFT::build_ontop_pair_density_gradients(const arma::mat &D2ab) {
      int nbfs = get_nbfs();
      size_t npts = get_npts();
      int p{0},
	  mu{0}, nu{0},
	  lambda{0}, sigma{0};
      arma::mat phi(get_phi());
      arma::mat phi_x(get_phi_x());
      arma::mat phi_y(get_phi_y());
      arma::mat phi_z(get_phi_z());
      arma::vec pi_x(npts, arma::fill::zeros);
      arma::vec pi_y(npts, arma::fill::zeros);
      arma::vec pi_z(npts, arma::fill::zeros);
      #pragma omp parallel default(shared) \
                           private(p, mu, nu, lambda, sigma)
      {
         #pragma omp for schedule(static)
            for (int p = 0; p < npts; p++) {
                double dum_x = 0.0;
                double dum_y = 0.0;
                double dum_z = 0.0;
                // pi(r) = D(mu,nu; lambda,sigma) * phi(r,mu) * phi(r,nu) * phi(r,lambda) * phi(r,sigma)
                #pragma omp parallel for schedule(static) \
                  	                 reduction(+:dum_x, dum_y, dum_z) \
    		                         num_threads(2) \
                                         collapse(4)
                for (int mu = 0; mu < nbfs; mu++) {
                    for (int nu = 0; nu < nbfs; nu++) {
                        for (int lambda = 0; lambda < nbfs; lambda++) {
                            for (int sigma = 0; sigma < nbfs; sigma++) {
                                dum_x += ( phi_x(p, mu) * phi(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi_x(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi_x(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi(p, sigma) * phi_x(p, nu) ) * D2ab(nu*nbfs+mu, sigma*nbfs+lambda);

                                dum_y += ( phi_y(p, mu) * phi(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi_y(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi_y(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi(p, sigma) * phi_y(p, nu) ) * D2ab(nu*nbfs+mu, sigma*nbfs+lambda);

                                dum_z += ( phi_z(p, mu) * phi(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi_z(p, lambda) * phi(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi_z(p, sigma) * phi(p, nu) +
                                           phi(p, mu) * phi(p, lambda) * phi(p, sigma) * phi_z(p, nu) ) * D2ab(nu*nbfs+mu, sigma*nbfs+lambda);
                            }
                        }
                    }
                }
                pi_x(p) = dum_x;
                pi_y(p) = dum_y;
                pi_z(p) = dum_z;
            }
      }
      set_pi_x(pi_x);
      set_pi_z(pi_y);
      set_pi_z(pi_z);
   }

   void MCPDFT::build_R() {
        double tol = 1.0e-20;
        size_t npts = get_npts();
        arma::vec temp(npts);
        arma::vec pi(get_pi());
        arma::vec rho(get_rho());
        #pragma parallel for schedule(static) \
	                     private(p) \
                             shared(npts, pi, rho, temp)
           for (int p = 0; p < npts; p++) {
               temp(p) = 4.0 * pi(p) / ( rho(p) * rho(p) );
           }
        set_R(temp);
   }

   void MCPDFT::translate() {
     translate_density();
     if ( is_gga_ ) {
        translate_density_gradients();
     }
   }

   void MCPDFT::translate_density() {
      double tol = 1.0e-20;
      size_t npts = get_npts();
      arma::vec rho_vec(get_rho());
      arma::vec pi_vec(get_pi());
      arma::vec R_vec(get_R());
      arma::vec W(get_w());
      arma::vec tr_rhoa(npts);
      arma::vec tr_rhob(npts);
      double dum_a = 0.0;
      double dum_b = 0.0;
      double dum_tot = 0.0;
      double rho = 0.0;
      double pi = 0.0;
      double zeta = 0.0;
      double R = 0.0;
      #pragma parallel for schedule(static) \
                           default(shared) \
                           private(p, zeta, pi, rho, R) \
                           reduction(+:dum_a, dum_b, dum_tot)
         for (int p = 0; p < npts; p++) {
             rho = rho_vec(p);
             pi = pi_vec(p);
             zeta = 0.0;
             R = 0.0;
             if ( !(rho < tol) && !(pi < 0.0) ) {
                R = R_vec(p);
                if ( (1.0 - R) > tol ) {
                   zeta = sqrt(1.0 - R);
                }else{
                     zeta = 0.0;
                }
                tr_rhoa(p) = (1.0 + zeta) * (rho/2.0);
                tr_rhob(p) = (1.0 - zeta) * (rho/2.0);
             }else {
                    tr_rhoa(p) = 0.0;
                    tr_rhob(p) = 0.0;
             }
             dum_a += tr_rhoa(p) * W(p);
             dum_b += tr_rhob(p) * W(p);
             dum_tot += ( tr_rhoa(p) + tr_rhob(p) ) * W(p) ;
         }
      set_tr_rhoa(tr_rhoa);
      set_tr_rhob(tr_rhob);

      printf("\n");
      printf("  Integrated translated total density = %20.12lf\n",dum_tot);
      printf("  Integrated translated alpha density = %20.12lf\n",dum_a);
      printf("  Integrated translated beta density  = %20.12lf\n",dum_b);
      printf("\n");
   }

   void MCPDFT::translate_density_gradients() {
     double tol = 1.0e-20;
     size_t npts = get_npts();
     arma::vec rho_vec(get_rho());
     arma::vec pi_vec(get_pi());
     arma::vec R_vec(get_R());
     double rho = 0.0;
     double pi = 0.0;
     double zeta = 0.0;
     double R = 0.0;
     double rho_x = 0.0;
     double rho_y = 0.0;
     double rho_z = 0.0;
     arma::vec rho_a_x(get_rhoa_x());
     arma::vec rho_a_y(get_rhoa_y());
     arma::vec rho_a_z(get_rhoa_z());
     arma::vec rho_b_x(get_rhob_x());
     arma::vec rho_b_y(get_rhob_y());
     arma::vec rho_b_z(get_rhob_z());
     arma::vec tr_rho_a_x(npts,  arma::fill::zeros);
     arma::vec tr_rho_b_x(npts,  arma::fill::zeros);
     arma::vec tr_rho_a_y(npts,  arma::fill::zeros);
     arma::vec tr_rho_b_y(npts,  arma::fill::zeros);
     arma::vec tr_rho_a_z(npts,  arma::fill::zeros);
     arma::vec tr_rho_b_z(npts,  arma::fill::zeros);
     arma::vec tr_sigma_aa(npts, arma::fill::zeros);
     arma::vec tr_sigma_ab(npts, arma::fill::zeros);
     arma::vec tr_sigma_bb(npts, arma::fill::zeros);
     #pragma parallel for schedule(static) \
                          default(shared) \
                          private(p, zeta, pi, rho, R,\
		                  rho_x, rho_y, rho_z)
        for (int p = 0; p < npts; p++) {
            rho = rho_vec(p);
            pi = pi_vec(p);
            rho_x = rho_a_x(p) + rho_b_x(p);
            rho_y = rho_a_y(p) + rho_b_y(p);
            rho_z = rho_a_z(p) + rho_b_z(p);
            zeta = 0.0;
            R = 0.0;
            if ( !(rho < tol) && !(pi < 0.0) ) {
               R = R_vec(p);
               if ( (1.0 - R) > tol )  {
                  zeta = sqrt(1.0 - R);
               }else{
                    zeta = 0.0;
               }
               tr_rho_a_x(p) = (1.0 + zeta) * (rho_x/2.0);
               tr_rho_b_x(p) = (1.0 - zeta) * (rho_x/2.0);

               tr_rho_a_y(p) = (1.0 + zeta) * (rho_y/2.0);
               tr_rho_b_y(p) = (1.0 - zeta) * (rho_y/2.0);

               tr_rho_a_z(p) = (1.0 + zeta) * (rho_z/2.0);
               tr_rho_b_z(p) = (1.0 - zeta) * (rho_z/2.0);
            }else {
                  tr_rho_a_x(p) = 0.0;
                  tr_rho_b_x(p) = 0.0;
                  tr_rho_a_y(p) = 0.0;
                  tr_rho_b_y(p) = 0.0;
                  tr_rho_a_z(p) = 0.0;
                  tr_rho_b_z(p) = 0.0;
            }
            tr_sigma_aa(p) = (tr_rho_a_x(p) * tr_rho_a_x(p)) + (tr_rho_a_y(p) * tr_rho_a_y(p)) + (tr_rho_a_z(p) * tr_rho_a_z(p));
            tr_sigma_ab(p) = (tr_rho_a_x(p) * tr_rho_b_x(p)) + (tr_rho_a_y(p) * tr_rho_b_y(p)) + (tr_rho_a_z(p) * tr_rho_b_z(p));
            tr_sigma_bb(p) = (tr_rho_b_x(p) * tr_rho_b_x(p)) + (tr_rho_b_y(p) * tr_rho_b_y(p)) + (tr_rho_b_z(p) * tr_rho_b_z(p));
        }
     set_tr_sigma_aa(tr_sigma_aa);
     set_tr_sigma_ab(tr_sigma_ab);
     set_tr_sigma_bb(tr_sigma_bb);
   }

   void MCPDFT::fully_translate() {
     fully_translate_density();
     if ( is_gga_ ) {
        fully_translate_density_gradients();
     }
   }

   void MCPDFT::fully_translate_density(){
        double tol = 1.0e-20;
        double const R0 = 0.9;
        double const R1 = 1.15;
        double const A = -475.60656009;
        double const B = -379.47331922;
        double const C = -85.38149682;

        size_t npts = get_npts();
        arma::vec rho_vec(get_rho());
        arma::vec pi_vec(get_pi());
        arma::vec R_vec(get_R());
        arma::vec W(get_w());
        arma::vec tr_rhoa(npts);
        arma::vec tr_rhob(npts);
        arma::vec tr_rho(npts);

        double rho = 0.0;
        double pi = 0.0;
        double zeta = 0.0;
        double R = 0.0;
        double DelR = 0.0;
        double temp_tot = 0.0;
        double temp_a = 0.0;
        double temp_b = 0.0;
        #pragma parallel for schedule(static) \
                             default(shared) \
                             private(p, zeta, pi, rho) \
                             reduction(+:dum_a, dum_b, dum_tot)
        for (int p = 0; p < npts; p++) {
            zeta = 0.0;
            R = 0.0;
            rho = rho_vec(p);
            pi = pi_vec(p);
            DelR = R_vec(p) - R1;
            if ( !(rho < tol) && !(pi < 0.0) ) {
               R = R_vec(p);
               if ( ((1.0 - R) > tol) && ( R < R0 ) ) {
                  zeta = sqrt(1.0 - R);
               }else if( !(R < R0) && !(R > R1) ) {
                       zeta = A * pow(DelR, 5.0) + B * pow(DelR, 4.0) + C * pow(DelR, 3.0);
               }else if( R > R1 ) {
                       zeta = 0.0;
               }
               tr_rhoa(p) = (1.0 + zeta) * (rho/2.0);
               tr_rhob(p) = (1.0 - zeta) * (rho/2.0);
            }else{
                 tr_rhoa(p) = 0.0;
                 tr_rhob(p) = 0.0;
            }
            temp_a += tr_rhoa(p) * W(p);
            temp_b += tr_rhob(p) * W(p);
            temp_tot += ( tr_rhob(p) + tr_rhoa(p) ) * W(p);
        }
        set_tr_rhoa(tr_rhoa);
        set_tr_rhob(tr_rhob);

        printf("\n");
        printf("      Integrated fully translated total density = %20.12lf\n",temp_tot);
        printf("      Integrated fully translated alpha density = %20.12lf\n",temp_a);
        printf("      Integrated fully translated beta density  = %20.12lf\n",temp_b);
        printf("\n");
   }

   void MCPDFT::fully_translate_density_gradients(){
        double tol = 1.0e-20;
        size_t npts = get_npts();

        arma::vec rho_vec(get_rho());
        arma::vec pi_vec(get_pi());
        arma::vec R_vec(get_R());
        arma::vec W(get_w());
        arma::vec tr_rhoa(npts);
        arma::vec tr_rhob(npts);
        arma::vec tr_rho(npts);

        double const R0 = 0.9;
        double const R1 = 1.15;
        double const A = -475.60656009;
        double const B = -379.47331922;
        double const C = -85.38149682;

        arma::vec tr_rho_a_x(npts,  arma::fill::zeros);
        arma::vec tr_rho_b_x(npts,  arma::fill::zeros);
        arma::vec tr_rho_a_y(npts,  arma::fill::zeros);
        arma::vec tr_rho_b_y(npts,  arma::fill::zeros);
        arma::vec tr_rho_a_z(npts,  arma::fill::zeros);
        arma::vec tr_rho_b_z(npts,  arma::fill::zeros);
        arma::vec tr_sigma_aa(npts, arma::fill::zeros);
        arma::vec tr_sigma_ab(npts, arma::fill::zeros);
        arma::vec tr_sigma_bb(npts, arma::fill::zeros);
        double rho_x = 0.0;
        double rho_y = 0.0;
        double rho_z = 0.0;
        double rho = 0.0;
        double pi = 0.0;
        double DelR = 0.0;
        double zeta = 0.0;
        double R = 0.0;
        double temp_tot = 0.0;
        double temp_a = 0.0;
        double temp_b = 0.0;
        #pragma parallel for schedule(static) \
                             default(shared) \
                             private(p, zeta, pi, rho, R, DelR,\
                                     rho_x, rho_y, rho_z)
           for (int p = 0; p < npts; p++) {
               rho_x = rho_a_x_(p) + rho_b_x_(p);
               rho_y = rho_a_y_(p) + rho_b_y_(p);
               rho_z = rho_a_z_(p) + rho_b_z_(p);
               rho = rho_vec(p);
               pi = pi_vec(p);
               DelR = R_vec(p) - R1;
               zeta = 0.0;
               R = 0.0;
               if ( !(rho < tol) && !(pi < 0.0) ) {
                   R = R_vec(p);
                   if ( ((1.0 - R) > tol) && ( R < R0 ) ) {
                      zeta = sqrt(1.0 - R);
                      tr_rho_a_x(p) = (1.0 + zeta) * (rho_x/2.0) + (R * rho_x) / (2.0*zeta) - pi_x_(p) / (rho*zeta);
                      tr_rho_b_x(p) = (1.0 - zeta) * (rho_x/2.0) - (R * rho_x) / (2.0*zeta) + pi_x_(p) / (rho*zeta);
                      tr_rho_a_y(p) = (1.0 + zeta) * (rho_y/2.0) + (R * rho_y) / (2.0*zeta) - pi_y_(p) / (rho*zeta);
                      tr_rho_b_y(p) = (1.0 - zeta) * (rho_y/2.0) - (R * rho_y) / (2.0*zeta) + pi_y_(p) / (rho*zeta);
                      tr_rho_a_z(p) = (1.0 + zeta) * (rho_z/2.0) + (R * rho_z) / (2.0*zeta) - pi_z_(p) / (rho*zeta);
                      tr_rho_b_z(p) = (1.0 - zeta) * (rho_z/2.0) - (R * rho_z) / (2.0*zeta) + pi_z_(p) / (rho*zeta);
                   }else if( !(R < R0) && !(R > R1) ) {
                           zeta = A * pow(DelR, 5.0) + B * pow(DelR, 4.0) + C * pow(DelR, 3.0);

                           tr_rho_a_x(p) = (1.0 + zeta) * (rho_x/2.0)
                                          + (A * pow(DelR, 4.0)) * ( (10.0 * pi_x_(p) / rho) - (5.0 * R * rho_x) )
                                          + (B * pow(DelR, 3.0)) * ( (8.0  * pi_x_(p) / rho) - (4.0 * R * rho_x) )
                                          + (C * pow(DelR, 2.0)) * ( (6.0  * pi_x_(p) / rho) - (3.0 * R * rho_x) );

                           tr_rho_b_x(p) = (1.0 - zeta) * (rho_x/2.0)
                                          + (A * pow(DelR, 4.0)) * (-(10.0 * pi_x_(p) / rho) + (5.0 * R * rho_x) )
                                          + (B * pow(DelR, 3.0)) * (-(8.0  * pi_x_(p) / rho) + (4.0 * R * rho_x) )
                                          + (C * pow(DelR, 2.0)) * (-(6.0  * pi_x_(p) / rho) + (3.0 * R * rho_x) );

                           tr_rho_a_y(p) = (1.0 + zeta) * (rho_y/2.0)
                                          + (A * pow(DelR, 4.0)) * ( (10.0 * pi_y_(p) / rho) - (5.0 * R * rho_y) )
                                          + (B * pow(DelR, 3.0)) * ( (8.0  * pi_y_(p) / rho) - (4.0 * R * rho_y) )
                                          + (C * pow(DelR, 2.0)) * ( (6.0  * pi_y_(p) / rho) - (3.0 * R * rho_y) );

                           tr_rho_b_y(p) = (1.0 - zeta) * (rho_y/2.0)
                                          + (A * pow(DelR, 4.0)) * (-(10.0 * pi_y_(p) / rho) + (5.0 * R * rho_y) )
                                          + (B * pow(DelR, 3.0)) * (-(8.0  * pi_y_(p) / rho) + (4.0 * R * rho_y) )
                                          + (C * pow(DelR, 2.0)) * (-(6.0  * pi_y_(p) / rho) + (3.0 * R * rho_y) );

                           tr_rho_a_z(p) = (1.0 + zeta) * (rho_z/2.0)
                                          + (A * pow(DelR, 4.0)) * ( (10.0 * pi_z_(p) / rho) - (5.0 * R * rho_z) )
                                          + (B * pow(DelR, 3.0)) * ( (8.0  * pi_z_(p) / rho) - (4.0 * R * rho_z) )
                                          + (C * pow(DelR, 2.0)) * ( (6.0  * pi_z_(p) / rho) - (3.0 * R * rho_z) );

                           tr_rho_b_z(p) = (1.0 - zeta) * (rho_z/2.0)
                                           + (A * pow(DelR, 4.0)) * (-(10.0 * pi_z_(p) / rho) + (5.0 * R * rho_z) )
                                           + (B * pow(DelR, 3.0)) * (-(8.0  * pi_z_(p) / rho) + (4.0 * R * rho_z) )
                                           + (C * pow(DelR, 2.0)) * (-(6.0  * pi_z_(p) / rho) + (3.0 * R * rho_z) );
                   }else if( R > R1 ) {
                           zeta = 0.0;
                           tr_rho_a_x(p) = (1.0 + zeta) * (rho_x/2.0);
                           tr_rho_b_x(p) = (1.0 - zeta) * (rho_x/2.0);
                           tr_rho_a_y(p) = (1.0 + zeta) * (rho_y/2.0);
                           tr_rho_b_y(p) = (1.0 - zeta) * (rho_y/2.0);
                           tr_rho_a_z(p) = (1.0 + zeta) * (rho_z/2.0);
                           tr_rho_b_z(p) = (1.0 - zeta) * (rho_z/2.0);
                   }
               }else{
                   tr_rho_a_x(p) = 0.0;
                   tr_rho_b_x(p) = 0.0;
                   tr_rho_a_y(p) = 0.0;
                   tr_rho_b_y(p) = 0.0;
                   tr_rho_a_z(p) = 0.0;
                   tr_rho_b_z(p) = 0.0;
               }
               tr_sigma_aa(p) = (tr_rho_a_x(p) * tr_rho_a_x(p)) + (tr_rho_a_y(p) * tr_rho_a_y(p)) + (tr_rho_a_z(p) * tr_rho_a_z(p));
               tr_sigma_ab(p) = (tr_rho_a_x(p) * tr_rho_b_x(p)) + (tr_rho_a_y(p) * tr_rho_b_y(p)) + (tr_rho_a_z(p) * tr_rho_b_z(p));
               tr_sigma_bb(p) = (tr_rho_b_x(p) * tr_rho_b_x(p)) + (tr_rho_b_y(p) * tr_rho_b_y(p)) + (tr_rho_b_z(p) * tr_rho_b_z(p));
           }
        set_tr_sigma_aa(tr_sigma_aa);
        set_tr_sigma_ab(tr_sigma_ab);
        set_tr_sigma_bb(tr_sigma_bb);
   }

   void MCPDFT::build_opdm() {
      // fetching the number of basis functions
      int nbfs;
      nbfs = get_nbfs();

      // getting the AO->MO transformation matrix C
      arma::mat ca(get_cmat());
      arma::mat cb(get_cmat());

      // building the 1-electron reduced density matrices (1RDMs)
      arma::mat D1a(nbfs, nbfs, arma::fill::zeros);
      arma::mat D1b(nbfs, nbfs, arma::fill::zeros);
      for (int mu = 0; mu < nbfs; mu++) {
          for (int nu = 0; nu < nbfs; nu++) {
              double duma = 0.0;
              double dumb = 0.0;
              for (int i = 0; i < nbfs/2; i++) {
                   duma += ca(mu, i) * ca(nu, i);
                   dumb += cb(mu, i) * cb(nu, i);
              }
              D1a(mu, nu) = duma;
              D1b(mu, nu) = dumb;
         }
      }
      // D1a(0,0) = 1.0;
      // D1b(0,0) = 1.0;
      D1a.print("D1a = ");
      D1b.print("D1b = ");
      set_D1a(D1a);
      set_D1b(D1b);
   }

   void MCPDFT::build_tpdm() {
      // fetching the number of basis functions
      int nbfs  = get_nbfs();
      int nbfs2 = nbfs * nbfs;

      arma::mat D1a(get_D1a());
      arma::mat D1b(get_D1b());

      arma::mat D2ab(nbfs2, nbfs2, arma::fill::zeros);
      D2ab = arma::kron(D1a,D1b);
      // D2ab.print("D2ab = ");
      set_D2ab(D2ab);
   }

   bool MCPDFT::is_gga() const { return is_gga_; }
   size_t MCPDFT::get_npts() const { return npts_; }
   int    MCPDFT::get_nbfs() const { return nbfs_; }
   arma::vec MCPDFT::get_w() const { return w_; }
   arma::vec MCPDFT::get_x() const { return x_; }
   arma::vec MCPDFT::get_y() const { return y_; }
   arma::vec MCPDFT::get_z() const { return z_; }
   arma::mat MCPDFT::get_phi() const { return phi_; }
   arma::mat MCPDFT::get_phi_x() const { return phi_x_; }
   arma::mat MCPDFT::get_phi_y() const { return phi_y_; }
   arma::mat MCPDFT::get_phi_z() const { return phi_z_; }
   double MCPDFT::get_eref() const { return eref_; }
   double MCPDFT::get_eclass()  const { return eclass_; }
   arma::mat MCPDFT::get_cmat() const { return cmat_; }
   arma::mat MCPDFT::get_D1a()  const { return D1a_ ; }
   arma::mat MCPDFT::get_D1b()  const { return D1b_ ; }
   arma::mat MCPDFT::get_D2ab()  const { return D2ab_ ; }
   arma::vec MCPDFT::get_rhoa() const { return rho_a_; }
   arma::vec MCPDFT::get_rhoa_x() const { return rho_a_x_; }
   arma::vec MCPDFT::get_rhoa_y() const { return rho_a_y_; }
   arma::vec MCPDFT::get_rhoa_z() const { return rho_a_z_; }
   arma::vec MCPDFT::get_rhob() const { return rho_b_; }
   arma::vec MCPDFT::get_rhob_x() const { return rho_b_x_; }
   arma::vec MCPDFT::get_rhob_y() const { return rho_b_y_; }
   arma::vec MCPDFT::get_rhob_z() const { return rho_b_z_; }
   arma::vec MCPDFT::get_rho() const { return rho_; }
   arma::vec MCPDFT::get_tr_rhoa() const { return tr_rho_a_; }
   arma::vec MCPDFT::get_tr_rhob() const { return tr_rho_b_; }
   arma::vec MCPDFT::get_tr_rho() const { return tr_rho_; }
   arma::vec MCPDFT::get_pi() const { return pi_; }
   arma::vec MCPDFT::get_R() const { return R_; }
   arma::vec MCPDFT::get_sigma_aa() const { return sigma_aa_; }
   arma::vec MCPDFT::get_sigma_ab() const { return sigma_ab_; }
   arma::vec MCPDFT::get_sigma_bb() const { return sigma_bb_; }
   arma::vec MCPDFT::get_tr_sigma_aa() const { return tr_sigma_aa_; }
   arma::vec MCPDFT::get_tr_sigma_ab() const { return tr_sigma_ab_; }
   arma::vec MCPDFT::get_tr_sigma_bb() const { return tr_sigma_bb_; }

   void MCPDFT::set_npts(const size_t npts) { npts_ = npts; }
   void MCPDFT::set_nbfs(const int nbfs)    { nbfs_ = nbfs; }
   void MCPDFT::set_w(const arma::vec &w) { w_ = w; }
   void MCPDFT::set_x(const arma::vec &x) { x_ = x; }
   void MCPDFT::set_y(const arma::vec &y) { y_ = y; }
   void MCPDFT::set_z(const arma::vec &z) { z_ = z; }
   void MCPDFT::set_phi(const arma::mat &phi) { phi_ = phi; }
   void MCPDFT::set_phi_x(const arma::mat &phi_x) { phi_x_ = phi_x; }
   void MCPDFT::set_phi_y(const arma::mat &phi_y) { phi_y_ = phi_y; }
   void MCPDFT::set_phi_z(const arma::mat &phi_z) { phi_z_ = phi_z; }
   void MCPDFT::set_eref(const double eref) { eref_ = eref; }
   void MCPDFT::set_eclass(const double eclass) { eclass_ = eclass; }
   void MCPDFT::set_cmat(const arma::mat &cmat) { cmat_ = cmat; }
   void MCPDFT::set_D1a(const arma::mat &D1a) { D1a_ = D1a; }
   void MCPDFT::set_D1b(const arma::mat &D1b) { D1b_ = D1b; }
   void MCPDFT::set_D2ab(const arma::mat &D2ab) { D2ab_ = D2ab; }
   void MCPDFT::set_rhoa(const arma::vec &rhoa) { rho_a_ = rhoa; }
   void MCPDFT::set_rhoa_x(const arma::vec &rhoa_x) { rho_a_x_ = rhoa_x; }
   void MCPDFT::set_rhoa_y(const arma::vec &rhoa_y) { rho_a_y_ = rhoa_y; }
   void MCPDFT::set_rhoa_z(const arma::vec &rhoa_z) { rho_a_z_ = rhoa_z; }
   void MCPDFT::set_rhob(const arma::vec &rhob) { rho_b_ = rhob; }
   void MCPDFT::set_rhob_x(const arma::vec &rhob_x) { rho_b_x_ = rhob_x; }
   void MCPDFT::set_rhob_y(const arma::vec &rhob_y) { rho_b_y_ = rhob_y; }
   void MCPDFT::set_rhob_z(const arma::vec &rhob_z) { rho_b_z_ = rhob_z; }
   void MCPDFT::set_rho(const arma::vec &rho) { rho_ = rho; }
   void MCPDFT::set_tr_rhoa(const arma::vec &tr_rhoa) { tr_rho_a_ = tr_rhoa; }
   void MCPDFT::set_tr_rhoa_x(const arma::vec &tr_rhoa_x) { tr_rho_a_x_ = tr_rhoa_x; }
   void MCPDFT::set_tr_rhoa_y(const arma::vec &tr_rhoa_y) { tr_rho_a_y_ = tr_rhoa_y; }
   void MCPDFT::set_tr_rhoa_z(const arma::vec &tr_rhoa_z) { tr_rho_a_z_ = tr_rhoa_z; }
   void MCPDFT::set_tr_rhob(const arma::vec &tr_rhob) { tr_rho_b_ = tr_rhob; }
   void MCPDFT::set_tr_rhob_x(const arma::vec &tr_rhob_x) { tr_rho_b_x_ = tr_rhob_x; }
   void MCPDFT::set_tr_rhob_y(const arma::vec &tr_rhob_y) { tr_rho_b_y_ = tr_rhob_y; }
   void MCPDFT::set_tr_rhob_z(const arma::vec &tr_rhob_z) { tr_rho_b_z_ = tr_rhob_z; }
   void MCPDFT::set_tr_rho(const arma::vec &tr_rho) { tr_rho_ = tr_rho; }
   void MCPDFT::set_pi(const arma::vec &pi) { pi_ = pi; }
   void MCPDFT::set_pi_x(const arma::vec &pi_x) { pi_x_ = pi_x; }
   void MCPDFT::set_pi_y(const arma::vec &pi_y) { pi_y_ = pi_y; }
   void MCPDFT::set_pi_z(const arma::vec &pi_z) { pi_z_ = pi_z; }
   void MCPDFT::set_R(const arma::vec &R) { R_ = R; }
   void MCPDFT::set_sigma_aa(const arma::vec &sigma_aa) { sigma_aa_ = sigma_aa; }
   void MCPDFT::set_sigma_ab(const arma::vec &sigma_ab) { sigma_ab_ = sigma_ab; }
   void MCPDFT::set_sigma_bb(const arma::vec &sigma_bb) { sigma_bb_ = sigma_bb; }
   void MCPDFT::set_tr_sigma_aa(const arma::vec &tr_sigma_aa) { tr_sigma_aa_ = tr_sigma_aa; }
   void MCPDFT::set_tr_sigma_ab(const arma::vec &tr_sigma_ab) { tr_sigma_ab_ = tr_sigma_ab; }
   void MCPDFT::set_tr_sigma_bb(const arma::vec &tr_sigma_bb) { tr_sigma_bb_ = tr_sigma_bb; }

   void MCPDFT::print_banner() const {
      printf("\n******************************************************************\n");
      printf("*                                                                *\n");
      printf("*                           OpenRDM:                             *\n");
      printf("*                                                                *\n");
      printf("*                 An open-source library for                     *\n");
      printf("*    reduced-density matrix-based analysis and computation       *\n");
      printf("*                                                                *\n");
      printf("*                     Mohammad Mostafanejad                      *\n");
      printf("*                   Florida State University                     *\n");
      printf("*                                                                *\n");
      printf("******************************************************************\n");

      printf("\n           Please cite the following article(s):\n\n");

      printf("    # M. Mostafanejad and A. E. DePrince III\n");
      printf("      J. Chem. Theory Comput. 15, 290-302 (2019).\n");
   }
}


#include "extent.hpp"

#include <algorithm>
#include <boost/algorithm/clamp.hpp>
#include <boost/math/constants/constants.hpp>
#include <cmath>
#include <utility>
#include "../common/geom.hpp"
#include "../common/helpers/eigen_helpers.hpp"

const double PI = boost::math::constants::pi<double>();

namespace ear {

  /** @brief Normalise position or return {0,1,0}.
   *
   * @param position  Position to normalise.
   *
   * @returns normalised position
   */
  Eigen::Vector3d safeNormPosition(Eigen::Vector3d position) {
    double norm = position.norm();
    if (norm < 1e-10) {
      return Eigen::Vector3d{0.0, 1.0, 0.0};
    } else {
      return position / norm;
    }
  }

  double extentMod(double extent, double distance) {
    double minSize = 0.2;
    double size = interp(extent, Eigen::Vector2d(0.0, 360.0),
                         Eigen::Vector2d(minSize, 1.0));
    double extent1 = 4.0 * degrees(atan2(size, 1.0));
    return interp(4.0 * degrees(atan2(size, distance)),
                  Eigen::Vector3d(0.0, extent1, 360.0),
                  Eigen::Vector3d(0.0, extent, 360.0));
  }

  Eigen::Matrix3d calcBasis(Eigen::Vector3d position) {
    position = safeNormPosition(position);
    double az = azimuth(position);
    double el = elevation(position);

    // points near the poles have indeterminate azimuth; assume 0
    if (std::abs(el) > (90.0 - 1e-5)) {
      az = 0.0;
    }
    return localCoordinateSystem(az, el);
  }

  Eigen::Vector3d cartOnBasis(Eigen::Matrix3d basis, double azimuth,
                              double elevation) {
    Eigen::RowVector3d cartPosRel{sin(azimuth) * cos(elevation),
                                  cos(azimuth) * cos(elevation),
                                  sin(elevation)};
    return cartPosRel * basis;
  }

  std::pair<double, double> azimuthElevationOnBasis(
      Eigen::Matrix3d basis, Eigen::RowVector3d position) {
    // project onto each basis, and clip components to keep asin happy
    Eigen::Vector3d components =
        (position * basis.transpose()).cwiseMin(1.0).cwiseMax(-1.0);

    double azimuth = atan2(components(0), components(1));
    double elevation = asin(components(2));

    return std::make_pair(azimuth, elevation);
  }

  WeightingFunction::WeightingFunction(Eigen::Vector3d position, double width,
                                       double height) {
    _width = radians(width) / 2;
    _height = radians(height) / 2;

    // basis vectors to rotate the vsource positions towards position
    Eigen::Matrix3d basises = calcBasis(position);

    _circleRadius = std::min(_width, _height);

    // Flip the width and the height such that it is always wider than it is
    // high from here in.
    if (_height > _width) {
      std::swap(_height, _width);
      _flippedBasis = basises.colwise().reverse();
    } else {
      _flippedBasis = basises;
    }

    // modify the width to make it meet at the back.
    double widthFull = PI + _height;
    // interpolate to this from a width of pi/2 to pi
    double widthMod = interp(_width, Eigen::Vector3d{0.0, PI / 2.0, PI},
                             Eigen::Vector3d{0.0, PI / 2.0, widthFull});
    // apply this fully for a height of less than pi/4; tail off until pi/2
    _width = interp(_height, Eigen::Vector4d{0, PI / 4.0, PI / 2.0, PI},  //
                    Eigen::Vector4d{widthMod, widthMod, _width, _width});  //

    // angle of the circle centres from the source position; width is to the
    // end of the rectangle.
    _circlePos = _width - _circleRadius;

    // Cartesian circle centres
    _circlePositions << cartOnBasis(_flippedBasis, -_circlePos, 0.0),
        cartOnBasis(_flippedBasis, _circlePos, 0.0);
  }

  double WeightingFunction::operator()(Eigen::Vector3d position) const {
    // Flipped azimuths and elevations; the straight edges are always along
    // azimuth lines.
    double azimuth, elevation;
    std::tie(azimuth, elevation) =
        azimuthElevationOnBasis(_flippedBasis, position);

    // The distance is the angle away from the defined shape; 0 or negative is
    // inside.
    double distance = 0.0;

    // for the straight lines
    if (std::abs(azimuth) <= _circlePos) {
      distance = std::abs(elevation) - _circleRadius;
    } else {
      // distance from the closest circle centre
      size_t nearest_circle = azimuth < 0 ? 0 : 1;
      double angle =
          position.transpose() * _circlePositions.col(nearest_circle);
      double circleDistance = acos(boost::algorithm::clamp(angle, -1.0, 1.0));
      distance = circleDistance - _circleRadius;
    }
    // fade the weight from one to zero over fadeWidth
    return interp(distance, Eigen::Vector2d{0.0, radians(_fadeWidth)},
                  Eigen::Vector2d{1.0, 0.0});
  }

  SpreadingPanner::SpreadingPanner(std::shared_ptr<PointSourcePanner> psp,
                                   int nRows)
      : _psp(psp), _nRows(nRows) {
    _panningPositions = _generatePanningPositionsEven();
    _panningPositionsResults = _generatePanningPositionsResults();
  }

  Eigen::VectorXd SpreadingPanner::panningValuesForWeight(
      const WeightingFunction& weightFunc) {
    Eigen::VectorXd weights(_panningPositions.rows());
    for (int i = 0; i < _panningPositions.rows(); ++i) {
      weights(i) = weightFunc(_panningPositions.row(i));
    }
    Eigen::VectorXd totalPv = weights.transpose() * _panningPositionsResults;
    return totalPv / totalPv.norm();
  }

  Eigen::MatrixXd SpreadingPanner::_generatePanningPositionsEven() {
    Eigen::VectorXd elevations =
        Eigen::VectorXd::LinSpaced(_nRows, -90.0, 90.0);
    Eigen::MatrixXd positions(0, 3);

    for (double el : elevations) {
      double radius = cos(radians(el));
      double perimiter = 2 * PI * radius;
      double perimiter_centre = 2 * PI;

      int nPoints = static_cast<int>(
          std::round((perimiter / perimiter_centre) * 2 * (_nRows - 1)));
      if (nPoints == 0) {
        nPoints = 1;
      }
      Eigen::VectorXd azimuths =
          Eigen::VectorXd::LinSpaced(nPoints + 1, 0.0, 360.0);
      for (int i = 0; i < azimuths.size() - 1; ++i) {
        double az = azimuths(i);
        positions.conservativeResize(positions.rows() + 1, Eigen::NoChange);
        positions.row(positions.rows() - 1) = cart(az, el, 1.0);
      }
    }
    return positions;
  }

  Eigen::MatrixXd SpreadingPanner::_generatePanningPositionsResults() {
    Eigen::MatrixXd results(_panningPositions.rows(),
                            _psp->numberOfOutputChannels());
    for (int i = 0; i < _panningPositions.rows(); ++i) {
      results.row(i) = _psp->handle(_panningPositions.row(i)).get();
    }
    return results;
  }

  PolarExtentPanner::PolarExtentPanner(std::shared_ptr<PointSourcePanner> psp)
      : _psp(psp), _spreadingPanner(SpreadingPanner(psp, _nRows)){};

  Eigen::VectorXd PolarExtentPanner::calcPvSpread(Eigen::Vector3d position,
                                                  double width, double height) {
    // When calculating the spread panning values the width and height are
    // set to at least fade_width. For sizes where any of the dimensions is
    // less than this, interpolate linearly between the point and spread
    // panning values.
    double ammount_spread =
        interp(std::max(width, height), Eigen::Vector2d(0.0, _fadeWidth),
               Eigen::Vector2d(0.0, 1.0));
    double ammountPoint = 1.0 - ammount_spread;
    Eigen::ArrayXd pv = Eigen::ArrayXd::Zero(_psp->numberOfOutputChannels());
    if (ammountPoint > 1e-10) {
      pv += ammountPoint * _psp->handle(position).get().array().square();
    }
    if (ammount_spread > 1e-10) {
      // minimum width and height as above
      width = std::max(width, _fadeWidth / 2.0);
      height = std::max(height, _fadeWidth / 2.0);

      WeightingFunction weightingFunction(position, width, height);
      Eigen::VectorXd panning_values =
          _spreadingPanner.panningValuesForWeight(weightingFunction);
      pv += ammount_spread * panning_values.array().square();
    }
    return pv.sqrt().matrix();
  }

  Eigen::VectorXd PolarExtentPanner::handle(Eigen::Vector3d position,
                                            double width, double height,
                                            double depth) {
    double distance = position.norm();

    if (depth != 0.0) {
      double distanceMin = distance - depth / 2.0;
      double distanceMax = distance + depth / 2.0;
      distanceMin = (distanceMin < 0) ? 0.0 : distanceMin;
      distanceMax = (distanceMax < 0) ? 0.0 : distanceMax;
      Eigen::VectorXd pvsMin =
          calcPvSpread(position, extentMod(width, distanceMin),
                       extentMod(height, distanceMin));
      Eigen::VectorXd pvsMax =
          calcPvSpread(position, extentMod(width, distanceMax),
                       extentMod(height, distanceMax));
      return ((pvsMin.array().square() + pvsMax.array().square()) / 2.0).sqrt();
    } else {
      Eigen::VectorXd pvs = calcPvSpread(position, extentMod(width, distance),
                                         extentMod(height, distance));
      return pvs;
    }
  }

}  // namespace ear


// C++ standard headers
#include <iostream>

// C++ third party headers
#include <Eigen/Core>

// User defined headers
#include "Header.h"
#include "RandomSampling.h"

using namespace std;
int main() {
	cout << "hello world" << endl;
	cout << "Eigen version : " << EIGEN_MAJOR_VERSION << " . " << EIGEN_MINOR_VERSION << endl;
	cout << sum(2, 3) << endl;

	Eigen::MatrixXd var(2, 2);
	var << 1, 0,
				0, 1;

	cout << var << "\n";

	RandomSampling* rs;
	GaussianRandomSampling g;
	rs = &g;

	// virtual random generator function binded at runtime
	cout << rs->generate_random_seq(1) << endl;
}


//==================================================================================================
/*!
  @file

  @copyright 2015 NumScale SAS
  @copyright 2015 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_MOD_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_MOD_HPP_INCLUDED

#include <boost/simd/function/div.hpp>
#include <boost/simd/function/idiv.hpp>
#include <boost/simd/function/is_nez.hpp>
#include <boost/simd/function/multiplies.hpp>
#include <boost/simd/function/selsub.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/config.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  BOOST_DISPATCH_OVERLOAD ( mod_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_<bd::floating_<A0> >
                          , bd::generic_<bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0, A0  a1) const BOOST_NOEXCEPT
    {
      return selsub(is_nez(a1),a0,div(a0,a1,floor)*a1 );
    }
  };

  BOOST_DISPATCH_OVERLOAD ( mod_
                          , (typename A0)
                          , bd::cpu_
                          , bd::generic_<bd::arithmetic_<A0> >
                          , bd::generic_<bd::arithmetic_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0, A0 a1) const BOOST_NOEXCEPT
    {
      return selsub(is_nez(a1),a0,simd::multiplies(idiv(a0,a1,floor),a1));
    }
  };
} } }

#endif


#ifndef STAN_MATH_PRIM_SCAL_FUN_BESSEL_SECOND_KIND_HPP
#define STAN_MATH_PRIM_SCAL_FUN_BESSEL_SECOND_KIND_HPP

#include <boost/math/special_functions/bessel.hpp>

namespace stan {
namespace math {

/**
 *
   \f[
   \mbox{bessel\_second\_kind}(v, x) =
   \begin{cases}
     \textrm{error} & \mbox{if } x \leq 0 \\
     Y_v(x) & \mbox{if } x > 0 \\[6pt]
     \textrm{NaN} & \mbox{if } x = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   \frac{\partial\, \mbox{bessel\_second\_kind}(v, x)}{\partial x} =
   \begin{cases}
     \textrm{error} & \mbox{if } x \leq 0 \\
     \frac{\partial\, Y_v(x)}{\partial x} & \mbox{if } x > 0 \\[6pt]
     \textrm{NaN} & \mbox{if } x = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   Y_v(x)=\frac{J_v(x)\cos(v\pi)-J_{-v}(x)}{\sin(v\pi)}
   \f]

   \f[
   \frac{\partial \, Y_v(x)}{\partial x} = \frac{v}{x}Y_v(x)-Y_{v+1}(x)
   \f]
 *
 */
template <typename T2>
inline T2 bessel_second_kind(int v, const T2 z) {
  return boost::math::cyl_neumann(v, z);
}

}  // namespace math
}  // namespace stan

#endif


#include "heuristics/null.hpp"
#include "heuristics/shortest_path.hpp"
#include "heuristics/mst.hpp"
#include "mapgen/mapgen.hpp"
#include "solver/tspsolver.hpp"

#include <fstream>

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_ALTERNATIVE_INIT_API
#include <boost/test/unit_test.hpp>

constexpr char csv1[] = "a,b,c,d,e\n"
                        "0,10,4,9,10\n"
                        "10,0,3,2,4\n"
                        "4,3,0,6,8\n"
                        "9,2,6,0,3\n"
                        "10,4,8,3,0\n";

// Example from "http://www.ams.org/samplings/feature-column/fcarc-tsp"
constexpr char csv2[] = "a,b,c,d,e\n"
                        "0,6,9,inf,10\n"
                        "6,0,11,inf,inf\n"
                        "9,11,0,5,14\n"
                        "inf,inf,5,0,8\n"
                        "10,inf,14,8,0\n";

bool init_unit_test() {
	std::ofstream file1("solver_1.csv");
	file1 << csv1;

	std::ofstream file2("solver_2.csv");
	file2 << csv2;

	return true;
}

int main(int argc, char* argv[]) {
	return ::boost::unit_test::unit_test_main(&init_unit_test, argc, argv);
}

BOOST_AUTO_TEST_CASE(dumb_null) {
	using TSPSolver        = awesome::TSPSolver<awesome::NullHeuristic>;
	using MapGraph         = TSPSolver::MapGraph;
	using ConstNode        = MapGraph::ConstNode_t;
	using WeightedProperty = graph::WeightedProperty;

	MapGraph map;

	map.addEdges({{"a", "b", WeightedProperty{1}},
	              {"b", "c", WeightedProperty{1}},
	              {"c", "d", WeightedProperty{1}},
	              {"d", "e", WeightedProperty{1}},
	              {"e", "a", WeightedProperty{1}}});

	TSPSolver solver(map);

	std::list<ConstNode> candidate{map["a"], map["b"], map["c"], map["d"], map["e"], map["a"]};

	BOOST_CHECK(solver.goForIt() == candidate);
}

BOOST_AUTO_TEST_CASE(dumb_shortest) {
	using TSPSolver        = awesome::TSPSolver<awesome::ShortestPathHeuristic>;
	using MapGraph         = TSPSolver::MapGraph;
	using ConstNode        = MapGraph::ConstNode_t;
	using WeightedProperty = graph::WeightedProperty;

	MapGraph map;

	map.addEdges({{"a", "b", WeightedProperty{1}},
	              {"b", "c", WeightedProperty{1}},
	              {"c", "d", WeightedProperty{1}},
	              {"d", "e", WeightedProperty{1}},
	              {"e", "a", WeightedProperty{1}}});

	TSPSolver solver(map);

	std::list<ConstNode> candidate{map["a"], map["b"], map["c"], map["d"], map["e"], map["a"]};

	BOOST_CHECK(solver.goForIt() == candidate);
}

BOOST_AUTO_TEST_CASE(dumb_mst) {
	using TSPSolver        = awesome::TSPSolver<awesome::MSTHeuristic>;
	using MapGraph         = TSPSolver::MapGraph;
	using ConstNode        = MapGraph::ConstNode_t;
	using WeightedProperty = graph::WeightedProperty;

	MapGraph map;

	map.addEdges({{"a", "b", WeightedProperty{1}},
	              {"b", "c", WeightedProperty{1}},
	              {"c", "d", WeightedProperty{1}},
	              {"d", "e", WeightedProperty{1}},
	              {"e", "a", WeightedProperty{1}}});

	TSPSolver solver(map);

	std::list<ConstNode> candidate{map["a"], map["b"], map["c"], map["d"], map["e"], map["a"]};

	BOOST_CHECK(solver.goForIt() == candidate);
}

BOOST_AUTO_TEST_CASE(solver_insolvable_null) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::NullHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	map.removeEdge(map["d"], map["c"]);
	map.removeEdge(map["c"], map["d"]);
	TSPSolver solver(map);

	BOOST_CHECK_THROW(solver.goForIt(), std::logic_error);
}

BOOST_AUTO_TEST_CASE(solver_insolvable_shortest) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::ShortestPathHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	map.removeEdge(map["d"], map["c"]);
	map.removeEdge(map["c"], map["d"]);
	TSPSolver solver(map);

	BOOST_CHECK_THROW(solver.goForIt(), std::logic_error);
}

BOOST_AUTO_TEST_CASE(solver_insolvable_mst) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::MSTHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	map.removeEdge(map["d"], map["c"]);
	map.removeEdge(map["c"], map["d"]);
	TSPSolver solver(map);

	BOOST_CHECK_THROW(solver.goForIt(), std::logic_error);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_1_null) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::NullHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_1.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["e"], map["d"], map["b"], map["c"], map["a"], map["e"]};

	BOOST_CHECK(result == expected);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_1_shortest) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::ShortestPathHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_1.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["a"], map["c"], map["b"], map["d"], map["e"], map["a"]};

	BOOST_CHECK(result == expected);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_1_mst) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::MSTHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_1.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["b"], map["c"], map["a"], map["e"], map["d"], map["b"]};

	BOOST_CHECK(result == expected);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_2_null) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::NullHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["c"], map["d"], map["e"], map["a"], map["b"], map["c"]};

	BOOST_CHECK(result == expected);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_2_shortest) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::ShortestPathHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["b"], map["a"], map["e"], map["d"], map["c"], map["b"]};

	BOOST_CHECK(result == expected);
}

BOOST_AUTO_TEST_CASE(solver_sample_problem_2_mst) {
	using awesome::MapGen;

	using TSPSolver = awesome::TSPSolver<awesome::MSTHeuristic>;
	using MapGraph  = awesome::MapGen::MapGraph;
	using ConstNode = MapGraph::ConstNode_t;

	MapGraph map = MapGen::fromFile("solver_2.csv");
	TSPSolver solver(map);

	std::list<ConstNode> result   = solver.goForIt(),
	                     expected = {map["d"], map["c"], map["b"], map["a"], map["e"], map["d"]};

	BOOST_CHECK(result == expected);
}


#include <boost/foreach.hpp>

#include "rb_tracker/MultiRBTracker.h"

#include "omip_common/OMIPUtils.h"

#include <cmath>

#include <pcl/correspondence.h>
#include <pcl/registration/correspondence_rejection_sample_consensus.h>

#include <pcl_conversions/pcl_conversions.h>

#if ROS_VERSION_MINIMUM(1, 11, 1) // if current ros version is >= 1.11.1 (indigo)
#else
#include "rb_tracker/pcl_conversions_indigo.h"
#endif

#include <pcl/console/print.h>

using namespace omip;

inline void estimateRigidTransformation(
        const FeaturesDataBase::MapOfFeatureLocationPairs& last2locations,
        Eigen::Matrix4d &transformation_matrix)
{
    // Convert to Eigen format
    const int npts = static_cast<int>(last2locations.size());

    Eigen::Matrix<double, 3, Eigen::Dynamic> cloud_src(3, npts);
    Eigen::Matrix<double, 3, Eigen::Dynamic> cloud_tgt(3, npts);

    FeaturesDataBase::MapOfFeatureLocationPairs::const_iterator locations_it =
            last2locations.begin();
    FeaturesDataBase::MapOfFeatureLocationPairs::const_iterator locations_it_end =
            last2locations.end();
    for (int i = 0; locations_it != locations_it_end; locations_it++, i++)
    {
        cloud_src(0, i) = locations_it->second.second.get<0>();
        cloud_src(1, i) = locations_it->second.second.get<1>();
        cloud_src(2, i) = locations_it->second.second.get<2>();

        cloud_tgt(0, i) = locations_it->second.first.get<0>();
        cloud_tgt(1, i) = locations_it->second.first.get<1>();
        cloud_tgt(2, i) = locations_it->second.first.get<2>();
    }

    // Call Umeyama directly from Eigen (PCL patched version until Eigen is released)
    transformation_matrix = Eigen::umeyama(cloud_src, cloud_tgt, false);
}

MultiRBTracker::MultiRBTracker(int max_num_rb,
                               double loop_period_ns,
                               int ransac_iterations,
                               double estimation_error_threshold,
                               double static_motion_threshold,
                               double new_rbm_error_threshold,
                               double max_error_to_reassign_feats,
                               int supporting_features_threshold,
                               int min_num_feats_for_new_rb,
                               int min_num_frames_for_new_rb,
                               int initial_cam_motion_constraint,
                               static_environment_tracker_t static_environment_tracker_type) :
    RecursiveEstimatorFilterInterface(loop_period_ns),
    _max_num_rb(max_num_rb),
    _ransac_iterations(ransac_iterations),
    _estimation_error_threshold(estimation_error_threshold),
    _static_motion_threshold(static_motion_threshold),
    _new_rbm_error_threshold(new_rbm_error_threshold),
    _max_error_to_reassign_feats(max_error_to_reassign_feats),
    _supporting_features_threshold(supporting_features_threshold),
    _min_num_feats_for_new_rb(min_num_feats_for_new_rb),
    _min_num_frames_for_new_rb(min_num_frames_for_new_rb),
    _min_num_points_in_segment(0),
    _min_probabilistic_value(0.0),
    _max_fitness_score(0.0),
    _min_num_supporting_feats_to_correct(7),
    _max_distance_ee(-1),
    _max_distance_irb(-1)
{

    //_really_free_feats_file.open("really_free_feats_file.txt");
    this->_filter_name = "RBMFilter";

    if (this->_features_db)
        this->_features_db.reset();

    this->_features_db = FeaturesDataBase::Ptr(new FeaturesDataBase());

    this->_predicted_measurement = rbt_measurement_t(new FeatureCloudPCLwc());

    // The static environment exists always as initial hypothesis
    // We can track the static enviroment (visual odometry) using different computation methods: ICP or EKF
    _static_environment_filter = StaticEnvironmentFilter::Ptr(new StaticEnvironmentFilter(this->_loop_period_ns,
                                                                                this->_features_db,
                                                                                this->_static_motion_threshold));
    _static_environment_filter->setComputationType(static_environment_tracker_type);
    _static_environment_filter->setFeaturesDatabase(this->_features_db);
    _static_environment_filter->setMotionConstraint(initial_cam_motion_constraint);

    this->_kalman_filters.push_back(_static_environment_filter);
}

void MultiRBTracker::Init()
{
    _static_environment_filter->setMeasurementDepthFactor(this->_meas_depth_factor);
    _static_environment_filter->setMinCovarianceMeasurementX(this->_min_cov_meas_x);
    _static_environment_filter->setMinCovarianceMeasurementY(this->_min_cov_meas_y);
    _static_environment_filter->setMinCovarianceMeasurementZ(this->_min_cov_meas_z);
    _static_environment_filter->setPriorCovariancePose(this->_prior_cov_pose);
    _static_environment_filter->setPriorCovarianceVelocity(this->_prior_cov_vel);
    _static_environment_filter->setCovarianceSystemAccelerationTx(this->_cov_sys_acc_tx);
    _static_environment_filter->setCovarianceSystemAccelerationTy(this->_cov_sys_acc_ty);
    _static_environment_filter->setCovarianceSystemAccelerationTz(this->_cov_sys_acc_tz);
    _static_environment_filter->setCovarianceSystemAccelerationRx(this->_cov_sys_acc_rx);
    _static_environment_filter->setCovarianceSystemAccelerationRy(this->_cov_sys_acc_ry);
    _static_environment_filter->setCovarianceSystemAccelerationRz(this->_cov_sys_acc_rz);
    _static_environment_filter->setNumberOfTrackedFeatures(this->_num_tracked_feats);
    _static_environment_filter->setMinNumberOfSupportingFeaturesToCorrectPredictedState(this->_min_num_supporting_feats_to_correct);
    _static_environment_filter->Init();
}

MultiRBTracker::~MultiRBTracker()
{

}

void MultiRBTracker::setMeasurement(rbt_measurement_t acquired_measurement, const double& measurement_timestamp)
{
    this->_previous_measurement_timestamp_ns = this->_measurement_timestamp_ns;

    {
        boost::mutex::scoped_lock(this->_measurement_timestamp_ns_mutex);
        this->_measurement_timestamp_ns = measurement_timestamp;
    }

    this->_features_db->clearListOfAliveFeatureIds();
    this->_measurement = acquired_measurement;

    int num_tracked_features = acquired_measurement->points.size();

    int tracked_features_index = 0;
    uint32_t feature_id;
    for (; tracked_features_index < num_tracked_features; tracked_features_index++)
    {
        feature_id = acquired_measurement->points[tracked_features_index].label;
        if(feature_id != 0)
        {
            Feature::Location feature_location = Feature::Location(acquired_measurement->points[tracked_features_index].x, acquired_measurement->points[tracked_features_index].y,
                                                                   acquired_measurement->points[tracked_features_index].z);
            this->addFeatureLocation(feature_id, feature_location);
        }
    }
    this->_features_db->step();
}

void MultiRBTracker::predictState(double time_interval_ns)
{
    //First we update the filters using the internal state (system update)
    for (std::vector<RBFilter::Ptr>::iterator filter_it = this->_kalman_filters.begin(); filter_it != this->_kalman_filters.end(); filter_it++)
    {
        (*filter_it)->predictState(time_interval_ns);
        (*filter_it)->doNotUsePredictionFromHigherLevel();
    }
}

// NOTE: The location of the Features predicted by the static RBFilter is the initial location of these Features
// When these predictions are used for the visual tracking the searching area for the Features is static until these Features can
// be assigned to an existing RB (or a group of them create a new one)
// The consequence is that the Features could go out of the searching area before they are detected as moving
// There are two solutions:
//      1- Use the static_motion_threshold in the estimation of the searching area. The searching area should cover motions of static_motion_threshold
//         at the depth of the Feature (COMPLEX)
//      2- Only for the static RBFilter the predicted locations retrieved here are the latest locations of the Features. This is ugly (treats the static
//         RBFilter as an special one) but easier
void MultiRBTracker::predictMeasurement()
{
    std::vector<RBFilter::Ptr>::iterator filter_it = this->_kalman_filters.begin();
    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        (*filter_it)->predictMeasurement();
    }

    this->_predicted_measurement->points.clear();

    std::vector<Feature::Id> feature_ids;
    FeatureCloudPCLwc one_filter_predicted_measurement;
    filter_it = this->_kalman_filters.begin();
    filter_it++;
    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        one_filter_predicted_measurement = (*filter_it)->getPredictedMeasurement();
        for (size_t feat_idx = 0; feat_idx < one_filter_predicted_measurement.points.size(); feat_idx++)
        {
            this->_predicted_measurement->points.push_back(one_filter_predicted_measurement.points.at(feat_idx));
            feature_ids.push_back(one_filter_predicted_measurement.points.at(feat_idx).label);
        }
    }

    // The free Features are not in any RBFilter. We collect their latest locations
    std::vector<Feature::Id> free_features_ids = this->EstimateFreeFeatures(feature_ids);
    FeaturePCLwc free_feat;
    for (size_t feat_idx = 0; feat_idx < free_features_ids.size(); feat_idx++)
    {
        free_feat.x = this->_features_db->getFeatureLastX(free_features_ids.at(feat_idx));
        free_feat.y = this->_features_db->getFeatureLastY(free_features_ids.at(feat_idx));
        free_feat.z = this->_features_db->getFeatureLastZ(free_features_ids.at(feat_idx));
        free_feat.label = free_features_ids.at(feat_idx);
        free_feat.covariance[0] = 20;
        this->_predicted_measurement->points.push_back(free_feat);
    }
}

rbt_measurement_t MultiRBTracker::getPredictedMeasurement() const
{
    this->_predicted_measurement->header.stamp = pcl_conversions::toPCL(ros::Time(this->_measurement_timestamp_ns/1e9+this->_loop_period_ns/1e9));
    return this->_predicted_measurement;
}

std::vector<Feature::Id> MultiRBTracker::estimateBestPredictionsAndSupportingFeatures()
{
    std::vector<Feature::Id> all_filters_supporting_features;
    for (std::vector<RBFilter::Ptr>::iterator filter_it = this->_kalman_filters.begin(); filter_it != this->_kalman_filters.end();)
    {
        (*filter_it)->estimateBestPredictionAndSupportingFeatures();
        std::vector<Feature::Id> this_filter_supporting_features = (*filter_it)->getSupportingFeatures();

        // If there are no features supporting to the EKF we delete it (no re-finding of RBs/EKFs possible)
        // Don't delete the filter of the environment!!! (filter_it == this->_kalman_filters.begin())
        if (this_filter_supporting_features.size()  < (size_t)this->_supporting_features_threshold
                && *filter_it != this->_static_environment_filter )
        {
            Eigen::Twistd last_pose_ec;
            TransformMatrix2Twist((*filter_it)->getPose(), last_pose_ec);
            ROS_ERROR_NAMED( "MultiRBTracker.estimateBestPredictionsAndSupportingFeatures",
                             "RBFilter %2d: Supporting feats:%3d (<%3d), Pose (vx,vy,vz,rx,ry,rz): % 2.2f,% 2.2f,% 2.2f,% 2.2f,% 2.2f,% 2.2f. REMOVED",
                             (int)(*filter_it)->getId(),
                             (int)this_filter_supporting_features.size() ,
                             (int)(this->_supporting_features_threshold),
                             last_pose_ec.vx(),
                             last_pose_ec.vy(),
                             last_pose_ec.vz(),
                             last_pose_ec.rx(),
                             last_pose_ec.ry(),
                             last_pose_ec.rz());
            filter_it = this->_kalman_filters.erase(filter_it);
        }
        else
        {
            all_filters_supporting_features.insert(all_filters_supporting_features.end(),
                                                   this_filter_supporting_features.begin(),
                                                   this_filter_supporting_features.end());
            filter_it++;
        }
    }
    return all_filters_supporting_features;
}

void MultiRBTracker::correctState()
{
    // 1: Estimate the supporting features of each filter
    std::vector<Feature::Id> all_filters_supporting_features = this->estimateBestPredictionsAndSupportingFeatures();

    // 2: Estimate the free features
    std::vector<Feature::Id> free_feat_ids = this->EstimateFreeFeatures(all_filters_supporting_features);

    // 3: Try to assign free features to existing RBs
    std::vector<Feature::Id> really_free_feat_ids = this->ReassignFreeFeatures(free_feat_ids);

    ROS_INFO_NAMED("MultiRBTracker::correctState","Moving bodies: %3d, free features: %3d", (int)this->_kalman_filters.size(), (int)really_free_feat_ids.size());

    // 4: Correct the predicted state of each RBFilter using the last acquired Measurement
    for (std::vector<RBFilter::Ptr>::iterator filter_it = this->_kalman_filters.begin(); filter_it != this->_kalman_filters.end(); filter_it++)
    {
        (*filter_it)->correctState();
    }

    // 5: Try to create new filters for the "really free" features
    // NOTE: Keep this step after the correctState of the RBFilters, otherwise you will create a filter and correct it in the first iteration!
    if(this->_kalman_filters.size() < this->_max_num_rb)
    {
        this->TryToCreateNewFilter(really_free_feat_ids);
    }
}

void MultiRBTracker::ReflectState()
{
    this->_state.rb_poses_and_vels.clear();

    std::vector<geometry_msgs::TwistWithCovariancePtr> tracked_poses_twist_with_cov = this->getPosesECWithCovariance();
    std::vector<geometry_msgs::TwistWithCovariancePtr> tracked_vels_twist_with_cov = this->getVelocitiesWithCovariance();
    std::vector<Eigen::Vector3d> centroids = this->getCentroids();

    RB_id_t RB_id;

    for (size_t poses_idx = 0; poses_idx < tracked_poses_twist_with_cov.size(); poses_idx++)
    {
        RB_id = this->getRBId(poses_idx);

        geometry_msgs::Point centroid;
        centroid.x = centroids.at(poses_idx).x();
        centroid.y = centroids.at(poses_idx).y();
        centroid.z = centroids.at(poses_idx).z();

        omip_msgs::RigidBodyPoseAndVelMsg rbpose_and_vel;
        rbpose_and_vel.rb_id = RB_id;
        rbpose_and_vel.pose_wc = *(tracked_poses_twist_with_cov.at(poses_idx));
        rbpose_and_vel.velocity_wc = *(tracked_vels_twist_with_cov.at(poses_idx));
        rbpose_and_vel.centroid = centroid;
        this->_state.rb_poses_and_vels.push_back(rbpose_and_vel);
    }
}

rbt_state_t MultiRBTracker::getState() const
{
    return this->_state;
}

void MultiRBTracker::setDynamicReconfigureValues(rb_tracker::RBTrackerDynReconfConfig &config)
{
    _static_environment_filter->setMotionConstraint(config.cam_motion_constraint);
    this->_min_num_feats_for_new_rb = config.min_feats_new_rb;
    this->_supporting_features_threshold = config.min_feats_survive_rb;
}

void MultiRBTracker::processMeasurementFromShapeTracker(const omip_msgs::ShapeTrackerStates &meas_from_st)
{
    for(int idx = 0; idx < meas_from_st.shape_tracker_states.size(); idx++)
    {
        BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
        {
            if (filter->getId() == meas_from_st.shape_tracker_states.at(idx).rb_id )
            {
                //Check if the flag indicates that this refinement is reliable
                if((meas_from_st.shape_tracker_states.at(idx).number_of_points_of_model > _min_num_points_in_segment) &&
                        (meas_from_st.shape_tracker_states.at(idx).number_of_points_of_current_pc > _min_num_points_in_segment) &&
                        (meas_from_st.shape_tracker_states.at(idx).probabilistic_value > _min_probabilistic_value) &&
                        (meas_from_st.shape_tracker_states.at(idx).fitness_score < _max_fitness_score))
                {
                    ROS_ERROR_STREAM("Integrating shape-based tracked pose RB " << meas_from_st.shape_tracker_states.at(idx).rb_id);
                    filter->integrateShapeBasedPose(meas_from_st.shape_tracker_states.at(idx).pose_wc, this->_measurement_timestamp_ns - meas_from_st.header.stamp.toNSec() );
                }else{
                    ROS_ERROR_STREAM("Ignoring shape-based tracked pose RB " << meas_from_st.shape_tracker_states.at(idx).rb_id);
                }
            }
        }
    }
}

void MultiRBTracker::CreateNewFilter(const std::vector<Eigen::Matrix4d>& initial_trajectory,
                                     const Eigen::Twistd& initial_velocity,
                                     const std::vector<Feature::Id>& initial_supporting_feats)
{
    RBFilter::Ptr new_kf = RBFilter::Ptr(new RBFilter(this->_loop_period_ns,
                                                      initial_trajectory,
                                                      initial_velocity,
                                                      this->_features_db,
                                                      this->_estimation_error_threshold));

    new_kf->setMeasurementDepthFactor(this->_meas_depth_factor);
    new_kf->setMinCovarianceMeasurementX(this->_min_cov_meas_x);
    new_kf->setMinCovarianceMeasurementY(this->_min_cov_meas_y);
    new_kf->setMinCovarianceMeasurementZ(this->_min_cov_meas_z);
    new_kf->setPriorCovariancePose(this->_prior_cov_pose);
    new_kf->setPriorCovarianceVelocity(this->_prior_cov_vel);
    new_kf->setCovarianceSystemAccelerationTx(this->_cov_sys_acc_tx);
    new_kf->setCovarianceSystemAccelerationTy(this->_cov_sys_acc_ty);
    new_kf->setCovarianceSystemAccelerationTz(this->_cov_sys_acc_tz);
    new_kf->setCovarianceSystemAccelerationRx(this->_cov_sys_acc_rx);
    new_kf->setCovarianceSystemAccelerationRy(this->_cov_sys_acc_ry);
    new_kf->setCovarianceSystemAccelerationRz(this->_cov_sys_acc_rz);
    new_kf->setNumberOfTrackedFeatures(this->_num_tracked_feats);
    new_kf->setMinNumberOfSupportingFeaturesToCorrectPredictedState(this->_min_num_supporting_feats_to_correct);

    new_kf->Init();

    this->_kalman_filters.push_back(new_kf);
    BOOST_FOREACH(Feature::Id supporting_feat_id, initial_supporting_feats)
    {
        new_kf->addSupportingFeature(supporting_feat_id);
    }
}

std::vector<Feature::Id> MultiRBTracker::EstimateFreeFeatures(const std::vector<Feature::Id>& supporting_features)
{
    // We check the set of free features = alive features - supporting features
    std::vector<Feature::Id> alive_feat_ids = this->_features_db->getListOfAliveFeatureIds();
    std::vector<Feature::Id> free_feat_ids;
    BOOST_FOREACH(Feature::Id alive_feat_id, alive_feat_ids)
    {
        // If the feature ID is not in the list of supporting features -> Free feature
        if (std::find(supporting_features.begin(), supporting_features.end(),alive_feat_id) == supporting_features.end())
        {
            free_feat_ids.push_back(alive_feat_id);
        }
    }
    _free_feat_ids = free_feat_ids;
    return free_feat_ids;
}

// NOTE: Here the features have an "extra" location compared to the moment when we predict the measurement for each filter (at the end of the iteration) (predictMeasurement)
// Therefore, the "age of the feature" is +1 the age it would have when predicting the measurement
// EXCEPT if the measurement prediction has to be called again in the loop!!!!!
std::vector<Feature::Id> MultiRBTracker::ReassignFreeFeatures(const std::vector<Feature::Id>& free_feat_ids)
{
    std::vector<Feature::Id> really_free_feat_ids;

    Feature::Ptr free_feat;
    Feature::Location free_feature_last_location;
    Feature::Location free_feature_pre_last_location;
    Feature::Location free_feature_predicted_location;
    double feat_error = 0.0;
    double prediction_error = 0.0;
    int best_ekf_idx = -1;
    Eigen::Matrix4d predicted_pose;

    // Check for each free feature if we can add it to one of the existing vision-based rigid bodies
    BOOST_FOREACH(Feature::Id free_feat_id, free_feat_ids)
    {
        free_feature_last_location = this->_features_db->getFeatureLastLocation(free_feat_id);
        free_feature_pre_last_location = this->_features_db->getFeatureNextToLastLocation(free_feat_id);
        free_feat = this->_features_db->getFeature(free_feat_id);
        feat_error = (double)this->_max_error_to_reassign_feats;
        best_ekf_idx = -1;
        Feature::Location best_predicted_location;

        // Get the prediction from each of the existing vision-based rigid bodies
        // TODO: Use the Mahalanobis distance between the uncertain feature location and its predicted location
        for (size_t filter_idx = 0; filter_idx < this->_kalman_filters.size();filter_idx++)
        {
            // Get the most likely prediction about the next feature state (the likelihood is estimated in estimateSupportingFeatures and copied to the others
            // so it is equivalent to get the belief pose from any hypothesis)
            //prediction_error = this->_kalman_filters.at(filter_idx)->PredictFeatureLocationNew(free_feat, false, true);

            predicted_pose = this->_kalman_filters.at(filter_idx)->getPose();
            this->_kalman_filters.at(filter_idx)->PredictFeatureLocation(free_feat, predicted_pose, false, free_feature_predicted_location,true);
            prediction_error = L2Distance(free_feature_last_location, free_feature_predicted_location);

            if (prediction_error < feat_error)
            {
                feat_error = prediction_error;
                best_ekf_idx = filter_idx;
                best_predicted_location = free_feature_predicted_location;
            }
        }

        if (best_ekf_idx != -1)
        {
            this->_kalman_filters.at(best_ekf_idx)->addSupportingFeature(free_feat_id);
            this->_kalman_filters.at(best_ekf_idx)->addPredictedFeatureLocation(best_predicted_location,free_feature_pre_last_location, free_feat_id);
        }
        else
        {
            really_free_feat_ids.push_back(free_feat_id);
        }
    }
    _really_free_feat_ids = really_free_feat_ids;
    return really_free_feat_ids;
}

void MultiRBTracker::TryToCreateNewFilter(const std::vector<Feature::Id>& really_free_feat_ids)
{
    // First check that the number of free features is over the minimum number to estimate motion
    if(really_free_feat_ids.size() > this->_min_num_feats_for_new_rb)
    {
        pcl::PointCloud<pcl::PointXYZ>::Ptr source = pcl::PointCloud<pcl::PointXYZ>::Ptr(new pcl::PointCloud<pcl::PointXYZ>());
        pcl::PointCloud<pcl::PointXYZ>::Ptr target = pcl::PointCloud<pcl::PointXYZ>::Ptr(new pcl::PointCloud<pcl::PointXYZ>());
        pcl::Correspondences initial_correspondeces, filtered_correspondences;
        int index = 0;
        std::map<int, int> index_to_feat_id;
        PointPCL point_pcl;

        //Iterate over the free features
        BOOST_FOREACH(Feature::Id free_feat_id, really_free_feat_ids)
        {
            // If the free feature is old enough, we add it to the source-target point clouds and add a correspondence
            if ((int)this->_features_db->getFeatureAge(free_feat_id) > this->_min_num_frames_for_new_rb)
            {
                Feature::Location first = this->_features_db->getFeatureNToLastLocation(free_feat_id, this->_min_num_frames_for_new_rb);
                Feature::Location second = this->_features_db->getFeatureLastLocation(free_feat_id);
                Location2PointPCL(first, point_pcl);
                source->push_back(point_pcl);
                Location2PointPCL(second, point_pcl);
                target->push_back(point_pcl);
                initial_correspondeces.push_back(pcl::Correspondence(index, index, 0.0));
                index_to_feat_id[index] = free_feat_id;
                index++;
            }
        }

        // If enough of the free features are old enough, we try to estimate the motion
        if(index > this->_min_num_feats_for_new_rb)
        {
            //pcl::console::setVerbosityLevel(pcl::console::L_DEBUG);
            // PCL RANSAC
            pcl::registration::CorrespondenceRejectorSampleConsensus<pcl::PointXYZ> rejector;
            rejector.setInputSource(source);
            rejector.setInputTarget(target);
            rejector.setInlierThreshold(this->_new_rbm_error_threshold);
            rejector.setMaximumIterations(this->_ransac_iterations);
            rejector.getRemainingCorrespondences(initial_correspondeces, filtered_correspondences);

            Eigen::Matrix4f best_transformation = rejector.getBestTransformation();

            // If the number of free features that are inliers of the best transformation is over the minimum number to estimate motion
            // AND if the best transformation is not just noise:
            //      - Translation is larger than min_amount_translation meters
            //      - Rotation around the axis is larger than min_amount_rotation degrees
            Eigen::Transform<float, 3, Eigen::Affine> best_transformation2(best_transformation);
            double amount_translation = best_transformation2.translation().norm();
            Eigen::AngleAxisf best_axis_angle;
            double amount_rotation = best_axis_angle.fromRotationMatrix(best_transformation2.rotation()).angle();

            if ((int)filtered_correspondences.size() > this->_min_num_feats_for_new_rb && (
                        amount_translation > _min_amount_translation_for_new_rb ||
                        amount_rotation > _min_amount_rotation_for_new_rb*(M_PI/180.0) ) )
            {
                std::vector<Eigen::Matrix4d> trajectory;
                Feature::Location centroid_in_cf;

                //Collect the ids of the matched features
                //and compute the centroid of the supporting features N frames ago
                //The centroid will be used to "place" the frame of the rigid body -> Initial transformation of the rigid body
                //is a pure translation to the location of the centroid
                std::vector<Feature::Id> ransac_best_trial_inliers;
                BOOST_FOREACH(pcl::Correspondence filter_correspondence, filtered_correspondences)
                {
                    int feature_id = index_to_feat_id.at(filter_correspondence.index_match);
                    ransac_best_trial_inliers.push_back(feature_id);
                    Feature::Location feat_loc = this->_features_db->getFeatureNToLastLocation(feature_id, this->_min_num_frames_for_new_rb-1);
                    boost::tuples::get<0>(centroid_in_cf) = boost::tuples::get<0>(centroid_in_cf) + boost::tuples::get<0>(feat_loc);
                    boost::tuples::get<1>(centroid_in_cf) = boost::tuples::get<1>(centroid_in_cf) + boost::tuples::get<1>(feat_loc);
                    boost::tuples::get<2>(centroid_in_cf) = boost::tuples::get<2>(centroid_in_cf) + boost::tuples::get<2>(feat_loc);
                }

                boost::tuples::get<0>(centroid_in_cf) = boost::tuples::get<0>(centroid_in_cf)/(double)filtered_correspondences.size();
                boost::tuples::get<1>(centroid_in_cf) = boost::tuples::get<1>(centroid_in_cf)/(double)filtered_correspondences.size();
                boost::tuples::get<2>(centroid_in_cf) = boost::tuples::get<2>(centroid_in_cf)/(double)filtered_correspondences.size();

                Eigen::Matrix4d pre_translation = Eigen::Matrix4d::Identity();
                pre_translation(0,3) = boost::tuples::get<0>(centroid_in_cf);
                pre_translation(1,3) = boost::tuples::get<1>(centroid_in_cf);
                pre_translation(2,3) = boost::tuples::get<2>(centroid_in_cf);

                trajectory.push_back(pre_translation);

                // Estimate the trajectory in the last _min_num_frames_for_new_rb frames
                for(int frame = this->_min_num_frames_for_new_rb-2; frame >= 0; frame--)
                {
                    FeaturesDataBase::MapOfFeatureLocationPairs one_frame_best_subset;
                    BOOST_FOREACH(Feature::Id matched_feat_id, ransac_best_trial_inliers)
                    {
                        one_frame_best_subset[matched_feat_id] = Feature::LocationPair(this->_features_db->getFeatureNToLastLocation(matched_feat_id, frame),
                                                                                this->_features_db->getFeatureNToLastLocation(matched_feat_id, this->_min_num_frames_for_new_rb-1)
                                                                                - centroid_in_cf);
                    }
                    Eigen::Matrix4d one_frame_best_transformation_matrix;
                    estimateRigidTransformation(one_frame_best_subset,one_frame_best_transformation_matrix);
                    trajectory.push_back(one_frame_best_transformation_matrix);
                }

                Eigen::Matrix4d last_delta = trajectory.at(trajectory.size()-1)*(trajectory.at(trajectory.size()-2).inverse());
                Eigen::Twistd initial_velocity;
                TransformMatrix2Twist(last_delta, initial_velocity);

                double measured_loop_period_ns= this->_measurement_timestamp_ns - this->_previous_measurement_timestamp_ns;
                initial_velocity /=(measured_loop_period_ns/1e9);

                this->CreateNewFilter(trajectory, initial_velocity, ransac_best_trial_inliers);
            }else
            {
                ROS_INFO_STREAM_NAMED("RBMTracker._createNewFilters", "Not enough Consensus to create a new Filter with the free Features ("<<
                                      filtered_correspondences.size() <<" inliers are too few, or the the amount of translation " <<
                                      amount_translation << " < " << _min_amount_translation_for_new_rb << " and the amount of rotation " <<
                                      amount_rotation << " < " << _min_amount_rotation_for_new_rb*(M_PI/180.0));
                return;
            }
        }
    }
}

std::vector<Eigen::Matrix4d> MultiRBTracker::getPoses() const
{
    std::vector<Eigen::Matrix4d> return_filter_results;

    BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
    {
        return_filter_results.push_back(filter->getPose());
    }
    return return_filter_results;
}

std::vector<geometry_msgs::PoseWithCovariancePtr> MultiRBTracker::getPosesWithCovariance() const
{
    std::vector<geometry_msgs::PoseWithCovariancePtr> return_filter_results;
    BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
    {
        return_filter_results.push_back(filter->getPoseWithCovariance());
    }
    return return_filter_results;
}

std::vector<geometry_msgs::TwistWithCovariancePtr> MultiRBTracker::getPosesECWithCovariance() const
{
    std::vector<geometry_msgs::TwistWithCovariancePtr> return_filter_results;
    BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
    {
        return_filter_results.push_back(filter->getPoseECWithCovariance());
    }
    return return_filter_results;
}

std::vector<geometry_msgs::TwistWithCovariancePtr> MultiRBTracker::getVelocitiesWithCovariance() const
{
    std::vector<geometry_msgs::TwistWithCovariancePtr> return_filter_results;
    BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
    {
        return_filter_results.push_back(filter->getVelocityWithCovariance());
    }
    return return_filter_results;
}

void MultiRBTracker::addFeatureLocation(Feature::Id f_id, Feature::Location f_loc)
{
    if (this->_features_db->addFeatureLocation(f_id, f_loc))
    {
        // addFeature::LocationOfFeature returns true if the feature is new
        _static_environment_filter->addSupportingFeature(f_id);
    }
}

RB_id_t MultiRBTracker::getRBId(int n) const
{
    return this->_kalman_filters.at(n)->getId();
}

int MultiRBTracker::getNumberSupportingFeatures(int n) const
{
    return this->_kalman_filters.at(n)->getNumberSupportingFeatures();
}

FeatureCloudPCL MultiRBTracker::getLabelledSupportingFeatures()
{
    FeatureCloudPCL colored_pc;
    std::vector<RBFilter::Ptr>::iterator filter_it =this->_kalman_filters.begin();
    // TODO: The next line is to not show the supporting features of the static rigid body.
    // I commented it out for the ICRA paper: we want to plot all points

    filter_it++;
    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        std::vector<Feature::Id> this_filter_supporting_features = (*filter_it)->getSupportingFeatures();
        for (size_t feat_idx = 0; feat_idx < this_filter_supporting_features.size();feat_idx++)
        {
            if (this->_features_db->isFeatureStored(this_filter_supporting_features.at(feat_idx)))
            {
                FeaturePCL temp_point;
                RB_id_t label = (*filter_it)->getId();
                temp_point.x = this->_features_db->getFeatureLastX(this_filter_supporting_features.at(feat_idx));
                temp_point.y = this->_features_db->getFeatureLastY(this_filter_supporting_features.at(feat_idx));
                temp_point.z = this->_features_db->getFeatureLastZ(this_filter_supporting_features.at(feat_idx));

                temp_point.label = label;
                colored_pc.points.push_back(temp_point);
            }
        }
    }
    return colored_pc;
}

FeatureCloudPCL MultiRBTracker::getFreeFeatures()
{
    FeatureCloudPCL colored_pc;
    for (size_t feat_idx = 0; feat_idx < _really_free_feat_ids.size();feat_idx++)
    {
        if (this->_features_db->isFeatureStored(_really_free_feat_ids.at(feat_idx)))
        {
            FeaturePCL temp_point;
            temp_point.x = this->_features_db->getFeatureLastX(_really_free_feat_ids.at(feat_idx));
            temp_point.y = this->_features_db->getFeatureLastY(_really_free_feat_ids.at(feat_idx));
            temp_point.z = this->_features_db->getFeatureLastZ(_really_free_feat_ids.at(feat_idx));

            temp_point.label = _really_free_feat_ids.at(feat_idx);
            colored_pc.points.push_back(temp_point);
        }
    }
    return colored_pc;
}

FeatureCloudPCL MultiRBTracker::getPredictedFeatures()
{
    FeatureCloudPCL colored_pc;
    std::vector<RBFilter::Ptr>::iterator filter_it =this->_kalman_filters.begin();
    filter_it++;
    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        FeatureCloudPCLwc::Ptr predicted_locations_pc = (*filter_it)->getFeaturesPredicted();
        for (size_t feat_idx = 0; feat_idx < predicted_locations_pc->points.size();feat_idx++)
        {
            FeaturePCL temp_point;
            temp_point.x = predicted_locations_pc->points[feat_idx].x;
            temp_point.y = predicted_locations_pc->points[feat_idx].y;
            temp_point.z = predicted_locations_pc->points[feat_idx].z;
            temp_point.label = predicted_locations_pc->points[feat_idx].label;
            colored_pc.points.push_back(temp_point);
        }
        (*filter_it)->resetFeaturesPredicted();
    }
    return colored_pc;
}

FeatureCloudPCL MultiRBTracker::getAtBirthFeatures()
{
    FeatureCloudPCL colored_pc;
    std::vector<RBFilter::Ptr>::iterator filter_it =this->_kalman_filters.begin();
    filter_it++;
    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        FeatureCloudPCLwc::Ptr at_birth_locations_pc = (*filter_it)->getFeaturesAtBirth();
        for (size_t feat_idx = 0; feat_idx < at_birth_locations_pc->points.size();feat_idx++)
        {
            FeaturePCL temp_point;
            temp_point.x = at_birth_locations_pc->points[feat_idx].x;
            temp_point.y = at_birth_locations_pc->points[feat_idx].y;
            temp_point.z = at_birth_locations_pc->points[feat_idx].z;
            temp_point.label = at_birth_locations_pc->points[feat_idx].label;
            colored_pc.points.push_back(temp_point);
        }
        (*filter_it)->resetFeaturesAtBirth();
    }
    return colored_pc;
}

std::vector<Eigen::Vector3d> MultiRBTracker::getCentroids() const
{
    std::vector<Eigen::Vector3d> centroids;
    centroids.push_back(Eigen::Vector3d(0., 0., 0.));
    std::vector<RBFilter::Ptr>::const_iterator filter_it = this->_kalman_filters.begin();
    filter_it++;    //Ignore the static environment

    for (; filter_it != this->_kalman_filters.end(); filter_it++)
    {
        //NEW: the first time we send as centroid the centroid we used for the trajectory estimation
        //In this way we can inform the next level of the initial pose of the rigid body
        if((*filter_it)->getTrajectory().size() < this->_min_num_frames_for_new_rb)
        {
            Feature::Location initial_location = (*filter_it)->getIntialLocationOfCentroid();
            centroids.push_back(Eigen::Vector3d(boost::tuples::get<0>(initial_location),
                                                boost::tuples::get<1>(initial_location),
                                                boost::tuples::get<2>(initial_location)));
        }else{
            Eigen::Vector3d temp_centroid = Eigen::Vector3d(0., 0., 0.);
            int supporting_feats_ctr = 0;
            std::vector<Feature::Id> this_filter_supporting_features = (*filter_it)->getSupportingFeatures();
            for (size_t feat_idx = 0; feat_idx < this_filter_supporting_features.size();feat_idx++)
            {
                if (this->_features_db->isFeatureStored(this_filter_supporting_features.at(feat_idx)))
                {
                    temp_centroid.x() += this->_features_db->getFeatureLastX(this_filter_supporting_features.at(feat_idx));
                    temp_centroid.y() += this->_features_db->getFeatureLastY(this_filter_supporting_features.at(feat_idx));
                    temp_centroid.z() += this->_features_db->getFeatureLastZ(this_filter_supporting_features.at(feat_idx));
                    supporting_feats_ctr++;
                }
            }
            centroids.push_back(temp_centroid / (double)supporting_feats_ctr);
        }
    }
    return centroids;
}

void MultiRBTracker::addPredictedState(const rbt_state_t& predicted_state, const double &predicted_state_timestamp_ns)
{
    BOOST_FOREACH(omip_msgs::RigidBodyPoseAndVelMsg predicted_pose_and_vel, predicted_state.rb_poses_and_vels)
    {
        BOOST_FOREACH(RBFilter::Ptr filter, this->_kalman_filters)
        {
            if (filter->getId() == predicted_pose_and_vel.rb_id)
            {
                //ROS_WARN_STREAM_NAMED("MultiRBTracker.addPredictedState", "Prediction for body " << predicted_pose_and_vel.rb_id);
                filter->setPredictedState(predicted_pose_and_vel);
            }
        }
    }
}

FeaturesDataBase::Ptr MultiRBTracker::getFeaturesDatabase()
{
    return this->_features_db;
}


#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

#include <boost/range/algorithm_ext/push_back.hpp>
#include <boost/range/irange.hpp>

#include <catch2/catch.hpp>
#include <cpp_algs.hpp>

TEST_CASE("Searching Algorithms test") {

    std::vector<int> v_int;
    boost::push_back(v_int, boost::irange<int>(1, 1001));
    REQUIRE(v_int.size() == 1000);

    std::vector<char> v_char;
    for (int i = 33; i < 128; i++) {
        char c = i;
        if (c != 'z')
            v_char.push_back(c);
    }
    REQUIRE(v_char.size() == 94);

    std::vector<double> v_double;
    boost::push_back(v_double, boost::irange<double>(1.0, 1001.0));
    REQUIRE(v_double.size() == 1000);

    SECTION("Linear Search") {
        REQUIRE(al::linearSearch<int>(v_int, 500) == 499);
        REQUIRE(al::linearSearch<int>(v_int, 0) == -1);

        REQUIRE(al::linearSearch<char>(v_char, 'a') == 64);
        REQUIRE(al::linearSearch<char>(v_char, 'z') == -1);

        REQUIRE(al::linearSearch<double>(v_double, 150.0) == 149);
        REQUIRE(al::linearSearch<double>(v_double, 0.0) == -1);
    }

    SECTION("Binary Search") {
        REQUIRE(al::binarySearch<int>(v_int, 500) == 499);
        REQUIRE(al::binarySearch<int>(v_int, 0) == -1);

        REQUIRE(al::binarySearch<char>(v_char, 'a') == 64);
        REQUIRE(al::binarySearch<char>(v_char, 'z') == -1);

        REQUIRE(al::binarySearch<double>(v_double, 150.0) == 149);
        REQUIRE(al::binarySearch<double>(v_double, 0.0) == -1);
    }

    SECTION("Ternary Search") {
        REQUIRE(al::ternarySearch<int>(v_int, 500) == 499);
        REQUIRE(al::ternarySearch<int>(v_int, 0) == -1);

        REQUIRE(al::ternarySearch<char>(v_char, 'a') == 64);
        REQUIRE(al::ternarySearch<char>(v_char, 'z') == -1);

        REQUIRE(al::ternarySearch<double>(v_double, 150.0) == 149);
        REQUIRE(al::ternarySearch<double>(v_double, 0.0) == -1);
    }
}


#include <geometry_msgs/TransformStamped.h>
#include <ros/ros.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2_ros/transform_broadcaster.h>
#include <tf_conversions/tf_eigen.h>

#include <Eigen/Dense>

geometry_msgs::TransformStamped get_static_transform(std::string parent_frame, std::string child_frame,
                                                     Eigen::Matrix<double, 4, 4, Eigen::RowMajor> &mat);

int main(int argc, char **argv)
{
  // init node
  ros::init(argc, argv, "static_transforms");
  ros::NodeHandle nh;
  ros::NodeHandle priv_nh_("~");

  // Setup Parameters
  std::vector<double> trans_world_param, trans12_param, trans13_param, trans14_param;

  if (!priv_nh_.getParam("world_sensor1_transform", trans_world_param))
    ROS_ERROR("Node: static_transforms:: Failed to get parameter sensor12_transform from server.");

  if (!priv_nh_.getParam("sensor1_sensor2_transform", trans12_param))
    ROS_ERROR("Node: static_transforms:: Failed to get parameter sensor12_transform from server.");

  if (!priv_nh_.getParam("sensor1_sensor3_transform", trans13_param))
    ROS_ERROR("Node: static_transforms:: Failed to get parameter sensor13_transform from server.");

  if (!priv_nh_.getParam("sensor1_sensor4_transform", trans14_param))
    ROS_ERROR("Node: static_transforms:: Failed to get parameter sensor14_transform from server.");

  // Populate transform matrices from parameters
  Eigen::Matrix<double, 4, 4, Eigen::RowMajor> transworld(trans_world_param.data());
  Eigen::Matrix<double, 4, 4, Eigen::RowMajor> trans12(trans12_param.data());
  Eigen::Matrix<double, 4, 4, Eigen::RowMajor> trans13(trans13_param.data());
  Eigen::Matrix<double, 4, 4, Eigen::RowMajor> trans14(trans14_param.data());

  //std::cout << trans12 << std::endl << trans13 << std::endl << trans14 << std::endl;

  // Define static tf publisher
  static tf2_ros::StaticTransformBroadcaster static_transform_broadcaster;

  geometry_msgs::TransformStamped world_transformStamped, sensor12_transformStamped, sensor13_transformStamped, sensor14_transformStamped;

  // Get tf messages built for all the sensors
  world_transformStamped = get_static_transform("sensor1_frame", "world_frame", transworld);
  sensor12_transformStamped = get_static_transform("sensor1_frame", "sensor2_frame", trans12);
  sensor13_transformStamped = get_static_transform("sensor1_frame", "sensor3_frame", trans13);
  sensor14_transformStamped = get_static_transform("sensor1_frame", "sensor4_frame", trans14);

  // Publish all tf messeges to tf_static
  const std::vector<geometry_msgs::TransformStamped> alltransforms{world_transformStamped,
                                                                   sensor12_transformStamped,
                                                                   sensor13_transformStamped,
                                                                   sensor14_transformStamped};
  static_transform_broadcaster.sendTransform(alltransforms);

  ROS_INFO_ONCE("Node: static_transforms:: Successfully published all static transforms.");

  ros::spin();

  return 0;
}

// TF messege builder
geometry_msgs::TransformStamped get_static_transform(std::string parent_frame, std::string child_frame,
                                                     Eigen::Matrix<double, 4, 4, Eigen::RowMajor> &mat)
{
  geometry_msgs::TransformStamped stat_trans;

  stat_trans.header.stamp = ros::Time::now();
  stat_trans.header.frame_id = parent_frame;
  stat_trans.child_frame_id = child_frame;

  stat_trans.transform.translation.x = mat(0, 3);
  stat_trans.transform.translation.y = mat(1, 3);
  stat_trans.transform.translation.z = mat(2, 3);

  Eigen::Matrix3d rot;
  rot = mat.block<3, 3>(0, 0);

  //std::cout << "rot: " << rot << std::endl;

  const Eigen::Quaterniond q(rot);

  stat_trans.transform.rotation.x = q.x();
  stat_trans.transform.rotation.y = q.y();
  stat_trans.transform.rotation.z = q.z();
  stat_trans.transform.rotation.w = q.w();

  ROS_INFO_ONCE("Transformation calculated : %s to %s", child_frame.c_str(), parent_frame.c_str());

  return stat_trans;
}