// Copyright (C) 2014 The Regents of the University of California (Regents).
// 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 Regents or University of California 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 HOLDERS 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.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)

// BSD 3-Clause License

// Copyright (c) 2020, Chenyu
// 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 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.

#include "solver/qp_solver.h"

#include <glog/logging.h>

#include <Eigen/Core>
#include <algorithm>
#include <limits>
#include <string>

#include "math/sparse_cholesky_llt.h"
#include "util/stringprintf.h"

namespace DAGSfM {

QPSolver::QPSolver(const Options& options, const Eigen::SparseMatrix<double>& P,
                   const Eigen::VectorXd& q, const double r)
    : options_(options), P_(P), q_(q), r_(r) {
  CHECK_EQ(P_.rows(), P_.cols()) << "P must be a symmetric matrix.";
  CHECK_EQ(P_.cols(), q_.size())
      << "The dimensions of P and q must be consistent.";

  // Set the lower and upper bounds to be negative and positive infinity (i.e
  // no bounds).
  lb_.setConstant(P_.cols(), -std::numeric_limits<double>::infinity());
  ub_.setConstant(P_.cols(), std::numeric_limits<double>::infinity());

  // Set up the linear solver to compute the cholesky decomposition of:
  //     P_ + rho * eye(N)
  Eigen::SparseMatrix<double> spd_mat(P_.rows(), P_.cols());
  spd_mat.setIdentity();
  spd_mat *= options_.rho;
  spd_mat += P_;

  linear_solver_.Compute(spd_mat);
  CHECK_EQ(linear_solver_.Info(), Eigen::Success);
}

void QPSolver::SetMaxIterations(const int max_iterations) {
  options_.max_num_iterations = max_iterations;
}

void QPSolver::SetUpperBound(const Eigen::VectorXd& ub) { ub_ = ub.array(); }

void QPSolver::SetLowerBound(const Eigen::VectorXd& lb) { lb_ = lb.array(); }

// Solve the quadratic program.
bool QPSolver::Solve(Eigen::VectorXd* solution) {
  // Ensure the bounds are valid. If there are any invalid bounds then the
  // difference between the bounds would be a negative value.
  int coeff_index = -1;
  if ((ub_ - lb_).minCoeff(&coeff_index) < 0) {
    LOG(WARNING) << "You specified invalid lower or upper bounds for the "
                    "problem. lower_bound["
                 << coeff_index << "] = " << lb_[coeff_index]
                 << " but upper_bound[" << coeff_index
                 << "] = " << ub_[coeff_index];
    return false;
  }

  CHECK_NOTNULL(solution)->setZero(q_.size());
  Eigen::VectorXd& x = *solution;
  Eigen::VectorXd z(P_.rows()), u(P_.rows());
  z.setZero();
  u.setZero();

  Eigen::VectorXd z_old(z.size()), x_hat(P_.rows());

  // Precompute some convergence terms.
  const double primal_abs_tolerance_eps =
      std::sqrt(P_.rows()) * options_.absolute_tolerance;
  const double dual_abs_tolerance_eps =
      std::sqrt(P_.rows()) * options_.absolute_tolerance;
  VLOG(2) << "Iteration   Residual         R norm          S norm          "
             "Primal eps      Dual eps";
  const std::string row_format =
      "  % 4d      % 4.4e     % 4.4e     % 4.4e     % 4.4e     % 4.4e";

  // Run the iterations.
  for (int i = 0; i < options_.max_num_iterations; i++) {
    // Update x.
    x.noalias() = linear_solver_.Solve(options_.rho * (z - u) - q_);
    if (linear_solver_.Info() != Eigen::Success) {
      return false;
    }

    // Update x_hat.
    std::swap(z, z_old);
    x_hat.noalias() = options_.alpha * x + (1.0 - options_.alpha) * z_old;

    // Update z.
    z.noalias() = ub_.min(lb_.max((x_hat + u).array())).matrix();

    // Update u.
    u.noalias() += x_hat - z;

    // Compute the convergence terms.
    const double objval = 0.5 * x.dot(P_ * x) + q_.dot(x) + r_;
    const double r_norm = (x - z).norm();
    const double s_norm = (-options_.rho * (z - z_old)).norm();
    const double max_norm = std::max({x.norm(), z.norm()});
    const double primal_eps =
        primal_abs_tolerance_eps + options_.relative_tolerance * max_norm;
    const double dual_eps =
        dual_abs_tolerance_eps +
        options_.relative_tolerance * (options_.rho * u).norm();

    // Log the result to the screen.
    VLOG(2) << StringPrintf(row_format.c_str(), objval, i, r_norm, s_norm,
                            primal_eps, dual_eps);
    // Determine if the minimizer has converged.
    if (r_norm < primal_eps && s_norm < dual_eps) {
      break;
    }
  }
  return true;
}

}  // namespace DAGSfM


#ifndef OFFLINECACLAAG_HPP
#define OFFLINECACLAAG_HPP

#include <vector>
#include <string>
#include <boost/serialization/list.hpp>
#include <boost/serialization/set.hpp>
#include <boost/serialization/vector.hpp>

#include "arch/AACAgent.hpp"
#include "bib/Seed.hpp"
#include "bib/Utils.hpp"
#include "bib/OrnsteinUhlenbeckNoise.hpp"
#include <bib/MetropolisHasting.hpp>
#include <bib/XMLEngine.hpp>
#include "nn/MLP.hpp"
#include "nn/DODevMLP.hpp"

typedef struct _sample {
  std::vector<double> s;
  std::vector<double> pure_a;
  std::vector<double> a;
  std::vector<double> next_s;
  double r;
  bool goal_reached;

  friend class boost::serialization::access;
  template <typename Archive>
  void serialize(Archive& ar, const unsigned int) {
    ar& BOOST_SERIALIZATION_NVP(s);
    ar& BOOST_SERIALIZATION_NVP(pure_a);
    ar& BOOST_SERIALIZATION_NVP(a);
    ar& BOOST_SERIALIZATION_NVP(next_s);
    ar& BOOST_SERIALIZATION_NVP(r);
    ar& BOOST_SERIALIZATION_NVP(goal_reached);
  }

  bool operator< (const _sample& b) const {
    for (uint i = 0; i < s.size(); i++) {
      if(s[i] != b.s[i])
        return s[i] < b.s[i];
    }

    for (uint i = 0; i < a.size(); i++) {
      if(a[i] != b.a[i])
        return a[i] < b.a[i];
    }

    return false;
  }

} sample;

template<typename NN = MLP>
class OfflineCaclaAg : public arch::AACAgent<NN, arch::AgentProgOptions> {
 public:
  typedef NN PolicyImpl;

  OfflineCaclaAg(unsigned int _nb_motors, unsigned int _nb_sensors)
    : arch::AACAgent<NN, arch::AgentProgOptions>(_nb_motors, _nb_sensors), nb_sensors(_nb_sensors), empty_action(0) {

  }

  virtual ~OfflineCaclaAg() {
    delete vnn;
    delete ann;

    delete ann_testing;
    if(batch_norm_critic != 0)
      delete vnn_testing;

    delete hidden_unit_v;
    delete hidden_unit_a;

    if(oun == nullptr)
      delete oun;
  }

  const std::vector<double>& _run(double reward, const std::vector<double>& sensors,
                                  bool learning, bool goal_reached, bool) override {

    // protect batch norm from testing data and poor data
    vector<double>* next_action = ann_testing->computeOut(sensors);
    if (last_action.get() != nullptr && learning)
      trajectory.push_back( {last_state, *last_pure_action, *last_action, sensors, reward, goal_reached});

    last_pure_action.reset(new vector<double>(*next_action));
    if(learning) {
      if(gaussian_policy == 1) {
        vector<double>* randomized_action = bib::Proba<double>::multidimentionnalTruncatedGaussian(*next_action, noise);
        delete next_action;
        next_action = randomized_action;
      } else if(gaussian_policy == 2) {
        oun->step(*next_action);
      } else if(gaussian_policy == 3 && bib::Utils::rand01() < noise2) {
        vector<double>* randomized_action = bib::Proba<double>::multidimentionnalTruncatedGaussian(*next_action, noise);
        delete next_action;
        next_action = randomized_action;
      } else if(gaussian_policy == 4) {
        vector<double>* randomized_action = bib::Proba<double>::multidimentionnalTruncatedGaussian(*next_action, noise * pow(noise2, noise3 - ((double) step)));
        delete next_action;
        next_action = randomized_action;
      } else if(bib::Utils::rand01() < noise) { //e-greedy
        for (uint i = 0; i < next_action->size(); i++)
          next_action->at(i) = bib::Utils::randin(-1.f, 1.f);
      }
    }
    last_action.reset(next_action);

    ann_testing->neutral_action(sensors, next_action);

    last_state.clear();
    for (uint i = 0; i < sensors.size(); i++)
      last_state.push_back(sensors[i]);

    step++;

    return *next_action;
  }

