// This file is part of snark, a generic and flexible library for robotics research
// Copyright (c) 2014 The University of Sydney
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the name of the University of Sydney nor the
//    names of its contributors may be used to endorse or promote products
//    derived from this software without specific prior written permission.
//
// NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE
// GRANTED BY THIS LICENSE.  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT
// HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED
// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
// BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
// OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
// IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include <cmath>
#include <Eigen/Eigen>
#include <comma/base/exception.h>
#include "../rotation_matrix.h"
#include "polytope.h"

namespace snark { namespace geometry {

convex_polytope::convex_polytope( const std::vector< Eigen::VectorXd >& normals, const std::vector< double >& distances )
    : normals_( normals.size(), normals[0].size() )
    , distances_( distances.size() )
{
    if( !normals.size() || !normals[0].size() ) { COMMA_THROW( comma::exception, "normals cannot be empty or have zero dimensions" ); }
    if( normals.size() != distances.size() ) { COMMA_THROW( comma::exception, "normals and distances should be of same size, got "<< normals.size()<<" normals and " << distances.size() << " distances" ); }
    for( unsigned int i = 0; i < normals.size(); ++i )
    {
        normals_.row( i ) = normals[i].transpose().normalized();
        distances_( i ) = distances[i];
    }
}

convex_polytope::convex_polytope( const Eigen::MatrixXd& normals, const Eigen::VectorXd& distances )
    : normals_( normals )
    , distances_( distances )
{
    if( !normals.size() ) { COMMA_THROW( comma::exception, "normals cannot be empty" ); }
    if( normals.rows() != distances.size() ) { COMMA_THROW( comma::exception, "normals and distances should be of same size, got "<< normals.rows() << " normals and " << distances.size() << " distances" ); }
    for( unsigned int i = 0; i < normals.rows(); ++i ) { normals_.row( i ).normalize(); }
}

convex_polytope convex_polytope::transformed( const Eigen::Vector3d& translation, const snark::roll_pitch_yaw& orientation ) const
{
    convex_polytope p( *this );
    const Eigen::MatrixXd& rotation_matrix = snark::rotation_matrix::rotation( orientation );
    for( unsigned int i = 0; i < normals_.rows(); ++i ) // todo: quick and dirty; can it be done in a single operation?
    {
        p.normals_.row( i ) = rotation_matrix * normals_.row( i ).transpose();
        const Eigen::VectorXd& n = p.normals_.row( i );
        p.distances_[i] = n.dot( n * distances_[i] + translation );
    }
    return p;
}

bool convex_polytope::has( const Eigen::VectorXd &rhs, double epsilon ) const
{
    if( rhs.size() != normals_.cols() ) { COMMA_THROW( comma::exception, "expected same dimension as polytope: "<< normals_.cols()<<"; got dimension: " << rhs.size() ); }
    return ( ( normals_ * rhs - distances_ ).array() <= epsilon ).all();
}

} } // namespace snark{ { namespace geometry {


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

  @copyright 2016 NumScale SAS

  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_FUNCTION_DEFINITION_SHR_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_DEFINITION_SHR_HPP_INCLUDED

#include <boost/simd/config.hpp>
#include <boost/simd/detail/dispatch/function/make_callable.hpp>
#include <boost/simd/detail/dispatch/hierarchy/functions.hpp>
#include <boost/simd/detail/dispatch.hpp>

namespace boost { namespace simd
{
  namespace tag
  {
    BOOST_DISPATCH_MAKE_TAG(ext, shr_, boost::dispatch::elementwise_<shr_>);
  }

  namespace ext
  {
    BOOST_DISPATCH_FUNCTION_DECLARATION(tag, shr_);
  }

  BOOST_DISPATCH_CALLABLE_DEFINITION(tag::shr_,shr);

} }

#endif


//
// Copyright (c) 2017 CNRS
//
// This file is part of tsid
// tsid is free software: you can redistribute it
// and/or modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation, either version
// 3 of the License, or (at your option) any later version.
// tsid is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// General Lesser Public License for more details. You should have
// received a copy of the GNU Lesser General Public License along with
// tsid If not, see
// <http://www.gnu.org/licenses/>.
//

#include <iostream>

#include <boost/test/unit_test.hpp>
#include <boost/utility/binary.hpp>

#include <tsid/math/utils.hpp>
#include "tsid/robots/robot-wrapper.hpp"

using namespace tsid;
using namespace tsid::robots;
using namespace std;
using namespace pinocchio;

typedef pinocchio::Motion Motion;

BOOST_AUTO_TEST_SUITE ( BOOST_TEST_MODULE )

BOOST_AUTO_TEST_CASE ( test_set_gravity )
{

  const string romeo_model_path = TSID_SOURCE_DIR"/models/romeo";

  vector<string> package_dirs;
  package_dirs.push_back(romeo_model_path);
  string urdfFileName = package_dirs[0] + "/urdf/romeo.urdf";
  RobotWrapper robot(urdfFileName,
                     package_dirs,
                     pinocchio::JointModelFreeFlyer(),
                     false);

  Motion g = pinocchio::Motion::Zero();
  Motion init_gravity = robot.model().gravity;

  robot.setGravity(g);

  Motion no_gravity = robot.model().gravity;

  BOOST_CHECK(no_gravity != init_gravity);

}

BOOST_AUTO_TEST_SUITE_END ()


#include <iostream>

#include <vector>
using namespace std;

#include <boost/timer.hpp>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "DM_NRG.hpp"

void DM_NRG_CalcRhoN(vector<double> ParamsTemp,
		     bool UseFDM,CNRGbasisarray* pAcutN,
		     CNRGbasisarray* pAcutNp1,
		     CNRGbasisarray* pAbasisNp1,
		     CNRGmatrix* pRhoN,
		     CNRGmatrix* pRhoNp1){

  // New routine, which rotates the matrix first with uBLAS (!)
  // FASTER... like a Ferrari compared to an old Volkswagen beetle.

  vector< vector<int> > ChildSt_SameType;
  vector< vector<int> >::iterator arrayit;
  vector<int>::iterator iit;

  //Trace...
  double trace=0.0;
  double traceKK=0.0;
  double traceDD=0.0;

  // Check time
  boost::timer MyTime;
  double time_elapsed;

  bool debug=false;

  CNRGmatrix RhoNp1basis;

  double TempBar=ParamsTemp[0];
  double ZNdisc=ParamsTemp[1];
  double betabar=1/TempBar;

  cout << " DM-NRG: Calculating Rho N ..." << endl;

  cout << " Rotating Rho N... " << endl;

  MyTime.restart();
  RotateMatrix(pRhoNp1,pAcutNp1,(&RhoNp1basis));
  time_elapsed=MyTime.elapsed();

  cout << " ... done in " << time_elapsed << " secs. " << endl;

  int ThisBl=0;
//   if (pAcutN->Nshell==4){
//     cout << " Rho_N=5: " << endl;
//     pRhoNp1->PrintAllBlocks();
// //     pRhoNp1->PrintMatBlock(11,11);
// //     RhoNp1basis.PrintMatBlock(11,11);
//     ThisBl=11;
//     cout << " ibl = " << ThisBl << endl;
//     cout << " pAcutNp1 block structure: " << endl;
//     pAcutNp1->PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << pAcutNp1->GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << pAcutNp1->GetBlockLimit(ThisBl,1)
// 	 << endl;
//     cout << " RhoNp1basis block structure: " << endl;
//     RhoNp1basis.PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << RhoNp1basis.GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << RhoNp1basis.GetBlockLimit(ThisBl,1)
// 	 << endl;
//     ThisBl=18;
//     cout << " pAbasisNp1 block structure: " << endl;
// //     pAbasisNp1->PrintBlockQNumbers(ThisBl);
// //     cout << " Block " << ThisBl << " goes from ist = "
// // 	 << pAbasisNp1->GetBlockLimit(ThisBl,0)
// // 	 << " to "
// // 	 << pAbasisNp1->GetBlockLimit(ThisBl,1)
// // 	 << endl;
// //     pAbasisNp1->PrintBlockBasis(ThisBl);
//   }

  // Initial set up
  pRhoN->ClearAll();
  pRhoN->SyncNRGarray(*pAcutN);
  pRhoN->UpperTriangular=true;
  pRhoN->IsComplex=pRhoNp1->IsComplex;

  // Calculate matrix elements
  int i1=0;
  for (int ibl=0;ibl<pAcutN->NumBlocks();ibl++){

    // Rho is block diagonal
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockBegEnd.push_back(i1);

    int ist0=pAcutN->GetBlockLimit(ibl,0);
    int ist1=pAcutN->GetBlockLimit(ibl,1);

    // testing...

    // Get Sold from ibl

    int BlSize=pAcutN->GetBlockSize(ibl);
    if (debug)
      cout << " Setting Block " << ibl << " of " << pAcutN->NumBlocks()-1
	   << "  Size : " << BlSize << " ... ";

    MyTime.restart();
    // Loop in block states (and calculate rhoN(ist,jst))
    for (int ist=ist0;ist<=ist1;ist++){
      for (int jst=ist;jst<=ist1;jst++){

	// Ok, watch out for the " type" thing for SU(2) symmetry

	DM_NRG_SetChildSameType(ist,jst,pAcutN,pAbasisNp1,ChildSt_SameType);

	double auxMatEl=0.0;
	complex<double> cMatEl=ZeroC;
	// Loop over blocks in ChildSt_SameType (N+1)
	for (arrayit=ChildSt_SameType.begin();
	     arrayit<ChildSt_SameType.end();arrayit++){
	  iit=(*arrayit).begin();
	  int iblbasis_Np1=(*iit);
	  int istbasis_Np1=(*(iit+1));
	  int jstbasis_Np1=(*(iit+2));

	  // Get S_Np1 from iblbasis_Np1
	  // Get S_type= from istbasis_Np1
	  // Calculate Clebsh-Gordan coefs.

	  //
	  // No need to find corresponding block in AcutNp1:
	  //
	  // if block is not there, then
	  // RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1)
	  // will return 0.0 (in principle!)
	  //

	  if (!(RhoNp1basis.IsComplex))
	    auxMatEl+=RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1);
	  else
	    cMatEl+=RhoNp1basis.cGetMatEl(istbasis_Np1,jstbasis_Np1);

//  	  if ( (debug)&&(iblbasis_Np1==18)&&(ist==jst) ){
// 	    cout << " ibl = " << ibl
// 		 << " ist = " << ist
// 		 << " jst = " << jst
// 		 << " iblbasis_Np1 = " << iblbasis_Np1
// 		 << " istbasis_Np1 = " << istbasis_Np1
// 		 << " jstbasis_Np1 = " << jstbasis_Np1
// 		 << " mat el = " << RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1)
// 		 << " auxMatEl = " << auxMatEl
// 		 << endl;
// 	  }

	}
	// end loop in Child states
	// What if states are NOT kept
	// // Diagonal only
	if ( (UseFDM)&&(jst==ist)&&(pAcutN->CheckKept(ist,false)) ){
// 	  cout << " Disc-disc : ist=" << ist <<  endl;
// 	  cout << " size ChildSt[ist]= " << pAcutN->ChildStates[ist].size();
	  if ( dNEqual(ZNdisc,0.0) )
	    auxMatEl=exp(-betabar*(pAcutN->dEn[ist]))/ZNdisc;
// 	  cout << " auxMatEl= " << auxMatEl;
//	  cout << endl;
	}
	// end if discarded, discarded

	if (ist==jst){
	  trace+=auxMatEl;
	  if (pAcutN->Kept[ist]) traceKK+=auxMatEl;
	  else traceDD+=auxMatEl;

	}
	// diagonal terms
	if (!(pRhoN->IsComplex))
	  pRhoN->MatEl.push_back(auxMatEl);
	else
	  pRhoN->MatElCplx.push_back(cMatEl);

	i1++;
      }
      // end loop in jst
    }
    // end loop in ist
    pRhoN->MatBlockBegEnd.push_back(i1-1);

    time_elapsed=MyTime.elapsed();
    if (debug)
     cout << " ... done in " << time_elapsed << " secs." << endl;

  }
  // end loop in pAcutN blocks

  cout << " ... Rho N  done. " << endl;
  cout << " Trace = " << trace
       << " TraceKK = " << traceKK
       << " TraceDD = " << traceDD << endl;

