// g2o - General Graph Optimization
// Copyright (C) 2011 R. Kuemmerle, G. Grisetti, H. Strasdat, W. Burgard
// 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.
//
// 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.

#include "simple_star_ops.h"
#include "backbone_tree_action.h"
#include "edge_types_cost_function.h"
#include "g2o/core/optimization_algorithm_with_hessian.h"
#include <iostream>
#include <Eigen/LU>
#include <Eigen/Cholesky>
#include <Eigen/Eigenvalues>

namespace g2o{

  using namespace std;
  using namespace Eigen;

  double activeVertexChi(const OptimizableGraph::Vertex* v){
    const SparseOptimizer* s = dynamic_cast<const SparseOptimizer*>(v->graph());
    const OptimizableGraph::EdgeContainer& av = s->activeEdges();
    double chi = 0;
    int ne =0;
    for (HyperGraph::EdgeSet::iterator it = v->edges().begin(); it!=v->edges().end(); it++){
      OptimizableGraph::Edge* e = dynamic_cast <OptimizableGraph::Edge*> (*it);
      if (!e)
	continue;
      if (s->findActiveEdge(e)!=av.end()) {
	chi +=e->chi2();
	ne++;
      }
    }
    if (! ne)
      return -1;
    return chi/ne;
  }

void constructEdgeStarMap(EdgeStarMap& esmap, StarSet& stars, bool low){
  esmap.clear();
  for (StarSet::iterator it=stars.begin(); it!=stars.end(); it++){
    Star* s = *it;
    if (low) {
      for (HyperGraph::EdgeSet::iterator it = s->lowLevelEdges().begin();
          it!=s->lowLevelEdges().end(); it++){
        HyperGraph::Edge* e=*it;
        esmap.insert(make_pair(e,s));
      }
    } else {
      for (HyperGraph::EdgeSet::iterator it = s->starEdges().begin();
          it!=s->starEdges().end(); it++){
        HyperGraph::Edge* e=*it;
        esmap.insert(make_pair(e,s));
      }
    }
  }
}

size_t vertexEdgesInStar(HyperGraph::EdgeSet& eset, HyperGraph::Vertex* v, Star* s, EdgeStarMap& esmap){
  eset.clear();
  for (HyperGraph::EdgeSet::iterator it=v->edges().begin(); it!=v->edges().end(); it++){
    HyperGraph::Edge* e=*it;
    EdgeStarMap::iterator eit=esmap.find(e);
    if (eit!=esmap.end() && eit->second == s)
      eset.insert(e);
  }
  return eset.size();
}

void starsInVertex(StarSet& stars, HyperGraph::Vertex* v, EdgeStarMap& esmap){
  for (HyperGraph::EdgeSet::iterator it=v->edges().begin(); it!=v->edges().end(); it++){
    HyperGraph::Edge* e=*it;
    EdgeStarMap::iterator eit=esmap.find(e);
    if (eit!=esmap.end())
      stars.insert(eit->second);
  }
}

void starsInEdge(StarSet& stars, HyperGraph::Edge* e, EdgeStarMap& esmap, HyperGraph::VertexSet& gauge){
  for (size_t i=0; i<e->vertices().size(); i++){
    OptimizableGraph::Vertex* v=(OptimizableGraph::Vertex*)e->vertices()[i];
    if (gauge.find(v)==gauge.end())
      starsInVertex(stars, v, esmap);
  }
}

void assignHierarchicalEdges(StarSet& stars, EdgeStarMap& esmap, EdgeLabeler* labeler, EdgeCreator* creator, SparseOptimizer* optimizer, int minNumEdges, int maxIterations){
  // now construct the hierarchical edges for all the stars
  int starNum=0;
  for (StarSet::iterator it=stars.begin(); it!=stars.end(); it++){
    cerr << "STAR# " << starNum << endl;
    Star* s=*it;
    std::vector<OptimizableGraph::Vertex*> vertices(2);
    vertices[0]= (OptimizableGraph::Vertex*) *s->_gauge.begin();
    cerr << "eIs"  << endl;
    HyperGraph::VertexSet vNew =s->lowLevelVertices();
    for (HyperGraph::VertexSet::iterator vit=s->_lowLevelVertices.begin(); vit!=s->_lowLevelVertices.end(); vit++){
      OptimizableGraph::Vertex* v=(OptimizableGraph::Vertex*)*vit;
      vertices[1]=v;
      if (v==vertices[0])
        continue;
      HyperGraph::EdgeSet eInSt;
      int numEdges = vertexEdgesInStar(eInSt, v, s, esmap);
      if (Factory::instance()->tag(v)==Factory::instance()->tag(vertices[0]) || numEdges>minNumEdges) {
        OptimizableGraph::Edge* e=creator->createEdge(vertices);
        //cerr << "creating edge" << e << endl;
        if (e) {
          e->setLevel(1);
          optimizer->addEdge(e);
          s->_starEdges.insert(e);
        } else {
          cerr << "THERE" << endl;
          cerr << "FATAL, cannot create edge" << endl;
        }
      } else {
        vNew.erase(v);
        // cerr << numEdges << " ";
        // cerr << "r " <<  v-> id() << endl;
        // remove from the star all edges that are not sufficiently connected
        for (HyperGraph::EdgeSet::iterator it=eInSt.begin(); it!=eInSt.end(); it++){
          HyperGraph::Edge* e=*it;
          s->lowLevelEdges().erase(e);
        }
      }
    }
    s->lowLevelVertices()=vNew;
    //cerr << endl;
    cerr <<  "gauge: " << (*s->_gauge.begin())->id()
      << " edges:" << s->_lowLevelEdges.size()
      << " hedges" << s->_starEdges.size() << endl;

    const bool debug = false;
    if (debug){
      char starLowName[100];
      sprintf(starLowName, "star-%04d-low.g2o", starNum);
      ofstream starLowStream(starLowName);
      optimizer->saveSubset(starLowStream, s->_lowLevelEdges);
    }
    bool labelOk=s->labelStarEdges(maxIterations, labeler);
    if (labelOk) {
      if (debug) {
        char starHighName[100];
        sprintf(starHighName, "star-%04d-high.g2o", starNum);
        ofstream starHighStream(starHighName);
        optimizer->saveSubset(starHighStream, s->_starEdges);
      }
    } else {
      cerr << "FAILURE" << endl;
    }
    starNum++;
  }
}

void computeBorder(StarSet& stars, EdgeStarMap& hesmap){
  cerr << "computing edges on the border" << endl;
  for (StarSet::iterator it=stars.begin(); it!=stars.end(); it++){
    Star* s=*it;
    for (HyperGraph::EdgeSet::iterator iit=s->_starEdges.begin(); iit!=s->_starEdges.end(); iit++){
      OptimizableGraph::Edge* e= (OptimizableGraph::Edge*) *iit;
      StarSet sset;
      starsInEdge(sset, e, hesmap, s->gauge());
      //cerr << "e: " << e << " l:" << e->level() << " sset.size()=" << sset.size() << endl;
      if (sset.size()>1){
        s->starFrontierEdges().insert(e);
      }
    }
  }
}