  void _unique_invoke(boost::property_tree::ptree* pt, boost::program_options::variables_map*) override {
//     bib::Seed::setFixedSeedUTest();
    hidden_unit_v           = bib::to_array<uint>(pt->get<std::string>("agent.hidden_unit_v"));
    hidden_unit_a           = bib::to_array<uint>(pt->get<std::string>("agent.hidden_unit_a"));
    noise                   = pt->get<double>("agent.noise");
    gaussian_policy         = pt->get<uint>("agent.gaussian_policy");
    update_delta_neg        = pt->get<bool>("agent.update_delta_neg");
    vnn_from_scratch        = pt->get<bool>("agent.vnn_from_scratch");
    update_critic_first     = pt->get<bool>("agent.update_critic_first");
    number_fitted_iteration = pt->get<uint>("agent.number_fitted_iteration");
    stoch_iter_actor        = pt->get<uint>("agent.stoch_iter_actor");
    stoch_iter_critic       = pt->get<uint>("agent.stoch_iter_critic");
    batch_norm_actor        = pt->get<uint>("agent.batch_norm_actor");
    batch_norm_critic       = pt->get<uint>("agent.batch_norm_critic");
    actor_output_layer_type = pt->get<uint>("agent.actor_output_layer_type");
    hidden_layer_type       = pt->get<uint>("agent.hidden_layer_type");
    alpha_a                 = pt->get<double>("agent.alpha_a");
    alpha_v                 = pt->get<double>("agent.alpha_v");
    lambda                  = pt->get<double>("agent.lambda");
    momentum                = pt->get<uint>("agent.momentum");
    corrected_update_ac     = false;
    gae                     = false;
    inverting_gradient      = false;
    update_each_episode = 1;

    if(gaussian_policy == 2){
      double oun_theta = pt->get<double>("agent.noise2");
      double oun_dt = pt->get<double>("agent.noise3");
      oun = new bib::OrnsteinUhlenbeckNoise<double>(this->nb_motors, noise, oun_theta, oun_dt);
    } else if (gaussian_policy == 3){
      noise2 = pt->get<double>("agent.noise2");
    } else if (gaussian_policy == 4){
      noise2 = pt->get<double>("agent.noise2");
      noise3 = pt->get<double>("agent.noise3");
    }

    try {
      update_each_episode     = pt->get<uint>("agent.update_each_episode");
    } catch(boost::exception const& ) {
    }

    try {
      corrected_update_ac   = pt->get<bool>("agent.corrected_update_ac");
    } catch(boost::exception const& ) {
    }
    try {
      inverting_gradient   = pt->get<bool>("agent.inverting_gradient");
    } catch(boost::exception const& ) {
    }
    if(corrected_update_ac){
      try {
        corrected_update_ac_factor   = pt->get<double>("agent.corrected_update_ac_factor");
      } catch(boost::exception const& ) {
      }
    }

    if(lambda >= 0.)
      gae = pt->get<bool>("agent.gae");

    if(lambda >=0. && batch_norm_critic != 0){
      LOG_DEBUG("to be done!");
      exit(1);
    }

    ann = new NN(nb_sensors, *hidden_unit_a, this->nb_motors, alpha_a, 1, hidden_layer_type, actor_output_layer_type, batch_norm_actor, true, momentum);
    if(std::is_same<NN, DODevMLP>::value)
      ann->exploit(pt, nullptr);

    vnn = new NN(nb_sensors, nb_sensors, *hidden_unit_v, alpha_v, 1, -1, hidden_layer_type, batch_norm_critic, false, momentum);
    if(std::is_same<NN, DODevMLP>::value)
      vnn->exploit(pt, ann);

    ann_testing = new NN(*ann, false, ::caffe::Phase::TEST);
    if(batch_norm_critic != 0)
      vnn_testing = new NN(*vnn, false, ::caffe::Phase::TEST);

    if(std::is_same<NN, DODevMLP>::value){
      try {
        if(pt->get<bool>("devnn.reset_learning_algo")){
          LOG_ERROR("NFAC cannot reset anything with DODevMLP");
          exit(1);
        }
      } catch(boost::exception const& ) {
      }
    }

    bestever_score = std::numeric_limits<double>::lowest();
  }

  void _start_episode(const std::vector<double>& sensors, bool learning) override {
    last_state.clear();
    for (uint i = 0; i < sensors.size(); i++)
      last_state.push_back(sensors[i]);

    last_action = nullptr;
    last_pure_action = nullptr;

    step = 0;
    if(gaussian_policy == 2)
      oun->reset();

    ratio_valid_advantage = -1;

    if(std::is_same<NN, DODevMLP>::value && learning){
      DODevMLP * ann_cast = static_cast<DODevMLP *>(ann);
      bool changed_ann = std::get<1>(ann_cast->inform(episode, this->last_sum_weighted_reward));
//       don't need to inform vnn because they share parameters
//       static_cast<DODevMLP *>(vnn)->inform(episode, this->last_sum_weighted_reward);
      if(changed_ann && ann_cast->ewc_enabled() && ann_cast->ewc_force_constraint()){
        static_cast<DODevMLP *>(vnn)->ewc_setup();
//         else if(changed_vnn && !changed_ann) //impossible cause of ann structure
      }
    }

    double* weights = new double[ann->number_of_parameters(false)];
    ann->copyWeightsTo(weights, false);
    ann_testing->copyWeightsFrom(weights, false);
    delete[] weights;
  }

  double update_critic() {
    double V_pi_s0 = 0.f;
    if (trajectory.size() > 0) {
      //remove trace of old policy
      auto iter = [&]() {
        std::vector<double> all_states(trajectory.size() * nb_sensors);
        std::vector<double> all_next_states(trajectory.size() * nb_sensors);
        std::vector<double> v_target(trajectory.size());
        int li=0;
        for (auto it : trajectory) {
          std::copy(it.s.begin(), it.s.end(), all_states.begin() + li * nb_sensors);
          std::copy(it.next_s.begin(), it.next_s.end(), all_next_states.begin() + li * nb_sensors);
          li++;
        }

        decltype(vnn_testing->computeOutVFBatch(all_next_states, empty_action)) all_nextV;
        if(batch_norm_critic != 0)
        {
          double* weights = new double[vnn->number_of_parameters(false)];
          vnn->copyWeightsTo(weights, false);
          vnn_testing->copyWeightsFrom(weights, false);
          delete[] weights;
          all_nextV = vnn_testing->computeOutVFBatch(all_next_states, empty_action);
        } else
          all_nextV = vnn->computeOutVFBatch(all_next_states, empty_action);

        li=0;
        for (auto it : trajectory) {
          double target = it.r;
          if (!it.goal_reached) {
            double nextV = all_nextV->at(li);
            target += this->gamma * nextV;
          }

          v_target[li] = target;
          li++;
        }

        ASSERT((uint)li == trajectory.size(), "");
        if(vnn_from_scratch){
          delete vnn;
          vnn = new NN(nb_sensors, nb_sensors, *hidden_unit_v, alpha_v, trajectory.size(), -1, hidden_layer_type, batch_norm_critic, false, momentum);
        }
        if(lambda < 0.f && batch_norm_critic == 0)
          vnn->learn_batch(all_states, empty_action, v_target, stoch_iter_critic);
        else if(lambda < 0.f){
          for(uint sia = 0; sia < stoch_iter_critic; sia++){
            delete vnn->computeOutVFBatch(all_states, empty_action);
            {
              double* weights = new double[vnn->number_of_parameters(false)];
              vnn->copyWeightsTo(weights, false);
              vnn_testing->copyWeightsFrom(weights, false);
              delete[] weights;
            }
            auto all_V = vnn_testing->computeOutVFBatch(all_states, empty_action);

            const auto q_values_blob = vnn->getNN()->blob_by_name(MLP::q_values_blob_name);
            double* q_values_diff = q_values_blob->mutable_cpu_diff();
            uint i=0;
            double s = trajectory.size();
            for (auto it : trajectory){
              q_values_diff[i] = (all_V->at(i)-v_target[i])/s;
              i++;
            }
            vnn->critic_backward();
            vnn->updateFisher(trajectory.size());
            vnn->regularize();
            vnn->getSolver()->ApplyUpdate();
            vnn->getSolver()->set_iter(vnn->getSolver()->iter() + 1);
            delete all_V;
          }
        }
        else {
          auto all_V = vnn->computeOutVFBatch(all_states, empty_action);
          std::vector<double> deltas(trajectory.size());
//
//        Simple computation for lambda return
//
//          bib::Logger::PRINT_ELEMENTS(*all_V, "V0 ");
          li=0;
          for (auto it : trajectory){
            deltas[li] = v_target[li] - all_V->at(li);
            ++li;
          }
//          bib::Logger::PRINT_ELEMENTS(deltas, "deltas ");

          std::vector<double> diff(trajectory.size());
          li=trajectory.size() - 1;
          double prev_delta = 0.;
          int index_ep = trajectory_end_points.size() - 1;
          for (auto it : trajectory) {
            if (index_ep >= 0 && trajectory_end_points[index_ep] - 1 == li){
                prev_delta = 0.;
                index_ep--;
            }
            diff[li] = deltas[li] + prev_delta;
            prev_delta = this->gamma * lambda * diff[li];
            --li;
          }
          ASSERT(diff[trajectory.size() -1] == deltas[trajectory.size() -1], "pb lambda");

// //           comment following lines to compare with the other formula
          li=0;
          for (auto it : trajectory){
            diff[li] = diff[li] + all_V->at(li);
            ++li;
          }

//          bib::Logger::PRINT_ELEMENTS(diff, "target ");
          V_pi_s0 = diff[0];
          vnn->learn_batch(all_states, empty_action, diff, stoch_iter_critic);
// //
// //        The mechanic formula
// //
//           std::vector<double> diff2(trajectory.size());
//           li=0;
//           for (auto it : trajectory){
//             diff2[li] = 0;
//             double sum_n = 0.f;
//             for(int n=1;n<=((int)trajectory.size()) - li - 1;n++){
//               double sum_i = 0.f;
//               for(int i=li;i<=li+n-1;i++)
//                 sum_i += std::pow(this->gamma, i-li) * trajectory[i].r;
//               sum_i += std::pow(this->gamma, n) * all_nextV->at(li+n-1);
//               sum_i *= pow(lambda, n-1);
//               sum_n += sum_i;
//             }
//             sum_n *= (1.f-lambda);
//
//             double sum_L = 0.f;
//             for(int i=li;i<(int)trajectory.size();i++)
//               sum_L += std::pow(this->gamma, i-li) * trajectory[i].r;
//             if(trajectory[trajectory.size()-1].goal_reached)
//               sum_n += std::pow(lambda, trajectory.size() - li - 1) * sum_L;
//             else {
//               sum_L += std::pow(this->gamma, ((int)trajectory.size()) - li) * all_nextV->at(trajectory.size()-1);
//               sum_n += std::pow(lambda, trajectory.size() - li - 1) * sum_L;
//             }
//
//             sum_n -= all_V->at(li);
//
//             diff2[li] = sum_n;
//             ++li;
//           }
//           bib::Logger::PRINT_ELEMENTS(diff2, "mech form ");
//
//           if(trajectory[trajectory.size()-1].goal_reached)
//             exit(1);
          delete all_V;
        }

        delete all_nextV;
      };

      for(uint i=0; i<number_fitted_iteration; i++)
        iter();
    }

    return V_pi_s0;
  }