//   if (pAcutN->Nshell==4){
//     cout << " Rho_N=4: " << endl;
//     //pRhoN->PrintMatBlock(10,10);
//     ThisBl=10;
//     cout << " ibl = " << ThisBl << endl;
//     cout << " pRhoN block structure: " << endl;
//     pRhoN->PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << pRhoN->GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << pRhoN->GetBlockLimit(ThisBl,1)
// 	 << endl;
//     //pRhoN->PrintMatBlock(13,13);
//     ThisBl=13;
//     cout << " ibl = " << ThisBl << endl;
//     cout << " pRhoN block structure: " << endl;
//     pRhoN->PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << pRhoN->GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << pRhoN->GetBlockLimit(ThisBl,1)
// 	 << endl;
//   }

}
// end set RhoN new

/////////////////////////////////////////
/////   Version with SU(2) symmetry /////
/////////////////////////////////////////

void DM_NRG_CalcRhoN_withSU2(vector<double> ParamsTemp,
			     bool UseFDM,
			     CNRGbasisarray* pAcutN,
			     CNRGbasisarray* pAcutNp1,
			     CNRGbasisarray* pAbasisNp1,
			     CNRGbasisarray* pSingleSite,
			     CNRGmatrix* pRhoN,
			     CNRGmatrix* pRhoNp1){

  // New routine, which rotates the matrix first with uBLAS (!)
  // FASTER... like a Ferrari compared to an old Volkswagen beetle.

  vector< vector<int> > ChildSt_SameType;
  vector< vector<int> >::iterator arrayit;
  vector<int>::iterator iit;

  // Check time
  boost::timer MyTime;
  double time_elapsed;

  //Trace...
  double trace=0.0;

  bool debug=false;

  CNRGmatrix RhoNp1basis;

  double TempBar=ParamsTemp[0];
  double ZNdisc=ParamsTemp[1];
  double betabar=1/TempBar;

  cout << " DM-NRG: Calculating Rho N (w SU(2))..." << endl;

  cout << " Rotating Rho N... " << endl;

  MyTime.restart();
  RotateMatrix(pRhoNp1,pAcutNp1,(&RhoNp1basis));
  time_elapsed=MyTime.elapsed();

  cout << " ... done in " << time_elapsed << " secs. " << endl;

  int ThisBl=0;

  // Initial set up
  pRhoN->ClearAll();
  pRhoN->SyncNRGarray(*pAcutN);
  pRhoN->UpperTriangular=true;

  // Calculate matrix elements
  int i1=0;
  for (int ibl=0;ibl<pAcutN->NumBlocks();ibl++){

    // Rho is block diagonal
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockBegEnd.push_back(i1);

    int ist0=pAcutN->GetBlockLimit(ibl,0);
    int ist1=pAcutN->GetBlockLimit(ibl,1);

    // testing...

    // Get Sold from ibl
    double Sold=0.0;
    if (pAcutN->totalS){
      Sold=pAcutN->GetQNumber(ibl,pAcutN->Sqnumbers[0]);
      // only a single SU(2) for now
    }

    int BlSize=pAcutN->GetBlockSize(ibl);
    if (debug)
      cout << " Setting Block " << ibl << " of " << pAcutN->NumBlocks()-1
	   << "  Size : " << BlSize << " ... ";

    MyTime.restart();
    // Loop in block states (and calculate rhoN(ist,jst))
    for (int ist=ist0;ist<=ist1;ist++){
      for (int jst=ist;jst<=ist1;jst++){

	// Ok, watch out for the " type" thing for SU(2) symmetry
	DM_NRG_SetChildSameType(ist,jst,pAcutN,pAbasisNp1,ChildSt_SameType);

	double auxMatEl=0.0;
	complex<double> cMatEl=ZeroC;
	// Loop over blocks in ChildSt_SameType (N+1)
	for (arrayit=ChildSt_SameType.begin();
	     arrayit<ChildSt_SameType.end();arrayit++){
	  iit=(*arrayit).begin();
	  int iblbasis_Np1=(*iit);
	  int istbasis_Np1=(*(iit+1));
	  int jstbasis_Np1=(*(iit+2));

	  // Get S_Np1 from iblbasis_Np1
	  double Si=0.0;
	  if (pAbasisNp1->totalS){
	    Si=pAbasisNp1->GetQNumber(iblbasis_Np1,pAbasisNp1->Sqnumbers[0]);
	    // only a single SU(2) for now
	  }
	  // Get S_type= from istbasis_Np1
	  int itype_Np1=pAbasisNp1->iType[istbasis_Np1];
	  double Stilde=0.0;
	  if (pSingleSite->totalS){
	    //Stilde=pSingleSite->GetQNumber(ibls,pSingleSite->Sqnumbers[0]);
	    Stilde=pSingleSite->GetQNumberFromSt(itype_Np1,pSingleSite->Sqnumbers[0]);
	    // only a single SU(2) for now
	  }

	  // Calculate Clebsh-Gordan coefs.

	  double CGfactor=0.0;
	  for (double Sz=-Si;Sz<=Si;Sz+=1.0){
	    double auxCG=CGordan(Sold,Sold,Stilde,Sz-Sold,Si,Sz);
	    CGfactor+=auxCG*auxCG; // Square it!! (oct 2010)
	  }
	  // end calcCG

	  //
	  // No need to find corresponding block in AcutNp1:
	  //
	  // if block is not there, then
	  // RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1)
	  // will return 0.0 (in principle!)
	  //
	  //auxMatEl+=CGfactor*RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1);
	  if (!(RhoNp1basis.IsComplex))
	    auxMatEl+=CGfactor*RhoNp1basis.GetMatEl(istbasis_Np1,jstbasis_Np1);
	  else
	    cMatEl+=CGfactor*RhoNp1basis.cGetMatEl(istbasis_Np1,jstbasis_Np1);

	}
	// end loop in Child states
	// What if states are NOT kept
	// // Diagonal only
	// What about the CB coefficient? Should we do (???)
	//auxExp=(2.0*Sold+1.0)*exp(-betabar*pAcutNp1->dEn[ist]);
	// NO! IT GIVES WRONG RESULTS...

	if ( (UseFDM)&&(jst==ist)&&(pAcutN->CheckKept(ist,false)) ){
	  //cout << " Disc-disc : ist=" << ist <<  endl;
	  //cout << " size ChildSt[ist]= " << pAcutN->ChildStates[ist].size();
	  if ( dNEqual(ZNdisc,0.0) )
	    auxMatEl=exp(-betabar*(pAcutN->dEn[ist]))/ZNdisc;
	}
	// end if discarded, discarded

	if (ist==jst) trace+=(2.0*Sold+1.0)*auxMatEl;

	if (pRhoN->IsComplex)
	  pRhoN->MatElCplx.push_back(cMatEl);
	else
	  pRhoN->MatEl.push_back(auxMatEl);

	i1++;
      }
      // end loop in jst
    }
    // end loop in ist
    pRhoN->MatBlockBegEnd.push_back(i1-1);

    time_elapsed=MyTime.elapsed();
    if (debug)
     cout << " ... done in " << time_elapsed << " secs." << endl;

  }
  // end loop in pAcutN blocks

  cout << " ... Rho N  done. Trace(2Si+1) = " << trace << endl;

//   if (pAcutN->Nshell==4){
//     cout << " Rho_N=4: " << endl;
//     //pRhoN->PrintMatBlock(10,10);
//     ThisBl=10;
//     cout << " ibl = " << ThisBl << endl;
//     cout << " pRhoN block structure: " << endl;
//     pRhoN->PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << pRhoN->GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << pRhoN->GetBlockLimit(ThisBl,1)
// 	 << endl;
//     //pRhoN->PrintMatBlock(13,13);
//     ThisBl=13;
//     cout << " ibl = " << ThisBl << endl;
//     cout << " pRhoN block structure: " << endl;
//     pRhoN->PrintBlockQNumbers(ThisBl);
//     cout << " Block " << ThisBl << " goes from ist = "
// 	 << pRhoN->GetBlockLimit(ThisBl,0)
// 	 << " to "
// 	 << pRhoN->GetBlockLimit(ThisBl,1)
// 	 << endl;
//   }

}
// end set RhoN new

/////////////////////////////////////////
/////////////////////////////////////////
/////////////////////////////////////////

void DM_NRG_CalcRhoN_old(CNRGbasisarray* pAcutN,
		     CNRGbasisarray* pAcutNp1,
		     CNRGbasisarray* pAbasisNp1,
		     CNRGmatrix* pRhoN,
		     CNRGmatrix* pRhoNp1){

  // Poor old routine, which does not use uBLAS... (slow old man)

  vector< vector<int> > ChildSt_SameType;
  vector< vector<int> >::iterator arrayit;
  vector<int>::iterator iit;

  // Check time
  boost::timer MyTime;
  double time_elapsed;

  cout << " DM-NRG: Calculating Rho N ..." << endl;

  // Initial set up
  pRhoN->ClearAll();
  pRhoN->SyncNRGarray(*pAcutN);
  pRhoN->UpperTriangular=true;

  // Calculate matrix elements
  int i1=0;
  for (int ibl=0;ibl<pAcutN->NumBlocks();ibl++){
  // Testing in a few blocks
  //for (int ibl=0;ibl<=8;ibl+=8){
    // Rho is block diagonal
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockMap.push_back(ibl);
    pRhoN->MatBlockBegEnd.push_back(i1);

    int ist0=pAcutN->GetBlockLimit(ibl,0);
    int ist1=pAcutN->GetBlockLimit(ibl,1);

    // testing...

    int BlSize=pAcutN->GetBlockSize(ibl);
    cout << " Setting Block " << ibl << " of " << pAcutN->NumBlocks()-1
	 << "  Size : " << BlSize << " ... ";

    MyTime.restart();
    // Loop in block states (and calculate rhoN(ist,jst))
    for (int ist=ist0;ist<=ist1;ist++){
      for (int jst=ist;jst<=ist1;jst++){

	double auxMatEl=0.0;
	// Debugging
// 	if ( (ist==61)&&(jst==61) ){
// 	  cout << " ist = " << ist << " jst = " << jst << endl;
// 	}

	DM_NRG_SetChildSameType(ist,jst,pAcutN,pAbasisNp1,ChildSt_SameType);

// 	if ( (ist==61)&&(jst==61) ){
// 	  cout << "bl  ist_Np1  jst_Np1 " << endl;
// 	  for (arrayit=ChildSt_SameType.begin();
// 	       arrayit<ChildSt_SameType.end();arrayit++){
// 	    for(iit=(*arrayit).begin();iit<(*arrayit).end();iit++){
// 	      cout << (*iit) << " ";
// 	    }
// 	    cout << endl;
// 	  }
// 	}
	// debug

	// Loop over blocks in ChildSt_SameType (N+1)

	for (arrayit=ChildSt_SameType.begin();
	     arrayit<ChildSt_SameType.end();arrayit++){
	  iit=(*arrayit).begin();
	  int iblbasis_Np1=(*iit);
	  int istbasis_Np1=(*(iit+1));
	  int jstbasis_Np1=(*(iit+2));

// 	  if ( (ist==61)&&(jst==61) ){
// 	    cout << " iblbasis_Np1 = " << iblbasis_Np1
// 		 << " istbasis_Np1 = "<< istbasis_Np1
// 		 << " jstbasis_Np1 = "<< jstbasis_Np1
// 		 << endl;
// 	  }
	  //
	  // Find corresponding block in AcutNp1 (compare QNumbers)
	  //
	  double* qnums=new double [pAbasisNp1->NQNumbers];
	  for (int iqn=0;iqn<pAbasisNp1->NQNumbers;iqn++){
	    qnums[iqn]=pAbasisNp1->GetQNumber(iblbasis_Np1,iqn);
	  }
	  int ibl_Np1=pAcutNp1->GetBlockFromQNumbers(qnums);
// 	  if (ibl_Np1==-1){
// 	    cout << " Block  not found in Acut N+1: ";
// 	    cout << " QNs = " ;
// 	    for (int iqn=0;iqn<pAbasisNp1->NQNumbers;iqn++){
// 	      cout << qnums[iqn] << " ";
// 	    }
// 	    cout << endl;
// 	  }
	  // if block not found
	  delete [] qnums;

// 	  if ( (ist==61)&&(jst==61) ){
// 	    cout << " ibl_Np1 = " << ibl_Np1 << endl;
// 	    if (ibl_Np1>0)
// 	      cout << " states: from " << pAcutNp1->GetBlockLimit(ibl_Np1,0)
// 		   << " to " << pAcutNp1->GetBlockLimit(ibl_Np1,1) << endl;
// 	  }
	  // debug
	  // Loop over states in the block
	  if (ibl_Np1>0){ // Block exists in cut basis (kept)
	    for(int ist_Np1=pAcutNp1->GetBlockLimit(ibl_Np1,0);
		ist_Np1<=pAcutNp1->GetBlockLimit(ibl_Np1,1);
		ist_Np1++){
	      for(int jst_Np1=pAcutNp1->GetBlockLimit(ibl_Np1,0);
		  jst_Np1<=pAcutNp1->GetBlockLimit(ibl_Np1,1);
		  jst_Np1++){
	       double U_wi=pAcutNp1->GetEigVecComponent(ist_Np1,istbasis_Np1);
	       double U_wpj=pAcutNp1->GetEigVecComponent(jst_Np1,jstbasis_Np1);
	       double rho_wwp=pRhoNp1->GetMatEl(ist_Np1,jst_Np1);

	       auxMatEl+=rho_wwp*U_wi*U_wpj;

// 	      if ( (ist==61)&&(jst==61) ){
// 		cout << " ist_Np1 = " << ist_Np1
// 		     << " jst_Np1 = " << jst_Np1
// 		     << endl;
// 		cout << " U_wi= " << U_wi
// 		     << " U_wpj= " << U_wpj
// 		     << " rho_wwp= " << rho_wwp
// 		     << endl;
// 		cout << " auxMatEl = " << auxMatEl << endl;
// 	      }
	      // debug
	      }
	      // end loop in j_states in ibl_Np1
	    }
	    // end loop in i_states in ibl_Np1
	  }
	  // end if block exists in Acut

// 	  if ( (ist==61)&&(jst==61) ){
// 	    cout << " auxMatEl = " << auxMatEl << endl;
// 	  }

	}
	// end loop in ChildSt_SameType

	pRhoN->MatEl.push_back(auxMatEl);
	i1++;
      }
      // loop over jst
    }
    // loop over ist
    pRhoN->MatBlockBegEnd.push_back(i1-1);

    time_elapsed=MyTime.elapsed();
    cout << " ... done in " << time_elapsed << " secs." << endl;

  }
  // end loop in blocks

  cout << " ... Rho N done." << endl;

}
// end subroutine