  void computeSimpleStars(StarSet& stars,
			  SparseOptimizer* optimizer,
			  EdgeLabeler* labeler,
			  EdgeCreator* creator,
			  OptimizableGraph::Vertex* gauge_,
			  std::string edgeTag,
			  std::string vertexTag,
			  int level,
			  int step,
			  int backboneIterations,
			  int starIterations,
			  double rejectionThreshold,
			  bool debug){

    cerr << "preforming the tree actions" << endl;
    HyperDijkstra d(optimizer);
    // compute a spanning tree based on the types of edges and vertices in the pool
    EdgeTypesCostFunction f(edgeTag, vertexTag, level);
    d.shortestPaths(gauge_,
        &f,
        std::numeric_limits< double >::max(),
        1e-6,
        false,
        std::numeric_limits< double >::max()/2);

    HyperDijkstra::computeTree(d.adjacencyMap());
    // constructs the stars on the backbone

    BackBoneTreeAction bact(optimizer, vertexTag, level, step);
    bact.init();

    cerr << "free edges size " << bact.freeEdges().size() << endl;

    // perform breadth-first visit of the visit tree and create the stars on the backbone
    d.visitAdjacencyMap(d.adjacencyMap(),&bact,true);
    stars.clear();

    for (VertexStarMultimap::iterator it=bact.vertexStarMultiMap().begin();
        it!=bact.vertexStarMultiMap().end(); it++){
      stars.insert(it->second);
    }
    cerr << "stars.size: " << stars.size() << endl;
    cerr << "size: " << bact.vertexStarMultiMap().size() << endl;

    //  for each star

    //    for all vertices in the backbone, select all edges leading/leaving from that vertex
    //    that are contained in freeEdges.

    //      mark the corresponding "open" vertices and add them to a multimap (vertex->star)

    //    select a gauge in the backbone

    //    push all vertices on the backbone

    //    compute an initial guess on the backbone

    //    one round of optimization backbone

    //    lock all vertices in the backbone

    //    push all "open" vertices

    //    for each open vertex,
    //      compute an initial guess given the backbone
    //      do some rounds of solveDirect
    //      if (fail)
    //        - remove the vertex and the edges in that vertex from the star
    //   - make the structures consistent

    //    pop all "open" vertices
    //    pop all "vertices" in the backbone
    //    unfix the vertices in the backbone

    int starNum=0;
    for (StarSet::iterator it=stars.begin(); it!=stars.end(); it++){
      Star* s =*it;
      HyperGraph::VertexSet backboneVertices = s->_lowLevelVertices;
      HyperGraph::EdgeSet backboneEdges = s->_lowLevelEdges;
      if (backboneEdges.empty())
	continue;

      // cerr << "optimizing backbone" << endl;
      // one of these  should be the gauge, to be simple we select the fisrt one in the backbone
      OptimizableGraph::VertexSet gauge;
      gauge.insert(*backboneVertices.begin());
      s->gauge()=gauge;
      s->optimizer()->push(backboneVertices);
      s->optimizer()->setFixed(gauge,true);
      s->optimizer()->initializeOptimization(backboneEdges);
      s->optimizer()->computeInitialGuess();
      s->optimizer()->optimize(backboneIterations);
      s->optimizer()->setFixed(backboneVertices, true);

      // cerr << "assignind edges.vertices not in bbone" << endl;
      HyperGraph::EdgeSet otherEdges;
      HyperGraph::VertexSet otherVertices;
      std::multimap<HyperGraph::Vertex*, HyperGraph::Edge*> vemap;
      for (HyperGraph::VertexSet::iterator bit=backboneVertices.begin(); bit!=backboneVertices.end(); bit++){
	HyperGraph::Vertex* v=*bit;
	for (HyperGraph::EdgeSet::iterator eit=v->edges().begin(); eit!=v->edges().end(); eit++){
	  OptimizableGraph::Edge* e = (OptimizableGraph::Edge*) *eit;
	  HyperGraph::EdgeSet::iterator feit=bact.freeEdges().find(e);
	  if (feit!=bact.freeEdges().end()){ // edge is admissible
	    otherEdges.insert(e);
	    bact.freeEdges().erase(feit);
	    for (size_t i=0; i<e->vertices().size(); i++){
	      OptimizableGraph::Vertex* ve= (OptimizableGraph::Vertex*)e->vertices()[i];
	      if (backboneVertices.find(ve)==backboneVertices.end()){
		otherVertices.insert(ve);
		vemap.insert(make_pair(ve,e));
	      }
	    }
	  }
	}
      }

      // RAINER TODO maybe need a better solution than dynamic casting here??
      OptimizationAlgorithmWithHessian* solverWithHessian = dynamic_cast<OptimizationAlgorithmWithHessian*>(s->optimizer()->solver());
      if (solverWithHessian) {
        s->optimizer()->push(otherVertices);
        // cerr << "optimizing vertices out of bbone" << endl;
        // cerr << "push" << endl;
        // cerr << "init" << endl;
        s->optimizer()->initializeOptimization(otherEdges);
        // cerr << "guess" << endl;
        s->optimizer()->computeInitialGuess();
        // cerr << "solver init" << endl;
        s->optimizer()->solver()->init();
        // cerr << "structure" << endl;
        if (!solverWithHessian->buildLinearStructure())
          cerr << "FATAL: failure while building linear structure" << endl;
        // cerr << "errors" << endl;
        s->optimizer()->computeActiveErrors();
        // cerr << "system" << endl;
        solverWithHessian->updateLinearSystem();
        // cerr << "directSolove" << endl;
      } else {
        cerr << "FATAL: hierarchical thing cannot be used with a solver that does not support the system structure construction" << endl;
      }

      // // then optimize the vertices one at a time to check if a solution is good
      for (HyperGraph::VertexSet::iterator vit=otherVertices.begin(); vit!=otherVertices.end(); vit++){
        OptimizableGraph::Vertex* v=(OptimizableGraph::Vertex*)(*vit);
        v->solveDirect();
        // cerr << " " << d;
        // if  a solution is found, add a vertex and all the edges in
        //othervertices that are pointing to that edge to the star
        s->_lowLevelVertices.insert(v);
        for (HyperGraph::EdgeSet::iterator eit=v->edges().begin(); eit!=v->edges().end(); eit++){
          OptimizableGraph::Edge* e = (OptimizableGraph::Edge*) *eit;
          if (otherEdges.find(e)!=otherEdges.end())
            s->_lowLevelEdges.insert(e);
        }
      }
      //cerr <<  endl;

      // relax the backbone and optimize it all
      // cerr << "relax bbone" << endl;
      s->optimizer()->setFixed(backboneVertices, false);
      //cerr << "fox gauge bbone" << endl;
      s->optimizer()->setFixed(s->gauge(),true);

      //cerr << "opt init" << endl;
      s->optimizer()->initializeOptimization(s->_lowLevelEdges);
      optimizer->computeActiveErrors();
      double initialChi = optimizer->activeChi2();
      int starOptResult = s->optimizer()->optimize(starIterations);
      //cerr << starOptResult << "(" << starIterations << ")  " << endl;
      double finalchi=-1.;

      cerr <<  "computing star: " << starNum << endl;

      int vKept=0, vDropped=0;
      if (!starIterations || starOptResult > 0  ){
	optimizer->computeActiveErrors();
	finalchi = optimizer->activeChi2();

#if 1

        s->optimizer()->computeActiveErrors();
        // cerr << "system" << endl;
        solverWithHessian->updateLinearSystem();
        HyperGraph::EdgeSet prunedStarEdges = backboneEdges;
        HyperGraph::VertexSet prunedStarVertices = backboneVertices;
        for (HyperGraph::VertexSet::iterator vit=otherVertices.begin(); vit!=otherVertices.end(); vit++){

	  //discard the vertices whose error is too big

          OptimizableGraph::Vertex* v=(OptimizableGraph::Vertex*)(*vit);
          MatrixXd h(v->dimension(), v->dimension());
          for (int i=0; i<v->dimension(); i++){
            for (int j=0; j<v->dimension(); j++)
              h(i,j)=v->hessian(i,j);
          }
          EigenSolver<Eigen::MatrixXd> esolver;
          esolver.compute(h);
          VectorXcd ev= esolver.eigenvalues();
          double emin = std::numeric_limits<double>::max();
          double emax = -std::numeric_limits<double>::max();
          for (int i=0; i<ev.size(); i++){
            emin = ev(i).real()>emin ? emin : ev(i).real();
            emax = ev(i).real()<emax ? emax : ev(i).real();
          }

          double d=emin/emax;

          // cerr << " " << d;
          if (d>rejectionThreshold){
	  // if  a solution is found, add a vertex and all the edges in
	  //othervertices that are pointing to that edge to the star
            prunedStarVertices.insert(v);
            for (HyperGraph::EdgeSet::iterator eit=v->edges().begin(); eit!=v->edges().end(); eit++){
              OptimizableGraph::Edge* e = (OptimizableGraph::Edge*) *eit;
              if (otherEdges.find(e)!=otherEdges.end())
                prunedStarEdges.insert(e);
            }
            //cerr << "K( " << v->id() << "," << d << ")" ;
            vKept ++;
          } else {
            vDropped++;
            //cerr << "R( " << v->id() << "," << d << ")" ;
          }
        }
        s->_lowLevelEdges=prunedStarEdges;
        s->_lowLevelVertices=prunedStarVertices;

#endif
	//cerr << "addHedges" << endl;
	//now add to the star the hierarchical edges
	std::vector<OptimizableGraph::Vertex*> vertices(2);
	vertices[0]= (OptimizableGraph::Vertex*) *s->_gauge.begin();

	for (HyperGraph::VertexSet::iterator vit=s->_lowLevelVertices.begin(); vit!=s->_lowLevelVertices.end(); vit++){
	  OptimizableGraph::Vertex* v=(OptimizableGraph::Vertex*)*vit;
	  vertices[1]=v;
	  if (v==vertices[0])
	    continue;
	  OptimizableGraph::Edge* e=creator->createEdge(vertices);
	  //rr << "creating edge" << e <<  Factory::instance()->tag(vertices[0]) << "->" <<  Factory::instance()->tag(v) <endl;
	  if (e) {
	    e->setLevel(level+1);
	    optimizer->addEdge(e);
	    s->_starEdges.insert(e);
	  } else {
            cerr << "HERE" << endl;
	    cerr << "FATAL, cannot create edge" << endl;
	  }
	}
      }

      cerr << " gauge: " << (*s->_gauge.begin())->id()
           << " kept: " << vKept
           << " dropped: " << vDropped
	   << " edges:" << s->_lowLevelEdges.size()
	   << " hedges" << s->_starEdges.size()
	   << " initial chi " << initialChi
	   << " final chi " << finalchi << endl;

      if (debug) {
	char starLowName[100];
	sprintf(starLowName, "star-%04d-low.g2o", starNum);
	ofstream starLowStream(starLowName);
	optimizer->saveSubset(starLowStream, s->_lowLevelEdges);
      }
      bool labelOk=false;
      if (!starIterations || starOptResult > 0)
        labelOk = s->labelStarEdges(0, labeler);
      if (labelOk) {
        if (debug) {
          char starHighName[100];
          sprintf(starHighName, "star-%04d-high.g2o", starNum);
          ofstream starHighStream(starHighName);
          optimizer->saveSubset(starHighStream, s->_starEdges);
        }
      } else {

        cerr << "FAILURE: " << starOptResult << endl;
      }
      starNum++;

      //label each hierarchical edge
      s->optimizer()->pop(otherVertices);
      s->optimizer()->pop(backboneVertices);
      s->optimizer()->setFixed(s->gauge(),false);
    }

    StarSet stars2;
    // now erase the stars that have 0 edges. They r useless
    for (StarSet::iterator it=stars.begin(); it!=stars.end(); it++){
      Star* s=*it;
      if (s->lowLevelEdges().size()==0) {
        delete s;
      } else
        stars2.insert(s);
    }
    stars=stars2;
  }

}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG, www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also tools/license/license.mtl.txt in the distribution.