  void end_episode(bool learning) override {
//     LOG_FILE("policy_exploration", ann->hash());
    if(!learning){
      if(ann->ewc_best_method() >= 4){
        ann->update_best_param_previous_task(this->sum_weighted_reward);
        vnn->update_best_param_previous_task(this->sum_weighted_reward);
      }
      return;
    }

    //learning phase
    if(ann->ewc_best_method() <= 3){
      ann->update_best_param_previous_task(this->sum_weighted_reward);
      vnn->update_best_param_previous_task(this->sum_weighted_reward);
    }

    trajectory_end_points.push_back(trajectory.size());
    if (episode % update_each_episode != 0)
      return;

    if(trajectory.size() > 0){
      vnn->increase_batchsize(trajectory.size());
      if(batch_norm_critic != 0)
        vnn_testing->increase_batchsize(trajectory.size());
    }

    double V_pi_s0 = 0;
    if(update_critic_first)
      V_pi_s0 = update_critic();

    if (trajectory.size() > 0) {
      std::vector<double> sensors(trajectory.size() * nb_sensors);
      std::vector<double> actions(trajectory.size() * this->nb_motors);
      std::vector<bool> disable_back(trajectory.size() * this->nb_motors, false);
      const std::vector<bool> disable_back_ac(this->nb_motors, true);
//       std::vector<double> deltas_blob(trajectory.size() * this->nb_motors);
      std::vector<double> deltas(trajectory.size());

      std::vector<double> all_states(trajectory.size() * nb_sensors);
      std::vector<double> all_next_states(trajectory.size() * nb_sensors);
      int li=0;
      for (auto it : trajectory) {
        std::copy(it.s.begin(), it.s.end(), all_states.begin() + li * nb_sensors);
        std::copy(it.next_s.begin(), it.next_s.end(), all_next_states.begin() + li * nb_sensors);
        li++;
      }

      decltype(vnn->computeOutVFBatch(all_next_states, empty_action)) all_nextV, all_mine;
      if(batch_norm_critic != 0)
      {
        double* weights = new double[vnn->number_of_parameters(false)];
        vnn->copyWeightsTo(weights, false);
        vnn_testing->copyWeightsFrom(weights, false);
        delete[] weights;
        all_nextV = vnn_testing->computeOutVFBatch(all_next_states, empty_action);
        all_mine = vnn_testing->computeOutVFBatch(all_states, empty_action);
      } else {
        all_nextV = vnn->computeOutVFBatch(all_next_states, empty_action);
        all_mine = vnn->computeOutVFBatch(all_states, empty_action);
      }

      li=0;
      for (auto it : trajectory){
        sample sm = it;
        double v_target = sm.r;
        if (!sm.goal_reached) {
          double nextV = all_nextV->at(li);
          v_target += this->gamma * nextV;
        }

        deltas[li] = v_target - all_mine->at(li);
        ++li;
      }

      if(gae){
        //
        //        Simple computation for lambda return
        //
        std::vector<double> diff(trajectory.size());
        li=trajectory.size() - 1;
        double prev_delta = 0.;
        int index_ep = trajectory_end_points.size() - 1;
        for (auto it : trajectory) {
          if (index_ep >= 0 && trajectory_end_points[index_ep] - 1 == li){
              prev_delta = 0.;
              index_ep--;
          }
          diff[li] = deltas[li] + prev_delta;
          prev_delta = this->gamma * lambda * diff[li];
          --li;
        }
        ASSERT(diff[trajectory.size() -1] == deltas[trajectory.size() -1], "pb lambda");

        li=0;
        for (auto it : trajectory){
//           diff[li] = diff[li] + all_V->at(li);
          deltas[li] = diff[li];
          ++li;
        }
      }

      uint n=0;
      li=0;
//       double dmax = std::max((double)0.f, *std::max_element(deltas.begin(), deltas.end()));
      for(auto it = trajectory.begin(); it != trajectory.end() ; ++it) {
        sample sm = *it;

        std::copy(it->s.begin(), it->s.end(), sensors.begin() + li * nb_sensors);
        if(deltas[li] > 0.) {
          std::copy(it->a.begin(), it->a.end(), actions.begin() + li * this->nb_motors);
          n++;
        } else if(update_delta_neg) {
          std::copy(it->pure_a.begin(), it->pure_a.end(), actions.begin() + li * this->nb_motors);
        } else {
          std::copy(it->a.begin(), it->a.end(), actions.begin() + li * this->nb_motors);
          std::copy(disable_back_ac.begin(), disable_back_ac.end(), disable_back.begin() + li * this->nb_motors);
        }
//         std::fill(deltas_blob.begin() + li * this->nb_motors, deltas_blob.begin() + (li+1) * this->nb_motors, deltas[li]);
        li++;
      }

      ratio_valid_advantage = ((float)n) / ((float) trajectory.size());

      if(n > 0) {
        for(uint sia = 0; sia < stoch_iter_actor; sia++){
          ann->increase_batchsize(trajectory.size());
          //learn BN
          auto ac_out = ann->computeOutBatch(sensors);
          if(batch_norm_actor != 0) {
            //re-compute ac_out with BN as testing
            double* weights = new double[ann->number_of_parameters(false)];
            ann->copyWeightsTo(weights, false);
            ann_testing->copyWeightsFrom(weights, false);
            delete[] weights;
            delete ac_out;
            ann_testing->increase_batchsize(trajectory.size());
            ac_out = ann_testing->computeOutBatch(sensors);
          }
          ann->ZeroGradParameters();

          const auto actor_actions_blob = ann->getNN()->blob_by_name(MLP::actions_blob_name);
          auto ac_diff = actor_actions_blob->mutable_cpu_diff();
          for(int i=0; i<actor_actions_blob->count(); i++){
            if(disable_back[i]){
              ac_diff[i] = 0.00000000f;
            } else {
              double x = actions[i] - ac_out->at(i);
              if(!corrected_update_ac){
                ac_diff[i] = -x;
                if(inverting_gradient){
                  const double min_ = -1.0;
                  const double max_ = 1.0;

                  if (ac_diff[i] < 0)
                    ac_diff[i] *= (max_ - ac_out->at(i)) / (max_ - min_);
                  else if (ac_diff[i] > 0)
                    ac_diff[i] *= (ac_out->at(i) - min_) / (max_ - min_);
                }
              } else {
                ac_diff[i] = -x * corrected_update_ac_factor;
              }
            }
          }
          ann->actor_backward();
          ann->updateFisher(n);
          ann->regularize();
          ann->getSolver()->ApplyUpdate();
          ann->getSolver()->set_iter(ann->getSolver()->iter() + 1);
          delete ac_out;
        }
      } else if(batch_norm_actor != 0){
        ann->increase_batchsize(trajectory.size());
        delete ann->computeOutBatch(sensors);
      }

      delete all_nextV;
      delete all_mine;

      if(batch_norm_actor != 0)
        ann_testing->increase_batchsize(1);
    }

    if(!update_critic_first){
      update_critic();
    }

    nb_sample_update= trajectory.size();
    trajectory.clear();
    trajectory_end_points.clear();

    ann->ewc_decay_update();
    vnn->ewc_decay_update();
  }

  void end_instance(bool learning) override {
    if(learning)
      episode++;
  }

  void save(const std::string& path, bool savebest, bool learning) override {
    if(savebest) {
      if(!learning && this->sum_weighted_reward >= bestever_score) {
        bestever_score = this->sum_weighted_reward;
        ann->save(path+".actor");
      }
    } else {
      ann->save(path+".actor");
      vnn->save(path+".critic");
    }
  }

  void save_run() override {
    ann->save("continue.actor");
    vnn->save("continue.critic");
    struct algo_state st = {episode};
    bib::XMLEngine::save(st, "algo_state", "continue.algo_state.data");
  }

  void load(const std::string& path) override {
    ann->load(path+".actor");
    vnn->load(path+".critic");
  }

  void load_previous_run() override {
    ann->load("continue.actor");
    vnn->load("continue.critic");
    auto p3 = bib::XMLEngine::load<struct algo_state>("algo_state", "continue.algo_state.data");
    episode = p3->episode;
    delete p3;
  }

  double criticEval(const std::vector<double>&, const std::vector<double>&) override {
    LOG_INFO("not implemented");
    return 0;
  }

  arch::Policy<NN>* getCopyCurrentPolicy() override {
//         return new arch::Policy<MLP>(new MLP(*ann) , gaussian_policy ? arch::policy_type::GAUSSIAN : arch::policy_type::GREEDY, noise, decision_each);
    return nullptr;
  }

 protected:
  void _display(std::ostream& out) const override {
    out << std::setw(12) << std::fixed << std::setprecision(10) << this->sum_weighted_reward/this->gamma << " " << this->sum_reward <<
        " " << std::setw(8) << std::fixed << std::setprecision(5) << vnn->error() << " " << noise << " " << nb_sample_update <<
          " " << std::setprecision(3) << ratio_valid_advantage << " " << vnn->weight_l1_norm() << " " << ann->weight_l1_norm();
  }

  void _dump(std::ostream& out) const override {
    out << std::setw(25) << std::fixed << std::setprecision(22) <<
    this->sum_weighted_reward/this->gamma << " " << this->sum_reward << " " << std::setw(8) << std::fixed <<
        std::setprecision(5) << vnn->error() << " " << nb_sample_update <<
        " " << std::setprecision(3) << ratio_valid_advantage ;
  }

  double sign(double x){
    if(x>=0)
      return 1.f;
    return -1.f;
  }

 private:
  uint nb_sensors;
  uint episode = 1;
  uint step = 0;

  double noise, noise2, noise3;
  uint gaussian_policy;
  bool vnn_from_scratch, update_critic_first,
        update_delta_neg, corrected_update_ac, gae;
  bool inverting_gradient;
  uint number_fitted_iteration, stoch_iter_actor, stoch_iter_critic;
  uint batch_norm_actor, batch_norm_critic, actor_output_layer_type, hidden_layer_type, momentum;
  double lambda, corrected_update_ac_factor;

  std::shared_ptr<std::vector<double>> last_action;
  std::shared_ptr<std::vector<double>> last_pure_action;
  std::vector<double> last_state;
  double alpha_v, alpha_a;

  std::deque<sample> trajectory;
  std::deque<int> trajectory_end_points;

  NN* ann;
  NN* vnn;
  NN* ann_testing;
  NN* vnn_testing;

  std::vector<uint>* hidden_unit_v;
  std::vector<uint>* hidden_unit_a;
  std::vector<double> empty_action; //dummy action cause c++ cannot accept null reference
  double bestever_score;
  int update_each_episode;
  bib::OrnsteinUhlenbeckNoise<double>* oun = nullptr;
  float ratio_valid_advantage=0;
  int nb_sample_update = 0;

  struct algo_state {
    uint episode;

    friend class boost::serialization::access;
    template <typename Archive>
    void serialize(Archive& ar, const unsigned int) {
      ar& BOOST_SERIALIZATION_NVP(episode);
    }
  };
};

#endif


//
// SPDX-License-Identifier: BSD-3-Clause
// Copyright Contributors to the OpenEXR Project.
//

// clang-format off

#include <Python.h>
#include <boost/python.hpp>
#include "PyImath.h"
#include "PyImathExport.h"
#include "PyImathBasicTypes.h"
#include "PyImathFixedArray.h"
#include "PyImathFixedVArray.h"
#include "PyImathBufferProtocol.h"

using namespace boost::python;