/////////////////////////////////////////
/////////////////////////////////////////
/////////////////////////////////////////


/**
BSD 3-Clause License

This file is part of the code accompanying the paper
Gradient-SDF: A Semi-Implicit Surface Representation for 3D Reconstruction
by Christiane Sommer*, Lu Sang*, David Schubert, and Daniel Cremers (* denotes equal contribution).

Copyright (c) 2021, Christiane Sommer and Lu Sang.
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

* Neither the name of the copyright holder nor the names of its
  contributors may be used to endorse or promote products derived from
  this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

// standard includes
#include <iostream>
#include <fstream>
#include <vector>
// library includes
#include <cstdlib>
#include <Eigen/Dense>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <CLI/CLI.hpp>
// class includes
#include "Timer.h"
#include "normals/NormalEstimator.h"
#include "sdf_tracker/MapGradPixelSdf.h"
#include "sdf_tracker/RigidPointOptimizer.h"
#include "img_loader/img_loader.h"
#include "ps_optimizer/PhotometricOptimizer.h"
#include "ps_optimizer/ColorUpsampler.h"
#include "ps_optimizer/SharpDetector.h"
// own includes
#include "mat.h"

void sampleKeyFrame(std::vector<int>& key_frames, std::vector<std::string>& key_stamps, std::vector<std::shared_ptr<cv::Mat>>& key_images, std::vector<Mat4f, Eigen::aligned_allocator<Mat4f>>& key_poses, int max_num);
/**
 * main function
 */
int main(int argc, char *argv[]) {

    Timer T;
    // Default input sequence in folder
    std::string input = "";
    std::string output = "../results/";
    std::string stype = "map-gp";
    std::string dtype = "";
    size_t first = 0;
    size_t last = 300;
    float voxel_size = 0.01; // Voxel size in m
    float z_max = 3.5; // maximal depth to take into account in m
    float truncation_factor = 5; // truncation in voxels
    float sharp_threshold = 0.0001;
    int num_frame = 30;

    CLI::App app{"Hash Table-Based 3D Scanning and Texture Optimization"};
    app.add_option("--input", input, "folder of input sequence");
    app.add_option("--results", output, "folder to store results in");
    app.add_option("--first", first, "number of first frame to be processed (default: 0)");
    app.add_option("--last", last, "number of last frame (default: all)");
    app.add_option("--data-type", dtype, "type of dataset");
    app.add_option("--voxel-size", voxel_size, "voxel size in meters (default: 0.01)");
    app.add_option("--trunc", truncation_factor, "truncation in multiples of voxels (default: 5)");
    app.add_option("--key-frame", num_frame, "number of the key frames to be sampled (default: 20)");

    // parse input arguments
    try {
        app.parse(argc, argv);
    } catch (const CLI::ParseError& e) {
        return app.exit(e);
    }

    // parse dataset type
    enum class DataType {
        TUM_RGBD,
        SYNTH,
        PRINTED_3D,
        REDWOOD,
    };
    DataType DT;
    if (dtype == "tum" || dtype == "tumrgbd") {
        DT = DataType::TUM_RGBD;
        sharp_threshold = 0.026;
    }
    else if (dtype == "synth") {
        DT = DataType::SYNTH;
    }
    else if (dtype == "printed") {
        DT = DataType::PRINTED_3D;
        sharp_threshold = 0.026;
    }
    else if (dtype == "rw" || dtype == "redwood") {
        DT = DataType::REDWOOD;
        sharp_threshold = 0.033;
    }
    else {
        std::cerr << "Your specified dataset type is not supported (yet)." << std::endl;
        return 1;
    }

    // create image loader
    ImageLoader* loader;
    switch (DT) {
    case DataType::TUM_RGBD :
        loader = new TumrgbdLoader(input);
        break;
    case DataType::SYNTH :
        loader = new SynthLoader(input);
        break;
    case DataType::PRINTED_3D :
        loader = new Printed3dLoader(input);
        break;
    case DataType::REDWOOD :
        loader = new RedwoodLoader(input);
        break;
    default:
        std::cerr << "Specified dataset type not recognized, return" << std::endl;
        return 1;
    }

    // Load camera intrinsics
    if (!loader->load_intrinsics("intrinsics.txt")) {
        std::cerr << "No intrinsics file found in " << input << "!" << std::endl;
        return 1;
    }
    const Mat3f K = loader->K();
    std::cout << "K: " << std::endl << K << std::endl;

    // create normal estimator
    T.tic();
    cv::NormalEstimator<float>* NEst;
    NEst = new cv::NormalEstimator<float>(640, 480, K, cv::Size(2*5+1, 2*5+1));
    T.toc("Init normal estimation");

    std::stringstream stream;
    stream << std::fixed << std::setprecision(3) << voxel_size;
    std::string voxel_size_str = stream.str();

    // Trunction distance for the tSDF
    const float truncation = truncation_factor * voxel_size;

    Sdf* tSDF;
    RigidPointOptimizer* pOpt;
    PhotometricOptimizer* cOpt;

    std::ofstream pose_file(output + stype + "_" + voxel_size_str + "m_T" + std::to_string(int(truncation_factor)) + "_poses.txt");

    std::vector<Mat4f, Eigen::aligned_allocator<Mat4f>> poses;
    // poses.push_back(Mat4f::Identity());
    std::vector<int> valid_frames;
    std::vector<int> invalid_frames;
    std::vector<int> keyframes;
    valid_frames.push_back(0);
    keyframes.push_back(0);
    // vector of sampled poses and frames
    std::vector<std::string> key_stamps;
    std::vector<std::shared_ptr<cv::Mat>> key_images;
    std::vector<Mat4f, Eigen::aligned_allocator<Mat4f>> key_poses;
    key_poses.push_back(Mat4f::Identity());

    // Frames to be processed
    cv::Mat color, depth;
    int dist_to_last_keyframe = 0;

    // Proceed until first frame
    for (size_t i = 0; i < first; ++i) {
        loader->load_next(color, depth);
    }

    // Actual scanning loop
    for (size_t i = first; i <= last; ++i) {
        std::cout << "Working on frame: " << i << std::endl;

        // Load data
        T.tic();
        if (!loader->load_next(color, depth)) {
            std::cerr << " -> Frame " << i << " could not be loaded!" << std::endl;
            T.toc("Load data");
            break;
        }
        T.toc("Load data");

        // Get initial volume pose from centroid of first depth map
        if (i == first) {
            // create SDF data
            T.tic();
            tSDF = new MapGradPixelSdf(voxel_size, truncation);
            T.toc("Create Sdf");
            // Initialize tSDF
            T.tic();
            tSDF->setup(color, depth, K, NEst);
            T.toc("Integrate depth data into Sdf");

			// Initialize optimizers
			T.tic();
			pOpt = new RigidPointOptimizer(tSDF);
			T.toc("Create RigidOptimizer");
            T.tic();
            cOpt = new PhotometricOptimizer(static_cast<MapGradPixelSdf*>(tSDF), voxel_size, K, output);
            T.toc("Create PhotometricOptimizer");
            key_stamps.push_back(loader->rgb_timestamp());
            cv::Mat new_color;
            color.copyTo(new_color);
            key_images.push_back(std::make_shared<cv::Mat>(new_color));

        }
		else {
            // Perform optimization
            T.tic();
            bool conv = pOpt->optimize(depth, K);
            T.toc("Point optimization");
            // Integrate data into model
            if (conv) {
                valid_frames.push_back(i-first);
                T.tic();
                tSDF->update(color, depth, K, pOpt->pose(), NEst);
                T.toc("Integrate depth data into Sdf");

                if (sharpDetector(color, sharp_threshold) || dist_to_last_keyframe > 5)
                {
                    dist_to_last_keyframe = 0;
                    keyframes.push_back(i-first);
                    key_stamps.push_back(loader->rgb_timestamp());
                    key_poses.push_back(pOpt->pose().matrix());
                    cv::Mat new_color;
                    color.copyTo(new_color);
                    key_images.push_back(std::make_shared<cv::Mat>(new_color));
                }
                else{

                    dist_to_last_keyframe++;
                }
            }
            else {
                invalid_frames.push_back(i-first);
            }
		}
		// write timestamp + pose in tx ty tz qx qy qz qw format
		std::cout << "Current pose:" << std::endl
		          << pOpt->pose().matrix() << std::endl;
        poses.push_back(pOpt->pose().matrix());
		Vec3f t(pOpt->pose().translation());
		Eigen::Quaternion<float> q(pOpt->pose().rotationMatrix());
		pose_file << loader->depth_timestamp() << " "
		          << t[0] << " " << t[1] << " " << t[2] << " "
		          << q.x() << " " << q.y() << " " << q.z() << " " << q.w() << "\n";
    }
    pose_file.close();

    // extract mesh and write to file
    T.tic();
    if (!tSDF->extract_mesh(output + "mesh_lr.ply")) {
        std::cerr << "Could not save mesh!" << std::endl;
    }
    T.toc("Save mesh to disk");

    // extract point cloud and write to file
    T.tic();
    if (!tSDF->extract_pc(output + "cloud_lr.ply")) {
        std::cerr << "Could not save point cloud!" << std::endl;
    }
    T.toc("Save point cloud to disk");

    sampleKeyFrame(keyframes, key_stamps, key_images, key_poses, num_frame);

    cOpt->setImages(key_images);
    cOpt->setKeyframes(keyframes);
    cOpt->setPoses(key_poses);
    cOpt->setKeytimestamps(key_stamps);
    cOpt->optimize();

    // up sampling
    ColorUpsampler tOpt((static_cast<MapGradPixelSdf*>(tSDF))->get_tsdf(),
                    (static_cast<MapGradPixelSdf*>(tSDF))->get_vis(),
                    key_images,
                    key_poses,
                    keyframes,
                    voxel_size,
                    K);

    std::cout << "Color optimizer constructed" << std::endl;
    tOpt.computeColor(); // just to be able to compare better
    tOpt.extractMesh(output + "coarse_BA_mesh_after_upsample");
    tOpt.extractCloud(output + "coarse_BA_cloud_after_upsample");

    return 0;
}