#ifndef MTL_VECTOR_VEC_SCAL_MIXED_EXPR_INCLUDE
#define MTL_VECTOR_VEC_SCAL_MIXED_EXPR_INCLUDE

#include <boost/numeric/mtl/utility/enable_if.hpp>
#include <boost/numeric/mtl/vector/map_view.hpp>
#include <boost/numeric/mtl/operation/tfunctor_mixed.hpp>

namespace mtl { namespace vector {

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E1, map_view<tfunctor::left_plus<E1, typename E2::value_type>, E2> >::type
inline operator+(const E1& e1, const vec_expr<E2>& e2)
{
    typedef tfunctor::left_plus<E1, typename E2::value_type> ftype;
    typedef map_view<ftype, E2>                             type;

    return type(ftype(e1), static_cast<const E2&>(e2));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E2, map_view<tfunctor::right_plus<typename E1::value_type, E2>, E1> >::type
inline operator+(const vec_expr<E1>& e1, const E2& e2)
{
    typedef tfunctor::right_plus<typename E1::value_type, E2>ftype;
    typedef map_view<ftype, E1>                             type;

    return type(ftype(e2), static_cast<const E1&>(e1));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E1, map_view<tfunctor::left_minus<E1, typename E2::value_type>, E2> >::type
inline operator-(const E1& e1, const vec_expr<E2>& e2)
{
    typedef tfunctor::left_minus<E1, typename E2::value_type> ftype;
    typedef map_view<ftype, E2>                             type;

    return type(ftype(e1), static_cast<const E2&>(e2));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E2, map_view<tfunctor::right_minus<typename E1::value_type, E2>, E1> >::type
inline operator-(const vec_expr<E1>& e1, const E2& e2)
{
    typedef tfunctor::right_minus<typename E1::value_type, E2>ftype;
    typedef map_view<ftype, E1>                             type;

    return type(ftype(e2), static_cast<const E1&>(e1));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E1, map_view<tfunctor::left_min<E1, typename E2::value_type>, E2> >::type
inline min(const E1& e1, const vec_expr<E2>& e2)
{
    typedef tfunctor::left_min<E1, typename E2::value_type> ftype;
    typedef map_view<ftype, E2>                             type;

    return type(ftype(e1), static_cast<const E2&>(e2));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E2, map_view<tfunctor::right_min<typename E1::value_type, E2>, E1> >::type
inline min(const vec_expr<E1>& e1, const E2& e2)
{
    typedef tfunctor::right_min<typename E1::value_type, E2>ftype;
    typedef map_view<ftype, E1>                             type;

    return type(ftype(e2), static_cast<const E1&>(e1));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E1, map_view<tfunctor::left_max<E1, typename E2::value_type>, E2> >::type
inline max(const E1& e1, const vec_expr<E2>& e2)
{
    typedef tfunctor::left_max<E1, typename E2::value_type> ftype;
    typedef map_view<ftype, E2>                             type;

    return type(ftype(e1), static_cast<const E2&>(e2));
}

template <typename E1, typename E2>
typename mtl::traits::enable_if_scalar<E2, map_view<tfunctor::right_max<typename E1::value_type, E2>, E1> >::type
inline max(const vec_expr<E1>& e1, const E2& e2)
{
    typedef tfunctor::right_max<typename E1::value_type, E2>ftype;
    typedef map_view<ftype, E1>                             type;

    return type(ftype(e2), static_cast<const E1&>(e1));
}

}} // namespace mtl::vector

#endif // MTL_VECTOR_VEC_SCAL_MIXED_EXPR_INCLUDE


//////////////////////////////////////////////////////////////////////////////////////
// This file is distributed under the University of Illinois/NCSA Open Source License.
// See LICENSE file in top directory for details.
//
// Copyright (c) 2016 Jeongnim Kim and QMCPACK developers.
//
// File developed by: Ken Esler, kpesler@gmail.com, University of Illinois at Urbana-Champaign
//                    Jeremy McMinnis, jmcminis@gmail.com, University of Illinois at Urbana-Champaign
//                    Luke Shulenburger, lshulen@sandia.gov, Sandia National Laboratories
//                    Jeongnim Kim, jeongnim.kim@gmail.com, University of Illinois at Urbana-Champaign
//                    Raymond Clay III, j.k.rofling@gmail.com, Lawrence Livermore National Laboratory
//                    Ye Luo, yeluo@anl.gov, Argonne National Laboratory
//                    Mark A. Berrill, berrillma@ornl.gov, Oak Ridge National Laboratory
//
// File created by: Jeongnim Kim, jeongnim.kim@gmail.com, University of Illinois at Urbana-Champaign
//////////////////////////////////////////////////////////////////////////////////////

#include "QMCDrivers/SimpleFixedNodeBranch.h"
#include "QMCDrivers/DMC/WalkerControlFactory.h"
#include <numeric>
#include "OhmmsData/FileUtility.h"
#include "Utilities/RandomGenerator.h"
#include "QMCDrivers/WalkerControlBase.h"
#include "Estimators/EstimatorManagerBase.h"
#include "QMCDrivers/BranchIO.h"
#include "Particle/Reptile.h"
#ifdef HAVE_ADIOS
#include <adios.h>
#endif

//#include <boost/archive/text_oarchive.hpp>

namespace qmcplusplus
{

///enum to yes/no options saved in sParam
enum {COMBOPT, USETAUOPT, MIXDMCOPT, DUMMYOPT};

SimpleFixedNodeBranch::SimpleFixedNodeBranch(RealType tau, int nideal)
  : vParam(1.0), WalkerController(0), BackupWalkerController(0),
    MyEstimator(0)//, PopHist(5), DMCEnergyHist(5)
{
  BranchMode.set(B_DMCSTAGE,0); //warmup stage
  BranchMode.set(B_POPCONTROL,1); //use standard DMC
  BranchMode.set(B_USETAUEFF,1); //use taueff
  BranchMode.set(B_CLEARHISTORY,0); //clear history and start with the current average
  BranchMode.set(B_KILLNODES,0); //when killing walkers at nodes etrial is updated differently
  vParam[B_TAU]=tau;
  vParam[B_TAUEFF]=tau;
  vParam[B_FEEDBACK]=1.0;
  vParam[B_FILTERSCALE]=10;
  R2Accepted(1.0e-10);
  R2Proposed(1.0e-10);
  //set the default values for integer parameters
  iParam[B_WARMUPSTEPS]=200;
  iParam[B_ENERGYUPDATEINTERVAL]=1;
  iParam[B_BRANCHINTERVAL]=1;
  iParam[B_TARGETWALKERS]=0;
  iParam[B_MAXWALKERS]=nideal;
  iParam[B_MINWALKERS]=nideal;
  iParam[B_COUNTER]=-1;
  //default is no
  sParam.resize(DUMMYOPT,"no");
  //default is classic
  branching_cutoff_scheme="classic";
  registerParameters();
  reset();
}

/** copy constructor
 *
 * Copy only selected data members and WalkerController is never copied.
 */
SimpleFixedNodeBranch::SimpleFixedNodeBranch(const SimpleFixedNodeBranch& abranch):
  BranchMode(abranch.BranchMode),
  iParam(abranch.iParam),
  vParam(abranch.vParam),
  WalkerController(0), MyEstimator(0),
  sParam(abranch.sParam),
  branching_cutoff_scheme(abranch.branching_cutoff_scheme)
{
  registerParameters();
  reset();
}

void SimpleFixedNodeBranch::registerParameters()
{
  m_param.add(iParam[B_WARMUPSTEPS],"warmupSteps","int");
  m_param.add(iParam[B_WARMUPSTEPS],"warmupsteps","int");
  m_param.add(iParam[B_ENERGYUPDATEINTERVAL],"energyUpdateInterval","int");
  m_param.add(iParam[B_BRANCHINTERVAL],"branchInterval","int");
  m_param.add(iParam[B_TARGETWALKERS],"targetWalkers","int");
  m_param.add(iParam[B_TARGETWALKERS],"targetwalkers","int");
  m_param.add(iParam[B_TARGETWALKERS],"target_walkers","int");
  //trial energy
  m_param.add(vParam[B_EREF],"refEnergy","AU");
  m_param.add(vParam[B_EREF],"ref_energy","AU");
  m_param.add(vParam[B_EREF],"en_ref","AU");
  m_param.add(vParam[B_TAU],"tau","AU");
  m_param.add(vParam[B_TAU],"timestep","AU");
  m_param.add(vParam[B_TAU],"timeStep","AU");
  m_param.add(vParam[B_TAU],"TimeStep","AU");
  //filterscale:  sets the filtercutoff to sigma*filterscale
  m_param.add(vParam[B_FILTERSCALE],"filterscale","double");
  //feed back parameter for population control
  m_param.add(vParam[B_FEEDBACK],"feedback","double");
  //turn on/off effective tau onl for time-step error comparisons
  m_param.add(sParam[USETAUOPT],"useBareTau","option");
  m_param.add(sParam[MIXDMCOPT],"warmupByReconfiguration","opt");
  m_param.add(branching_cutoff_scheme,"branching_cutoff_scheme","option");
}

void SimpleFixedNodeBranch::start(const std::string& froot, bool append)
{
  RootName=froot;
  MyEstimator->RootName=froot;
  MyEstimator->reset();
}

int SimpleFixedNodeBranch::initWalkerController(MCWalkerConfiguration& walkers, bool fixW, bool killwalker)
{
  BranchMode.set(B_DMC,1);//set DMC
  BranchMode.set(B_DMCSTAGE,iParam[B_WARMUPSTEPS]==0);//use warmup
  //this is not necessary
  //check if tau is different and set the initial values
  //vParam[B_TAU]=tau;
  bool fromscratch=false;
  EstimatorRealType tau=vParam[B_TAU];

  int nwtot_now=walkers.getGlobalNumWalkers();

  //this is the first time DMC is used
  if(WalkerController == 0)
  {
    if(iParam[B_TARGETWALKERS]==0)
    {
      Communicate* acomm=MyEstimator->getCommunicator();
      int ncontexts=acomm->size();
      std::vector<int> nw(ncontexts,0),nwoff(ncontexts+1,0);
      nw[acomm->rank()]=walkers.getActiveWalkers();
      acomm->allreduce(nw);
      for(int ip=0; ip<ncontexts; ++ip)
        nwoff[ip+1]=nwoff[ip]+nw[ip];
      walkers.setGlobalNumWalkers(nwoff[ncontexts]);
      walkers.setWalkerOffsets(nwoff);
      iParam[B_TARGETWALKERS]=nwoff[ncontexts];
    }
    WalkerController = createWalkerController(iParam[B_TARGETWALKERS], MyEstimator->getCommunicator(), myNode);
    if(!BranchMode[B_RESTART])
    {
      fromscratch=true;
      app_log() << "  START ALL OVER " << std::endl;
      vParam[B_TAUEFF]=tau;
      BranchMode.set(B_POPCONTROL,!fixW);//fixW -> 0
      BranchMode.set(B_KILLNODES,killwalker);
      iParam[B_MAXWALKERS]=WalkerController->Nmax;
      iParam[B_MINWALKERS]=WalkerController->Nmin;
      if(!fixW && sParam[MIXDMCOPT]=="yes")
      {
        app_log() << "Warmup DMC is done with a fixed population " << iParam[B_TARGETWALKERS] << std::endl;
        BackupWalkerController=WalkerController; //save the main controller
        WalkerController=createWalkerController(iParam[B_TARGETWALKERS],MyEstimator->getCommunicator(), myNode,true);
        BranchMode.set(B_POPCONTROL,0);
      }
      //PopHist.clear();
      //PopHist.reserve(std::max(iParam[B_ENERGYUPDATEINTERVAL],5));
    }
    WalkerController->setWalkerID(walkers);
  }
  //else
  //{
  //  BranchMode.set(B_DMCSTAGE,0);//always reset warmup
  //}
  MyEstimator->reset();
  //update the simulation parameters
  WalkerController->put(myNode);
  //assign current Eref and a large number for variance
  WalkerController->setTrialEnergy(vParam[B_ETRIAL]);
  this->reset();
  if(fromscratch)
  {
    //determine the branch cutoff to limit wild weights based on the sigma and sigmaBound
    setBranchCutoff(vParam[B_SIGMA2],WalkerController->targetSigma,50,walkers.R.size());
    vParam[B_TAUEFF]=tau*R2Accepted.result()/R2Proposed.result();
  }
  //reset controller
  WalkerController->reset();
  if(BackupWalkerController)
    BackupWalkerController->reset();
  app_log() << "  QMC counter      = " << iParam[B_COUNTER] << std::endl;
  app_log() << "  time step        = " << vParam[B_TAU] << std::endl;
  app_log() << "  effective time step = " << vParam[B_TAUEFF] << std::endl;
  app_log() << "  trial energy     = " << vParam[B_ETRIAL] << std::endl;
  app_log() << "  reference energy = " << vParam[B_EREF] << std::endl;
  app_log() << "  Feedback = " << vParam[B_FEEDBACK] <<  std::endl;
  app_log() << "  reference variance = " << vParam[B_SIGMA2] << std::endl;
  app_log() << "  target walkers = " << iParam[B_TARGETWALKERS] << std::endl;
  app_log() << "  branching cutoff scheme " << branching_cutoff_scheme << std::endl;
  app_log() << "  branch cutoff = " <<  vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX] << std::endl;
  app_log() << "  Max and minimum walkers per node= " << iParam[B_MAXWALKERS] << " " << iParam[B_MINWALKERS] << std::endl;
  app_log() << "  QMC Status (BranchMode) = " << BranchMode << std::endl;

  return int(round(double(iParam[B_TARGETWALKERS])/double(nwtot_now)));
}

void SimpleFixedNodeBranch::initReptile(MCWalkerConfiguration& W)
{
  RealType allowedFlux=50.0;
  BranchMode.set(B_RMC,1);//set RMC
  BranchMode.set(B_RMCSTAGE,iParam[B_WARMUPSTEPS]==0);//use warmup
  //this is not necessary
  //check if tau is different and set the initial values
  //vParam[B_TAU]=tau;
  bool fromscratch=false;
  EstimatorRealType tau=vParam[B_TAU];
  //this is the first time DMC is used
  if(WalkerController == 0)
  {
    //  if(iParam[B_TARGETWALKERS]==0)
    //  {
    //    Communicate* acomm=MyEstimator->getCommunicator();
    //    int ncontexts=acomm->size();
    //    std::vector<int> nw(ncontexts,0),nwoff(ncontexts+1,0);
    //    nw[acomm->rank()]=W.getActiveWalkers();
    //   acomm->allreduce(nw);
    //    for(int ip=0; ip<ncontexts; ++ip)
    //      nwoff[ip+1]=nwoff[ip]+nw[ip];
    //    W.setGlobalNumWalkers(nwoff[ncontexts]);
    //    W.setWalkerOffsets(nwoff);
    //    iParam[B_TARGETWALKERS]=nwoff[ncontexts];
    //  }
    if(!BranchMode[B_RESTART])
    {
      fromscratch=true;
      app_log() << "  START ALL OVER " << std::endl;
      vParam[B_TAUEFF]=tau;
      //PopHist.clear();
      //PopHist.reserve(std::max(iParam[B_ENERGYUPDATEINTERVAL],5));
    }
  }
  //else
  //{
  //  BranchMode.set(B_DMCSTAGE,0);//always reset warmup
  //}
  MyEstimator->reset();
  this->reset();
  if(fromscratch)
  {
    //determine the branch cutoff to limit wild weights based on the sigma and sigmaBound
    setBranchCutoff(vParam[B_SIGMA2],allowedFlux,50,W.R.size());
    vParam[B_TAUEFF]=tau*R2Accepted.result()/R2Proposed.result();
  }
  //reset controller
  app_log() << "  QMC counter      = " << iParam[B_COUNTER] << std::endl;
  app_log() << "  time step        = " << vParam[B_TAU] << std::endl;
  app_log() << "  effective time step = " << vParam[B_TAUEFF] << std::endl;
  app_log() << "  reference energy = " << vParam[B_EREF] << std::endl;
  app_log() << "  Feedback = " << vParam[B_FEEDBACK] <<  std::endl;
  app_log() << "  reference variance = " << vParam[B_SIGMA2] << std::endl;
  app_log() << "  branching cutoff scheme " << branching_cutoff_scheme << std::endl;
  app_log() << "  branch cutoff = " <<  vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX] << std::endl;
  app_log() << "  QMC Status (BranchMode) = " << BranchMode << std::endl;
}

void SimpleFixedNodeBranch::flush(int counter)
{
  if(counter==0 && WalkerController)
    WalkerController->reset();
}

void SimpleFixedNodeBranch::branch(int iter, MCWalkerConfiguration& walkers)
{
  //collect the total weights and redistribute the walkers accordingly, using a fixed tolerance
  //RealType pop_now= WalkerController->branch(iter,walkers,0.1);
  EstimatorRealType pop_now;
  if(BranchMode[B_DMCSTAGE]||iter)
    pop_now= WalkerController->branch(iter,walkers,0.1);
  else
    pop_now = WalkerController->doNotBranch(iter,walkers);//do not branch for the first step of a warmup
  //population for trial energy modification should not include any released node walkers.
  pop_now -= WalkerController->EnsembleProperty.RNSamples;
  //current energy
  vParam[B_ENOW]=WalkerController->EnsembleProperty.Energy;
  VarianceHist(WalkerController->EnsembleProperty.Variance);
  R2Accepted(WalkerController->EnsembleProperty.R2Accepted);
  R2Proposed(WalkerController->EnsembleProperty.R2Proposed);
  //PopHist(pop_now);
  vParam[B_EREF]=EnergyHist.mean();//current mean
  if(BranchMode[B_USETAUEFF])
    vParam[B_TAUEFF]=vParam[B_TAU]*R2Accepted.result()/R2Proposed.result();
  if(BranchMode[B_KILLNODES])
    EnergyHist(vParam[B_ENOW]-std::log(WalkerController->EnsembleProperty.LivingFraction)/vParam[B_TAUEFF]);
  else
    EnergyHist(vParam[B_ENOW]);
  if(BranchMode[B_DMCSTAGE]) // main stage
  {
    if(BranchMode[B_POPCONTROL])
    {
      if(ToDoSteps>0)
        --ToDoSteps;
      else
      {
        vParam[B_ETRIAL]=vParam[B_EREF]+vParam[B_FEEDBACK]*(logN-std::log(pop_now));
        ToDoSteps=iParam[B_ENERGYUPDATEINTERVAL]-1;
      }
    }
    else
      vParam[B_ETRIAL]=vParam[B_EREF];
  }
  else//warmup
  {
    if(BranchMode[B_USETAUEFF])
      vParam[B_TAUEFF]=vParam[B_TAU]*R2Accepted.result()/R2Proposed.result();
    if(BranchMode[B_POPCONTROL])
    {
      //RealType emix=((iParam[B_WARMUPSTEPS]-ToDoSteps)<100)?(0.25*vParam[B_EREF]+0.75*vParam[B_ENOW]):vParam[B_EREF];
      //vParam[B_ETRIAL]=emix+Feedback*(logN-std::log(pop_now));
      //vParam[B_ETRIAL]=vParam[B_EREF]+Feedback*(logN-std::log(pop_now));
      if(BranchMode[B_KILLNODES])
        vParam[B_ETRIAL]=(0.00*vParam[B_EREF]+1.0*vParam[B_ENOW])
                         +vParam[B_FEEDBACK]*(logN-std::log(pop_now))-std::log(WalkerController->EnsembleProperty.LivingFraction)/vParam[B_TAU];
      else
        vParam[B_ETRIAL]=vParam[B_ENOW]+(logN-std::log(pop_now))/vParam[B_TAU];
    }
    else
      vParam[B_ETRIAL]=vParam[B_EREF];
    --ToDoSteps;
    if(ToDoSteps==0)  //warmup is done
    {
      vParam[B_SIGMA2]=VarianceHist.mean();
      setBranchCutoff(vParam[B_SIGMA2],WalkerController->targetSigma,10,walkers.R.size());
      app_log() << "\n Warmup is completed after " << iParam[B_WARMUPSTEPS] << std::endl;
      if(BranchMode[B_USETAUEFF])
        app_log() << "\n  TauEff     = " << vParam[B_TAUEFF] << "\n TauEff/Tau = " << vParam[B_TAUEFF]/vParam[B_TAU];
      else
        app_log() << "\n  TauEff proposed   = " << vParam[B_TAUEFF]*R2Accepted.result()/R2Proposed.result();
      app_log()  << "\n  Etrial     = " << vParam[B_ETRIAL] << std::endl;
      app_log() << " Running average of energy = " << EnergyHist.mean() << std::endl;
      app_log() << "                  Variance = " << vParam[B_SIGMA2] << std::endl;
      app_log() << "branch cutoff = " <<  vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX] << std::endl;
      ToDoSteps = iParam[B_ENERGYUPDATEINTERVAL]-1;
      iParam[B_WARMUPSTEPS]=0;
      BranchMode.set(B_DMCSTAGE,1); //set BranchModex to main stage
      //reset the histogram
      EnergyHist.clear();
      EnergyHist(vParam[B_ENOW]);
      if(sParam[MIXDMCOPT]=="yes")
      {
        app_log() << "Switching to DMC with fluctuating populations" << std::endl;
        BranchMode.set(B_POPCONTROL,1); //use standard DMC
        delete WalkerController;
        WalkerController=BackupWalkerController;
        BackupWalkerController=0;
        vParam[B_ETRIAL]=vParam[B_EREF];
        app_log()  << "  Etrial     = " << vParam[B_ETRIAL] << std::endl;
        WalkerController->start();
      }
      //This is not necessary
      //EnergyHist(DMCEnergyHist.mean());
    }
  }
  WalkerController->setTrialEnergy(vParam[B_ETRIAL]);
  //accumulate collectables and energies for scalar.dat
  EstimatorRealType wgt_inv=WalkerController->NumContexts/WalkerController->EnsembleProperty.Weight;
  walkers.Collectables *= wgt_inv;
  MyEstimator->accumulate(walkers);
}

/**
 *
 */
void SimpleFixedNodeBranch::collect(int iter, MCWalkerConfiguration& W)
{
  //Update the current energy and accumulate.
  MCWalkerConfiguration::Walker_t& head = W.reptile->getHead();
  MCWalkerConfiguration::Walker_t& tail = W.reptile->getTail();
  vParam[B_ENOW]=0.5*(head.Properties(LOCALENERGY)+tail.Properties(LOCALENERGY));
// app_log()<<"IN SimpleFixedNodeBranch::collect\n";
// app_log()<<"\tvParam[B_ENOW]="<<vParam[B_ENOW]<< std::endl;
  EnergyHist(vParam[B_ENOW]);
  vParam[B_EREF]=EnergyHist.mean();
// app_log()<<"\tvParam[B_EREF]="<<vParam[B_EREF]<< std::endl;
  //Update the energy variance and R2 for effective timestep and filtering.
  VarianceHist(std::pow(vParam[B_ENOW]-vParam[B_EREF],2));
  R2Accepted(head.Properties(R2ACCEPTED));
  R2Proposed(head.Properties(R2PROPOSED));
// app_log()<<"\thead.Properties(R2ACCEPTED)="<<head.Properties(R2ACCEPTED)<< std::endl;
// app_log()<<"\thead.Properties(R2PROPOSED)="<<head.Properties(R2PROPOSED)<< std::endl;
//  app_log()<<"\tR2Accepted="<<R2Accepted.result()<< std::endl;
// app_log()<<"\tR2Proposed="<<R2Proposed.result()<< std::endl;
//  app_log()<<"\tR2Accept/R2Prop="<<R2Accepted.result()/R2Proposed.result()<< std::endl;
// app_log()<<"\t <E^2> = "<<VarianceHist.mean()<< std::endl;
// app_log()<<"\t <E>   = "<<EnergyHist.mean()<< std::endl;
//  app_log()<<"\t <E>^2 = "<<std::pow(EnergyHist.mean(),2)<< std::endl;
// app_log()<<"\t var = "<<VarianceHist.mean()-pow(EnergyHist.mean(),2)<< std::endl;
// app_log()<<"--------------\n";
  //current mean
  if(BranchMode[B_USETAUEFF])
  {
    //app_log()<<" BRANCHMODE = "<<BranchMode[B_USETAUEFF]<< std::endl;
    vParam[B_TAUEFF]=vParam[B_TAU]*R2Accepted.result()/R2Proposed.result();
    //  app_log()<<"\tvParam[B_TAU]="<<vParam[B_TAU]<<" "<<vParam[B_TAUEFF]<< std::endl;
  }
  /*
  if(BranchMode[B_RMCSTAGE]) // main stage
  {
    if(BranchMode[B_POPCONTROL])
    {
      if(ToDoSteps>0)
        --ToDoSteps;
      else
      {
        vParam[B_ETRIAL]=vParam[B_EREF]+vParam[B_FEEDBACK]*(logN-std::log(pop_now));
        ToDoSteps=iParam[B_ENERGYUPDATEINTERVAL]-1;
      }
    }
    else
      vParam[B_ETRIAL]=vParam[B_EREF];
  }*/
  //app_log()<<"BranchMode[B_RMCSTAGE]="<<BranchMode[B_RMCSTAGE]<< std::endl;
  if(!BranchMode[B_RMCSTAGE])//warmup
  {
    if(BranchMode[B_USETAUEFF])
      vParam[B_TAUEFF]=vParam[B_TAU]*R2Accepted.result()/R2Proposed.result();
    // app_log()<<"\t <E^2> = "<<VarianceHist.mean()<< std::endl;
// app_log()<<"\t <E>   = "<<EnergyHist.mean()<< std::endl;
// app_log()<<"\t <E>^2 = "<<std::pow(EnergyHist.mean(),2)<< std::endl;
    //app_log()<<"\t var = "<<VarianceHist.mean()-std::pow(EnergyHist.mean(),2)<< std::endl;
// app_log()<<"\t var = "<<VarianceHist.mean()<< std::endl;
//  app_log()<<"----\n";
    //app_log()<<"ToDoSteps="<<ToDoSteps<< std::endl;
    vParam[B_ETRIAL]=vParam[B_EREF];
    --ToDoSteps;
    if(ToDoSteps==0)  //warmup is done
    {
      vParam[B_TAUEFF]=vParam[B_TAU]*R2Accepted.result()/R2Proposed.result();
      vParam[B_SIGMA2]=VarianceHist.mean();
      setBranchCutoff(vParam[B_SIGMA2],vParam[B_FILTERSCALE],vParam[B_FILTERSCALE],W.R.size());
      app_log() << "\n Warmup is completed after " << iParam[B_WARMUPSTEPS] << " steps."<< std::endl;
      if(BranchMode[B_USETAUEFF])
        app_log() << "\n  TauEff     = " << vParam[B_TAUEFF] << "\n TauEff/Tau = " << vParam[B_TAUEFF]/vParam[B_TAU];
      else
        app_log() << "\n  TauEff proposed   = " << vParam[B_TAUEFF]*R2Accepted.result()/R2Proposed.result();
      app_log() << "\n Running average of energy = " << EnergyHist.mean() << std::endl;
      app_log() << "\n                  Variance = " << vParam[B_SIGMA2] << std::endl;
      app_log() << "\nbranch cutoff = " <<  vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX] << std::endl;
      ToDoSteps = iParam[B_ENERGYUPDATEINTERVAL]-1;
      iParam[B_WARMUPSTEPS]=0;
      BranchMode.set(B_RMCSTAGE,1); //set BranchModex to main stage
      //reset the histogram
      EnergyHist.clear();
      EnergyHist(vParam[B_ENOW]);
    }
  }
  //accumulate collectables and energies for scalar.dat
  MyEstimator->accumulate(W);
}

/** Calculates and saves various action components, also does necessary updates for running averages.
 *
 */
void SimpleFixedNodeBranch::reset()
{
  //use effective time step of BranchInterval*Tau
  //Feed = 1.0/(static_cast<RealType>(NumGeneration*BranchInterval)*Tau);
  //logN = Feed*std::log(static_cast<RealType>(Nideal));
  //JNKIM passive
  //BranchMode.set(B_DMC,1);//set DMC
  //BranchMode.set(B_DMCSTAGE,0);//set warmup
  if(WalkerController)
  {
    //this is to compare the time step errors
    BranchMode.set(B_USETAUEFF,sParam[USETAUOPT]=="no");
    if(BranchMode[B_DMCSTAGE]) //
      ToDoSteps = iParam[B_ENERGYUPDATEINTERVAL]-1;
    else
      ToDoSteps = iParam[B_WARMUPSTEPS];
    if(BranchMode[B_POPCONTROL])
    {
      //logN = Feedback*std::log(static_cast<RealType>(iParam[B_TARGETWALKERS]));
      logN = std::log(static_cast<EstimatorRealType>(iParam[B_TARGETWALKERS]));
      if (vParam[B_FEEDBACK]==0.0) vParam[B_FEEDBACK] = 1.0;
    }
    else
    {
      //may set Eref to a safe value
      //if(EnergyHistory.count()<5) Eref -= vParam[EnergyWindowIndex];
      vParam[B_ETRIAL]=vParam[B_EREF];
      vParam[B_FEEDBACK]=0.0;
      logN=0.0;
    }
//       vParam(abranch.vParam)
    WalkerController->start();
  }
  if(BranchMode[B_RMC])
  {
    //this is to compare the time step errors
    // BranchMode.set(B_USETAUEFF,sParam[USETAUOPT]=="no");
    if(BranchMode[B_RMCSTAGE]) //
      ToDoSteps = iParam[B_ENERGYUPDATEINTERVAL]-1;
    else
      ToDoSteps = iParam[B_WARMUPSTEPS];
  }
}

void SimpleFixedNodeBranch::setRN (bool rn)
{
  RN=rn;
  WalkerController->WriteRN = rn;
  WalkerController->start();
}

int SimpleFixedNodeBranch::resetRun(xmlNodePtr cur)
{
  app_log() << "BRANCH resetRun" << std::endl;
  //estimator is always reset
  MyEstimator->reset();
  MyEstimator->setCollectionMode(true);
  std::bitset<B_MODE_MAX> bmode(BranchMode);
  IParamType iparam_old(iParam);
  VParamType vparam_old(vParam);
  myNode=cur;
  //store old target
  int nw_target=iParam[B_TARGETWALKERS];
  m_param.put(cur);

  int target_min=-1;
  ParameterSet p;
  p.add(target_min,"minimumtargetwalkers","int"); //p.add(target_min,"minimumTargetWalkers","int");
  p.add(target_min,"minimumsamples","int"); //p.add(target_min,"minimumSamples","int");
  p.put(cur);

  if (iParam[B_TARGETWALKERS] < target_min) {
    iParam[B_TARGETWALKERS] = target_min;
  }

  bool same_wc=true;
  if(BranchMode[B_DMC] && WalkerController)
  {
    std::string reconfig("no");
    if(WalkerController->MyMethod) reconfig="yes";
    std::string reconfig_prev(reconfig);
    ParameterSet p;
    p.add(reconfig,"reconfiguration","string");
    p.put(cur);
    same_wc=(reconfig==reconfig_prev);
  }

  //everything is the same, do nothing
  if(same_wc && bmode==BranchMode
      && std::equal(iParam.begin(),iParam.end(),iparam_old.begin())
      && std::equal(vParam.begin(),vParam.end(),vparam_old.begin())
    )
  {
    app_log() << "  Continue with the same input as the previous block." << std::endl;
    app_log().flush();
    //return 1;
  }
  app_log() << " SimpleFixedNodeBranch::resetRun detected changes in <parameter>'s " << std::endl;
  app_log() << " BranchMode : " << bmode << " " << BranchMode << std::endl;

  //vmc does not need to do anything with WalkerController
  if(!BranchMode[B_DMC])
  {
    app_log() << " iParam (old): "  <<iparam_old << std::endl;
    app_log() << " iParam (new): "  <<iParam << std::endl;
    app_log() << " vParam (old): "  <<vparam_old << std::endl;
    app_log() << " vParam (new): "  <<vParam << std::endl;
    app_log().flush();
    return 1;
  }

  if(WalkerController==0)
  {
    APP_ABORT("SimpleFixedNodeBranch::resetRun cannot initialize WalkerController");
  }

  if(!same_wc)
  {
    app_log() << "Destroy WalkerController. Existing method " << WalkerController->MyMethod << std::endl;;
    delete WalkerController;
    WalkerController = createWalkerController(iParam[B_TARGETWALKERS], MyEstimator->getCommunicator(), myNode);
    app_log().flush();

    BranchMode[B_POPCONTROL]= (WalkerController->MyMethod == 0);
    if(BranchMode[B_POPCONTROL])
    {
      vParam[B_ETRIAL]=vParam[B_EREF];
      if (vParam[B_FEEDBACK]==0.0) vParam[B_FEEDBACK]=1.0;
    }
  }

  //always add a warmup step using default 10 steps
  R2Accepted.clear();
  R2Proposed.clear();
  //R2Accepted(1.0e-12);
  //R2Proposed(1.0e-12);
  BranchMode[B_DMCSTAGE]=0;
  WalkerController->put(myNode);
  ToDoSteps=iParam[B_WARMUPSTEPS]=(iParam[B_WARMUPSTEPS])?iParam[B_WARMUPSTEPS]:10;
  setBranchCutoff(vParam[B_SIGMA2],WalkerController->targetSigma,10);
  WalkerController->reset();
#ifdef QMC_CUDA
  reset(); // needed. Ye
#endif
  if(BackupWalkerController)
    BackupWalkerController->reset();

  iParam[B_MAXWALKERS]=WalkerController->Nmax;
  iParam[B_MINWALKERS]=WalkerController->Nmin;

  app_log() << " iParam (old): "  <<iparam_old << std::endl;
  app_log() << " iParam (new): "  <<iParam << std::endl;
  app_log() << " vParam (old): "  <<vparam_old << std::endl;
  app_log() << " vParam (new): "  <<vParam << std::endl;

  app_log() << std::endl << " Using branching cutoff scheme " << branching_cutoff_scheme << std::endl;

  app_log().flush();

  //  return static_cast<int>(iParam[B_TARGETWALKERS]*1.01/static_cast<double>(nw_target));
  return static_cast<int>(round(static_cast<double>(iParam[B_TARGETWALKERS]/static_cast<double>(nw_target))));
}

void SimpleFixedNodeBranch::checkParameters(MCWalkerConfiguration& w)
{
  std::ostringstream o;
  if(!BranchMode[B_DMCSTAGE])
  {
    EstimatorRealType e, sigma2;
    MyEstimator->getCurrentStatistics(w,e,sigma2);
    vParam[B_ETRIAL]=vParam[B_EREF]=e;
    vParam[B_SIGMA2]=sigma2;
    EnergyHist.clear();
    VarianceHist.clear();
    //DMCEnergyHist.clear();
    EnergyHist(vParam[B_EREF]);
    VarianceHist(vParam[B_SIGMA2]);
    //DMCEnergyHist(vParam[B_EREF]);
    o << "SimpleFixedNodeBranch::checkParameters " << std::endl;
    o << "  Average Energy of a population  = " << e << std::endl;
    o << "  Energy Variance = " << vParam[B_SIGMA2] << std::endl;
  }
  app_log() << o.str() << std::endl;
  app_log().flush();
}

void SimpleFixedNodeBranch::finalize(MCWalkerConfiguration& w)
{
  std::ostringstream o;
  if(WalkerController)
  {
    o << "====================================================";
    o << "\n  SimpleFixedNodeBranch::finalize after a DMC block" ;
    o << "\n    QMC counter                   = " << iParam[B_COUNTER];
    o << "\n    time step                     = " << vParam[B_TAU];
    o << "\n    effective time step           = " << vParam[B_TAUEFF];
    o << "\n    trial energy                  = " << vParam[B_ETRIAL];
    o << "\n    reference energy              = " << vParam[B_EREF];
    o << "\n    reference variance            = " << vParam[B_SIGMA2];
    o << "\n    target walkers                = " << iParam[B_TARGETWALKERS];
    o << "\n    branch cutoff                 = " << vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX];
    o << "\n    Max and minimum walkers per node= " << iParam[B_MAXWALKERS] << " " << iParam[B_MINWALKERS];
    o << "\n    Feedback                      = " << vParam[B_FEEDBACK];
    o << "\n    QMC Status (BranchMode)       = " << BranchMode;
    o << "\n====================================================";
  }
  //running RMC
  else if (BranchMode[B_RMC])
  {
    o << "====================================================";
    o << "\n  SimpleFixedNodeBranch::finalize after a RMC block" ;
    o << "\n    QMC counter                   = " << iParam[B_COUNTER];
    o << "\n    time step                     = " << vParam[B_TAU];
    o << "\n    effective time step           = " << vParam[B_TAUEFF];
    o << "\n    reference energy              = " << vParam[B_EREF];
    o << "\n    reference variance            = " << vParam[B_SIGMA2];
    o << "\n    cutoff energy                 = " << vParam[B_BRANCHCUTOFF] << " " << vParam[B_BRANCHMAX];
    o << "\n    QMC Status (BranchMode)       = " << BranchMode;
    o << "\n====================================================";
  }
  else
  {
    //running VMC
    EstimatorRealType e, sigma2;
    //MyEstimator->getEnergyAndWeight(e,w,sigma2);
    MyEstimator->getCurrentStatistics(w,e,sigma2);
    vParam[B_ETRIAL]=vParam[B_EREF]=e;
    vParam[B_SIGMA2]=sigma2;
    //this is just to avoid diving by n-1  == 0
    EnergyHist(vParam[B_EREF]);
    //add Eref to the DMCEnergyHistory
    //DMCEnergyHist(vParam[B_EREF]);
    o << "====================================================";
    o << "\n  SimpleFixedNodeBranch::finalize after a VMC block" ;
    o << "\n    QMC counter        = " << iParam[B_COUNTER];
    o << "\n    time step          = " << vParam[B_TAU];
    o << "\n    reference energy   = " << vParam[B_EREF];
    o << "\n    reference variance = " << vParam[B_SIGMA2];
    o << "\n====================================================";
  }
  app_log() << o.str() << std::endl;
  write(RootName,true);
}

/**  Parse the xml file for parameters
 *@param cur current xmlNode
 *@param LogOut std::ostream to which the run-time report is sent
 *
 * Few important parameters are:
 * <ul>
 * <li> en_ref: a reference energy
 * <li> num_gen: number of generations \f$N_G\f$ to reach  equilibrium, used in the feedback parameter
 * \f$ feed = \frac{1}{N_G \tau} \f$
 * </ul>
 */
bool SimpleFixedNodeBranch::put(xmlNodePtr cur)
{
  //save it
  myNode=cur;
  //check dmc/vmc and decide to create WalkerControllerBase
  m_param.put(cur);
  reset();
  MyEstimator->setCollectionMode(true); //always collect
  return true;
}

void SimpleFixedNodeBranch::write(const std::string& fname, bool overwrite)
{
  RootName=fname;
  if(MyEstimator->is_manager())
  {
    //\since 2008-06-24
    vParam[B_ACC_ENERGY]=EnergyHist.result();
    vParam[B_ACC_SAMPLES]=EnergyHist.count();
    BranchIO hh(*this,MyEstimator->getCommunicator());
    bool success= hh.write(fname);
  }
}

#ifdef HAVE_ADIOS
void SimpleFixedNodeBranch::save_energy()
{
  if(MyEstimator->is_manager())
  {
    //\since 2008-06-24
    vParam[B_ACC_ENERGY]=EnergyHist.result();
    vParam[B_ACC_SAMPLES]=EnergyHist.count();
  }
}
#endif

void SimpleFixedNodeBranch::read(const std::string& fname)
{
  BranchMode.set(B_RESTART,0);
  if(fname.empty())
    return;
  vParam[B_ACC_ENERGY]=EnergyHist.result();
  vParam[B_ACC_SAMPLES]=EnergyHist.count();
  BranchIO hh(*this,MyEstimator->getCommunicator());
  BranchModeType bmode(BranchMode);
  bool success=hh.read(fname);
  if(success && R2Proposed.good() && bmode[B_POPCONTROL]==BranchMode[B_POPCONTROL])
  {
    BranchMode.set(B_RESTART,1);
    app_log() << "    Restarting, cummulative properties:"
              << "\n      energy     = " << EnergyHist.mean()
              << "\n      variance   = " << VarianceHist.mean()
              << "\n      r2accepted = " << R2Accepted.mean()
              << "\n      r2proposed = " << R2Proposed.mean()
              << std::endl;
  }
  else
  {
    if(BranchMode[B_POPCONTROL]!=bmode[B_POPCONTROL])
    {
      app_log() << "  Population control method has changed from " << BranchMode[B_POPCONTROL]
        << " to " << bmode[B_POPCONTROL] << std::endl;
      BranchMode[B_POPCONTROL]=bmode[B_POPCONTROL];
    }
  }

  //char fname2[128];
  //sprintf(fname2,"%s.p%03d.config",fname.c_str(),OHMMS::Controller->rank());
  //ofstream fout(fname2);
  //fout << "    Restarting, cummulative properties:"
  //          << "\n      energy     = " << EnergyHist.mean()
  //          << "\n      variance   = " << VarianceHist.mean()
  //          << "\n      r2accepted = " << R2Accepted.mean()
  //          << "\n      r2proposed = " << R2Proposed.mean()
  //          << std::endl;
  app_log().flush();
}

//   void SimpleFixedNodeBranch::storeConfigsForForwardWalking(MCWalkerConfiguration& w)
//   {
//     WalkerController->storeConfigsForForwardWalking(w);
//   }
//
//   void SimpleFixedNodeBranch::clearConfigsForForwardWalking( )
//   {
//     WalkerController->clearConfigsForForwardWalking( );
//   }
//
//   void SimpleFixedNodeBranch::debugFWconfig()
//   {
//     std::cout <<"FW size "<<WalkerController->sizeOfConfigsForForwardWalking()<< std::endl;
//     for(int i=0;i<WalkerController->ForwardWalkingHistory.size();i++) {
//       std::cout <<" Next Gen "<<i<< std::endl;
//       for(int j=0;j<WalkerController->ForwardWalkingHistory[i].size();j++)
//       {
//         std::cout <<j<<" "<<WalkerController->ForwardWalkingHistory[i][j].ID<<" "<<WalkerController->ForwardWalkingHistory[i][j].ParentID<< std::endl;
//       }
//     }
//   }

void SimpleFixedNodeBranch::setBranchCutoff(EstimatorRealType variance, EstimatorRealType targetSigma, EstimatorRealType maxSigma, int Nelec)
{
  if(branching_cutoff_scheme=="DRV")
  {
    // eq.(3), J. Chem. Phys. 89, 3629 (1988).
    // eq.(9), J. Chem. Phys. 99, 2865 (1993).
    vParam[B_BRANCHCUTOFF]=2.0/std::sqrt(vParam[B_TAU]);
  }
  else if(branching_cutoff_scheme=="ZSGMA")
  {
    // eq.(6), Phys. Rev. B 93, 241118(R) (2016)
    // do nothing if Nelec is not passed in.
    if(Nelec>0)
      vParam[B_BRANCHCUTOFF]=0.2*std::sqrt(Nelec/vParam[B_TAU]);
  }
  else if(branching_cutoff_scheme=="YL")
  {
    // a scheme from Ye Luo.
    vParam[B_BRANCHCUTOFF]=std::sqrt(variance)*std::min(targetSigma, std::sqrt(1.0/vParam[B_TAU]));
  }
  else if(branching_cutoff_scheme=="classic")
  {
    // default QMCPACK choice which is the same as v3.0.0 and before.
    vParam[B_BRANCHCUTOFF]=std::min(std::max(variance*targetSigma, maxSigma), 2.5/vParam[B_TAU]);
  }
  else
    APP_ABORT("SimpleFixedNodeBranch::setBranchCutoff unknown branching cutoff scheme "+branching_cutoff_scheme);

  vParam[B_BRANCHMAX]=vParam[B_BRANCHCUTOFF]*1.5;
  vParam[B_BRANCHFILTER]=1.0/(vParam[B_BRANCHMAX]-vParam[B_BRANCHCUTOFF]);
}

}


#pragma once

#include "multibrot_parallel_calculator.h"

#include <boost/compute.hpp>
#include <chrono>

#include "utils/utils.h"

template <typename P>
const Duration MultibrotParallelCalculator<P>::target_execution_time_ =
    Duration(std::chrono::seconds(1));

template <typename P>
MultibrotParallelCalculator<P>::MultibrotParallelCalculator(size_t width_pix, size_t height_pix)
    : partitioner_(width_pix, height_pix, fragment_width_pix_, fragment_height_pix_),
      width_pix_(width_pix),
      height_pix_(height_pix) {
    std::vector<boost::compute::device> devices;
    {
        std::vector<boost::compute::device> cpus, gpus;
        for (const auto& platform : boost::compute::system::platforms()) {
            if (use_cpus_) {
                Utils::AppendVectorToVector(cpus, platform.devices(CL_DEVICE_TYPE_CPU));
            }
            if (use_gpus_) {
                Utils::AppendVectorToVector(gpus, platform.devices(CL_DEVICE_TYPE_GPU));
            }
        }
        if (!cpus.empty())  // We have at least one CPU device
        {
            if (reserve_1_cpu_thread_) {
                // Split device into two - the first one contains all minus one compute units,
                // it is then added to devices list
                std::vector<boost::compute::device> split_cpu =
                    cpus[0].partition_by_counts({cpus[0].compute_units() - 1, 1});
                EXCEPTION_ASSERT(split_cpu.size() == 2);
                devices.push_back(split_cpu[0]);
            } else {
                devices.push_back(cpus[0]);  // Add first CPU device as is
            }
            // Add all other CPU devices
            devices.insert(devices.end(), cpus.begin() + 1, cpus.end());
        }
        // Add all GPUs
        Utils::AppendVectorToVector(devices, gpus);
    }

    for (auto& device : devices) {
        // TODO implement automatic resizing of memory buffers?
        device_states_.emplace(device, DeviceState(device));
    }
}

template <typename P>
std::complex<double> MultibrotParallelCalculator<P>::CalcComplexVal(
    std::complex<double> input_min, std::complex<double> input_max, size_t x, size_t y) {
    return input_min + std::complex<double>{
        (input_max.real() - input_min.real()) * x / width_pix_,
        (input_max.imag() - input_min.imag()) * y / height_pix_,
    };
}

template <typename P>
void MultibrotParallelCalculator<P>::Calculate(
    std::complex<double> input_min, std::complex<double> input_max, double power,
    int max_iterations, Callback cb) {
    partitioner_.Reset();

    CalculateFirstPhase(input_min, input_max, power, max_iterations, cb);

    // Process all remaining operations
    for (auto& device_state : device_states_) {
        if (device_state.second.prev_operation_info) {
            ProcessOperationResults(device_state, cb);
        }
    }
}

template <typename P>
void MultibrotParallelCalculator<P>::CalculateFirstPhase(
    std::complex<double> input_min, std::complex<double> input_max, double power,
    int max_iterations, Callback cb) {
    // TODO implement some more reliable check to prevent infinite loops?
    while (1) {
        for (auto& device_state : device_states_) {
            if (device_state.second.prev_operation_info) {
                if (device_state.second.prev_operation_info.value().copy_event.status() ==
                    CL_COMPLETE) {
                    ProcessOperationResults(device_state, cb);
                }
            } else {
                size_t fragment_count = starting_fragment_count_;
                if (device_state.second.prev_operations_duration_sum > Duration()) {
                    fragment_count =
                        (device_state.second.processed_pixels / fragment_size_pix_) *
                        (target_execution_time_ / device_state.second.prev_operations_duration_sum);
                    fragment_count =
                        std::min(fragment_count, max_segment_size_pix_ / fragment_size_pix_);
                }
                ImagePartitioner::Segment segment = partitioner_.Partition(fragment_count);
                if (segment.IsEmpty()) {
                    // No more work to assign, first phase is finished.
                    return;
                }

                device_state.second.prev_operation_info =
                    PrevOperationInfo{boost::compute::event(), boost::compute::event(), segment};
                std::complex<double> min =
                    CalcComplexVal(input_min, input_max, segment.x, segment.y);
                std::complex<double> max = CalcComplexVal(
                    input_min, input_max, segment.x + segment.width_pix,
                    segment.y + segment.height_pix);
                auto future = device_state.second.calculator.Calculate(
                    min, max, segment.width_pix, segment.height_pix, power, max_iterations,
                    device_state.second.output_vector.begin(),
                    &device_state.second.prev_operation_info.value().calc_event);
                device_state.second.prev_operation_info.value().copy_event = future.get_event();
            }
        }
    }
}

template <typename P>
void MultibrotParallelCalculator<P>::ProcessOperationResults(
    std::pair<const boost::compute::device, DeviceState>& device_info, Callback cb) {
    PrevOperationInfo& info = device_info.second.prev_operation_info.value();
    // Checking event status is not a synchronization point (OpenCL 2.2. Reference p.196),
    // so wait until it finishes completely
    info.copy_event.wait();

    cb(device_info.first, info.segment, device_info.second.output_vector.data());

    // Collect statistics for previous operation
    device_info.second.prev_operations_duration_sum +=
        Duration(info.calc_event) + Duration(info.copy_event);
    device_info.second.processed_pixels += info.segment.width_pix * info.segment.height_pix;
    device_info.second.prev_operation_info = boost::optional<PrevOperationInfo>();
}


/*!
* \file Particle.hpp
* \brief Particle structure
* \author Tycho Tatitscheff
* \version 1.0
* \date 16 février 2017
*
* Define a structure containing all field that a fluid particle may contain.
 */

#pragma once

#include <map>
#include <memory>
#include <vector>
#include <armadillo>

#include "utils/macros.hpp"
#include "kernel/Kernel.hpp"
#include "kernel/KernelPoly6.hpp"
#include "kernel/KernelSpiky.hpp"
#include "kernel/KernelViscosity.hpp"

#include "common/Fluid.hpp"
#include "common/Force.hpp"

using namespace arma;
using namespace std;

/**
 * \struct Particle
 * \brief Particle structure
 *
 * This structure contains all parameter carried by the fluid :
 *  - the position
 *  - the speed
 *  - the accelaration
 *  - the force applied
 *  - the mass
 *  - a pointer to the fluid
 */
class Particle {
  public:
    Particle(double r,
             shared_ptr<Fluid> flu,
             shared_ptr<listForces> f,
             shared_ptr<KernelPoly6> fKern,
             shared_ptr<KernelSpiky> pKern,
             shared_ptr<KernelViscosity> vKern,
             shared_ptr<KernelPoly6> sKern,
             int id,
             VEC pos,
             VEC spe,
             VEC acc);
    Particle(double r,
             shared_ptr<Fluid> flu,
             shared_ptr<listForces> f,
             shared_ptr<KernelPoly6> fKern,
             shared_ptr<KernelSpiky> pKern,
             shared_ptr<KernelViscosity> vKern,
             shared_ptr<KernelPoly6> sKern,
             int id);