namespace PyImath {

void
register_basicTypes()
{
    class_<BoolArray> bclass = BoolArray::register_("Fixed length array of bool");
    add_comparison_functions(bclass);

    class_<SignedCharArray> scclass = SignedCharArray::register_("Fixed length array of signed chars");
    add_arithmetic_math_functions(scclass);
    add_mod_math_functions(scclass);
    add_comparison_functions(scclass);
    add_ordered_comparison_functions(scclass);

    class_<UnsignedCharArray> ucclass = UnsignedCharArray::register_("Fixed length array of unsigned chars");
    add_arithmetic_math_functions(ucclass);
    add_mod_math_functions(ucclass);
    add_comparison_functions(ucclass);
    add_ordered_comparison_functions(ucclass);
    add_buffer_protocol<UnsignedCharArray>(ucclass);

    class_<ShortArray> sclass = ShortArray::register_("Fixed length array of shorts");
    add_arithmetic_math_functions(sclass);
    add_mod_math_functions(sclass);
    add_comparison_functions(sclass);
    add_ordered_comparison_functions(sclass);

    class_<UnsignedShortArray> usclass = UnsignedShortArray::register_("Fixed length array of unsigned shorts");
    add_arithmetic_math_functions(usclass);
    add_mod_math_functions(usclass);
    add_comparison_functions(usclass);
    add_ordered_comparison_functions(usclass);

    class_<IntArray> iclass = IntArray::register_("Fixed length array of ints");
    add_arithmetic_math_functions(iclass);
    add_mod_math_functions(iclass);
    add_comparison_functions(iclass);
    add_ordered_comparison_functions(iclass);
    add_explicit_construction_from_type<float>(iclass);
    add_explicit_construction_from_type<double>(iclass);
    add_buffer_protocol<IntArray>(iclass);

    class_<UnsignedIntArray> uiclass = UnsignedIntArray::register_("Fixed length array of unsigned ints");
    add_arithmetic_math_functions(uiclass);
    add_mod_math_functions(uiclass);
    add_comparison_functions(uiclass);
    add_ordered_comparison_functions(uiclass);
    add_explicit_construction_from_type<float>(uiclass);
    add_explicit_construction_from_type<double>(uiclass);

    class_<FloatArray> fclass = FloatArray::register_("Fixed length array of floats");
    add_arithmetic_math_functions(fclass);
    add_pow_math_functions(fclass);
    add_comparison_functions(fclass);
    add_ordered_comparison_functions(fclass);
    add_explicit_construction_from_type<int>(fclass);
    add_explicit_construction_from_type<double>(fclass);
    add_buffer_protocol<FloatArray>(fclass);

    class_<DoubleArray> dclass = DoubleArray::register_("Fixed length array of doubles");
    add_arithmetic_math_functions(dclass);
    add_pow_math_functions(dclass);
    add_comparison_functions(dclass);
    add_ordered_comparison_functions(dclass);
    add_explicit_construction_from_type<int>(dclass);
    add_explicit_construction_from_type<float>(dclass);
    add_buffer_protocol<DoubleArray>(dclass);

    class_<VIntArray>   ivclass = VIntArray::register_("Variable fixed length array of ints");
    class_<VFloatArray> fvclass = VFloatArray::register_("Variable fixed length array of floats");
    class_<VV2iArray> v2ivclass = VV2iArray::register_("Variable fixed length array of V2i");
    class_<VV2fArray> v2fvclass = VV2fArray::register_("Variable fixed length array of V2f");
    // Don't add other functionality until its defined better.
}

} // namespace PyImath


/*******************************************************************
* Copyright (c) 2010-2013 MiGraNT - DTAI.cs.kuleuven.be
* License details in the root directory of this distribution in
* the LICENSE file
********************************************************************/

#include "treeDecompositionDAG.hpp"
#include <iostream>
#include <stack>
#include <queue>
#include <set>
#include <algorithm>
#include <fstream>
#include <iostream>

#include <boost/graph/topological_sort.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/lexical_cast.hpp>
#include "treeDecompositionDAGAndNode.hpp"

#include <boost/graph/graphviz.hpp>