//! To selected limited number of frames if there are too many input as key frames.
void sampleKeyFrame(std::vector<int>& key_frames, std::vector<std::string>& key_stamps, std::vector<std::shared_ptr<cv::Mat>>& key_images, std::vector<Mat4f, Eigen::aligned_allocator<Mat4f>>& key_poses, int max_num){
    if (key_frames.size() < max_num ){
        return;
    }
    max_num -= 1;
    float step = static_cast<float>(key_frames.size()) / static_cast<float>(max_num);
    std::vector<int> frames;
    std::vector<std::string> stamps;
    std::vector<std::shared_ptr<cv::Mat>> images;
    std::vector<Mat4f, Eigen::aligned_allocator<Mat4f>> poses;
    float idx = 0;
    for(int count = 0; count < max_num; count++){
        int i = static_cast<int>(idx);
        frames.push_back(key_frames[i]);
        stamps.push_back(key_stamps[i]);
        images.push_back(key_images[i]);
        poses.push_back(key_poses[i]);
        idx+=step;
    }
    frames.push_back(key_frames.back()); //we need the last frame for resize the visibility vector
    stamps.push_back(key_stamps.back());
    images.push_back(key_images.back());
    poses.push_back(key_poses.back());

    key_frames = frames;
    key_stamps = stamps;
    key_poses = poses;
    key_images = images;
}


//***************************************************************************
//* Copyright (c) 2015 Saint Petersburg State University
//* Copyright (c) 2011-2014 Saint Petersburg Academic University
//* All Rights Reserved
//* See file LICENSE for details.
//***************************************************************************

#pragma once

#include "adt/iterator_range.hpp"
#include "assembly_graph/core/action_handlers.hpp"
#include "utils/parallel/openmp_wrapper.h"
#include "paired_info_buffer.hpp"
#include <type_traits>
#include <boost/iterator/iterator_facade.hpp>
#include <btree/safe_btree_map.h>
#include <set>

namespace omnigraph {

namespace de {

template<typename G, typename Traits, template<typename, typename> class Container>
class PairedIndex : public PairedBuffer<G, Traits, Container> {
    typedef PairedIndex<G, Traits, Container> self;
    typedef PairedBuffer<G, Traits, Container> base;

    typedef typename base::InnerHistogram InnerHistogram;
    typedef typename base::InnerHistPtr InnerHistPtr;
    typedef typename base::InnerPoint InnerPoint;

    using typename base::EdgePair;

public:
    using typename base::Graph;
    using typename base::EdgeId;
    typedef typename base::InnerMap InnerMap;
    typedef typename base::StorageMap StorageMap;
    using typename base::Point;

    typedef omnigraph::de::Histogram<Point> Histogram;

    //--Data access types--

    typedef typename StorageMap::const_iterator ImplIterator;

    //---------------- Data accessing methods ----------------

    /**
     * @brief Underlying raw implementation data (for custom iterator helpers).
     */
    ImplIterator data_begin() const {
        return this->storage_.begin();
    }

    /**
     * @brief Underlying raw implementation data (for custom iterator helpers).
     */
    ImplIterator data_end() const {
        return this->storage_.end();
    }

    /**
     * @brief Smart proxy set representing a composite histogram of points between two edges.
     * @detail You can work with the proxy just like any constant set.
     *         The only major difference is that it returns all consisting points by value,
     *         because some of them don't exist in the underlying sets and are
     *         restored from the conjugate info on-the-fly.
     */
    class HistProxy {

    public:
        /**
         * @brief Iterator over a proxy set of points.
         */
        class Iterator: public boost::iterator_facade<Iterator, Point, boost::bidirectional_traversal_tag, Point> {

            typedef typename InnerHistogram::const_iterator InnerIterator;

        public:
            Iterator(InnerIterator iter, DEDistance offset, bool back = false)
                    : iter_(iter), offset_(offset), back_(back)
            {}

        private:
            friend class boost::iterator_core_access;

            Point dereference() const {
                auto i = iter_;
                if (back_) --i;
                Point result = Traits::Expand(*i, offset_);
                if (back_)
                    result.d = -result.d;
                return result;
            }

            void increment() {
                back_ ? --iter_ : ++iter_;
            }

            void decrement() {
                back_ ? ++iter_ : --iter_;
            }

            inline bool equal(const Iterator &other) const {
                return iter_ == other.iter_ && back_ == other.back_;
            }

            InnerIterator iter_; //current position
            DEDistance offset_; //edge length
            bool back_;
        };

        /**
         * @brief Returns a wrapper for a histogram.
         */
        HistProxy(const InnerHistogram& hist, DEDistance offset = 0, bool back = false)
            : hist_(hist), offset_(offset), back_(back)
        {}

        /**
         * @brief Returns an empty proxy (effectively a Null object pattern).
         */
        static const InnerHistogram& empty_hist() {
            static InnerHistogram res;
            return res;
        }

        Iterator begin() const {
            return Iterator(back_ ? hist_.end() : hist_.begin(), offset_, back_);
        }

        Iterator end() const {
            return Iterator(back_ ? hist_.begin() : hist_.end(), offset_, back_);
        }

        /**
         * @brief Finds the point with the minimal distance.
         */
        Point min() const {
            VERIFY(!empty());
            return *begin();
        }

        /**
         * @brief Finds the point with the maximal distance.
         */
        Point max() const {
            VERIFY(!empty());
            return *--end();
        }

        /**
         * @brief Returns the copy of all points in a simple flat histogram.
         */
        Histogram Unwrap() const {
            return Histogram(begin(), end());
        }

        size_t size() const {
            return hist_.size();
        }

        bool empty() const {
            return hist_.empty();
        }

    private:
        const InnerHistogram& hist_;
        DEDistance offset_;
        bool back_;
    };

    typedef typename HistProxy::Iterator HistIterator;

    //---- Traversing edge neighbours ----

    using EdgeHist = std::pair<EdgeId, HistProxy>;

    /**
     * @brief A proxy map representing neighbourhood of an edge,
     *        where `Key` is the graph edge ID and `Value` is the proxy histogram.
     * @detail You can work with the proxy just like with any constant map.
     *         The only major difference is that it returns all consisting pairs by value,
     *         because proxies are constructed on-the-fly.
     */
    class EdgeProxy {
    public:

        /**
         * @brief Iterator over a proxy map.
         * @detail For a full proxy, traverses both straight and conjugate pairs.
         *         For a half proxy, traverses only lesser pairs (i.e., (a,b) where (a,b)<=(b',a')) of edges.
         */
        class Iterator: public boost::iterator_facade<Iterator, EdgeHist, boost::forward_traversal_tag, EdgeHist> {

            typedef typename InnerMap::const_iterator InnerIterator;

            void Skip() { //For a half iterator, skip conjugate pairs
                while (half_ && iter_ != stop_ && !index_.IsCanonical(edge_, iter_->first))
                    ++iter_;
            }

        public:
            Iterator(const PairedIndex &index, InnerIterator iter, InnerIterator stop, EdgeId edge, bool half)
                    : index_ (index)
                    , iter_(iter)
                    , stop_(stop)
                    , edge_(edge)
                    , half_(half)
            {
                Skip();
            }

            void increment() {
                ++iter_;
                Skip();
            }

        private:
            friend class boost::iterator_core_access;

            bool equal(const Iterator &other) const {
                return iter_ == other.iter_;
            }

            EdgeHist dereference() const {
                const auto& hist = *iter_->second;
                return std::make_pair(iter_->first, HistProxy(hist, index_.CalcOffset(edge_)));
            }

        private:
            const PairedIndex &index_; //TODO: get rid of this somehow
            InnerIterator iter_, stop_;
            EdgeId edge_;
            bool half_;
        };

        EdgeProxy(const PairedIndex &index, const InnerMap& map, EdgeId edge, bool half = false)
            : index_(index), map_(map), edge_(edge), half_(half)
        {}

        Iterator begin() const {
            return Iterator(index_, map_.begin(), map_.end(), edge_, half_);
        }

        Iterator end() const {
            return Iterator(index_, map_.end(), map_.end(), edge_, half_);
        }

        HistProxy operator[](EdgeId e2) const {
            if (half_ && !index_.IsCanonical(edge_, e2))
                return HistProxy::empty_hist();
            return index_.Get(edge_, e2);
        }

        bool empty() const {
            return map_.empty();
        }

    private:
        const PairedIndex& index_;
        const InnerMap& map_;
        EdgeId edge_;
        //When false, represents all neighbours (consisting both of directly added data and "restored" conjugates).
        //When true, proxifies only half of the added edges.
        bool half_;
    };

    typedef typename EdgeProxy::Iterator EdgeIterator;

    //---------------- Constructor ----------------

    PairedIndex(const Graph &graph)
        : base(graph) {}

    /**
     * @brief Adds a lot of info from another index, using fast merging strategy.
     *        Should be used instead of point-by-point index merge.
     */
    template<class Buffer>
    void Merge(Buffer& index_to_add) {
        if (index_to_add.size() == 0)
            return;

        auto locked_table = index_to_add.lock_table();
        for (auto& kvpair : locked_table) {
            EdgeId e1_to_add = kvpair.first; auto& map_to_add = kvpair.second;

            for (auto& to_add : map_to_add) {
                EdgePair ep(e1_to_add, to_add.first), conj = this->ConjugatePair(e1_to_add, to_add.first);
                if (ep > conj)
                    continue;

                base::Merge(ep.first, ep.second, *to_add.second);
            }
        }
        VERIFY(this->size() >= index_to_add.size());
    }

    template<class Buffer>
    typename std::enable_if<std::is_convertible<typename Buffer::InnerMap, InnerMap>::value,
        void>::type MoveAssign(Buffer& from) {
        auto& base_index = this->storage_;
        base_index.clear();
        auto locked_table = from.lock_table();
        for (auto& kvpair : locked_table) {
            base_index[kvpair.first] = std::move(kvpair.second);
        }
        this->size_ = from.size();
    }

public:
    //---------------- Data deleting methods ----------------

    /**
     * @brief Removes the specific entry from the index, and its conjugate.
     * @warning Don't use it on unclustered index, because hashmaps require set_deleted_item
     * @return The number of deleted entries (0 if there wasn't such entry)
     */
    size_t Remove(EdgeId e1, EdgeId e2, Point p) {
        InnerPoint point = Traits::Shrink(p, this->graph_.length(e1));

        // We remove first "non-owning part"
        EdgePair minep, maxep;
        std::tie(minep, maxep) = this->MinMaxConjugatePair({ e1, e2 });

        size_t res = RemoveSingle(minep.first, minep.second, point);
        size_t removed = (this->IsSelfConj(e1, e2) ? res : 2 * res);
        this->size_ -= removed;

        Prune(maxep.first, maxep.second);
        Prune(minep.first, minep.second);

        return removed;
    }