    void updateField(vector<shared_ptr<Particle>> neighb);
    void updateForce(vector<shared_ptr<Particle>> neighb);

    shared_ptr<Fluid> flu; // réference vers un fluid
    shared_ptr<KernelPoly6> fieldKernel;
    shared_ptr<KernelSpiky> pressureKernel;
    shared_ptr<KernelPoly6> surfaceKernel;
    shared_ptr<KernelViscosity> viscosityKernel;

    // Particule caracteristic
    double rad; // rayon
    double mass; // masse de la particule

    // Pressure related field
    double density;
    double pressure;

    // Colour field
    double colour;
    double colourLaplacian;
    VEC colourDirection;

    // Position
    VEC curr_pos;
    VEC next_pos;

    // Speed
    // prev_half and next_half are for leapfrog
    VEC prev_half_spe;
    VEC curr_spe;
    VEC next_half_spe;
    VEC next_spe;

    // Accélération
    VEC curr_acc;
    VEC next_acc;

    // Force
    VEC result_force;
    shared_ptr<listForces> ext_forces;

    VEC pressure_force;
    VEC viscosity_force;
    VEC surfaceTension_force;

    int id;
};

inline bool operator<(const Particle& lhs, const Particle& rhs) {
    return lhs.id < rhs.id;
}

inline bool operator==(const Particle& lhs, const Particle& rhs) {
    return lhs.id == rhs.id;
}