namespace treeDAG {

TreeDecompositionDAG::AndNodeTransformer::AndNodeTransformer()
    : dag_(0)
{
}

TreeDecompositionDAG::AndNodeTransformer::AndNodeTransformer(bool isSource, const TreeDecompositionDAG * dag)
    : dag_(dag),
      isSource_(isSource)
{
}

TreeDecompositionDAGAndNode TreeDecompositionDAG::AndNodeTransformer::operator()(boost::graph_traits<Graph>::edge_descriptor edge) const
{
    return TreeDecompositionDAGAndNode(dag_, isSource_ ? boost::source(edge, dag_->structure_) : boost::target(edge, dag_->structure_) );
}

TreeDecompositionDAG::ComponentTransformer::ComponentTransformer(const Graph * graph)
    : graph_(graph)
{
}

const TreeDecompositionSubgraph & TreeDecompositionDAG::ComponentTransformer::operator()(boost::graph_traits<Graph>::vertex_descriptor vd) const
{
    return graph_->operator [](vd).vertices;
}

TreeDecompositionDAG::ComponentFilter::ComponentFilter(const Graph * graph)
    : graph_(graph)
{
}

bool TreeDecompositionDAG::ComponentFilter::operator()(const TreeDecompositionSubgraph & component) const
{
    return component.size() != 0;
}

TreeDecompositionDAG::TreeDecompositionDAG(std::size_t patternSize)
    : patternSize_(patternSize)
{

}

TreeDecompositionDAG::TreeDecompositionDAG(const TreeDecompositionDAG & rhs)
    : structure_(rhs.structure_),
      patternSize_(rhs.patternSize_)
{
    refillMap();
}

void TreeDecompositionDAG::reset(std::size_t patternSize)
{
    structure_.clear();
    map_.clear();
    patternSize_ = patternSize;
}

TreeDecompositionDAG & TreeDecompositionDAG::operator=(const TreeDecompositionDAG & rhs)
{
    if(&rhs == this)
        return *this;

    patternSize_ = rhs.patternSize();
    structure_ = rhs.structure_;

    refillMap();

    return *this;
}

std::size_t TreeDecompositionDAG::patternSize() const
{
    return patternSize_;
}

std::size_t TreeDecompositionDAG::numComponents() const
{
    return map_.size();
}

bool TreeDecompositionDAG::hasComponent(const TreeDecompositionSubgraph & component) const
{
    return map_.find(component) != map_.end();
}

void TreeDecompositionDAG::addComponent(const TreeDecompositionSubgraph & component)
{
    addOrFindNode(component);
}

std::pair<TreeDecompositionDAG::ChildIterator, TreeDecompositionDAG::ChildIterator> TreeDecompositionDAG::children(const TreeDecompositionSubgraph & component) const
{
    Node n = findExistingNode(component);
    if(n == invalidDescriptor())
        return std::make_pair(ChildIterator(), ChildIterator());

    typedef boost::graph_traits<Graph>::out_edge_iterator oei;
    std::pair<oei, oei> p = boost::out_edges(n, structure_);

    return std::make_pair(
                boost::make_transform_iterator(p.first, AndNodeTransformer(false, this)),
                boost::make_transform_iterator(p.second, AndNodeTransformer(false, this))
                );
}

std::pair<TreeDecompositionDAG::ParentIterator, TreeDecompositionDAG::ParentIterator> TreeDecompositionDAG::parents(const TreeDecompositionSubgraph &component) const
{
    Node n = findExistingNode(component);
    if(n == invalidDescriptor())
        return std::make_pair(ParentIterator(), ParentIterator());

    typedef boost::graph_traits<Graph>::in_edge_iterator iet;
    std::pair<iet, iet> p = boost::in_edges(n, structure_);

    return std::make_pair(
                boost::make_transform_iterator(p.first, AndNodeTransformer(true, this)),
                boost::make_transform_iterator(p.second, AndNodeTransformer(true, this))
                );
}

std::pair<TreeDecompositionDAG::ComponentIterator, TreeDecompositionDAG::ComponentIterator> TreeDecompositionDAG::components() const
{
    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    typedef boost::transform_iterator<ComponentTransformer, vit> transform_it;

    ComponentTransformer transformer(&structure_);
    ComponentFilter filter(&structure_);

    std::pair<vit, vit> p = boost::vertices(structure_);
    std::pair<transform_it, transform_it> p2(
                boost::make_transform_iterator(p.first, transformer),
                boost::make_transform_iterator(p.second, transformer)
                );

    return std::make_pair(
        boost::make_filter_iterator(filter, p2.first, p2.second),
        boost::make_filter_iterator(filter, p2.second, p2.second)
                );
}

TreeDecompositionDAG::Node TreeDecompositionDAG::findExistingNode(const TreeDecompositionSubgraph & component) const
{
    boost::unordered_map<TreeDecompositionSubgraph, Node>::const_iterator it = map_.find(component);
    if(it == map_.end())
        return invalidDescriptor();

    return it->second;
}

const std::pair<TreeDecompositionSubgraph, TreeDecompositionDAG::Node> TreeDecompositionDAG::addOrFindNode(const TreeDecompositionSubgraph & component)
{
    boost::unordered_map<TreeDecompositionSubgraph, Node>::const_iterator it = map_.find(component);

    // a new node
    if(it == map_.end())
    {
        TreeDecompositionDAGLabel lbl;
        lbl.vertices = component;
        lbl.type = TreeDecompositionDAGLabel::NODE_ChildrenOR;
        Node n = boost::add_vertex(lbl, structure_);

        it = map_.insert(std::make_pair(component, n)).first;
    }

    return *it;
}
void TreeDecompositionDAG::remapIndices(const std::vector<std::size_t> &newIndices)

{
    if(newIndices.size() != patternSize())
        throw std::logic_error("Not a valid indices map");

    // count them all
    bool valid = true;
    std::vector<std::size_t> cnt(patternSize());

    // set the new indices
    for(std::size_t i =0; i < newIndices.size() && valid; ++i)
        if(newIndices[i] >= patternSize())
            valid = false;
        else
            ++cnt[newIndices[i]];

    for(std::size_t i = 0; i < cnt.size(); ++i)
        if(cnt[i] != 1)
            valid = false;

    if(!valid)
        throw std::logic_error("Not a valid indices map");

    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    for(std::pair<vit, vit> p = boost::vertices(structure_); p.first != p.second; ++p.first)
    {
        TreeDecompositionDAGLabel & lbl = structure_[*p.first];
        if(lbl.type == TreeDecompositionDAGLabel::NODE_ChildrenOR)
            remapComponent(newIndices, lbl.vertices);
    }

    map_.clear();
    refillMap();
}

void TreeDecompositionDAG::remapComponent(const std::vector<std::size_t> & newIndices, TreeDecompositionSubgraph & subgraph)
{
    TreeDecompositionSubgraph newComp(subgraph.size());

    for(std::size_t i = 0; i < subgraph.size(); ++i)
        newComp[newIndices[i]] = subgraph[i];

    std::swap(subgraph, newComp);
}

TreeDecompositionSubgraph TreeDecompositionDAG::addUnaryAnd(const TreeDecompositionSubgraph & child, const TreeDecompositionSubgraph & parent, TreeDecompositionDAGLabel::NodeType type)
{
    const std::pair<TreeDecompositionSubgraph, Node> & c = addOrFindNode(child);
    const std::pair<TreeDecompositionSubgraph, Node> & p = addOrFindNode(parent);

    // check whether this link already exists
    for(std::pair<ChildIterator, ChildIterator> p1 = children(parent); p1.first != p1.second; ++p1.first)
    {
        TreeDecompositionDAGAndNode andNode = *p1.first;
        if(andNode.type() == type && andNode.childAt(0) == child)
            return p.first;
    }

    // add the link
    TreeDecompositionDAGLabel lbl;
    lbl.type = type;
    Node andNode = boost::add_vertex(lbl, structure_);

    boost::add_edge(p.second, andNode, structure_);
    boost::add_edge(andNode, c.second, structure_);

    return p.first;
}

TreeDecompositionSubgraph TreeDecompositionDAG::addNaryAnd(std::vector<TreeDecompositionSubgraph> childComponents, const TreeDecompositionSubgraph & parent, TreeDecompositionDAGLabel::NodeType type)
{
    // sort the children
    std::sort(childComponents.begin(), childComponents.end());

    // find the parent node
    const std::pair<TreeDecompositionSubgraph, Node> & p = addOrFindNode(parent);

    // check whether this link already exists
    for(std::pair<ChildIterator, ChildIterator> p1 = children(parent); p1.first != p1.second; ++p1.first)
    {
        const TreeDecompositionDAGAndNode & andNode = *p1.first;
        if(andNode.type() != type || andNode.numberOfChildren() != childComponents.size())
            continue;

        bool allSame = true;
        for(std::size_t i = 0; i < andNode.numberOfChildren() && allSame; ++i)
            if(andNode.childAt(i) != childComponents[i])
                allSame = false;

        if(allSame)
            return p.first;
    }

    // now find the child nodes
    std::vector<Node> childNodes(childComponents.size());
    for(std::size_t i = 0; i < childNodes.size(); ++i)
        childNodes[i] = addOrFindNode(childComponents[i]).second;

    // add the link
    TreeDecompositionDAGLabel lbl;
    lbl.type = type;
    Node andNode = boost::add_vertex(lbl, structure_);

    boost::add_edge(p.second, andNode, structure_);
    for(std::size_t i = 0; i < childNodes.size(); ++i)
        boost::add_edge(andNode, childNodes[i], structure_);

    return p.first;
}

TreeDecompositionSubgraph TreeDecompositionDAG::addBinaryAnd(const TreeDecompositionSubgraph & childA, const TreeDecompositionSubgraph & childB, const TreeDecompositionSubgraph & parent, TreeDecompositionDAGLabel::NodeType type)
{
    std::vector<TreeDecompositionSubgraph> nodes(2);
    nodes[0] = childA;
    nodes[1] = childB;

    return addNaryAnd(nodes, parent, type);
}

// addition methods
TreeDecompositionSubgraph TreeDecompositionDAG::addList(std::size_t vertex)
{
    if(vertex > patternSize())
        throw errOutOfBounds();

    TreeDecompositionSubgraph child(patternSize());
    TreeDecompositionSubgraph parent(patternSize());
    parent[vertex] = TreeDecompositionSubgraph::ActiveVertex;

    return addUnaryAnd(child, parent, TreeDecompositionDAGLabel::NODE_ChildrenList);
}

TreeDecompositionSubgraph TreeDecompositionDAG::addExtend(const TreeDecompositionSubgraph & child, std::size_t vertexToExtend)
{
    if(vertexToExtend > patternSize())
        throw errOutOfBounds();

    TreeDecompositionSubgraph parent(child);

    if(parent[vertexToExtend] != TreeDecompositionSubgraph::UnseenVertex)
        throw errNotUnseen();

    parent[vertexToExtend] = TreeDecompositionSubgraph::ActiveVertex;

    return addUnaryAnd(child, parent, TreeDecompositionDAGLabel::NODE_ChildrenExtend);
}

TreeDecompositionSubgraph TreeDecompositionDAG::addProject(const TreeDecompositionSubgraph & child, std::size_t vertexToProjectAway)
{
    if(vertexToProjectAway > patternSize())
        throw errOutOfBounds();

    TreeDecompositionSubgraph parent(child);
    if(parent[vertexToProjectAway] != TreeDecompositionSubgraph::ActiveVertex)
        throw errNotActive();

    parent[vertexToProjectAway] = TreeDecompositionSubgraph::ProjectedAwayVertex;

    return addUnaryAnd(child, parent, TreeDecompositionDAGLabel::NODE_ChildrenProject);
}

TreeDecompositionSubgraph TreeDecompositionDAG::addJoin(const TreeDecompositionSubgraph & lhsComponent, const TreeDecompositionSubgraph & rhsComponent)
{
    TreeDecompositionSubgraph parent(lhsComponent);
    parent += rhsComponent;

    return addBinaryAnd(lhsComponent, rhsComponent, parent, TreeDecompositionDAGLabel::NODE_ChildrenJoin);
}

// removal method
void TreeDecompositionDAG::removeNode(const TreeDecompositionSubgraph & component)
{
    boost::unordered_map<TreeDecompositionSubgraph, Node>::iterator it = map_.find(component);
    if(it == map_.end())
        return;

    Node n = it->second;

    typedef boost::graph_traits<Graph>::in_edge_iterator iet;
    typedef boost::graph_traits<Graph>::out_edge_iterator oei;

    std::set<Node> toRemove;
    for(std::pair<iet, iet> p = boost::in_edges(n, structure_); p.first != p.second; ++p.first)
        toRemove.insert(boost::source(*p.first, structure_));
    for(std::pair<oei, oei> p = boost::out_edges(n, structure_); p.first != p.second; ++p.first)
        toRemove.insert(boost::target(*p.first, structure_));

    for(std::set<Node>::const_iterator it = toRemove.begin(); it != toRemove.end(); ++it)
    {
        Node n2 = *it;
        boost::clear_vertex(n2, structure_);
        boost::remove_vertex(n2, structure_);
    }

    boost::remove_vertex(n, structure_);
    map_.erase(it);
}

void TreeDecompositionDAG::removeSubtree(const TreeDecompositionSubgraph & component)
{
    // find the root node
    Node root = findExistingNode(component);
    if(root == invalidDescriptor())
        return;

    // visit the subtree which belongs uniquely to this node
    boost::unordered_map<Node, std::size_t> degreeMap;
    boost::unordered_set<Node> toRemove;
    std::queue<Node> todo;

    degreeMap.insert(std::make_pair(root, boost::in_degree(root, structure_)));
    todo.push(root);

    while(!todo.empty())
    {
        Node cur = todo.front();
        todo.pop();

        toRemove.insert(cur);

        // check the children of the children
        typedef boost::graph_traits<Graph>::out_edge_iterator oei;
        for(std::pair<oei, oei> pAnd = boost::out_edges(cur, structure_); pAnd.first != pAnd.second; ++pAnd.first)
        {
            // add the and node
            Node andNode = boost::target(*pAnd.first, structure_);
            toRemove.insert(andNode);

            // and add the children which reside only in the subtree
            for(std::pair<oei, oei> pChild = boost::out_edges(andNode, structure_); pChild.first != pChild.second; ++pChild.first)
            {
                Node childNode = boost::target(*pChild.first, structure_);
                std::size_t & curD = degreeMap[childNode];
                ++curD;

                // have we seen al the in edges?
                if(boost::in_degree(childNode, structure_) == curD)
                    todo.push(childNode);
            }
        }
    }

    // add the and node parent of the component
    {
        typedef boost::graph_traits<Graph>::in_edge_iterator iet;
        for(std::pair<iet, iet> p = boost::in_edges(root, structure_); p.first != p.second; ++p.first)
            toRemove.insert(boost::source(*p.first, structure_));

    }

    // first clear all the nodes (edges)
    for(boost::unordered_set<Node>::const_iterator it = toRemove.begin(); it != toRemove.end(); ++it)
    {
        Node node = *it;
        boost::clear_vertex(node, structure_);

        // remove from the map if it is an a or node
        if(structure_[node].type == TreeDecompositionDAGLabel::NODE_ChildrenOR)
            map_.erase(structure_[node].vertices);

        // and remove from the graph
        boost::remove_vertex(node, structure_);
    }
}

std::logic_error TreeDecompositionDAG::errOutOfBounds() const
{
    return std::logic_error("TreeDecompositionDAG: The supplied vertex is out of bound");
}
std::logic_error TreeDecompositionDAG::errNotActive() const
{
    return std::logic_error("TreeDecompositionDAG: The supplied vertex is not active");
}
std::logic_error TreeDecompositionDAG::errAlreadyActive() const
{
    return std::logic_error("TreeDecompositionDAG: The supplied vertex is active");
}

std::logic_error TreeDecompositionDAG::errBadRange() const
{
    return std::logic_error("TreeDecompositionDAG: The supplied range is invalid");
}

std::logic_error TreeDecompositionDAG::errNotUnseen() const
{
    return std::logic_error("TreeDecompositionDAG: The supplied vertex is not unseen");
}

TreeDecompositionDAG::Node TreeDecompositionDAG::invalidDescriptor()
{
    return boost::graph_traits<Graph>::null_vertex();
}

TreeDecompositionDAG::InsertResult TreeDecompositionDAG::checkForInsertion(const TreeDecompositionSubgraph & component, const TreeDecompositionSubgraph & toInsert) const
{
    bool thisPossible = true;
    for(std::size_t i = 0; i < patternSize(); ++i)
    {
        if(toInsert[i] == TreeDecompositionSubgraph::ActiveVertex)
        {
            if(component[i] == TreeDecompositionSubgraph::UnseenVertex)
                return NotPossible;
            else if(component[i] == TreeDecompositionSubgraph::ProjectedAwayVertex)
                thisPossible = false;
        }
        else if(toInsert[i] == TreeDecompositionSubgraph::ProjectedAwayVertex)
        {
            if(component[i] != TreeDecompositionSubgraph::UnseenVertex)
                return NotPossible;
        }
    }

    return thisPossible ? CanInsert : MaybeIntoChildren;
}

bool TreeDecompositionDAG::constructInsertColorMap(Node rootNode, std::set<TreeDecompositionSubgraph> insertNodes, std::map<Node, InsertResult> & map) const
{
    map.clear();

    std::stack<Node> todo;
    todo.push(rootNode);
    std::set<Node> visitedNodes;
    std::stack<Node> toVisit;

    // depth first visiting of all nodes
    while(todo.empty())
    {
        Node curNode = todo.top();
        todo.pop();

        // bookkeeping
        if(visitedNodes.count(curNode) > 0)
            continue;
        visitedNodes.insert(curNode);

        // get the label
        const TreeDecompositionSubgraph & comp = structure_[curNode].vertices;

        // is this one of the nodes we are looking for?
        std::set<TreeDecompositionSubgraph>::iterator it = insertNodes.find(comp);
        if(it != insertNodes.end())
        {
            map.insert(std::make_pair(curNode, CanInsert));
            toVisit.push(curNode);
        }

        // and continue the of the children
        typedef boost::graph_traits<Graph>::in_edge_iterator iet;
        for(std::pair<iet, iet> p1 = boost::in_edges(curNode, structure_); p1.first != p1.second; ++p1.first)
        {
            Node andNode = boost::source(*p1.first, structure_);
            for(std::pair<iet, iet> p2 = boost::in_edges(andNode, structure_); p2.first != p2.second; ++p2.first)
            {
                Node childNode = boost::source(*p2.first, structure_);
                todo.push(childNode);
            }
        }
    }

    // check if all nodes found
    if(insertNodes.size() != map.size())
        return false;

    // no go bottom up, and tag all nodes (till rootnode)
    std::set<Node> visitedParents;
    while(!toVisit.empty())
    {
        Node curNode = toVisit.top();
        toVisit.pop();

        // bookkeeping
        if(!visitedParents.insert(curNode).second)
            continue;

        // if no color, assign intoChild color
        if(map.count(curNode) == 0)
            map.insert(std::make_pair(curNode, MaybeIntoChildren));

        // and go to the parents
        typedef boost::graph_traits<Graph>::out_edge_iterator oei;
        for(std::pair<oei, oei> p1 = boost::out_edges(curNode, structure_); p1.first != p1.second; ++p1.first)
        {
            Node andNode = boost::target(*p1.first, structure_);
            Node parentNode = boost::target(*boost::out_edges(andNode, structure_).first, structure_);

            if(visitedNodes.count(parentNode) > 0)
            {
                map.insert(std::make_pair(andNode, MaybeIntoChildren));
                toVisit.push(parentNode);
            }
        }
    }

    return true;
}

TreeDecompositionDAG::InsertResult TreeDecompositionDAG::findInMap(const TreeDecompositionSubgraph & component, const InsertionMap & map) const
{
    InsertionMap::const_iterator it = map.find(component);
    return it == map.end() ? NotPossible : it->second;
}

TreeDecompositionSubgraph TreeDecompositionDAG::insert(const TreeDecompositionSubgraph & rootComponent, const TreeDecompositionSubgraph & toInsertComponent, const InsertionMap & insertionMap)
{
    // check for validity
    if(!hasComponent(rootComponent))
        throw errBadRange();

    if(findInMap(rootComponent, insertionMap) == NotPossible)
        throw errBadRange();

    std::stack<TreeDecompositionSubgraph> todo;
    std::set<TreeDecompositionSubgraph> handledNodes;
    todo.push(rootComponent);

    while(!todo.empty())
    {
        TreeDecompositionSubgraph curComp = todo.top();
        todo.pop();

        // already handled?
        if(!handledNodes.insert(curComp).second)
            continue;

        InsertResult result = findInMap(curComp, insertionMap);
        if(result == NotPossible)
            continue;

        if(result == CanInsert)
        {
            if(!partial_order_compare(curComp, toInsertComponent))
                addJoin(curComp, toInsertComponent);
        }

        // loop over the children
        for(std::pair<ChildIterator, ChildIterator> p = children(curComp); p.first != p.second; ++p.first)
        {
            const TreeDecompositionDAGAndNode & andNode = *p.first;

            switch(andNode.type())
            {
            case TreeDecompositionDAGLabel::NODE_ChildrenExtend:
            {
                const TreeDecompositionSubgraph & child = andNode.childAt(0);
                if(findInMap(child, insertionMap) != NotPossible)
                {
                    todo.push(child);
                    addExtend(createInsertComponent(child, toInsertComponent), andNode.extendVertex());
                }

                break;
            }

            case TreeDecompositionDAGLabel::NODE_ChildrenProject:
            {
                const TreeDecompositionSubgraph & child = andNode.childAt(0);
                if(findInMap(child, insertionMap) != NotPossible)
                {
                    todo.push(child);
                    addProject(createInsertComponent(child, toInsertComponent), andNode.projectAwayVertices().begin(), andNode.projectAwayVertices().end());
                }

                break;
            }

            case TreeDecompositionDAGLabel::NODE_ChildrenJoin:
            {
                std::vector<TreeDecompositionSubgraph> children(andNode.numberOfChildren());

                for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
                    children[i] = andNode.childAt(i);

                for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
                {
                    if(findInMap(children[i], insertionMap) != NotPossible)
                    {
                        todo.push(children[i]);
                        children[i] = createInsertComponent(children[i], toInsertComponent);
                        addJoin(children.begin(), children.end());

                        children[i] = andNode.childAt(i);
                    }
                }
                break;
            }

            default:
                break;
            }

        }
    }

    return createInsertComponent(rootComponent, toInsertComponent);
}

TreeDecompositionSubgraph TreeDecompositionDAG::createInsertComponent(const TreeDecompositionSubgraph & node, const TreeDecompositionSubgraph & toInsertComponent) const
{
    TreeDecompositionSubgraph result(patternSize());

    for(std::size_t i = 0 ; i < patternSize(); ++i)
        result[i] = (node[i] == TreeDecompositionSubgraph::UnseenVertex ? toInsertComponent[i] : node[i]);

    return result;
}

namespace {

template <typename LabelMap>
struct TreeDecompositionDAGLabel_DOTWriter
{
    typedef typename boost::property_traits<LabelMap>::key_type KeyType;
    typedef std::set<TreeDecompositionSubgraph> Keys;