    /**
     * @brief Removes the whole histogram from the index, and its conjugate.
     * @warning Don't use it on unclustered index, because hashmaps require set_deleted_item
     * @return The number of deleted entries
     */
    size_t Remove(EdgeId e1, EdgeId e2) {
        EdgePair minep, maxep;
        std::tie(minep, maxep) = this->MinMaxConjugatePair({ e1, e2 });

        size_t removed = RemoveAll(maxep.first, maxep.second);
        removed += RemoveAll(minep.first, minep.second);
        this->size_ -= removed;

        return removed;
    }

  private:
    void Prune(EdgeId e1, EdgeId e2) {
        auto i1 = this->storage_.find(e1);
        if (i1 == this->storage_.end())
            return;

        auto& map = i1->second;
        auto i2 = map.find(e2);
        if (i2 == map.end())
            return;

        if (!i2->second->empty())
            return;

        map.erase(e2);
        if (map.empty())
            this->storage_.erase(e1);
    }

    size_t RemoveSingle(EdgeId e1, EdgeId e2, InnerPoint point) {
        auto i1 = this->storage_.find(e1);
        if (i1 == this->storage_.end())
            return 0;

        auto& map = i1->second;
        auto i2 = map.find(e2);
        if (i2 == map.end())
            return 0;

        if (!i2->second->erase(point))
           return 0;

        return 1;
    }

    size_t RemoveAll(EdgeId e1, EdgeId e2) {
        auto i1 = this->storage_.find(e1);
        if (i1 == this->storage_.end())
            return 0;
        auto& map = i1->second;
        auto i2 = map.find(e2);
        if (i2 == map.end())
            return 0;

        size_t size_decrease = i2->second->size();
        map.erase(i2);
        if (map.empty()) //Prune empty maps
            this->storage_.erase(i1);
        return size_decrease;
    }

public:

    /**
     * @brief Removes all neighbourhood of an edge (all edges referring to it, and their histograms)
     * @warning To keep the symmetricity, it also deletes all conjugates, so the actual complexity is O(size).
     * NB: size - not all index but size of conjugates!!!
     * @return The number of deleted entries
     */
    size_t Remove(EdgeId edge) {
        auto i = this->storage_.find(edge);
        if (i == this->storage_.end())
            return 0;
        InnerMap &inner_map = i->second;
        std::vector<EdgeId> to_remove;
        to_remove.reserve(inner_map.size());
        size_t old_size = this->size();
        for (const auto& ep : inner_map)
            to_remove.push_back(ep.first);
        for (auto e2 : to_remove)
            this->Remove(edge, e2);
        return old_size - this->size();
    }

private:
    //When there is no such edge, returns a fake empty map for safety
    const InnerMap& GetImpl(EdgeId e) const {
        auto i = this->storage_.find(e);
        if (i != this->storage_.end())
            return i->second;
        return empty_map_;
    }

    //When there is no such histogram, returns a fake empty histogram for safety
    const InnerHistogram& GetImpl(EdgeId e1, EdgeId e2) const {
        auto i = this->storage_.find(e1);
        if (i != this->storage_.end()) {
            auto j = i->second.find(e2);
            if (j != i->second.end())
                return *j->second;
        }
        return HistProxy::empty_hist();
    }

public:

    /**
     * @brief Returns a whole proxy map to the neighbourhood of some edge.
     * @param e ID of starting edge
     */
    EdgeProxy Get(EdgeId e) const {
        return EdgeProxy(*this, GetImpl(e), e);
    }

    /**
     * @brief Returns a half proxy map to the neighbourhood of some edge.
     * @param e ID of starting edge
     */
    EdgeProxy GetHalf(EdgeId e) const {
        return EdgeProxy(*this, GetImpl(e), e, true);
    }

    /**
     * @brief Operator alias of Get(id).
     */
    EdgeProxy operator[](EdgeId e) const {
        return Get(e);
    }

    /**
     * @brief Returns a histogram proxy for all points between two edges.
     */
    HistProxy Get(EdgeId e1, EdgeId e2) const {
        return HistProxy(GetImpl(e1, e2), this->CalcOffset(e1));
    }

    /**
     * @brief Operator alias of Get(e1, e2).
     */
    HistProxy operator[](EdgePair p) const {
        return Get(p.first, p.second);
    }

    /**
     * @brief Checks if an edge (or its conjugated twin) is consisted in the index.
     */
    bool contains(EdgeId edge) const {
        return this->storage_.count(edge) + this->storage_.count(this->graph_.conjugate(edge)) > 0;
    }

    /**
     * @brief Checks if there is a histogram for two points (or their conjugated pair).
     */
    bool contains(EdgeId e1, EdgeId e2) const {
        auto i1 = this->storage_.find(e1);
        if (i1 != this->storage_.end() && i1->second.count(e2))
            return true;
        return false;
    }

    /**
     * @brief Inits the index with graph data. For each edge, adds a loop with zero weight.
     * @warning Do not call this on non-empty indexes.
     */
    void Init() {
        //VERIFY(size() == 0);
        for (auto it = this->graph_.ConstEdgeBegin(); !it.IsEnd(); ++it)
            this->Add(*it, *it, Point());
    }

    void VerifyIndex() {
        std::set<EdgeId> edges_in_index;
        size_t sz = 0;
        for( auto iter = data_begin(); iter != data_end(); iter ++ ) {
            VERIFY_MSG(iter->second.size() > 0, " empty map! ");
            VERIFY_MSG(this->graph_.edges().find(iter->first) != this->graph_.edges().end(), " left edge wrong  " << iter->first);
            for (const auto &e_iter: iter->second) {
                VERIFY_MSG(this->graph_.edges().find(e_iter.first) != this->graph_.edges().end(), " right edge wrong  " << e_iter.first);
                sz += e_iter.second->size();

            }
        }
        DEBUG(sz << " " << this->size_);
        VERIFY_MSG(sz == this ->size_, "Different sizes");
    }
private:
    InnerMap empty_map_; //null object
};

template<class T>
class NoLockingAdapter : public T {
  public:
    class locked_table {
      public:
        using iterator = typename T::iterator;
        using const_iterator = typename T::const_iterator;

        locked_table(T& table)
                : table_(table) {}

        iterator begin() { return table_.begin();  }
        const_iterator begin() const { return table_.begin(); }
        const_iterator cbegin() const { return table_.begin(); }

        iterator end() { return table_.end(); }
        const_iterator end() const { return table_.end(); }
        const_iterator cend() const { return table_.end(); }

        size_t size() const { return table_.size(); }

      private:
        T& table_;
    };

    // Nothing to lock here
    locked_table lock_table() {
        return locked_table(*this);
    }
};

//Aliases for common graphs
template<typename K, typename V>
using safe_btree_map = NoLockingAdapter<btree::safe_btree_map<K, V>>; //Two-parameters wrapper
template<typename Graph>
using PairedInfoIndexT = PairedIndex<Graph, PointTraits, safe_btree_map>;

template<typename K, typename V>
using btree_map = NoLockingAdapter<btree::btree_map<K, V>>; //Two-parameters wrapper

template<typename Graph>
using UnclusteredPairedInfoIndexT = PairedIndex<Graph, RawPointTraits, btree_map>;

template<typename G, typename Traits, template<typename, typename> class Container>
class PairedIndexHandler : public omnigraph::GraphActionHandler<G> {

public:
    using BaseIndex = PairedIndex<G, Traits, Container>;

    typedef typename G::EdgeId EdgeId;

    BaseIndex &paired_index_;

    PairedIndexHandler(PairedIndex<G, Traits, Container>& p): GraphActionHandler<G>(p.graph(), "PairedIndexHandler"), paired_index_(p) {}

    virtual void HandleDelete(EdgeId e) override {
        if (e == paired_index_.graph().conjugate(e)) {
            DEBUG("removing self-conj");
        }
        DEBUG("Deleted:" << e << " " << omp_get_thread_num());
        paired_index_.Remove(e);
    }

    virtual void HandleMerge(const std::vector<EdgeId> &old_edges, EdgeId new_edge) override {
//        VerifyIndex();
//        return;
        size_t shift = 0;
        DEBUG("created: " << new_edge);
        if (new_edge == paired_index_.graph().conjugate(new_edge))
            DEBUG("merging self-conj");
        std::set<EdgeId> forbidden;
        std::map<EdgeId, size_t> shifts;
        for (EdgeId e : old_edges)  {
            shifts[e] = shift;
            shift += paired_index_.graph().length(e);
            forbidden.insert(paired_index_.graph().conjugate(e));
        }

        DEBUG_EXPR({
            for (EdgeId e : old_edges)
                DEBUG(e << " " << paired_index_.graph().length(e));
        });

        for (EdgeId e : old_edges) {
            DEBUG("trying " << e);
            typename BaseIndex::EdgeProxy old_e = paired_index_.Get(e);
            std::vector<std::pair<EdgeId, Point>> to_add;
            for (auto it : old_e) {
                EdgeId next = it.first;
                if (forbidden.find(next) != forbidden.end()) {
                    DEBUG("skpping self-conjugate merge");
                    continue;
                }
                size_t neg_shift = 0;
                if (shifts.find(next) != shifts.end()) {
                    DEBUG("old edge was: " << next );
                    neg_shift = shifts[next];
                    next = new_edge;
                }
                DEBUG(e <<" " << next << " "<< it.second.size());
                for (auto pp : it.second) {
                    Point point (pp);
                    point.d += (double) shifts[e] - double (neg_shift);
                    DEBUG("from "  <<new_edge << " to " << next  <<" old " << e << " old_dist " << pp.d << "new_dist " << point.d << " "  << point.weight);
                    if (paired_index_.graph().length(new_edge) > point.d && new_edge != next) {
                        DEBUG("not adding, assert failed");
                    } else {
                        to_add.push_back(std::make_pair(next, point));
                        DEBUG("added");
                    }
                }
            }
            for (const auto &p : to_add)  {
                paired_index_.Add(new_edge, p.first, p.second);
            }

        }
    }
};

/**
 * @brief A collection of paired indexes which can be manipulated as one.
 *        Used as a convenient wrapper in parallel index processing.
 */
template<class Index>
class PairedIndices {
    typedef std::vector<Index> Storage;
    Storage data_;

public:
    typedef Index value_type;

    PairedIndices() {}

    PairedIndices(const typename Index::Graph& graph, size_t lib_num) {
        for (size_t i = 0; i < lib_num; ++i)
            data_.emplace_back(graph);
    }

    /**
     * @brief Initializes all indexes with zero points.
     */
    void Init() { for (auto& it : data_) it.Init(); }

    /**
     * @brief Clears all indexes.
     */
    void Clear() { for (auto& it : data_) it.clear(); }

    Index& operator[](size_t i) { return data_[i]; }

    const Index& operator[](size_t i) const { return data_[i]; }

    size_t size() const { return data_.size(); }

    typename Storage::iterator begin() { return data_.begin(); }
    typename Storage::iterator end() { return data_.end(); }

    typename Storage::const_iterator begin() const { return data_.begin(); }
    typename Storage::const_iterator end() const { return data_.end(); }
};

template<class Graph>
using PairedInfoIndicesT = PairedIndices<PairedInfoIndexT<Graph>>;

template<typename Graph>
using PairedInfoIndexHandlerT = PairedIndexHandler<Graph, PointTraits, safe_btree_map>;

template<class Graph>
using PairedInfoIndicesHandlerT = std::vector<PairedInfoIndexHandlerT<Graph>>;

template<class Graph>
using UnclusteredPairedInfoIndicesT = PairedIndices<UnclusteredPairedInfoIndexT<Graph>>;

template<typename K, typename V>
using unordered_map = NoLockingAdapter<std::unordered_map<K, V>>; //Two-parameters wrapper
template<class Graph>
using PairedInfoBuffer = PairedBuffer<Graph, RawPointTraits, unordered_map>;

template<class Graph>
using PairedInfoBuffersT = PairedIndices<PairedInfoBuffer<Graph>>;

}

}