    TreeDecompositionDAGLabel_DOTWriter(const LabelMap & map, const Keys & selectedKeys = Keys())
        : map_(map),
          selectedKeys_(selectedKeys)
    {
    }

    template <typename VertexDescriptor>
    void operator()(std::ostream & str, const VertexDescriptor & vd) const
    {
        typedef typename boost::property_traits<LabelMap>::value_type LabelType;
        const LabelType & label = boost::get(map_, vd);

        str << "[label=\"";

        if((label.type & TreeDecompositionDAGLabel::NODE_ORANDMASK) == TreeDecompositionDAGLabel::NODE_ChildrenAND)
        {
            switch(label.type)
            {
            case TreeDecompositionDAGLabel::NODE_ChildrenExtend:
                str << "Extend";
                break;
            case TreeDecompositionDAGLabel::NODE_ChildrenJoin:
                str << "Join";
                break;
            case TreeDecompositionDAGLabel::NODE_ChildrenList:
                str << "List";
                break;
            case TreeDecompositionDAGLabel::NODE_ChildrenProject:
                str << "Project";
                break;
            default:
                str << "Unknown";
                break;
            }
        }
        else
        {
            str << "{";
            unsigned int written = 0;
            for(unsigned int i = 0; i < label.vertices.size(); ++i)
                if(label.vertices[i] == TreeDecompositionSubgraph::ActiveVertex)
                {
                    if(written++ != 0)
                        str << ", ";
                    str << "v" << i;
                }
            str << "}, *{";
            written = 0;
            for(unsigned int i = 0; i < label.vertices.size(); ++i)
                if(label.vertices[i] == TreeDecompositionSubgraph::ProjectedAwayVertex)
                {
                    if(written++ != 0)
                        str << ", ";
                    str << "v" << i;
                }
            str << "}";
        }

        str << "\"";

        if(selectedKeys_.count(label.vertices) > 0)
            str << "color=red; style=filled;";

        str << "]";
    }