//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE.md file in the project root for full license information.
//
#include "stdafx.h"
#include <algorithm>
#include <array>
#include <random>
#include <numeric>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include "../../../Source/Math/Matrix.h"
#include "../../../Source/Math/CPUMatrix.h"
#include "../../../Source/Math/GPUMatrix.h"
#include "../../../Source/Math/ConvolutionEngine.h"
#include "../../../Source/Math/CuDnnFactories.h"
#include "common.h"

namespace Microsoft { namespace MSR { namespace CNTK { namespace Test {

using vec = std::vector<float>;

using ConvEng = ConvolutionEngine<float>;

bool AreEqual(float a, float b, float maxRelError, float maxAbsError)
{
    float diff = std::abs(a - b);
    if (diff <= maxAbsError)
        return true;
    float largest = std::max(std::abs(a), std::abs(b));
    return diff < largest * maxRelError;
}
bool AreEqual(double a, double b, double maxRelError, double maxAbsError)
{
    double diff = std::abs(a - b);
    if (diff <= maxAbsError)
        return true;
    double largest = std::max(std::abs(a), std::abs(b));
    return diff < largest * maxRelError;
}

size_t CountNans(const SingleMatrix& src)
{
    size_t n = 0;
    foreach_coord (i, j, src)
    {
        n += std::isnan(src(i, j)) ? 1 : 0;
    }
    return n;
}

// Returns vector of engine config parameters: <kind, device, maxTempMemSizeInSamples>
std::vector<std::tuple<ConvolutionEngineKind, DEVICEID_TYPE, size_t>> GetTestEngineConfigs()
{
    std::vector<std::tuple<ConvolutionEngineKind, DEVICEID_TYPE, size_t>> res;
    // Reference engine. The engine does not use temp memory so safe to set it to 0.
    res.push_back(std::make_tuple(ConvolutionEngineKind::Reference, -1, 0));
    res.push_back(std::make_tuple(ConvolutionEngineKind::Reference, 0, 0));

    // Gemm engine. Implemented only for CPU for now. Uses temp memory.
    res.push_back(std::make_tuple(ConvolutionEngineKind::Gemm, -1, 0));
    res.push_back(std::make_tuple(ConvolutionEngineKind::Gemm, -1, 1));
    res.push_back(std::make_tuple(ConvolutionEngineKind::Gemm, -1, 3));
    return res;
}

std::vector<ConvolveGeometryPtr> GenerateConvTestConfigs()
{
    std::vector<ConvolveGeometryPtr> res;
    // REVIEW alexeyk: add test cases with even dimensions of a kernel. There are some corner cases which cuDNN does not support (which essentially require negative padding).
    for (size_t kW : {1, 3})
    {
        for (size_t kH : {1, 3})
        {
            for (size_t inW : {kW, 2 * kW, 2 * kW - 1})
            {
                for (size_t inC : {1, 3})
                {
                    for (size_t mapCount : {1, 5})
                    {
                        for (size_t stride : {1, min((int)kW, min((int)kH, 2))})
                        {
                            // Note: must use sharing=false in channel dimension otherwise geometry will not be cuDNN compatible but cuDNN won't fail.
                            res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(inW, max(kH, inW) + 1, inC),
                                TensorShape(kW, kH, inC), TensorShape(mapCount), TensorShape(stride, stride, inC),
                                ConvolveGeometry::BoolVec{true},
                                ConvolveGeometry::BoolVec{(kW & 1) != 0, (kH & 1) != 0, false},
                                TensorShape(0), TensorShape(0)));
                        }
                    }
                }
            }
        }
    }
    // For debugging.
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(3, 3, 1),
        TensorShape(3, 3, 1), TensorShape(2), TensorShape(1, 1, 1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{true, true, false},
        TensorShape(0), TensorShape(0)));

    // Simple 3D convolution.
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(5, 5, 5, 2),
        TensorShape(3, 3, 3, 2), TensorShape(2), TensorShape(1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{false},
        TensorShape(0), TensorShape(0)));
    // Example of 3D convolution that can be represented with 3D tensors in reference engine
    // but requires 4D tensors in other engines.
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(5, 5, 3, 1),
        TensorShape(3, 3, 2, 1), TensorShape(2), TensorShape(1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{false},
        TensorShape(0), TensorShape(0)));

    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(16, 16, 1),
        TensorShape(3, 3, 1), TensorShape(8), TensorShape(1, 2, 1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{true, true, false},
        TensorShape(0), TensorShape(0)));

    // 1x1 convolution (shortcuts in ResNet).
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(16, 16, 2),
        TensorShape(1, 1, 2), TensorShape(1), TensorShape(2, 2, 1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{false},
        TensorShape(0, 0, 0), TensorShape(0)));
    return res;
}

std::vector<ConvolveGeometryPtr> GeneratePoolTestConfigs()
{
    std::vector<ConvolveGeometryPtr> res;
    for (size_t kW : {1, 2, 3})
    {
        for (size_t kH : {1, 2, 3})
        {
            for (size_t inW : {kW, 2 * kW, 2 * kW - 1})
            {
                for (size_t inC : {1, 3})
                {
                    for (size_t stride : {1, min((int)kW, min((int)kH, 2))})
                    {
                        // Note: must always use autopadding otherwise there might be configurations that
                        // require negative padding that cuDNN does not support.
                        res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(inW, max(kH, inW) + 1, inC),
                            TensorShape(kW, kH, 1), TensorShape(1), TensorShape(stride, stride, 1),
                            ConvolveGeometry::BoolVec{true},
                            ConvolveGeometry::BoolVec{true, true, false},
                            TensorShape(0), TensorShape(0)));
                    }
                }
            }
        }
    }
    // For debugging.
    // Ordinary pooling.
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(4, 4, 1),
        TensorShape(2, 2, 1), TensorShape(1), TensorShape(2, 2, 1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{true, true, false},
        TensorShape(0), TensorShape(0)));
    // Overlapped with padding.
    res.push_back(std::make_shared<ConvolveGeometry>(TensorShape(4, 4, 1),
        TensorShape(3, 3, 1), TensorShape(1), TensorShape(2, 2, 1),
        ConvolveGeometry::BoolVec{true}, ConvolveGeometry::BoolVec{true, true, false},
        TensorShape(0), TensorShape(0)));
    return res;
}

BOOST_AUTO_TEST_SUITE(ConvolutionSuite)