    LabelMap map_;
    Keys selectedKeys_;

};

template <typename VertexLabelWriter, typename DAGType>
void write_dot_format_helper(
        const DAGType & dag,
        std::ostream & stream,
        const std::set<TreeDecompositionSubgraph> & chosenNodes = std::set<TreeDecompositionSubgraph>())
{
    typedef typename boost::property_map<DAGType, boost::vertex_bundle_t>::const_type DAG_Label;
    typedef typename boost::graph_traits<DAGType>::vertex_iterator vertexIterator;
    typedef typename boost::graph_traits<DAGType>::vertex_descriptor vertexDescriptor;
    typedef std::map<vertexDescriptor, std::size_t> VertexIDAssocContainer;
    typedef VertexLabelWriter LabelWriter;

    // create vertexID MAP
    VertexIDAssocContainer map;
    boost::associative_property_map<VertexIDAssocContainer> vertexIDMap(map);

    // fill the map
    unsigned int count = 0;
    for(std::pair<vertexIterator, vertexIterator> p = boost::vertices(dag); p.first != p.second; ++p.first)
        map[*p.first] = count++;

    DAG_Label lmap = boost::get(boost::vertex_bundle, dag);
    VertexLabelWriter w(lmap, chosenNodes);

    boost::write_graphviz(stream, dag, w, boost::default_writer(), boost::default_writer(), vertexIDMap);
}

} //  namespace

void TreeDecompositionDAG::writeGraphviz(std::ostream & stream, const std::set<TreeDecompositionSubgraph> & activeComponents) const
{
    typedef boost::property_map<Graph, boost::vertex_bundle_t>::const_type DAG_Label;
    typedef TreeDecompositionDAGLabel_DOTWriter<DAG_Label> LabelWriter;

    write_dot_format_helper<LabelWriter>(structure_, stream, activeComponents);
}

void TreeDecompositionDAG::cleanUp(const TreeDecompositionSubgraph & rootComponent)
{
    Node node = findExistingNode(rootComponent);
    if(node == invalidDescriptor())
    {
        structure_.clear();
        map_.clear();
    }

    // visit all child nodes
    std::set<Node> visitedNodes;
    std::stack<Node> todo;
    todo.push(node);

    while(!todo.empty())
    {
        Node curNode = todo.top();
        todo.pop();

        // already seen
        if(!visitedNodes.insert(curNode).second)
            continue;

        typedef boost::graph_traits<Graph>::out_edge_iterator oei;
        for(std::pair<oei, oei> p1 = boost::out_edges(curNode, structure_); p1.first != p1.second; ++p1.first)
        {
            Node andNode = boost::target(*p1.first, structure_);
            visitedNodes.insert(andNode);

            for(std::pair<oei, oei> p2 = boost::out_edges(andNode, structure_); p2.first != p2.second; ++p2.first)
            {
                todo.push(boost::target(*p2.first, structure_));
            }
        }
    }

    // list all nodes to remove
    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    std::list<TreeDecompositionSubgraph> toRemove;
    for(std::pair<vit, vit> p = boost::vertices(structure_); p.first != p.second; ++p.first)
    {
        const TreeDecompositionDAGLabel & lbl = structure_[*p.first];
        if(lbl.type == TreeDecompositionDAGLabel::NODE_ChildrenOR && visitedNodes.count(*p.first) == 0)
            toRemove.push_back(lbl.vertices);
    }

    for(std::list<TreeDecompositionSubgraph>::iterator it = toRemove.begin(); it != toRemove.end(); ++it)
        removeNode(*it);
}

namespace {

std::size_t convertToInt(const std::string & str)
{
    return boost::lexical_cast<std::size_t>(str);
}

std::string convertToString(const TreeDecompositionDAGLabel & label)
{
    std::stringstream str;
    str << static_cast<int>(label.type) << ',' << label.vertices.size() << ',';
    for(std::size_t i = 0; i < label.vertices.size(); ++i)
    {
        if(label.vertices[i] == TreeDecompositionSubgraph::ActiveVertex)
            str << "a";
        else if(label.vertices[i] == TreeDecompositionSubgraph::ProjectedAwayVertex)
            str << "p";
        else if(label.vertices[i] == TreeDecompositionSubgraph::UnseenVertex)
            str << "u";
        else
            throw std::logic_error("Bad formatted TreeDecompositionSubgraph");
    }

    return str.str();
}

TreeDecompositionDAGLabel convertFromString(const std::string & str)
{
    std::vector<std::string> vct;
    boost::split(vct, str, boost::is_any_of(","));

    if(vct.size() < 2)
        throw std::logic_error("Bad formatted TreeDecompositionSubgraph");

    TreeDecompositionDAGLabel lbl;
    lbl.type = static_cast<TreeDecompositionDAGLabel::NodeType>(convertToInt(vct[0]));
    lbl.vertices = TreeDecompositionSubgraph(convertToInt(vct[1]));

    if(lbl.vertices.size() != 0)
    {
        if(vct.size() < 3 || vct[2].size() != lbl.vertices.size())
            throw std::logic_error("Bad formatted TreeDecompositionSubgraph");

        for(std::size_t i = 0; i < lbl.vertices.size(); ++i)
            if(vct[2][i] == 'a')
                lbl.vertices[i] = TreeDecompositionSubgraph::ActiveVertex;
            else if(vct[2][i] == 'p')
                lbl.vertices[i] = TreeDecompositionSubgraph::ProjectedAwayVertex;
            else if(vct[2][i] == 'u')
                lbl.vertices[i] = TreeDecompositionSubgraph::UnseenVertex;
            else
                throw std::logic_error("Bad formatted TreeDecompositionSubgraph");

    }

    return lbl;
}

} //  namespace

void TreeDecompositionDAG::serialize(std::ostream & stream) const
{
    typedef boost::graph_traits<Graph>::vertex_iterator vit;

    stream << "t # " << patternSize() << std::endl;

    // write the vertices
    std::map<Node, std::size_t> positionMap;
    std::size_t count = 0;
    for(std::pair<vit, vit> p = vertices(structure_); p.first != p.second; ++p.first, ++count)
    {
        Node v = *p.first;
        stream << "v " << count << " " << convertToString(structure_[v]) << std::endl;
        positionMap.insert(std::make_pair(v, count));
    }

    // write the edges
    typedef boost::graph_traits<Graph>::edge_iterator eit;
    for(std::pair<eit, eit> p = boost::edges(structure_); p.first != p.second; ++p.first)
    {
        Node src = boost::source(*p.first, structure_);
        Node tgt = boost::target(*p.first, structure_);

        stream << "e " << positionMap.find(src)->second << " " << positionMap.find(tgt)->second << "  " << std::endl;
    }
}

bool TreeDecompositionDAG::deserialize(std::istream & stream)
{
    std::string line;
    std::size_t pos = 0;

    std::vector<Node> nodevct;

    try
    {
        while(std::getline(stream, line))
        {
            std::vector<std::string> vct;
            boost::split(vct, line, boost::is_any_of(" "));

            if(vct.empty())
                continue;

            if(pos == 0)
            {
                if(vct.size() != 3 && vct[0] != "t" && vct[1] != "#")
                    return false;

                patternSize_ = convertToInt(vct[2]);

                // initialise the graph
                structure_.clear();
                pos = 1;
            }
            else if(pos == 1)
            {
                if(vct[0]  != "v")
                {
                    pos = 2;
                }
                else
                {
                    nodevct.push_back(boost::add_vertex(convertFromString(vct[2]), structure_));
                }
            }
            if(pos == 2)
            {
                if(vct[0] != "e")
                    return false;

                std::size_t src = convertToInt(vct[1]);
                std::size_t dst = convertToInt(vct[2]);
                boost::add_edge(nodevct[src], nodevct[dst], structure_);
            }
        }
    }
    catch(...)
    {
        structure_.clear();
        patternSize_ = 0;
        map_.clear();
    }

    refillMap();

    return true;
}

void TreeDecompositionDAG::refillMap()
{
    map_.clear();

    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    for(std::pair<vit, vit> p = boost::vertices(structure_); p.first != p.second; ++p.first)
    {
        const TreeDecompositionDAGLabel & lbl = structure_[*p.first];
        if(lbl.type == TreeDecompositionDAGLabel::NODE_ChildrenOR)
            map_.insert(std::make_pair(lbl.vertices, *p.first));
    }
}

void TreeDecompositionDAG::writeGraphviz(std::ostream & stream) const
{
    writeGraphviz(stream, std::set<TreeDecompositionSubgraph>());
}

namespace {

typedef std::map<TreeDecompositionSubgraph, std::size_t> DegreeMap;
typedef std::set<TreeDecompositionSubgraph> VisitedNodeSet;

VisitedNodeSet visitSubdag(const TreeDecompositionDAG & dag, const TreeDecompositionSubgraph & rootComponent)
{
    VisitedNodeSet visitedSubgraph;
    std::stack<TreeDecompositionSubgraph> todo;
    todo.push(rootComponent);

    // loop over the children of the root component
    while(!todo.empty())
    {
        TreeDecompositionSubgraph comp = todo.top();
        todo.pop();

        // already seen?
        if(!visitedSubgraph.insert(comp).second)
            continue;

        // visit all the children
        for(std::pair<TreeDecompositionDAG::ChildIterator, TreeDecompositionDAG::ChildIterator> p = dag.children(comp); p.first != p.second; ++p.first)
        {
            const TreeDecompositionDAGAndNode & andNode = *p.first;
            for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
                todo.push(andNode.childAt(i));
        }
    }

    return visitedSubgraph;
}

DegreeMap fillDegreeMap(const VisitedNodeSet & visitedNodes, const TreeDecompositionDAG & dag)
{
    DegreeMap map;
    for(VisitedNodeSet::const_iterator it = visitedNodes.begin(); it != visitedNodes.end(); ++it)
    {
        const TreeDecompositionSubgraph & comp = *it;
        std::size_t count = 0;
        for(std::pair<TreeDecompositionDAG::ParentIterator, TreeDecompositionDAG::ParentIterator> p = dag.parents(comp); p.first != p.second; ++p.first)
        {
            const TreeDecompositionDAGAndNode & andNode = *p.first;
            count += visitedNodes.count(andNode.parent());
        }

        map.insert(std::make_pair(comp, count));
    }

    return map;
}

std::list<TreeDecompositionSubgraph> extractAndRemoveZeros(DegreeMap & map)
{
    std::list<TreeDecompositionSubgraph> zeros;

    for(DegreeMap::const_iterator it = map.begin(); it != map.end(); ++it)
        if(it->second == 0)
            zeros.push_back(it->first);

    for(std::list<TreeDecompositionSubgraph>::const_iterator it = zeros.begin(); it != zeros.end(); ++it)
        map.erase(*it);

    return zeros;
}

} //  namespace

std::vector<TreeDecompositionSubgraph> TreeDecompositionDAG::topologicalOrder(const TreeDecompositionSubgraph & rootComponent) const
{
    VisitedNodeSet nodes = visitSubdag(*this, rootComponent);
    DegreeMap degreeMap = fillDegreeMap(nodes, *this);

    std::vector<TreeDecompositionSubgraph> result;

    // perform the sorting
    while(!degreeMap.empty())
    {
        // extract the non-zero's
        const std::list<TreeDecompositionSubgraph> & zeros = extractAndRemoveZeros(degreeMap);
        if(zeros.empty())
            throw std::logic_error("TreeDecompositionDAG: Not in DAG order");

        // add them to the list
        result.insert(result.end(), zeros.begin(), zeros.end());

        // and process the children
        for(std::list<TreeDecompositionSubgraph>::const_iterator it = zeros.begin(); it != zeros.end(); ++it)
        {
            const TreeDecompositionSubgraph & comp = *it;
            for(std::pair<ChildIterator, ChildIterator> p = children(comp); p.first != p.second; ++p.first)
            {
                const TreeDecompositionDAGAndNode & andNode = *p.first;
                for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
                    --degreeMap[andNode.childAt(i)];
            }
        }
    }

    return result;
}

std::vector<TreeDecompositionSubgraph> TreeDecompositionDAG::topologicalOrder() const
{
    VisitedNodeSet nodes;

    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    for(std::pair<vit, vit> p = boost::vertices(structure_); p.first != p.second; ++p.first)
    {
        const TreeDecompositionDAGLabel & lbl = structure_[*p.first];
        if(lbl.type == TreeDecompositionDAGLabel::NODE_ChildrenOR)
            nodes.insert(lbl.vertices);
    }

    DegreeMap degreeMap = fillDegreeMap(nodes, *this);

    std::vector<TreeDecompositionSubgraph> result;

    // perform the sorting
    while(!degreeMap.empty())
    {
        // extract the non-zero's
        const std::list<TreeDecompositionSubgraph> & zeros = extractAndRemoveZeros(degreeMap);
        if(zeros.empty())
            throw std::logic_error("TreeDecompositionDAG: Not in DAG order");

        // add them to the list
        result.insert(result.end(), zeros.begin(), zeros.end());

        // and process the children
        for(std::list<TreeDecompositionSubgraph>::const_iterator it = zeros.begin(); it != zeros.end(); ++it)
        {
            const TreeDecompositionSubgraph & comp = *it;
            for(std::pair<ChildIterator, ChildIterator> p = children(comp); p.first != p.second; ++p.first)
            {
                const TreeDecompositionDAGAndNode & andNode = *p.first;
                for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
                    --degreeMap[andNode.childAt(i)];
            }
        }
    }

    return result;
}

namespace {

std::size_t childrenTreewidth(const TreeDecompositionDAGAndNode & andNode, const std::map<TreeDecompositionSubgraph, std::size_t> & sizeMap)
{
    std::size_t curMax = sizeMap.find(andNode.childAt(0))->second;
    if(andNode.numberOfChildren() == 2)
        curMax = std::max(curMax, sizeMap.find(andNode.childAt(1))->second);

    return curMax;
}

} //  namespace

void TreeDecompositionDAG::removeAndNode(const TreeDecompositionDAGAndNode & andNode)
{
    Node n = findExistingNode(andNode.parent());
    if(n == invalidDescriptor())
        return;

    // find the node
    typedef boost::graph_traits<Graph>::out_edge_iterator oei;

    for(std::pair<oei, oei> p = boost::out_edges(n, structure_); p.first != p.second; ++p.first)
    {
        Node an = boost::target(*p.first, structure_);

        if(TreeDecompositionDAGAndNode(this, an) == andNode)
        {
            boost::clear_vertex(an, structure_);
            boost::remove_vertex(an, structure_);

            return;
        }
    }
}

std::size_t TreeDecompositionDAG::calculateTreewidth(const TreeDecompositionSubgraph & rootComponent) const
{
    const std::vector<TreeDecompositionSubgraph> & order = topologicalOrder(rootComponent);
    std::map<TreeDecompositionSubgraph, std::size_t> sizeMap;

    std::size_t result = 0;

    for(std::vector<TreeDecompositionSubgraph>::const_reverse_iterator it = order.rbegin(); it != order.rend(); ++it)
    {
        const TreeDecompositionSubgraph & component = *it;
        std::size_t curSize = component.count(TreeDecompositionSubgraph::ActiveVertex);

        // find the minimum value over the children
        std::size_t minValue = std::numeric_limits<std::size_t>::max();
        for(std::pair<ChildIterator, ChildIterator> p = children(component); p.first != p.second; ++p.first)
            minValue = std::min(minValue, childrenTreewidth(*p.first, sizeMap));

        // update the value
        curSize = minValue == std::numeric_limits<std::size_t>::max() ? curSize : std::max(curSize, minValue);

        // insert in the map
        sizeMap.insert(std::make_pair(component, curSize));

        result = std::max(curSize, result);
    }

    return result - 1;
}

void TreeDecompositionDAG::makeMinimalTreewidth(const TreeDecompositionSubgraph & rootComponent)
{
    std::list<TreeDecompositionDAGAndNode> toRemove;

    const std::vector<TreeDecompositionSubgraph> & order = topologicalOrder(rootComponent);
    std::map<TreeDecompositionSubgraph, std::size_t> sizeMap;

    for(std::vector<TreeDecompositionSubgraph>::const_reverse_iterator it = order.rbegin(); it != order.rend(); ++it)
    {
        const TreeDecompositionSubgraph & component = *it;
        std::size_t curSize = component.count(TreeDecompositionSubgraph::ActiveVertex);

        // find the minimum value over the children
        std::size_t minValue = std::numeric_limits<std::size_t>::max();
        for(std::pair<ChildIterator, ChildIterator> p = children(component); p.first != p.second; ++p.first)
            minValue = std::min(minValue, childrenTreewidth(*p.first, sizeMap));

        // find the children which should be removed
        for(std::pair<ChildIterator, ChildIterator> p = children(component); p.first != p.second; ++p.first)
        {
            const TreeDecompositionDAGAndNode & andNode = *p.first;
            if(childrenTreewidth(andNode, sizeMap) > minValue)
                toRemove.push_back(andNode);
        }

        // update the value
        curSize = minValue == std::numeric_limits<std::size_t>::max() ? curSize : std::max(curSize, minValue);

        // insert in the map
        sizeMap.insert(std::make_pair(component, curSize));
    }

    // and now remove these links
    for(std::list<TreeDecompositionDAGAndNode>::const_iterator it = toRemove.begin(); it != toRemove.end(); ++it)
        removeAndNode(*it);
}

bool TreeDecompositionDAG::isBinary() const
{
    typedef boost::graph_traits<Graph>::vertex_iterator vit;
    for(std::pair<vit, vit> p = boost::vertices(structure_); p.first != p.second; ++p.first)
        if(structure_[*p.first].type != TreeDecompositionDAGLabel::NODE_ChildrenOR && boost::out_degree(*p.first, structure_) > 2)
            return false;

    return true;
}

TreeDecompositionSubgraph TreeDecompositionDAG::getRootNode() const
{
    TreeDecompositionSubgraph result;

    for(std::pair<ComponentIterator, ComponentIterator> p = components(); p.first != p.second; ++p.first)
    {
        std::pair<ParentIterator, ParentIterator> pit = parents(*p.first);
        if(pit.first == pit.second)
        {
            if(result.size() != 0)
                throw std::logic_error("TreeDecompositionDAG: The treedecomposition dag has multiple roots");

            result = *p.first;
        }
    }

    return result;
}

std::vector<TreeDecompositionSubgraph> TreeDecompositionDAG::getRootNodes() const
{
    std::vector<TreeDecompositionSubgraph> roots;
    for(std::pair<ComponentIterator, ComponentIterator> p = components(); p.first != p.second; ++p.first)
    {
        std::pair<ParentIterator, ParentIterator> pit = parents(*p.first);
        if(pit.first == pit.second)
            roots.push_back(*p.first);
    }

    return roots;
}

namespace {

void findCompleteJoinVertices(const TreeDecompositionDAGAndNode & andNode, std::vector<TreeDecompositionSubgraph> & completeNodes, std::vector<TreeDecompositionSubgraph> & nonCompleteNodes)
{
    assert(andNode.type() == TreeDecompositionDAGLabel::NODE_ChildrenJoin);
    const std::vector<std::size_t> & joinVertices = andNode.joinVertices();

    for(std::size_t i = 0; i < andNode.numberOfChildren(); ++i)
    {
        const TreeDecompositionSubgraph & cur = andNode.childAt(i);

        bool allSet = true;
        for(std::vector<std::size_t>::const_iterator it = joinVertices.begin(); it != joinVertices.end() && allSet; ++it)
            if(cur[*it] != TreeDecompositionSubgraph::ActiveVertex)
                allSet = false;

        if(allSet)
            completeNodes.push_back(cur);
        else
            nonCompleteNodes.push_back(cur);
    }
}

void performAllPermutationsJoin(std::vector<TreeDecompositionSubgraph> & completeChildren, std::vector<TreeDecompositionSubgraph> & nonCompleteChildren, TreeDecompositionDAG & dag)
{
    std::sort(completeChildren.begin(), completeChildren.end());
    std::sort(nonCompleteChildren.begin(), nonCompleteChildren.end());

    TreeDecompositionSubgraph result;
    do
    {
        TreeDecompositionSubgraph curResult = completeChildren[0];
        for(std::size_t i = 1; i < completeChildren.size(); ++i)
            curResult = dag.addJoin(curResult, completeChildren[i]);

        result = curResult;
    }
    while(std::next_permutation(completeChildren.begin(), completeChildren.end()));

    do
    {
        TreeDecompositionSubgraph curResult = result;
        for(std::size_t i = 0; i < nonCompleteChildren.size(); ++i)
            curResult = dag.addJoin(curResult, nonCompleteChildren[i]);
    }
    while(std::next_permutation(nonCompleteChildren.begin(), nonCompleteChildren.end()));
}

void performSimpleJoin(std::vector<TreeDecompositionSubgraph> & completeChildren, std::vector<TreeDecompositionSubgraph> & nonCompleteChildren, TreeDecompositionDAG & dag)
{
    TreeDecompositionSubgraph result = completeChildren[0];
    for(std::size_t i = 1; i < completeChildren.size(); ++i)
        result = dag.addJoin(result, completeChildren[i]);

    for(std::size_t i = 0; i < nonCompleteChildren.size(); ++i)
        result = dag.addJoin(result, nonCompleteChildren[i]);
}

void make_binary_node(const TreeDecompositionDAGAndNode & andNode, TreeDecompositionDAG & dag, bool usePermutations)
{
    TreeDecompositionSubgraph parent = andNode.parent();
    std::vector<TreeDecompositionSubgraph> completeChildren, nonCompleteChildren;

    // find the complete and the non-complete children
    findCompleteJoinVertices(andNode, completeChildren, nonCompleteChildren);
    assert(!completeChildren.empty());

    // remove the and node
    dag.removeAndNode(andNode);

    // and now join the children
    if(usePermutations)
        performAllPermutationsJoin(completeChildren, nonCompleteChildren, dag);
    else
        performSimpleJoin(completeChildren, nonCompleteChildren, dag);
}

} //  namespace

void make_binary_treedecomposition(TreeDecompositionDAG & dag, bool useAllPermutations)
{
    std::stack<TreeDecompositionDAGAndNode> nonBinaryAndNodes;

    typedef TreeDecompositionDAG::ComponentIterator cit;
    typedef TreeDecompositionDAG::ChildIterator chit;

    // find out all the non-binary and nodes
    for(std::pair<cit, cit> p = dag.components(); p.first != p.second; ++p.first)
    {
        for(std::pair<chit, chit> p2 = dag.children(*p.first); p2.first != p2.second; ++p2.first)
        {
            const TreeDecompositionDAGAndNode & andNode = *p2.first;
            if(andNode.type() == TreeDecompositionDAGLabel::NODE_ChildrenJoin && andNode.numberOfChildren() > 2)
                nonBinaryAndNodes.push(andNode);
        }
    }

    // and make them binary
    while(!nonBinaryAndNodes.empty())
    {
        TreeDecompositionDAGAndNode andNode = nonBinaryAndNodes.top();
        nonBinaryAndNodes.pop();
        make_binary_node(andNode, dag, useAllPermutations);
    }
}

void swap(TreeDecompositionDAG & lhs, TreeDecompositionDAG & rhs)
{
    using std::swap;

    swap(lhs.structure_, rhs.structure_);
    swap(lhs.map_, rhs.map_);
    swap(lhs.patternSize_, rhs.patternSize_);
}

} // treeDAG namespace