BOOST_AUTO_TEST_CASE(ConvolutionForward)
{
    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    boost::random::normal_distribution<float> nd;

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int baseDeviceId = 0;
    for (const auto& engCfg : GetTestEngineConfigs())
    {
        auto engKind = std::get<0>(engCfg);
        auto deviceId = std::get<1>(engCfg);
        auto maxTempMem = std::get<2>(engCfg);
        for (const auto& g : GenerateConvTestConfigs())
        {
            auto baseEng = ConvEng::Create(g, baseDeviceId, ImageLayoutKind::CHW, 0, PoolKind::None, ConvolutionEngineKind::CuDnn);
            auto testEng = ConvEng::Create(g, deviceId, ImageLayoutKind::CHW, maxTempMem, PoolKind::None, engKind);

            size_t n = batchSizeG(rng);
            vec buf;
            buf.resize(g->InputShape().GetNumElements() * n);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix in(g->InputShape().GetNumElements(), n, buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix inB(g->InputShape().GetNumElements(), n, buf.data(), baseDeviceId, matrixFlagNormal);

            size_t mapCount = g->GetMapCount(g->InputShape().GetRank() - 1);
            buf.resize(g->KernelShape().GetNumElements() * mapCount);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix kernel(mapCount, g->KernelShape().GetNumElements(), buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix kernelB(mapCount, g->KernelShape().GetNumElements(), buf.data(), baseDeviceId, matrixFlagNormal);

            size_t crowOut = g->OutputShape().GetNumElements();
            SingleMatrix outBuf(deviceId);
            SingleMatrix out = initMat(outBuf, crowOut, n, buf);
            SingleMatrix outB(out.DeepClone(), baseDeviceId);

            SingleMatrix workspace(deviceId);
            SingleMatrix workspaceB(baseDeviceId);

            testEng->Forward(in, kernel, out, workspace);
            baseEng->Forward(inB, kernelB, outB, workspaceB);

            std::stringstream tmsg;
            tmsg << "Geometry: " << (std::string)(*g) << ", Batch: " << n << ", Device: " << deviceId << ", MaxTempMem: " << maxTempMem;
            std::string msg = " are not equal, " + tmsg.str();
            std::string msgNan = " has NaNs, " + tmsg.str();
            std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

            float relErr = Err<float>::Rel;
            float absErr = Err<float>::Abs;
            std::string emsg;

            BOOST_REQUIRE_MESSAGE(!out.HasNan("out"), "out" << msgNan);
            BOOST_REQUIRE_MESSAGE(CheckEqual(out, outB, emsg, relErr * 4, absErr * 14), "out" << msg << ". " << emsg);
            BOOST_REQUIRE_MESSAGE(CountNans(outBuf) == crowOut * 2 * n, "out" << msgNotNan);
        }
    }
}

BOOST_AUTO_TEST_CASE(ConvolutionBackwardData)
{
    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    boost::random::normal_distribution<float> nd;

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int baseDeviceId = 0;
    for (const auto& engCfg : GetTestEngineConfigs())
    {
        auto engKind = std::get<0>(engCfg);
        auto deviceId = std::get<1>(engCfg);
        auto maxTempMem = std::get<2>(engCfg);
        for (const auto& g : GenerateConvTestConfigs())
        {
            auto baseEng = ConvEng::Create(g, baseDeviceId, ImageLayoutKind::CHW, 0, PoolKind::None, ConvolutionEngineKind::CuDnn);
            auto testEng = ConvEng::Create(g, deviceId, ImageLayoutKind::CHW, maxTempMem, PoolKind::None, engKind);

            size_t n = batchSizeG(rng);
            vec buf;
            buf.resize(g->OutputShape().GetNumElements() * n);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix srcGrad(g->OutputShape().GetNumElements(), n, buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix srcGradB(g->OutputShape().GetNumElements(), n, buf.data(), baseDeviceId, matrixFlagNormal);

            size_t mapCount = g->GetMapCount(g->InputShape().GetRank() - 1);
            buf.resize(g->KernelShape().GetNumElements() * mapCount);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix kernel(mapCount, g->KernelShape().GetNumElements(), buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix kernelB(mapCount, g->KernelShape().GetNumElements(), buf.data(), baseDeviceId, matrixFlagNormal);

            size_t crowGrad = g->InputShape().GetNumElements();
            SingleMatrix gradBuf(deviceId);
            SingleMatrix grad = initMat(gradBuf, crowGrad, n, buf);
            SingleMatrix gradB(grad.DeepClone(), baseDeviceId);

            SingleMatrix workspace(deviceId);
            SingleMatrix workspaceB(baseDeviceId);

            testEng->BackwardData(srcGrad, kernel, grad, true, workspace);
            baseEng->BackwardData(srcGradB, kernelB, gradB, true, workspaceB);

            std::stringstream tmsg;
            tmsg << "Geometry: " << (std::string)(*g) << ", Batch: " << n << ", Device: " << deviceId;
            std::string msg = " are not equal, " + tmsg.str();
            std::string msgNan = " has NaNs, " + tmsg.str();
            std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

            float relErr = Err<float>::Rel;
            float absErr = Err<float>::Abs;
            std::string emsg;

            BOOST_REQUIRE_MESSAGE(!grad.HasNan("grad"), "grad" << msgNan);
            BOOST_REQUIRE_MESSAGE(CheckEqual(grad, gradB, emsg, relErr * 16, absErr * 16), "grad" << msg << ". " << emsg);
            BOOST_REQUIRE_MESSAGE(CountNans(gradBuf) == crowGrad * 2 * n, "grad" << msgNotNan);
        }
    }
}

BOOST_AUTO_TEST_CASE(ConvolutionBackwardKernel)
{
    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    boost::random::normal_distribution<float> nd;

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int baseDeviceId = 0;
    for (const auto& engCfg : GetTestEngineConfigs())
    {
        auto engKind = std::get<0>(engCfg);
        auto deviceId = std::get<1>(engCfg);
        auto maxTempMem = std::get<2>(engCfg);
        for (const auto& g : GenerateConvTestConfigs())
        {
            auto baseEng = ConvEng::Create(g, baseDeviceId, ImageLayoutKind::CHW, 0, PoolKind::None, ConvolutionEngineKind::CuDnn);
            auto testEng = ConvEng::Create(g, deviceId, ImageLayoutKind::CHW, maxTempMem, PoolKind::None, engKind);

            size_t n = batchSizeG(rng);
            vec buf;
            buf.resize(g->InputShape().GetNumElements() * n);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix in(g->InputShape().GetNumElements(), n, buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix inB(g->InputShape().GetNumElements(), n, buf.data(), baseDeviceId, matrixFlagNormal);

            buf.resize(g->OutputShape().GetNumElements() * n);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix grad(g->OutputShape().GetNumElements(), n, buf.data(), deviceId, matrixFlagNormal);
            SingleMatrix gradB(g->OutputShape().GetNumElements(), n, buf.data(), baseDeviceId, matrixFlagNormal);

            size_t mapCount = g->GetMapCount(g->InputShape().GetRank() - 1);
            buf.resize(g->KernelShape().GetNumElements() * mapCount);
            std::generate(begin(buf), end(buf), [&] { return nd(rng); });
            SingleMatrix kernelBuf(deviceId);
            SingleMatrix kernel = initMat(kernelBuf, mapCount, g->KernelShape().GetNumElements(), buf);
            SingleMatrix kernelB(kernel.DeepClone(), baseDeviceId);

            SingleMatrix workspace(deviceId);
            SingleMatrix workspaceB(baseDeviceId);

            testEng->BackwardKernel(grad, in, kernel, true, false, workspace);
            baseEng->BackwardKernel(gradB, inB, kernelB, true, false, workspaceB);

            std::stringstream tmsg;
            tmsg << "Geometry: " << (std::string)(*g) << ", Batch: " << n << ", Device: " << deviceId;
            std::string msg = " are not equal, " + tmsg.str();
            std::string msgNan = " has NaNs, " + tmsg.str();
            std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

            float relErr = Err<float>::Rel;
            float absErr = Err<float>::Abs;
            std::string emsg;

            BOOST_REQUIRE_MESSAGE(!kernel.HasNan("kernel"), "kernel" << msgNan);
            // Todo: check the threashold value after we have setttings regard determinstics in place.
            BOOST_REQUIRE_MESSAGE(CheckEqual(kernel, kernelB, emsg, relErr * 192, absErr * 32), "kernel" << msg << ". " << emsg);
            BOOST_REQUIRE_MESSAGE(CountNans(kernelBuf) == kernel.GetNumElements() * 2, "kernel" << msgNotNan);
        }
    }
}

BOOST_AUTO_TEST_CASE(PoolingForward)
{
    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    boost::random::normal_distribution<float> nd;

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int baseDeviceId = 0;
    auto engKind = ConvolutionEngineKind::Reference;
    for (auto kind : {PoolKind::Max, PoolKind::Average})
    {
        for (int deviceId : {-1, 0})
        {
            for (const auto& g : GeneratePoolTestConfigs())
            {
                auto baseEng = ConvEng::Create(g, baseDeviceId, ImageLayoutKind::CHW, 0, kind, ConvolutionEngineKind::CuDnn);
                auto testEng = ConvEng::Create(g, deviceId, ImageLayoutKind::CHW, 0, kind, engKind);

                size_t n = batchSizeG(rng);
                vec buf;
                buf.resize(g->InputShape().GetNumElements() * n);
                std::generate(begin(buf), end(buf), [&] { return nd(rng); });
                SingleMatrix in(g->InputShape().GetNumElements(), n, buf.data(), deviceId, matrixFlagNormal);
                SingleMatrix inB(g->InputShape().GetNumElements(), n, buf.data(), baseDeviceId, matrixFlagNormal);

                size_t crowOut = g->OutputShape().GetNumElements();
                SingleMatrix outBuf(deviceId);
                SingleMatrix out = initMat(outBuf, crowOut, n, buf);
                SingleMatrix outB(out.DeepClone(), baseDeviceId);

                testEng->ForwardPooling(in, out);
                baseEng->ForwardPooling(inB, outB);

                std::stringstream tmsg;
                tmsg << "Geometry: " << (std::string)(*g) << ", Pool: " << (int)kind << ", Batch: " << n << ", Device: " << deviceId;
                std::string msg = " are not equal, " + tmsg.str();
                std::string msgNan = " has NaNs, " + tmsg.str();
                std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

                float relErr = Err<float>::Rel;
                float absErr = Err<float>::Abs;
                std::string emsg;

                BOOST_REQUIRE_MESSAGE(!out.HasNan("out"), "out" << msgNan);
                BOOST_REQUIRE_MESSAGE(CheckEqual(out, outB, emsg, relErr, absErr * 8), "out" << msg << ". " << emsg);
                BOOST_REQUIRE_MESSAGE(CountNans(outBuf) == crowOut * 2 * n, "out" << msgNotNan);
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(PoolingBackward)
{
    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    boost::random::normal_distribution<float> nd;

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int baseDeviceId = 0;
    auto engKind = ConvolutionEngineKind::Reference;
    for (auto kind : {PoolKind::Max, PoolKind::Average})
    {
        for (int deviceId : {-1, 0})
        {
            for (const auto& g : GeneratePoolTestConfigs())
            {
                auto baseEng = ConvEng::Create(g, baseDeviceId, ImageLayoutKind::CHW, 0, kind, ConvolutionEngineKind::CuDnn);
                auto testEng = ConvEng::Create(g, deviceId, ImageLayoutKind::CHW, 0, kind, engKind);

                size_t n = batchSizeG(rng);
                vec buf;
                size_t crowIn = g->InputShape().GetNumElements();
                buf.resize(crowIn * n);
                std::generate(begin(buf), end(buf), [&] { return nd(rng); });
                SingleMatrix inB(crowIn, n, buf.data(), baseDeviceId, matrixFlagNormal);
                SingleMatrix in(crowIn, n, buf.data(), deviceId, matrixFlagNormal);

                size_t crowOut = g->OutputShape().GetNumElements();
                buf.resize(crowOut * n);
                std::generate(begin(buf), end(buf), [&] { return nd(rng); });
                SingleMatrix srcGradB(crowOut, n, buf.data(), baseDeviceId, matrixFlagNormal);
                SingleMatrix srcGrad(crowOut, n, buf.data(), deviceId, matrixFlagNormal);
                // Do not generate for out as it will be replaced anyway.
                SingleMatrix outB(crowOut, n, buf.data(), baseDeviceId, matrixFlagNormal);
                SingleMatrix out(crowOut, n, buf.data(), deviceId, matrixFlagNormal);

                testEng->ForwardPooling(in, out);
                baseEng->ForwardPooling(inB, outB);

                SingleMatrix gradBuf(deviceId);
                SingleMatrix grad = initMat(gradBuf, crowIn, n, buf);
                SingleMatrix gradB(grad.DeepClone(), baseDeviceId);

                testEng->BackwardPooling(out, srcGrad, in, grad);
                baseEng->BackwardPooling(outB, srcGradB, inB, gradB);

                std::stringstream tmsg;
                tmsg << "Geometry: " << (std::string)(*g) << ", Pool: " << (int)kind << ", Batch: " << n << ", Device: " << deviceId;
                std::string msg = " are not equal, " + tmsg.str();
                std::string msgNan = " has NaNs, " + tmsg.str();
                std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

                float relErr = Err<float>::Rel;
                float absErr = Err<float>::Abs;
                std::string emsg;

                BOOST_REQUIRE_MESSAGE(!grad.HasNan("grad"), "grad" << msgNan);
                BOOST_REQUIRE_MESSAGE(CheckEqual(grad, gradB, emsg, relErr, absErr * 8), "grad" << msg << ". " << emsg);
                BOOST_REQUIRE_MESSAGE(CountNans(gradBuf) == crowIn * 2 * n, "grad" << msgNotNan);
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(MaxUnpooling)
{
    using IntMatrix = Matrix<int>;

    std::mt19937 rng(0);
    boost::random::uniform_int_distribution<> batchSizeG(1, 8);
    // Using uniform distribution with positive values to avoid issues with
    // unpooling negative values.
    boost::random::uniform_real_distribution<float> nd(0, 1);

    auto initMat = [&](SingleMatrix& buf, size_t r, size_t c, vec& data) -> SingleMatrix
    {
        data.resize(r * 3 * c);
        std::fill(begin(data), end(data), std::numeric_limits<float>::quiet_NaN());
        std::generate(begin(data) + r * c, begin(data) + 2 * r * c, [&] { return nd(rng); });
        buf.SetValue(r, 3 * c, buf.GetDeviceId(), data.data());
        // Get center slice.
        return buf.ColumnSlice(c, c);
    };

    int cpuDeviceId = -1;
    int gpuDeviceId = 0;

    for (const auto& g : GeneratePoolTestConfigs())
    {
        // cpuEng and gpuEng are used to compare results against each other.
        auto cpuEng = ConvEng::Create(g, cpuDeviceId, ImageLayoutKind::CHW, 0, PoolKind::Max, ConvolutionEngineKind::Reference);
        auto gpuEng = ConvEng::Create(g, gpuDeviceId, ImageLayoutKind::CHW, 0, PoolKind::Max, ConvolutionEngineKind::Reference);

        size_t n = batchSizeG(rng);
        vec buf;
        buf.resize(g->InputShape().GetNumElements() * n);
        std::generate(begin(buf), end(buf), [&] { return nd(rng); });
        SingleMatrix inC(g->InputShape().GetNumElements(), n, buf.data(), cpuDeviceId, matrixFlagNormal);
        SingleMatrix inG(g->InputShape().GetNumElements(), n, buf.data(), gpuDeviceId, matrixFlagNormal);

        // First, compute max pooling output and corresponding mask.
        SingleMatrix outC(g->OutputShape().GetNumElements(), n, cpuDeviceId);
        SingleMatrix outG(g->OutputShape().GetNumElements(), n, gpuDeviceId);

        cpuEng->ForwardPooling(inC, outC);
        gpuEng->ForwardPooling(inG, outG);

        // Second, do the unpooling.
        size_t crowIn = g->InputShape().GetNumElements();
        SingleMatrix inUBufC(cpuDeviceId);
        SingleMatrix inUC = initMat(inUBufC, crowIn, n, buf);
        SingleMatrix inUBufG(inUBufC.DeepClone(), gpuDeviceId);
        SingleMatrix inUG = initMat(inUBufG, crowIn, n, buf);

        cpuEng->MaxUnpooling(outC, inC, inUC);
        gpuEng->MaxUnpooling(outG, inG, inUG);

        // Check that CPU/GPU results are the same.
        std::stringstream tmsg;
        tmsg << "Geometry: " << (std::string)(*g) << ", Batch: " << n;
        std::string msg = " are not equal, " + tmsg.str();
        std::string msgNan = " has NaNs, " + tmsg.str();
        std::string msgNotNan = " has buffer overflow/underflow, " + tmsg.str();

        float relErr = 0;
        float absErr = 0;
        std::string emsg;

        BOOST_REQUIRE_MESSAGE(!inUC.HasNan("inUC"), "inUC" << msgNan);
        BOOST_REQUIRE_MESSAGE(!inUG.HasNan("inUG"), "inUG" << msgNan);
        BOOST_REQUIRE_MESSAGE(CheckEqual(inUC, inUG, emsg, relErr, absErr), "inU" << msg << ". " << emsg);
        BOOST_REQUIRE_MESSAGE(CountNans(inUBufC) == crowIn * 2 * n, "inUBufC" << msgNotNan);
        BOOST_REQUIRE_MESSAGE(CountNans(inUBufG) == crowIn * 2 * n, "inUBufG" << msgNotNan);

        // Now do the pooling from unpooled source and compare with original pooling.
        SingleMatrix outC_2(g->OutputShape().GetNumElements(), n, cpuDeviceId);
        SingleMatrix outG_2(g->OutputShape().GetNumElements(), n, gpuDeviceId);
        cpuEng->ForwardPooling(inUC, outC_2);
        gpuEng->ForwardPooling(inUG, outG_2);

        BOOST_REQUIRE_MESSAGE(CheckEqual(outC_2, outC, emsg, relErr, absErr), "outC_2" << msg << ". " << emsg);
        BOOST_REQUIRE_MESSAGE(CheckEqual(outG_2, outG, emsg, relErr, absErr), "outG_2" << msg << ". " << emsg);
    }
}

BOOST_AUTO_TEST_SUITE_END()

} } } }


#ifndef FAST_GICP_FAST_VGICP_CUDA_HPP
#define FAST_GICP_FAST_VGICP_CUDA_HPP

#include <Eigen/Core>
#include <Eigen/Geometry>

#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/search/kdtree.h>
#include <pcl/registration/registration.h>

#include <fast_gicp/gicp/gicp_settings.hpp>
#include <fast_gicp/gicp/lsq_registration.hpp>

namespace fast_gicp {

namespace cuda {
class FastVGICPCudaCore;
}

enum class NearestNeighborMethod { CPU_PARALLEL_KDTREE, GPU_BRUTEFORCE, GPU_RBF_KERNEL };

/**
 * @brief Fast Voxelized GICP algorithm boosted with CUDA
 */
template<typename PointSource, typename PointTarget>
class FastVGICPCuda : public LsqRegistration<PointSource, PointTarget> {
public:
  using Scalar = float;
  using Matrix4 = typename pcl::Registration<PointSource, PointTarget, Scalar>::Matrix4;

  using PointCloudSource = typename pcl::Registration<PointSource, PointTarget, Scalar>::PointCloudSource;
  using PointCloudSourcePtr = typename PointCloudSource::Ptr;
  using PointCloudSourceConstPtr = typename PointCloudSource::ConstPtr;

  using PointCloudTarget = typename pcl::Registration<PointSource, PointTarget, Scalar>::PointCloudTarget;
  using PointCloudTargetPtr = typename PointCloudTarget::Ptr;
  using PointCloudTargetConstPtr = typename PointCloudTarget::ConstPtr;

#if PCL_VERSION >= PCL_VERSION_CALC(1, 10, 0)
  using Ptr = pcl::shared_ptr<FastVGICPCuda<PointSource, PointTarget>>;
  using ConstPtr = pcl::shared_ptr<const FastVGICPCuda<PointSource, PointTarget>>;
#else
  using Ptr = boost::shared_ptr<FastVGICPCuda<PointSource, PointTarget>>;
  using ConstPtr = boost::shared_ptr<const FastVGICPCuda<PointSource, PointTarget>>;
#endif

protected:
  using pcl::Registration<PointSource, PointTarget, Scalar>::input_;
  using pcl::Registration<PointSource, PointTarget, Scalar>::target_;

public:
  FastVGICPCuda();
  virtual ~FastVGICPCuda() override;

  void setCorrespondenceRandomness(int k);
  void setResolution(double resolution);
  void setKernelWidth(double kernel_width, double max_dist = -1.0);
  void setRegularizationMethod(RegularizationMethod method);
  void setNeighborSearchMethod(NeighborSearchMethod method, double radius = -1.0);
  void setNearestNeighborSearchMethod(NearestNeighborMethod method);

  virtual void swapSourceAndTarget() override;
  virtual void clearSource() override;
  virtual void clearTarget() override;

  virtual void setInputSource(const PointCloudSourceConstPtr& cloud) override;
  virtual void setInputTarget(const PointCloudTargetConstPtr& cloud) override;

protected:
  virtual void computeTransformation(PointCloudSource& output, const Matrix4& guess) override;
  virtual double linearize(const Eigen::Isometry3d& trans, Eigen::Matrix<double, 6, 6>* H = nullptr, Eigen::Matrix<double, 6, 1>* b = nullptr) override;
  virtual double compute_error(const Eigen::Isometry3d& trans) override;

  template<typename PointT>
  std::vector<int> find_neighbors_parallel_kdtree(int k, typename pcl::PointCloud<PointT>::ConstPtr cloud) const;

private:
  int k_correspondences_;
  double voxel_resolution_;
  RegularizationMethod regularization_method_;
  NearestNeighborMethod neighbor_search_method_;

  std::unique_ptr<cuda::FastVGICPCudaCore> vgicp_cuda_;
};

}  // namespace fast_gicp

#endif


#pragma once
#include <algorithm>
#include <type_traits>
#include <functional>
#include <autoppl/util/var_traits.hpp>
#include <autoppl/util/model_expr_traits.hpp>
#include <autoppl/util/dist_expr_traits.hpp>
#include <type_traits>
#include <typeinfo>       // operator typeid
#include <boost/uuid/uuid.hpp>
#include <boost/uuid/uuid_generators.hpp>
#include <boost/uuid/uuid_io.hpp>
#include <autoppl/program_analysis/dependencies_intervals.hpp>
#include <autoppl/program_analysis/UncertainIntervals/IntervalAnalysis.hpp>

namespace ppl {
namespace expr {

struct EqNodeParent {};
struct EqNodeCont {};
struct EqNodeDisc {};

/**
 * This class represents a "node" in the model expression
 * that relates a var with a distribution.
 */
#if __cplusplus <= 201703L
template <class VarType, class DistType>
#else
template <util::var VarType, util::dist_expr DistType>
#endif
struct EqNode : util::ModelExpr<EqNode<VarType, DistType>>, EqNodeParent, std::conditional_t<std::is_base_of_v<util::continuous_dist_base, DistType>, EqNodeCont, EqNodeDisc>
{

#if __cplusplus <= 201703L
    static_assert(util::assert_is_var_v<VarType>);
    static_assert(util::assert_is_dist_expr_v<DistType>);
#endif

    using var_t = VarType;
    using dist_t = DistType;
    using dist_value_t = typename util::dist_expr_traits<dist_t>::dist_value_t;

    EqNode(var_t& var,
           const dist_t& dist) noexcept
        : orig_var_ref_{var}
        , dist_{dist}
    {}

    template <class EqNodeFunc>
    void traverse(EqNodeFunc&& eq_f)
    {
        using this_t = EqNode<VarType, DistType>;
        eq_f(static_cast<this_t&>(*this));
    }

    template <class EqNodeFunc>
    void traverse(EqNodeFunc&& eq_f) const
    {
        using this_t = EqNode<VarType, DistType>;
        eq_f(static_cast<const this_t&>(*this));
    }

    dist_value_t pdf() const {
        return dist_.pdf(get_variable());
    }

    dist_value_t log_pdf() const {
        return dist_.log_pdf(get_variable());
    }

    auto& get_variable() { return orig_var_ref_.get(); }
    const auto& get_variable() const { return orig_var_ref_.get(); }
    const auto& get_distribution() const { return dist_; }

    std::set<boost::uuids::uuid> getDeps(){
	return dist_.getDeps();
    }

    std::set<boost::uuids::uuid> getVarID(){
	return orig_var_ref_.get().getDeps();
    }

    void getInterval(){
	//get  distribution's interval
	DeterministicInterval dist_range = dist_.getInterval();
	//get  var's id
	boost::uuids::uuid var_id = orig_var_ref_.get().getUUID();
	std::cout << dist_range.upper << std::endl;
	//map them together
	globals::intervals.insert({var_id,dist_range});
    }

    void getLikelihoodInterval(){

	//check if my variable is a data or param
	if constexpr (util::is_data_v<var_t>){
		boost::uuids::uuid var_id = orig_var_ref_.get().getUUID();
		int NumDataPoints = orig_var_ref_.get().size();
		DeterministicInterval DataInterval = orig_var_ref_.get().getDataInterval();

		DeterministicInterval BiggestInterval(0.,0.);
		DeterministicInterval LikelihoodInterval = dist_.log_pdf_interval(DataInterval,BiggestInterval);

		globals::single_likelihood_intervals.insert({var_id,LikelihoodInterval});

		//account for the likelihood sum over all data points
		if (NumDataPoints > 0){
			LikelihoodInterval = scale(double(NumDataPoints)/double(4),LikelihoodInterval);
		}
		BiggestInterval = union_(BiggestInterval,LikelihoodInterval);

		globals::likelihood_intervals.insert({var_id,LikelihoodInterval});
		globals::biggest_intervals.insert({var_id,BiggestInterval});
	}

    }

private:
    using var_ref_t = std::reference_wrapper<var_t>;
    var_ref_t orig_var_ref_;
    dist_t dist_;

};

} // namespace expr
} // namespace ppl


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

#include "read_into_bdd.hpp"

#include <boost/algorithm/string/predicate.hpp>
#include <boost/range/algorithm.hpp>

#include <core/utils/range_utils.hpp>
#include <core/utils/timer.hpp>
#include <classical/aig.hpp>
#include <classical/dd/aig_to_cirkit_bdd.hpp>
#include <classical/io/read_aiger.hpp>
#include <classical/io/read_pla_to_cirkit_bdd.hpp>
#include <classical/utils/aig_utils.hpp>

namespace cirkit
{

/******************************************************************************
 * Types                                                                      *
 ******************************************************************************/

/******************************************************************************
 * Private functions                                                          *
 ******************************************************************************/

/******************************************************************************
 * Public functions                                                           *
 ******************************************************************************/

std::pair<bdd_manager_ptr, std::vector<bdd>> read_into_bdd( const std::string& filename,
                                                            const properties::ptr& settings,
                                                            const properties::ptr& statistics )
{
  /* settings */
  auto log_max_objs = get( settings, "log_max_objs", 24u );

  /* timing */
  properties_timer t( statistics );

  if ( boost::ends_with( filename, ".pla" ) )
  {
    auto function = read_pla_into_cirkit_bdd( filename, settings );
    std::vector<bdd> fs( function->num_outputs() );
    for ( auto i = 0u; i < function->num_outputs(); ++i )
    {
      fs[i] = function->lookupOutput( i );
    }

    set( statistics, "input_labels", function->input_labels() );
    set( statistics, "output_labels", function->output_labels() );

    return std::make_pair( function->manager(), fs );
  }
  else if ( boost::ends_with( filename, ".aag" ) )
  {
    aig_graph aig;
    read_aiger( aig, filename );

    std::vector<bdd> fs;
    cirkit_bdd_simulator sim( aig, log_max_objs );
    auto map = simulate_aig( aig, sim );

    for ( const auto& m : map )
    {
      fs.push_back( m.second );
    }

    if ( statistics )
    {
      const auto info = aig_info( aig );
      std::vector<std::string> input_labels( info.inputs.size() );
      boost::transform( info.inputs, input_labels.begin(), [&]( const aig_node& n ) { return info.node_names.at( n ); } );
      auto output_labels = get_map_values( info.outputs );
      set( statistics, "input_labels", input_labels );
      set( statistics, "output_labels", output_labels );
    }

    return std::make_pair( sim.mgr, fs );
  }

  std::cerr << "[e] unknown suffix" << std::endl;
  assert( false );

  return std::make_pair( bdd_manager_ptr(), std::vector<bdd>() );
}

}

// Local Variables:
// c-basic-offset: 2
// eval: (c-set-offset 'substatement-open 0)
// eval: (c-set-offset 'innamespace 0)
// End: