#ifndef STAN_MATH_TORSTEN_PKMODELONECPT_HPP
#define STAN_MATH_TORSTEN_PKMODELONECPT_HPP

#include <Eigen/Dense>
#include <boost/math/tools/promotion.hpp>
#include <stan/math/torsten/PKModel/PKModel.hpp>
#include <stan/math/torsten/PKModel/Pred/Pred1_oneCpt.hpp>
#include <stan/math/torsten/PKModel/Pred/PredSS_oneCpt.hpp>
#include <string>
#include <vector>

namespace torsten {

/**
 * Computes the predicted amounts in each compartment at each event
 * for a one compartment model with first oder absorption.
 *
 * @tparam T0 type of scalar for time of events.
 * @tparam T1 type of scalar for amount at each event.
 * @tparam T2 type of scalar for rate at each event.
 * @tparam T3 type of scalar for inter-dose inteveral at each event.
 * @tparam T4 type of scalars for the model parameters.
 * @tparam T5 type of scalar for bio-variability F.
 * @tparam T6 type of scalar for lag times.
 * @param[in] pMatrix parameters at each event
 * @param[in] time times of events
 * @param[in] amt amount at each event
 * @param[in] rate rate at each event
 * @param[in] ii inter-dose interval at each event
 * @param[in] evid event identity:
 *                    (0) observation
 *                    (1) dosing
 *                    (2) other
 *                    (3) reset
 *                    (4) reset AND dosing
 * @param[in] cmt compartment number at each event
 * @param[in] addl additional dosing at each event
 * @param[in] ss steady state approximation at each event (0: no, 1: yes)
 * @return a matrix with predicted amount in each compartment
 *         at each event.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<std::vector<T4> >& pMatrix,
              const std::vector<std::vector<T5> >& biovar,
              const std::vector<std::vector<T6> >& tlag) {
  using std::vector;
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using boost::math::tools::promote_args;
  using stan::math::check_positive_finite;

  int nCmt = 2;
  int nParm = 3;

  // Check arguments -- FIX ME: handle the new parameter arguments
  static const char* function("PKModelOneCpt");
  torsten::pmetricsCheck(time, amt, rate, ii, evid, cmt, addl, ss,
                pMatrix, biovar, tlag, function);
  for (size_t i = 0; i < pMatrix.size(); i++) {
    check_positive_finite(function, "PK parameter CL", pMatrix[i][0]);
    check_positive_finite(function, "PK parameter V2", pMatrix[i][1]);
  }

  // FIX ME - we want to check every array of pMatrix, not
  // just the first one (at index 0)
  std::string message4 = ", but must equal the number of parameters in the model: " // NOLINT
    + boost::lexical_cast<std::string>(nParm) + "!";
  const char* length_error4 = message4.c_str();
  if (!(pMatrix[0].size() == (size_t) nParm))
    stan::math::invalid_argument(function,
    "The number of parameters per event (length of a vector in the ninth argument) is", // NOLINT
    pMatrix[0].size(), "", length_error4);

  std::string message5 = ", but must equal the number of compartments in the model: " // NOLINT
    + boost::lexical_cast<std::string>(nCmt) + "!";
  const char* length_error5 = message5.c_str();
  if (!(biovar[0].size() == (size_t) nCmt))
    stan::math::invalid_argument(function,
    "The number of biovariability parameters per event (length of a vector in the tenth argument) is", // NOLINT
    biovar[0].size(), "", length_error5);

  if (!(tlag[0].size() == (size_t) nCmt))
    stan::math::invalid_argument(function,
                                 "The number of lag times parameters per event (length of a vector in the eleventh argument) is", // NOLINT
                                 tlag[0].size(), "", length_error5);

  // Construct dummy matrix for last argument of pred
  Eigen::Matrix<T4, Eigen::Dynamic, Eigen::Dynamic> dummy_system;
  std::vector<Eigen::Matrix<T4, Eigen::Dynamic, Eigen::Dynamic> >
    dummy_systems(1, dummy_system);

  return Pred(time, amt, rate, ii, evid, cmt, addl, ss,
              pMatrix, biovar, tlag,
              nCmt, dummy_systems,
              Pred1_oneCpt(), PredSS_oneCpt());
}

/**
 * Overload function to allow user to pass an std::vector for pMatrix.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<T4>& pMatrix,
              const std::vector<std::vector<T5> >& biovar,
              const std::vector<std::vector<T6> >& tlag) {
  std::vector<std::vector<T4> > vec_pMatrix(1, pMatrix);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       vec_pMatrix, biovar, tlag);
}

/**
 * Overload function to allow user to pass an std::vector for pMatrix,
 * and biovar.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<T4>& pMatrix,
              const std::vector<T5>& biovar,
              const std::vector<std::vector<T6> >& tlag) {
  std::vector<std::vector<T4> > vec_pMatrix(1, pMatrix);
  std::vector<std::vector<T5> > vec_biovar(1, biovar);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       vec_pMatrix, vec_biovar, tlag);
}

/**
 * Overload function to allow user to pass an std::vector for pMatrix,
 * and tlag.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<T4>& pMatrix,
              const std::vector<std::vector<T5> >& biovar,
              const std::vector<T6>& tlag) {
  std::vector<std::vector<T4> > vec_pMatrix(1, pMatrix);
  std::vector<std::vector<T6> > vec_tlag(1, tlag);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       vec_pMatrix, biovar, vec_tlag);
}

/**
 * Overload function to allow user to pass an std::vector for pMatrix,
 * biovar, and tlag.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<T4>& pMatrix,
              const std::vector<T5>& biovar,
              const std::vector<T6>& tlag) {
  std::vector<std::vector<T4> > vec_pMatrix(1, pMatrix);
  std::vector<std::vector<T5> > vec_biovar(1, biovar);
  std::vector<std::vector<T6> > vec_tlag(1, tlag);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       vec_pMatrix, vec_biovar, vec_tlag);
}

/**
 * Overload function to allow user to pass an std::vector for biovar.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<std::vector<T4> >& pMatrix,
              const std::vector<T5>& biovar,
              const std::vector<std::vector<T6> >& tlag) {
  std::vector<std::vector<T5> > vec_biovar(1, biovar);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       pMatrix, vec_biovar, tlag);
}

/**
 * Overload function to allow user to pass an std::vector for biovar,
 * and tlag.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<std::vector<T4> >& pMatrix,
              const std::vector<T5>& biovar,
              const std::vector<T6>& tlag) {
  std::vector<std::vector<T5> > vec_biovar(1, biovar);
  std::vector<std::vector<T6> > vec_tlag(1, tlag);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       pMatrix, vec_biovar, vec_tlag);
}

/**
 * Overload function to allow user to pass an std::vector for tlag.
 */
template <typename T0, typename T1, typename T2, typename T3, typename T4,
          typename T5, typename T6>
Eigen::Matrix <typename boost::math::tools::promote_args<T0, T1, T2, T3,
  typename boost::math::tools::promote_args<T4, T5, T6>::type>::type,
  Eigen::Dynamic, Eigen::Dynamic>
PKModelOneCpt(const std::vector<T0>& time,
              const std::vector<T1>& amt,
              const std::vector<T2>& rate,
              const std::vector<T3>& ii,
              const std::vector<int>& evid,
              const std::vector<int>& cmt,
              const std::vector<int>& addl,
              const std::vector<int>& ss,
              const std::vector<std::vector<T4> >& pMatrix,
              const std::vector<std::vector<T5> >& biovar,
              const std::vector<T6>& tlag) {
  std::vector<std::vector<T6> > vec_tlag(1, tlag);

  return PKModelOneCpt(time, amt, rate, ii, evid, cmt, addl, ss,
                       pMatrix, biovar, vec_tlag);
}

}
#endif


#ifndef OFFNFACAG_HPP
#define OFFNFACAG_HPP

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

#include "arch/ARLAgent.hpp"
#include "bib/Seed.hpp"
#include "bib/Utils.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;
  double dpmu;

  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;

typedef struct _trajectory {
  std::shared_ptr<std::deque<sample>> transitions;
  double rewards;
} trajectory;

template<typename NN = MLP>
class OffNFACAg : public arch::ARLAgent<arch::AgentProgOptions> {
 public:
  typedef NN PolicyImpl;

  OffNFACAg(unsigned int _nb_motors, unsigned int _nb_sensors)
    : arch::ARLAgent<arch::AgentProgOptions>(_nb_motors, _nb_sensors), empty_action(0) {

  }

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

    delete ann_testing;

    delete hidden_unit_v;
    delete hidden_unit_a;
  }

  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) {
      double p0 = 1.f;
      for(uint i=0; i < this->nb_motors; i++) {
        p0 *= bib::Proba<double>::truncatedGaussianDensity(last_action->at(i), last_pure_action->at(i), noise);
      }
      trajectories.back()->transitions->push_back( {last_state, *last_pure_action, *last_action, sensors, reward, goal_reached, p0});
    }

    last_pure_action.reset(new vector<double>(*next_action));
    if(learning) {
      if(gaussian_policy) {
        vector<double>* randomized_action = bib::Proba<double>::multidimentionnalTruncatedGaussian(*next_action, noise);
        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);

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

    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<bool>("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");
    max_trajectory          = pt->get<uint>("agent.max_trajectory");
    offpolicy_strategy      = pt->get<uint>("agent.offpolicy_strategy");
    add_v_corrector         = pt->get<bool>("agent.add_v_corrector");
    offpolicy_actor         = pt->get<bool>("agent.offpolicy_actor");
    number_global_fitted_iteration         = pt->get<uint>("agent.number_global_fitted_iteration");
    a3c                                    = pt->get<bool>("agent.a3c");
    offpolicy_critic                       = pt->get<bool>("agent.offpolicy_critic");
    shuffle_buffer                         = pt->get<bool>("agent.shuffle_buffer");
    testing_ann_offpolicy                  = pt->get<bool>("agent.testing_ann_offpolicy");
    uint momentum           = pt->get<uint>("agent.momentum");
    testing_ann_offpolicy = false;
//     testing_vnn                            = pt->get<bool>("agent.testing_vnn");
    gae                     = false;

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

    if(lambda < 0. && offpolicy_critic) {
      LOG_DEBUG("set lambda please!");
      exit(1);
    }

    if(offpolicy_actor && !gae) {
      LOG_DEBUG("to be done? (offpolicy_actor without gae)");
      exit(1);
    }

    if(offpolicy_actor && a3c) {
      LOG_DEBUG("a3c is on-policy");
      exit(1);
    }

    if(gae && a3c) {
      LOG_DEBUG("choose either gae or a3c");
      exit(1);
    }

    if(offpolicy_actor && gae && stoch_iter_actor > 1) {
      LOG_DEBUG("to be done!");
      exit(1);
    }

    ann = new NN(this->get_state_size(), *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(this->get_state_size(), this->get_state_size(), *hidden_unit_v, alpha_v, 1, -1, hidden_layer_type, batch_norm_critic,
                 add_v_corrector && offpolicy_critic && offpolicy_strategy != 0, momentum);
//     if(std::is_same<NN, DODevMLP>::value)
//       vnn->exploit(pt, ann);

    ann_testing = new NN(*ann, false, ::caffe::Phase::TEST);
    ann_testing->increase_batchsize(1);

//     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& ) {
//       }
//     }

    best_testing_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;

    if(learning) {
      if(trajectories.size() >= max_trajectory)
        trajectories.pop_front();

      trajectory* t = new trajectory;
      t->transitions.reset(new std::deque<sample>);
      t->rewards = 0;
      trajectories.push_back(std::shared_ptr<trajectory>(t));
    }

//     if(std::is_same<NN, DODevMLP>::value && learning) {
//       static_cast<DODevMLP *>(vnn)->inform(episode, last_sum_weighted_reward);
//       static_cast<DODevMLP *>(ann)->inform(episode, last_sum_weighted_reward);
//       static_cast<DODevMLP *>(ann_testing)->inform(episode, last_sum_weighted_reward);
//     }

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

  void update_critic() {
    int all_size = alltransitions();
    if (all_size > 0) {
      //remove trace of old policy
      auto iter = [&]() {
        std::vector<double> all_states(all_size * this->get_state_size());
        std::vector<double> all_next_states(all_size * this->get_state_size());
        std::vector<double> v_target(all_size);

        int li=0;
        for(auto one_trajectory : trajectories) {
          const std::deque<sample>& trajectory = *one_trajectory->transitions;

          for (auto it : trajectory) {
            std::copy(it.s.begin(), it.s.end(), all_states.begin() + li * this->get_state_size());
            std::copy(it.next_s.begin(), it.next_s.end(), all_next_states.begin() + li * this->get_state_size());
            li++;
          }
        }
        ASSERT(li == all_size, "pb");

        decltype(vnn->computeOutVFBatch(all_next_states, empty_action)) all_nextV;
        all_nextV = vnn->computeOutVFBatch(all_next_states, empty_action);

        li=0;
        for(auto one_trajectory : trajectories) {
          const std::deque<sample>& trajectory = *one_trajectory->transitions;
          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(li == all_size, "pb");
//         bib::Logger::PRINT_ELEMENTS(v_target, "v_target");

        if(vnn_from_scratch) {
          delete vnn;
          vnn = new NN(this->get_state_size(), this->get_state_size(), *hidden_unit_v, alpha_v, all_size, -1, hidden_layer_type,
                       batch_norm_critic, add_v_corrector);
        }
        if(lambda < 0.f)
          vnn->learn_batch(all_states, empty_action, v_target, stoch_iter_critic);
        else {
          auto all_V = vnn->computeOutVFBatch(all_states, empty_action);
          std::vector<double> deltas(all_size);
          //
          //        Simple computation for lambda return
          //
          for (int i=0; i<all_size; i++)
            deltas[i] = v_target[i] - all_V->at(i);

          std::vector<double>* all_pi;
          double* ptheta;
          double max_ptheta;
          std::vector<double>* sample_weight;
          if(offpolicy_strategy != 0 && offpolicy_critic) {
            ann->increase_batchsize(all_size);
            if(testing_ann_offpolicy){
              ann_testing->copyParametersFrom(ann);
              ann_testing->increase_batchsize(all_size);
              all_pi = ann_testing->computeOutBatch(all_states);
              ann_testing->increase_batchsize(1);
            }
            else
              all_pi = ann->computeOutBatch(all_states);
            ptheta = new double[all_size];

            uint i=0;
            if(offpolicy_strategy >= 1 && offpolicy_strategy <= 3) {
              for(auto one_trajectory : trajectories) {
                const std::deque<sample>& trajectory = *one_trajectory->transitions;
                for (auto it : trajectory) {
                  double p0 = 1.f;
                  for(uint j=0; j < this->nb_motors; j++)
                    p0 *= bib::Proba<double>::truncatedGaussianDensity(it.a[j], all_pi->at(i*this->nb_motors+j), noise);

                  ptheta[i] = p0;
                  i++;
                }
              }
              max_ptheta = *std::max_element(ptheta, ptheta+all_size);
            }

            if(add_v_corrector)
              sample_weight = new std::vector<double>(all_size);
          }

          std::vector<double> diff(all_size);
          int index_shift=0;
          for(auto one_trajectory : trajectories) {
            const std::deque<sample>& trajectory = *one_trajectory->transitions;
            li=0;
            for (auto it : trajectory) {
              diff[index_shift+li] = 0;
//               offpolicy_strategy
//               0 : lambda only
//               1 : lambda \pi (TB)
//               2 : \pi / \mu (IS)
//               3 : lambda min(1,\pi / \mu) (Retrace)
//               4 : lambda * (1-||a_t - \pi||) (ours)
//               5 : lambda * (1-min(||a_t - \pi||, ||u_t - \pi|| )) (ours)
              if(offpolicy_strategy == 0 || !offpolicy_critic) {
                for (uint n=li; n<trajectory.size(); n++)
                  diff[index_shift+li] += std::pow(this->gamma * lambda, n-li) * deltas[index_shift+n];
              } else if(offpolicy_strategy == 1) {
//                 for (int n=trajectory.size()-1; n>=li; n--){
//                   double ci = 1.f;
//                   for(int subn = 0;subn<n-li;subn++)
//                     ci = ci * lambda;
//                   diff[index_shift+li] += std::pow(this->gamma, n-li) * ci * deltas[index_shift+n];
//                 }
                double target_sum = 0;
                for (int n=trajectory.size()-1; n>=li; n--) {
                  target_sum += deltas[index_shift+n];
                  diff[index_shift+li] = target_sum;
                  target_sum *= this->gamma * lambda * (ptheta[index_shift+n]/max_ptheta);
                }
                if(add_v_corrector)
                  sample_weight->at(index_shift+li) = ptheta[index_shift+li]/max_ptheta;
              } else if(offpolicy_strategy == 2) {
                double target_sum = 0;
                for (int n=trajectory.size()-1; n>=li; n--) {
                  target_sum += deltas[index_shift+n];
                  diff[index_shift+li] = target_sum;
                  target_sum *= this->gamma * lambda * (ptheta[index_shift+n]/trajectory[n].dpmu);
                }
                if(add_v_corrector)
                  sample_weight->at(index_shift+li) = ptheta[index_shift+li]/trajectory[li].dpmu;
              } else if(offpolicy_strategy == 3) {
                double target_sum = 0;
                for (int n=trajectory.size()-1; n>=li; n--) {
                  target_sum += deltas[index_shift+n];
                  diff[index_shift+li] = target_sum;
                  target_sum *= this->gamma * lambda * std::min((double)1.f,ptheta[index_shift+n]/trajectory[n].dpmu);
                }
                if(add_v_corrector)
                  sample_weight->at(index_shift+li) = std::min((double)1.f,ptheta[index_shift+li]/trajectory[li].dpmu);
              } else if(offpolicy_strategy == 4) {
                double target_sum = 0;
                for (int n=trajectory.size()-1; n>=li; n--) {
                  target_sum += deltas[index_shift+n];
                  diff[index_shift+li] = target_sum;
                  target_sum *= this->gamma * lambda * (1.f - l2dist(trajectory[n].a, *all_pi, index_shift+n));
                }
                if(add_v_corrector)
                  sample_weight->at(index_shift+li) = 1.f - l2dist(trajectory[li].a, *all_pi, index_shift+li);
              } else if(offpolicy_strategy == 5) {
                double target_sum = 0;
                for (int n=trajectory.size()-1; n>=li; n--) {
                  target_sum += deltas[index_shift+n];
                  diff[index_shift+li] = target_sum;
                  target_sum *= this->gamma * lambda *
                                (1.f - std::min(l2dist(trajectory[n].a, *all_pi, index_shift+n),
                                                l2dist(trajectory[n].pure_a, *all_pi, index_shift+n)));
                }
                if(add_v_corrector)
                  sample_weight->at(index_shift+li) = 1.f - std::min(l2dist(trajectory[li].a, *all_pi, index_shift+li),
                                                      l2dist(trajectory[li].pure_a, *all_pi, index_shift+li));
              }

              li++;
            }
            ASSERT(diff[index_shift+trajectory.size() -1] == deltas[index_shift+trajectory.size() -1], "pb lambda ");
            index_shift += trajectory.size();
          }
          ASSERT(index_shift == all_size, "index pb");

          // comment following lines to compare with the other formula
          for (int i=0; i<all_size; i++)
            diff[i] = diff[i] + all_V->at(i);

//           bib::Logger::PRINT_ELEMENTS(*all_V, "all v ");
//           bib::Logger::PRINT_ELEMENTS(diff, "final diff ");
          if(!offpolicy_critic){
            int size_last_traj = trajectories.back()->transitions->size();
            vnn->increase_batchsize(size_last_traj);

            std::vector<double> subdiff(size_last_traj);
            std::vector<double> subsensors(size_last_traj * this->get_state_size());

            const std::deque<sample>& trajectory = *trajectories.back()->transitions;

            li=0;
            for (auto it : trajectory){
              std::copy(it.s.begin(), it.s.end(), subsensors.begin() + li * this->get_state_size());
              li++;
            }

            int index=size_last_traj-1;
            for(li=all_size - 1;li>=(all_size-size_last_traj);li--){
              subdiff[index]=diff[li];
              index--;
            }
            vnn->learn_batch(subsensors, empty_action, subdiff, stoch_iter_critic);

            vnn->increase_batchsize(all_size);
          } else if(offpolicy_strategy != 0 && add_v_corrector)
            vnn->learn_batch_lw(all_states, empty_action, diff, *sample_weight, stoch_iter_critic);
          else
            vnn->learn_batch(all_states, empty_action, diff, stoch_iter_critic);

          if(offpolicy_strategy != 0 && offpolicy_critic) {
            delete[] ptheta;
            delete all_pi;
            if(add_v_corrector)
              delete sample_weight;
          }

          delete all_V;
        }

        delete all_nextV;
      };

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

  void end_episode(bool learning) override {
    //     LOG_FILE("policy_exploration", ann->hash());
    if(!learning)
      return;

    ann->update_best_param_previous_task(sum_weighted_reward);
    vnn->update_best_param_previous_task(sum_weighted_reward);
    ann_testing->update_best_param_previous_task(sum_weighted_reward);

    uint all_size = alltransitions();
    if(all_size > 0) {
      vnn->increase_batchsize(all_size);
    }

    for(uint fi = 0 ; fi < number_global_fitted_iteration; fi++){
      if(update_critic_first)
        update_critic();

      if(!offpolicy_actor){
        if(a3c)
          actor_update_onpolicy_a3c();
        else
          actor_update_onpolicy();
      } else
        actor_update_offpolicy();

      if(!update_critic_first) {
        if(all_size > 0 && !offpolicy_actor) {
          vnn->increase_batchsize(all_size);
        }
        update_critic();
      }
    }

    if(shuffle_buffer){
      std::random_shuffle(trajectories.begin(), trajectories.end());
    }
  }

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

  void actor_update_onpolicy() {
    const std::deque<sample>& trajectory = *trajectories.back()->transitions;
    vnn->increase_batchsize(trajectory.size());
    if (trajectory.size() > 0) {
      std::vector<double> sensors(trajectory.size() * this->get_state_size());
      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() * this->get_state_size());
      std::vector<double> all_next_states(trajectory.size() * this->get_state_size());
      uint li=0;
      for (auto it : trajectory) {
        std::copy(it.s.begin(), it.s.end(), all_states.begin() + li * this->get_state_size());
        std::copy(it.next_s.begin(), it.next_s.end(), all_next_states.begin() + li * this->get_state_size());
        li++;
      }

      decltype(vnn->computeOutVFBatch(all_next_states, empty_action)) all_nextV, all_mine;
      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=0;
        for (auto it : trajectory) {
          diff[li] = 0;
          for (uint n=li; n<trajectory.size(); n++)
            diff[li] += std::pow(this->gamma * lambda, n-li) * deltas[n];
          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;
      for(auto it = trajectory.begin(); it != trajectory.end() ; ++it) {
        sample sm = *it;

        std::copy(it->s.begin(), it->s.end(), sensors.begin() + li * this->get_state_size());
        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++;
      }

      ann->increase_batchsize(trajectory.size());
      if(n > 0) {
        for(uint sia = 0; sia < stoch_iter_actor; sia++) {
          //learn BN
          std::vector<double>* ac_out = ann->computeOutBatch(sensors);
          if(batch_norm_actor != 0){
            //re-compute ac_out with BN as testing
            ann_testing->copyParametersFrom(ann);
            ann_testing->increase_batchsize(trajectory.size());
            delete ac_out;
            ac_out = ann_testing->computeOutBatch(sensors);
            if(sia + 1 == stoch_iter_actor )
              ann_testing->increase_batchsize(1);
          }

          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);
              ac_diff[i] = -x;
            }
          }
          ann->actor_backward();
          ann->regularize();
          ann->getSolver()->ApplyUpdate();
          ann->getSolver()->set_iter(ann->getSolver()->iter() + 1);
          delete ac_out;
        }
      }

      delete all_nextV;
      delete all_mine;
    }
  }

  void actor_update_onpolicy_a3c() {
    const std::deque<sample>& trajectory = *trajectories.back()->transitions;
    vnn->increase_batchsize(trajectory.size());
    if (trajectory.size() > 0) {
      std::vector<double> sensors(trajectory.size() * this->get_state_size());
      std::vector<double> actions(trajectory.size() * this->nb_motors);
      std::vector<double> deltas_blob(trajectory.size() * this->nb_motors);
      std::vector<double> deltas(trajectory.size());

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

      decltype(vnn->computeOutVFBatch(all_next_states, empty_action)) all_nextV, all_mine;
      all_nextV = vnn->computeOutVFBatch(all_next_states, empty_action);
      all_mine = vnn->computeOutVFBatch(all_states, empty_action);

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

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

      li=0;
      for(auto it = trajectory.begin(); it != trajectory.end() ; ++it) {
        sample sm = *it;

        std::copy(it->s.begin(), it->s.end(), sensors.begin() + li * this->get_state_size());
        std::copy(it->a.begin(), it->a.end(), actions.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++;
      }

      for(uint sia = 0; sia < stoch_iter_actor; sia++) {
        ann->increase_batchsize(trajectory.size());
        auto ac_out = ann->computeOutBatch(sensors);
        if(batch_norm_actor != 0){
          //re-compute ac_out with BN as testing
          ann_testing->copyParametersFrom(ann);
          ann_testing->increase_batchsize(trajectory.size());
          delete ac_out;
          ac_out = ann_testing->computeOutBatch(sensors);
          if(sia + 1 == stoch_iter_actor )
            ann_testing->increase_batchsize(1);
        }

        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++) {
            double x = actions[i] - ac_out->at(i);
            ac_diff[i] = -x*deltas_blob[i];
        }
        ann->actor_backward();
        ann->regularize();
        ann->getSolver()->ApplyUpdate();
        ann->getSolver()->set_iter(ann->getSolver()->iter() + 1);
        delete ac_out;
      }

      delete all_nextV;
      delete all_mine;
    }
  }

  void actor_update_offpolicy() {
    int all_size = alltransitions();
    ann->increase_batchsize(all_size);

    std::vector<double> sensors(all_size * this->get_state_size());
    std::vector<double> actions(all_size * this->nb_motors);
    std::vector<bool> disable_back(all_size * this->nb_motors, false);
    const std::vector<bool> disable_back_ac(this->nb_motors, true);
    std::vector<double> deltas_blob(all_size * this->nb_motors);
    std::vector<double> deltas(all_size);

    std::vector<double> all_states(all_size * this->get_state_size());
    std::vector<double> all_next_states(all_size * this->get_state_size());

    int li=0;
    for(auto one_trajectory : trajectories) {
      const std::deque<sample>& trajectory = *one_trajectory->transitions;
      for (auto it : trajectory) {
        std::copy(it.s.begin(), it.s.end(), all_states.begin() + li * this->get_state_size());
        std::copy(it.next_s.begin(), it.next_s.end(), all_next_states.begin() + li * this->get_state_size());
        li++;
      }
    }

    decltype(vnn->computeOutVFBatch(all_next_states, empty_action)) all_nextV, all_mine;
    all_nextV = vnn->computeOutVFBatch(all_next_states, empty_action);
    all_mine = vnn->computeOutVFBatch(all_states, empty_action);

    li=0;
    for(auto one_trajectory : trajectories) {
      const std::deque<sample>& trajectory = *one_trajectory->transitions;
      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>* all_pi;
      std::vector<double> diff(all_size);
      int index_shift=0;
      double* ptheta;
      double max_ptheta;
      if(offpolicy_strategy != 0) {
        if(testing_ann_offpolicy){
          ann_testing->copyParametersFrom(ann);
          ann_testing->increase_batchsize(all_size);
          all_pi = ann_testing->computeOutBatch(all_states);
          ann_testing->increase_batchsize(1);
        } else
          all_pi = ann->computeOutBatch(all_states);
        ptheta = new double[all_size];

        uint i=0;
        if(offpolicy_strategy >= 1 && offpolicy_strategy <= 3) {
          for(auto one_trajectory : trajectories) {
            const std::deque<sample>& trajectory = *one_trajectory->transitions;
            for (auto it : trajectory) {
              double p0 = 1.f;
              for(uint j=0; j < this->nb_motors; j++)
                p0 *= bib::Proba<double>::truncatedGaussianDensity(it.a[j], all_pi->at(i*this->nb_motors+j), noise);

              ptheta[i] = p0;
              i++;
            }
          }
          max_ptheta = *std::max_element(ptheta, ptheta+all_size);
        }
      }
      for(auto one_trajectory : trajectories) {
        li=0;
        const std::deque<sample>& trajectory = *one_trajectory->transitions;
        for (auto it : trajectory) {
          diff[index_shift+li] = 0;
          if(offpolicy_strategy == 0) {
            for (uint n=li; n<trajectory.size(); n++)
              diff[index_shift+li] += std::pow(this->gamma * lambda, n-li) * deltas[index_shift+n];
          } else if(offpolicy_strategy == 1) {
            double target_sum = 0;
            for (int n=trajectory.size()-1; n>=li; n--) {
              target_sum += deltas[index_shift+n];
              diff[index_shift+li] = target_sum;
              target_sum *= this->gamma * lambda * (ptheta[index_shift+n]/max_ptheta);
            }
            if(add_v_corrector)
              diff[index_shift+li] *= ptheta[index_shift+li]/max_ptheta;
          } else if(offpolicy_strategy == 2) {
            double target_sum = 0;
            for (int n=trajectory.size()-1; n>=li; n--) {
              target_sum += deltas[index_shift+n];
              diff[index_shift+li] = target_sum;
              target_sum *= this->gamma * lambda * (ptheta[index_shift+n]/trajectory[n].dpmu);
            }
            if(add_v_corrector)
              diff[index_shift+li] *= ptheta[index_shift+li]/trajectory[li].dpmu;
          } else if(offpolicy_strategy == 3) {
            double target_sum = 0;
            for (int n=trajectory.size()-1; n>=li; n--) {
              target_sum += deltas[index_shift+n];
              diff[index_shift+li] = target_sum;
              target_sum *= this->gamma * lambda * std::min((double)1.f,ptheta[index_shift+n]/trajectory[n].dpmu);
            }
            if(add_v_corrector)
              diff[index_shift+li] *= std::min((double)1.f,ptheta[index_shift+li]/trajectory[li].dpmu);
          } else if(offpolicy_strategy == 4) {
            double target_sum = 0;
            for (int n=trajectory.size()-1; n>=li; n--) {
              target_sum += deltas[index_shift+n];
              diff[index_shift+li] = target_sum;
              target_sum *= this->gamma * lambda * (1.f - l2dist(trajectory[n].a, *all_pi, index_shift+n));
            }
            if(add_v_corrector)
              diff[index_shift+li] *= 1.f - l2dist(trajectory[li].a, *all_pi, index_shift+li);
          } else if(offpolicy_strategy == 5) {
            double target_sum = 0;
            for (int n=trajectory.size()-1; n>=li; n--) {
              target_sum += deltas[index_shift+n];
              diff[index_shift+li] = target_sum;
              target_sum *= this->gamma * lambda *
                            (1.f - std::min(l2dist(trajectory[n].a, *all_pi, index_shift+n),
                                            l2dist(trajectory[n].pure_a, *all_pi, index_shift+n)));
            }
            if(add_v_corrector)
              diff[index_shift+li] *= 1.f - std::min(l2dist(trajectory[li].a, *all_pi, index_shift+li),
                                                     l2dist(trajectory[li].pure_a, *all_pi, index_shift+li));
          }
          li++;
        }
        ASSERT(diff[index_shift+trajectory.size() -1] == deltas[index_shift+trajectory.size() -1], "pb lambda");
        index_shift += trajectory.size();
      }

      if(offpolicy_strategy != 0) {
        delete all_pi;
        delete [] ptheta;
      }

      for (li=0; li < all_size; li++) {
        //           diff[li] = diff[li] + all_V->at(li);
        deltas[li] = diff[li];
      }
    }

    uint n=0;
    li=0;
    for(auto one_trajectory : trajectories) {
      const std::deque<sample>& trajectory = *one_trajectory->transitions;
      for(auto it = trajectory.begin(); it != trajectory.end() ; ++it) {
        std::copy(it->s.begin(), it->s.end(), sensors.begin() + li * this->get_state_size());
        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++;
      }
    }

    if(n > 0) {
      for(uint sia = 0; sia < stoch_iter_actor; sia++) {
        std::vector<double>* ac_out = ann->computeOutBatch(sensors);
        if(batch_norm_actor != 0){
          //re-compute ac_out with BN as testing
          ann_testing->copyParametersFrom(ann);
          ann_testing->increase_batchsize(all_size);
          delete ac_out;
          ac_out = ann_testing->computeOutBatch(sensors);
          if(sia + 1 == stoch_iter_actor )
            ann_testing->increase_batchsize(1);
        }

        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);
            ac_diff[i] = -x;
          }
        }
        ann->actor_backward();
        ann->regularize();
        ann->getSolver()->ApplyUpdate();
        ann->getSolver()->set_iter(ann->getSolver()->iter() + 1);
        delete ac_out;
      }
    }

    delete all_nextV;
    delete all_mine;
  }

  void save(const std::string& path, bool save_best, bool learning) override {
    if(save_best && !learning && best_testing_score < this->sum_weighted_reward){
      ann->save(path+".actor");
      best_testing_score = this->sum_weighted_reward;
    } else if(!save_best && learning){
      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;
  }

 protected:

  void _display(std::ostream& out) const override {
    out << std::setw(12) << std::fixed << std::setprecision(10) << this->sum_weighted_reward << " " << std::setw(
          8) << std::fixed << std::setprecision(5) << vnn->error() << " " << noise << " " << alltransitions();
  }

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

  uint alltransitions() const {
    uint count=0;
    for (auto it : trajectories)
      count += it->transitions->size();
    return count;
  }

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

  double l2dist(const std::vector<double>& a, const std::vector<double>& b, const uint start) const {
    double r = 0.f;
    for(uint i=0; i<a.size(); i++) {
      double d = (a[i] - b[start+i]);
      r += d*d;
    }
    return sqrt(r)/(2.f *((double) a.size()));
  }

 private:
  uint episode = 0;

  double noise;
  bool gaussian_policy, vnn_from_scratch, update_critic_first,
       update_delta_neg, gae, add_v_corrector, offpolicy_actor;
  uint number_fitted_iteration, stoch_iter_actor, stoch_iter_critic;
  uint batch_norm_actor, batch_norm_critic, actor_output_layer_type,
       hidden_layer_type, max_trajectory, offpolicy_strategy;
  double lambda;
  uint number_global_fitted_iteration;
  bool offpolicy_critic, a3c, shuffle_buffer, testing_ann_offpolicy;

  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<std::shared_ptr<trajectory>> trajectories;

  NN* ann;
  NN* vnn;
  NN* ann_testing;
  //to make video on determinist env
  double best_testing_score;

  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

  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


#include <Eigen/Dense>
#include "control.hpp"
#include "defines.hpp"
#include "fad.hpp"
#include "global_residual.hpp"
#include "J2.hpp"
#include "material_params.hpp"

namespace calibr8 {

static ParameterList get_valid_local_residual_params() {
  ParameterList p;
  p.set<std::string>("type", "J2");
  p.set<int>("nonlinear max iters", 0);
  p.set<double>("nonlinear absolute tol", 0.);
  p.set<double>("nonlinear relative tol", 0.);
  p.sublist("materials");
  return p;
}

static ParameterList get_valid_material_params() {
  ParameterList p;
  p.set<double>("E", 0.);
  p.set<double>("nu", 0.);
  p.set<double>("K", 0.);
  p.set<double>("Y", 0.);
  return p;
}

template <typename T>
J2<T>::J2(ParameterList const& inputs, int ndims) {

  this->m_params_list = inputs;
  this->m_params_list.validateParameters(get_valid_local_residual_params(), 0);

  int const num_residuals = 3;

  this->m_num_residuals = num_residuals;
  this->m_num_eqs.resize(num_residuals);
  this->m_var_types.resize(num_residuals);
  this->m_resid_names.resize(num_residuals);

  this->m_resid_names[0] = "zeta";
  this->m_var_types[0] = SYM_TENSOR;
  this->m_num_eqs[0] = get_num_eqs(SYM_TENSOR, ndims);

  this->m_resid_names[1] = "Ie";
  this->m_var_types[1] = SCALAR;
  this->m_num_eqs[1] = get_num_eqs(SCALAR, ndims);

  this->m_resid_names[2] = "alpha";
  this->m_var_types[2] = SCALAR;
  this->m_num_eqs[2] = get_num_eqs(SCALAR, ndims);

  m_max_iters = inputs.get<int>("nonlinear max iters");
  m_abs_tol = inputs.get<double>("nonlinear absolute tol");
  m_rel_tol = inputs.get<double>("nonlinear relative tol");

}

template <typename T>
J2<T>::~J2() {
}

template <typename T>
void J2<T>::init_params() {

  int const num_params = 4;
  this->m_params.resize(num_params);
  this->m_param_names.resize(num_params);

  this->m_param_names.resize(num_params);
  this->m_param_names[0] = "E";
  this->m_param_names[1] = "nu";
  this->m_param_names[2] = "K";
  this->m_param_names[3] = "Y";

  int const num_elem_sets = this->m_elem_set_names.size();
  resize(this->m_param_values, num_elem_sets, num_params);

  ParameterList& all_material_params =
      this->m_params_list.sublist("materials", true);

  for (int es = 0; es < num_elem_sets; ++es) {
    std::string const& elem_set_name = this->m_elem_set_names[es];
    ParameterList& material_params =
        all_material_params.sublist(elem_set_name, true);
    material_params.validateParameters(get_valid_material_params(), 0);
    this->m_param_values[es][0] = material_params.get<double>("E");
    this->m_param_values[es][1] = material_params.get<double>("nu");
    this->m_param_values[es][2] = material_params.get<double>("K");
    this->m_param_values[es][3] = material_params.get<double>("Y");
  }

  this->m_active_indices.resize(1);
  this->m_active_indices[0].resize(1);
  this->m_active_indices[0][0] = 0;
}

template <typename T>
void J2<T>::init_variables_impl() {

  int const ndims = this->m_num_dims;
  int const zeta_idx = 0;
  int const Ie_idx = 1;
  int const alpha_idx = 2;

  T const Ie = 1.0;
  T const alpha = 0.0;
  Tensor<T> const zeta = minitensor::zero<T>(ndims);

  this->set_scalar_xi(Ie_idx, Ie);
  this->set_scalar_xi(alpha_idx, alpha);
  this->set_sym_tensor_xi(zeta_idx, zeta);

}

template <typename T>
Tensor<T> eval_be_bar(
    RCP<GlobalResidual<T>> global,
    Tensor<T> const& zeta,
    T const& Ie) {
  int const ndims = global->num_dims();
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const grad_u = global->grad_vector_x(0);
  Tensor<T> const grad_u_prev = global->grad_vector_x_prev(0);
  Tensor<T> const F = grad_u + I;
  Tensor<T> const F_prev = grad_u_prev + I;
  Tensor<T> const rF = F * minitensor::inverse(F_prev);
  T const det_rF = minitensor::det(rF);
  T const det_rF_13 = cbrt(det_rF);
  Tensor<T> const rF_bar = rF / det_rF_13;
  Tensor<T> const rF_barT = minitensor::transpose(rF_bar);
  Tensor<T> const be_bar = rF_bar * (zeta + Ie * I) * rF_barT;
  return be_bar;
}

template <>
int J2<double>::solve_nonlinear(RCP<GlobalResidual<double>>) {
  return 0;
}

template <>
int J2<FADT>::solve_nonlinear(RCP<GlobalResidual<FADT>> global) {

  int path;

  // pick an initial guess for the local variables
  {
    Tensor<FADT> const zeta_old = this->sym_tensor_xi_prev(0);
    FADT const Ie_old = this->scalar_xi_prev(1);
    FADT const alpha_old = this->scalar_xi_prev(2);
    Tensor<FADT> const be_bar_trial = eval_be_bar(global, zeta_old, Ie_old);
    Tensor<FADT> const zeta = minitensor::dev(be_bar_trial);
    FADT const Ie = minitensor::trace(be_bar_trial) / 3.;
    FADT const alpha = alpha_old;
    this->set_sym_tensor_xi(0, zeta);
    this->set_scalar_xi(1, Ie);
    this->set_scalar_xi(2, alpha);
    path = ELASTIC;
  }

  // newton iteration until convergence

  int iter = 1;
  double R_norm_0 = 1.;
  bool converged = false;

  while ((iter <= m_max_iters) && (!converged)) {

    path = this->evaluate(global);

    double const R_norm = this->norm_residual();
    if (iter == 1) R_norm_0 = R_norm;
    double const R_norm_rel = R_norm / R_norm_0;
    if ((R_norm_rel < m_rel_tol) || (R_norm < m_abs_tol)) {
      converged = true;
      break;
    }

    EMatrix const J = this->eigen_jacobian();
    EVector const R = this->eigen_residual();
    EVector const dxi = J.fullPivLu().solve(-R);

    this->add_to_sym_tensor_xi(0, dxi);
    this->add_to_scalar_xi(1, dxi);
    this->add_to_scalar_xi(2, dxi);

    iter++;

  }

  // fail if convergence was not achieved
  if ((iter > m_max_iters) && (!converged)) {
    fail("J2:solve_nonlinear failed in %d iterations", m_max_iters);
  }

  return path;

}

template <typename T>
int J2<T>::evaluate(
    RCP<GlobalResidual<T>> global,
    bool force_path,
    int path_in) {

  int path = ELASTIC;
  int const ndims = this->m_num_dims;
  double const sqrt_23 = std::sqrt(2./3.);
  double const sqrt_32 = std::sqrt(3./2.);

  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const K = this->m_params[2];
  T const Y = this->m_params[3];
  T const mu = compute_mu(E, nu);

  Tensor<T> const zeta_old = this->sym_tensor_xi_prev(0);
  T const Ie_old = this->scalar_xi_prev(1);
  T const alpha_old = this->scalar_xi_prev(2);

  Tensor<T> const zeta = this->sym_tensor_xi(0);
  T const Ie = this->scalar_xi(1);
  T const alpha = this->scalar_xi(2);

  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const be_bar_trial = eval_be_bar(global, zeta_old, Ie_old);
  Tensor<T> const s = mu * zeta;
  T const s_mag = minitensor::norm(s);
  Tensor<T> const n = s / s_mag;
  T const sigma_yield = Y + K * alpha;
  T const f = s_mag - sqrt_23 * sigma_yield;

  Tensor<T> R_zeta;
  T R_Ie;
  T R_alpha;

  if (!force_path) {
    // plastic step
    if (f > m_abs_tol || std::abs(f) < m_abs_tol) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_zeta = zeta - minitensor::dev(be_bar_trial) + 2. * dgam * Ie * n;
      R_Ie = minitensor::det(zeta + Ie * I) - 1.;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu);
      path = PLASTIC;
    }
    // elastic step
    else {
      R_zeta = (0. * mu + 1.) * zeta - minitensor::dev(be_bar_trial);
      R_Ie = Ie - minitensor::trace(be_bar_trial) / 3. + 0. * mu;
      R_alpha = alpha - alpha_old + 0. * mu;
      path = ELASTIC;
    }
  }

  // force the path
  else {
    path = path_in;
    // plastic step
    if (path == PLASTIC) {
      T const dgam = sqrt_32 * (alpha - alpha_old);
      R_zeta = zeta - minitensor::dev(be_bar_trial) + 2. * dgam * Ie * n;
      R_Ie = minitensor::det(zeta + Ie * I) - 1.;
      R_alpha = (s_mag - sqrt_23 * sigma_yield) / val(mu);
    }
    // elastic step
    else {
      R_zeta = (0. * mu + 1.) * zeta - minitensor::dev(be_bar_trial);
      R_Ie = Ie - minitensor::trace(be_bar_trial) / 3. + 0. * mu;
      R_alpha = alpha - alpha_old + 0. * mu;
    }
  }

  this->set_sym_tensor_R(0, R_zeta);
  this->set_scalar_R(1, R_Ie);
  this->set_scalar_R(2, R_alpha);

  return path;

}

template <typename T>
Tensor<T> J2<T>::dev_cauchy(RCP<GlobalResidual<T>> global) {
  int const ndims = global->num_dims();
  T const E = this->m_params[0];
  T const nu = this->m_params[1];
  T const mu = E / (2. * (1. + nu));
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const grad_u = global->grad_vector_x(0);
  Tensor<T> const F = grad_u + I;
  Tensor<T> const zeta = this->sym_tensor_xi(0);
  T const J = minitensor::det(F);
  return mu * zeta / J;
}

template <typename T>
Tensor<T> J2<T>::cauchy(RCP<GlobalResidual<T>> global, T p) {
  int const ndims = global->num_dims();
  Tensor<T> const I = minitensor::eye<T>(ndims);
  Tensor<T> const dev_sigma = this->dev_cauchy(global);
  Tensor<T> const sigma = dev_sigma - p * I;
  return sigma;
}

template class J2<double>;
template class J2<FADT>;

}


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

#include "bempp/common/config_ahmed.hpp"
#include "bempp/common/config_trilinos.hpp"

#include "aca_global_assembler.hpp"

#include "assembly_options.hpp"
#include "cluster_construction_helper.hpp"
#include "evaluation_options.hpp"
#include "index_permutation.hpp"
#include "discrete_boundary_operator_composition.hpp"
#include "discrete_sparse_boundary_operator.hpp"

#include "../common/armadillo_fwd.hpp"
#include "../common/auto_timer.hpp"
#include "../common/boost_shared_array_fwd.hpp"
#include "../common/chunk_statistics.hpp"
#include "../common/to_string.hpp"
#include "../fiber/explicit_instantiation.hpp"
#include "../fiber/local_assembler_for_operators.hpp"
#include "../fiber/local_assembler_for_potential_operators.hpp"
#include "../fiber/serial_blas_region.hpp"
#include "../fiber/scalar_traits.hpp"
#include "../space/space.hpp"

#include <stdexcept>
#include <fstream>
#include <iostream>

#include <boost/type_traits/is_complex.hpp>

#include <tbb/atomic.h>
#include <tbb/parallel_for.h>
#include <tbb/task_scheduler_init.h>
#include <tbb/concurrent_queue.h>

#ifdef WITH_AHMED
#include "ahmed_aux.hpp"

#ifdef __INTEL_COMPILER
#pragma warning(disable:381)
#endif

#include <apprx.h>

#ifdef __INTEL_COMPILER
#pragma warning(default:381)
#endif

#include "discrete_aca_boundary_operator.hpp"
#include "modified_aca.hpp"
#include "potential_operator_aca_assembly_helper.hpp"
#include "scattered_range.hpp"
#include "weak_form_aca_assembly_helper.hpp"
#endif

// #define DUMP_DENSE_BLOCKS // if defined, contents and DOF lists of blocks
                             // stored as dense matrices will be printed to the
                             // screen

namespace Bempp
{

// Body of parallel loop
namespace
{

#ifdef WITH_AHMED
template <typename BasisFunctionType, typename ResultType,
          typename AcaAssemblyHelper>
class AcaAssemblerLoopBody
{
    typedef typename Fiber::ScalarTraits<ResultType>::RealType CoordinateType;
    typedef AhmedDofWrapper<CoordinateType> AhmedDofType;
    typedef bemblcluster<AhmedDofType, AhmedDofType> AhmedBemBlcluster;
    typedef mblock<typename AhmedTypeTraits<ResultType>::Type> AhmedMblock;
public:
    typedef tbb::concurrent_queue<size_t> LeafClusterIndexQueue;

    AcaAssemblerLoopBody(
            AcaAssemblyHelper& helper,
            AhmedLeafClusterArray& leafClusters,
            boost::shared_array<AhmedMblock*> blocks,
            const AcaOptions& options,
            tbb::atomic<size_t>& done,
            bool verbose,
            LeafClusterIndexQueue& leafClusterIndexQueue,
            bool symmetric,
            std::vector<ChunkStatistics>& stats) :
        m_helper(helper),
        m_leafClusters(leafClusters), m_blocks(blocks),
        m_options(options), m_done(done), m_verbose(verbose),
        m_leafClusterIndexQueue(leafClusterIndexQueue),
        m_symmetric(symmetric),
        m_stats(stats)
    {
    }

    template <typename Range>
    void operator() (const Range& r) const {
        const char* TEXT = "Approximating ... ";
        for (typename Range::const_iterator i = r.begin(); i != r.end(); ++i) {
            size_t leafClusterIndex = 0;
            if (!m_leafClusterIndexQueue.try_pop(leafClusterIndex)) {
                std::cerr << "AcaWeakFormAssemblerLoopBody::operator(): "
                             "Warning: try_pop failed; this shouldn't happen!"
                          << std::endl;
                continue;
            }
            m_stats[leafClusterIndex].valid = true;
            m_stats[leafClusterIndex].chunkStart = r.begin();
            m_stats[leafClusterIndex].chunkSize = r.size();
            m_stats[leafClusterIndex].startTime = tbb::tick_count::now();

            AhmedBemBlcluster* cluster =
                    dynamic_cast<AhmedBemBlcluster*>(m_leafClusters[leafClusterIndex]);
            if (m_symmetric)
                apprx_sym(m_helper, m_blocks[cluster->getidx()],
                          cluster, m_options.eps, m_options.maximumRank,
                          true /* complex_sym */);
            else {
                if (m_options.useAhmedAca)
                    apprx_unsym(m_helper, m_blocks[cluster->getidx()],
                                cluster, m_options.eps, m_options.maximumRank);
                else
                    apprx_unsym_shooting(
                                m_helper, m_blocks[cluster->getidx()],
                                cluster, m_options.eps, m_options.maximumRank);
            }
            m_stats[leafClusterIndex].endTime = tbb::tick_count::now();
            // TODO: recompress
            const int HASH_COUNT = 20;
            if (m_verbose)
                progressbar(std::cout, TEXT, (++m_done) - 1,
                            m_leafClusters.size(), HASH_COUNT, true);
        }

    }

private:
    AcaAssemblyHelper& m_helper;
    AhmedLeafClusterArray& m_leafClusters;
    boost::shared_array<AhmedMblock*> m_blocks;
    const AcaOptions& m_options;
    tbb::atomic<size_t>& m_done;
    bool m_verbose;
    LeafClusterIndexQueue& m_leafClusterIndexQueue;
    bool m_symmetric;
    std::vector<ChunkStatistics>& m_stats;
};

void reallyGetClusterIds(const cluster& clusterTree,
                         const std::vector<unsigned int>& p2oDofs,
                         std::vector<unsigned int>& clusterIds,
                         unsigned int& id)
{
    if (clusterTree.isleaf())
        for (unsigned int nDof = clusterTree.getnbeg(); nDof < clusterTree.getnend(); ++nDof)
            clusterIds[p2oDofs[nDof]] = id;
    else
        for (unsigned int nSon = 0; nSon < clusterTree.getns(); ++nSon)
            reallyGetClusterIds(*clusterTree.getson(nSon), p2oDofs, clusterIds, ++id);
}

void getClusterIds(const cluster& clusterTree,
                   const std::vector<unsigned int>& p2oDofs,
                   std::vector<unsigned int>& clusterIds)
{
    clusterIds.resize(p2oDofs.size());
    unsigned int id = 0;
    reallyGetClusterIds(clusterTree, p2oDofs, clusterIds, id);
}

template <typename T>
void save_arma_matrix(const arma::Mat<T>& a, const std::string& fname)
{
    std::ofstream out;
    out.precision(17);
    out.open(fname.c_str());
    arma::diskio::save_raw_ascii(a, out);
    out.close();
}
template <typename T>
void save_arma_matrix(const arma::Mat<std::complex<T> >& a, const std::string& fname)
{
    std::ofstream out;
    out.precision(17);
    out.open((fname + "-real.txt").c_str());
    arma::Mat<T> component = arma::real(a);
    arma::diskio::save_raw_ascii(component, out);
    out.close();
    out.open((fname + "-imag.txt").c_str());
    component = arma::imag(a);
    arma::diskio::save_raw_ascii(component, out);
    out.close();
}

template <typename ValueType>
void dumpDenseBlocks(
        typename DiscreteAcaBoundaryOperator<ValueType>::AhmedBemBlcluster* clusterTree,
        typename DiscreteAcaBoundaryOperator<ValueType>::AhmedMblockArray& blocks,
        const std::vector<unsigned int>& p2oRows,
        const std::vector<unsigned int>& p2oCols,
        const std::vector<Point3D<typename Fiber::ScalarTraits<ValueType>::RealType> >& rowDofs,
        const std::vector<Point3D<typename Fiber::ScalarTraits<ValueType>::RealType> >& colDofs)
{
    if (!clusterTree)
        return;
    typedef typename Fiber::ScalarTraits<ValueType>::RealType CoordinateType;
    typedef DiscreteAcaBoundaryOperator<ValueType> AcaOp;
    typedef typename AcaOp::AhmedDofType AhmedDofType;
    typedef typename AcaOp::AhmedMblock AhmedMblock;
    typedef bemcluster<AhmedDofType> Cluster;
    if (clusterTree->isleaf()) {
        unsigned int idx = clusterTree->getidx();
        std::cout << "LEAF; idx = " << idx << std::endl;
        if (clusterTree->isGeM(blocks.get()) && clusterTree->isadm()) {
            std::cout << "Dense block; " << clusterTree->getb1()
                      << " " << clusterTree->getb2()
                      << " " << clusterTree->getn1()
                      << " " << clusterTree->getn2() << "\n";
            // if (clusterTree->getn1() < 500 || clusterTree->getn2() < 500)
            //     return;
            Cluster* clRow = clusterTree->getcl1();
            assert(clRow);
            std::cout << "Row center of mass: ("
                      << clRow->getcom(0) << ", "
                      << clRow->getcom(1) << ", "
                      << clRow->getcom(2) << ")\n";
            std::cout << "Row icm: "
                      << clRow->geticom() - clusterTree->getb1() << std::endl;
            for (unsigned int nDof = clRow->getnbeg();
                 nDof < clRow->getnend(); ++nDof) {
                assert(nDof < p2oRows.size());
                assert(p2oRows[nDof] < rowDofs.size());
                const Point3D<CoordinateType> dofPos = rowDofs[p2oRows[nDof]];
                std::cout << "  Row dof #" << p2oRows[nDof] << " at "
                          << dofPos.x << ", " << dofPos.y << ", "
                          << dofPos.z << "\n";
            }
            Cluster* clCol = clusterTree->getcl2();
            assert(clCol);
            std::cout << "Col center of mass: ("
                      << clCol->getcom(0) << ", "
                      << clCol->getcom(1) << ", "
                      << clCol->getcom(2) << ")\n";
            std::cout << "Column icm: "
                      << clCol->geticom() - clusterTree->getb2() << std::endl;
            for (unsigned int nDof = clCol->getnbeg();
                 nDof < clCol->getnend(); ++nDof) {
                assert(nDof < p2oCols.size());
                assert(p2oCols[nDof] < colDofs.size());
                const Point3D<CoordinateType> dofPos = colDofs[p2oCols[nDof]];
                std::cout << "  Col dof #" << p2oCols[nDof] << " at "
                          << dofPos.x << ", " << dofPos.y << ", "
                          << dofPos.z << "\n";
            }
            AhmedMblock* block = blocks[idx];
            arma::Mat<ValueType> ablock(clusterTree->getn1(),
                                        clusterTree->getn2());
            for (size_t i = 0; i < block->nvals(); ++i)
                ablock[i] = block->getdata()[i];
            save_arma_matrix(ablock, "block-" + toString(idx) + ".txt");
        }
    }
    else
        for (unsigned int nRowSon = 0; nRowSon < clusterTree->getnrs(); ++nRowSon)
            for (unsigned int nColSon = 0; nColSon < clusterTree->getncs(); ++nColSon)
            dumpDenseBlocks<ValueType>(
                        dynamic_cast<typename AcaOp::AhmedBemBlcluster*>(
                            clusterTree->getson(nRowSon, nColSon)),
                        blocks,
                        p2oRows, p2oCols, rowDofs, colDofs);
}

template <typename AcaAssemblyHelper,
          typename BasisFunctionType, typename ResultType>
std::auto_ptr<DiscreteAcaBoundaryOperator<ResultType> >
assembleAcaOperator(
        AcaAssemblyHelper& helper,
        const shared_ptr<typename DiscreteAcaBoundaryOperator<ResultType>::
            AhmedBemBlcluster>& bemBlclusterTree,
        const ParallelizationOptions& parallelOptions,
        const AcaOptions& acaOptions,
        bool verbosityAtLeastDefault,
        bool symmetric,
        const shared_ptr<IndexPermutation>& test_o2pPermutation,
        const shared_ptr<IndexPermutation>& trial_o2pPermutation
#ifdef DUMP_DENSE_BLOCKS
        ,
        const shared_ptr<IndexPermutation>& test_p2oPermutation,
        const shared_ptr<IndexPermutation>& trial_p2oPermutation,
        const std::vector<Point3D<
            typename Fiber::ScalarTraits<ResultType>::RealType> >& testDofCenters,
        const std::vector<Point3D<
            typename Fiber::ScalarTraits<ResultType>::RealType> >& trialDofCenters
#endif // DUMP_DENSE_BLOCKS
        )
{
    typedef mblock<typename AhmedTypeTraits<ResultType>::Type> AhmedMblock;
    boost::shared_array<AhmedMblock*> blocks =
            allocateAhmedMblockArray<ResultType>(bemBlclusterTree.get());

    const size_t testDofCount = test_o2pPermutation->size();
    const size_t trialDofCount = trial_o2pPermutation->size();

    AhmedLeafClusterArray leafClusters(bemBlclusterTree.get());
    leafClusters.sortAccordingToClusterSize();
    const size_t leafClusterCount = leafClusters.size();

    int maxThreadCount = 1;
    if (!parallelOptions.isOpenClEnabled())
    {
        if (parallelOptions.maxThreadCount() == ParallelizationOptions::AUTO)
            maxThreadCount = tbb::task_scheduler_init::automatic;
        else
            maxThreadCount = parallelOptions.maxThreadCount();
    }
    tbb::task_scheduler_init scheduler(maxThreadCount);
    tbb::atomic<size_t> done;
    done = 0;

    std::vector<ChunkStatistics> chunkStats(leafClusterCount);

    typedef AcaAssemblerLoopBody<
            BasisFunctionType, ResultType, AcaAssemblyHelper> Body;
    typename Body::LeafClusterIndexQueue leafClusterIndexQueue;
    for (size_t i = 0; i < leafClusterCount; ++i)
        leafClusterIndexQueue.push(i);

    if (verbosityAtLeastDefault)
        std::cout << "About to start the ACA assembly loop" << std::endl;
    tbb::tick_count loopStart = tbb::tick_count::now();
    {
        Fiber::SerialBlasRegion region; // if possible, ensure that BLAS is single-threaded
        tbb::parallel_for(tbb::blocked_range<size_t>(0, leafClusterCount),
                          Body(helper, leafClusters, blocks, acaOptions, done,
                               verbosityAtLeastDefault,
                               leafClusterIndexQueue, symmetric, chunkStats));
    }
    tbb::tick_count loopEnd = tbb::tick_count::now();
    if (verbosityAtLeastDefault) {
        std::cout << "\n"; // the progress bar doesn't print the final \n
        std::cout << "ACA loop took " << (loopEnd - loopStart).seconds() << " s"
                  << std::endl;
    }

    // TODO: parallelise!
    if (acaOptions.recompress) {
        if (verbosityAtLeastDefault)
            std::cout << "About to start ACA agglomeration" << std::endl;
        agglH(bemBlclusterTree.get(), blocks.get(),
              acaOptions.eps, acaOptions.maximumRank);
        if (verbosityAtLeastDefault)
            std::cout << "Agglomeration finished" << std::endl;
    }

#ifdef DUMP_DENSE_BLOCKS
    dumpDenseBlocks<ResultType>(bemBlclusterTree.get(), blocks,
                                test_p2oPermutation->permutedIndices(),
                                trial_p2oPermutation->permutedIndices(),
                                testDofCenters, trialDofCenters);
#endif // DUMP_DENSE_BLOCKS

    //    dumpAhmedMblockArray<ResultType>(blocks, blockCount);

    // // Dump timing data of individual chunks
    //    std::cout << "\nChunks:\n";
    //    for (int i = 0; i < leafClusterCount; ++i)
    //        if (chunkStats[i].valid) {
    //            int blockIndex = leafClusters[i]->getidx();
    //            std::cout << chunkStats[i].chunkStart << "\t"
    //                      << chunkStats[i].chunkSize << "\t"
    //                      << (chunkStats[i].startTime - loopStart).seconds() << "\t"
    //                      << (chunkStats[i].endTime - loopStart).seconds() << "\t"
    //                      << (chunkStats[i].endTime - chunkStats[i].startTime).seconds() << "\t"
    //                      << blocks[blockIndex]->getn1() << "\t"
    //                      << blocks[blockIndex]->getn2() << "\t"
    //                      << blocks[blockIndex]->islwr() << "\t"
    //                      << (blocks[blockIndex]->islwr() ? blocks[blockIndex]->rank() : 0) << "\n";
    //        }

    if (verbosityAtLeastDefault) {
#ifdef CHECK_ACA_ERROR // a define from include/AHMED/apprx.h
        std::cout << "ACA finished. Max error: " << ACA_error_max << std::endl;
#endif
        size_t totalEntryCount = testDofCount * trialDofCount;
        size_t origMemory = sizeof(ResultType) * totalEntryCount;
        size_t ahmedMemory = sizeH(bemBlclusterTree.get(), blocks.get());
        int maximumRank = Hmax_rank(bemBlclusterTree.get(), blocks.get());
        size_t accessedEntryCount = helper.accessedEntryCount();
        double accessedFraction = double(accessedEntryCount) / totalEntryCount;
        std::cout << "\nNeeded storage: "
                  << ahmedMemory / 1024. / 1024. << " MB.\n"
                  << "Without approximation: "
                  << origMemory / 1024. / 1024. << " MB.\n"
                  << "Compressed to "
                  << (100. * ahmedMemory) / origMemory << "%.\n"
                  << "Maximum rank: " << maximumRank << ".\n"
                  << "Accessed "
                  << 100. * accessedFraction << "% matrix entries.\n"
                  << std::endl;
    }

    if (acaOptions.outputPostscript) {
        if (verbosityAtLeastDefault)
            std::cout << "Writing matrix partition ..." << std::flush;
        std::ofstream os(acaOptions.outputFname.c_str());
        if (symmetric)
            // psoutputHeH() seems to work also for symmetric matrices
            psoutputHeH(os, bemBlclusterTree.get(),
                        trialDofCount, blocks.get());
        else
            psoutputGeH(os, bemBlclusterTree.get(),
                        std::max(testDofCount, trialDofCount), blocks.get());
        os.close();
        if (verbosityAtLeastDefault)
            std::cout << " done." << std::endl;
    }

    int outSymmetry = NO_SYMMETRY;
    if (symmetric) {
        outSymmetry = SYMMETRIC;
        if (!boost::is_complex<ResultType>())
            outSymmetry |= HERMITIAN;
    }
    typedef DiscreteAcaBoundaryOperator<ResultType> DiscreteAcaLinOp;
    std::auto_ptr<DiscreteAcaLinOp> acaOp(
                new DiscreteAcaLinOp(testDofCount, trialDofCount,
                                     acaOptions.eps,
                                     acaOptions.maximumRank,
                                     outSymmetry,
                                     bemBlclusterTree, blocks,
                                     *trial_o2pPermutation, // domain
                                     *test_o2pPermutation, // range
                                     parallelOptions));
    return acaOp;
}

#endif

} // namespace

template <typename BasisFunctionType, typename ResultType>
std::auto_ptr<DiscreteBoundaryOperator<ResultType> >
AcaGlobalAssembler<BasisFunctionType, ResultType>::assembleDetachedWeakForm(
        const Space<BasisFunctionType>& testSpace,
        const Space<BasisFunctionType>& trialSpace,
        const std::vector<LocalAssemblerForBoundaryOperators*>& localAssemblers,
        const std::vector<const DiscreteBndOp*>& sparseTermsToAdd,
        const std::vector<ResultType>& denseTermMultipliers,
        const std::vector<ResultType>& sparseTermMultipliers,
        const AssemblyOptions& options,
        int symmetry)
{
#ifdef WITH_AHMED
    typedef AhmedDofWrapper<CoordinateType> AhmedDofType;
    typedef ExtendedBemCluster<AhmedDofType> AhmedBemCluster;
    typedef bemblcluster<AhmedDofType, AhmedDofType> AhmedBemBlcluster;
    typedef DiscreteAcaBoundaryOperator<ResultType> DiscreteAcaLinOp;

    const AcaOptions& acaOptions = options.acaOptions();
    const bool indexWithGlobalDofs = acaOptions.globalAssemblyBeforeCompression;
    const bool verbosityAtLeastDefault =
            (options.verbosityLevel() >= VerbosityLevel::DEFAULT);
    const bool verbosityAtLeastHigh =
            (options.verbosityLevel() >= VerbosityLevel::HIGH);

    // Currently we don't support Hermitian ACA operators. This is because we
    // don't have the means to really test them -- we would need complex-valued
    // basis functions for that. (Assembly of such a matrix would be very easy
    // -- just change complex_sym from true to false in the call to apprx_sym()
    // in AcaWeakFormAssemblerLoopBody::operator() -- but operations on
    // symmetric/Hermitian matrices are not always trivial and we do need to be
    // able to test them properly.)
    bool symmetric = symmetry & SYMMETRIC;
    if (symmetry & HERMITIAN && !(symmetry & SYMMETRIC) &&
            verbosityAtLeastDefault)
        std::cout << "Warning: assembly of non-symmetric Hermitian H-matrices "
                     "is not supported yet. A general H-matrix will be assembled"
                  << std::endl;

#ifndef WITH_TRILINOS
    if (!indexWithGlobalDofs)
        throw std::runtime_error("AcaGlobalAssembler::assembleDetachedWeakForm(): "
                                 "ACA assembly with globalAssemblyBeforeCompression "
                                 "set to false requires BEM++ to be linked with "
                                 "Trilinos");
#endif // WITH_TRILINOS

    const size_t testDofCount = indexWithGlobalDofs ?
                testSpace.globalDofCount() : testSpace.flatLocalDofCount();
    const size_t trialDofCount = indexWithGlobalDofs ?
                trialSpace.globalDofCount() : trialSpace.flatLocalDofCount();

    if (symmetric && testDofCount != trialDofCount)
        throw std::invalid_argument("AcaGlobalAssembler::assembleDetachedWeakForm(): "
                                    "you cannot generate a symmetric weak form "
                                    "using test and trial spaces with different "
                                    "numbers of DOFs");

    // o2p: map of original indices to permuted indices
    // p2o: map of permuted indices to original indices
    typedef ClusterConstructionHelper<BasisFunctionType> CCH;
    shared_ptr<AhmedBemCluster> testClusterTree;
    shared_ptr<IndexPermutation> test_o2pPermutation, test_p2oPermutation;
    CCH::constructBemCluster(testSpace, indexWithGlobalDofs, acaOptions,
                             testClusterTree,
                             test_o2pPermutation, test_p2oPermutation);
    shared_ptr<AhmedBemCluster> trialClusterTree;
    shared_ptr<IndexPermutation> trial_o2pPermutation, trial_p2oPermutation;
    if (symmetric || &testSpace == &trialSpace) {
        trialClusterTree = testClusterTree;
        trial_o2pPermutation = test_o2pPermutation;
        trial_p2oPermutation = test_p2oPermutation;
    } else
        CCH::constructBemCluster(trialSpace, indexWithGlobalDofs, acaOptions,
                                 trialClusterTree,
                                 trial_o2pPermutation, trial_p2oPermutation);

//    // Export VTK plots showing the disctribution of leaf cluster ids
//    std::vector<unsigned int> testClusterIds;
//    getClusterIds(*testClusterTree, test_p2oPermutation->permutedIndices(), testClusterIds);
//    testSpace.dumpClusterIds("testClusterIds", testClusterIds,
//                             indexWithGlobalDofs ? GLOBAL_DOFS : FLAT_LOCAL_DOFS);
//    std::vector<unsigned int> trialClusterIds;
//    getClusterIds(*trialClusterTree, trial_p2oPermutation->permutedIndices(), trialClusterIds);
//    trialSpace.dumpClusterIds("trialClusterIds", trialClusterIds,
//                              indexWithGlobalDofs ? GLOBAL_DOFS : FLAT_LOCAL_DOFS);

    if (verbosityAtLeastHigh)
        std::cout << "Test cluster count: " << testClusterTree->getncl()
                  << "\nTrial cluster count: " << trialClusterTree->getncl()
                  << std::endl;

    unsigned int blockCount = 0;
    shared_ptr<AhmedBemBlcluster> bemBlclusterTree(
                CCH::constructBemBlockCluster(acaOptions, symmetric,
                                              *testClusterTree, *trialClusterTree,
                                              blockCount).release());

    if (verbosityAtLeastHigh)
        std::cout << "Mblock count: " << blockCount << std::endl;

    std::vector<unsigned int> p2oTestDofs =
        test_p2oPermutation->permutedIndices();
    std::vector<unsigned int> p2oTrialDofs =
        trial_p2oPermutation->permutedIndices();
    assert(p2oTestDofs.size() == testDofCount);
    assert(p2oTrialDofs.size() == trialDofCount);

#ifdef DUMP_DENSE_BLOCKS
    std::vector<Point3D<CoordinateType> > testDofCenters, trialDofCenters;
    if (indexWithGlobalDofs) {
        testSpace.getGlobalDofPositions(testDofCenters);
        trialSpace.getGlobalDofPositions(trialDofCenters);
    } else {
        testSpace.getFlatLocalDofPositions(testDofCenters);
        trialSpace.getFlatLocalDofPositions(trialDofCenters);
    }
#endif // DUMP_DENSE_BLOCKS

    typedef WeakFormAcaAssemblyHelper<BasisFunctionType, ResultType>
            AcaAssemblyHelper;
    // TODO: It might be better (more efficient and elegant)
    // to pass p2oPermutation than p2oDofs.
    // Also, it might be more logical to rename IndexPermutation to IndexMapping
    // and permute/unpermute to map/unmap.
    AcaAssemblyHelper helper(
                testSpace, trialSpace, p2oTestDofs, p2oTrialDofs,
                localAssemblers, sparseTermsToAdd,
                denseTermMultipliers, sparseTermMultipliers, options);

    std::auto_ptr<DiscreteAcaBoundaryOperator<ResultType> > acaOp =
    assembleAcaOperator<AcaAssemblyHelper, BasisFunctionType, ResultType>(
                helper, bemBlclusterTree,
                options.parallelizationOptions(), options.acaOptions(),
                verbosityAtLeastDefault, symmetric,
                test_o2pPermutation, trial_o2pPermutation
#ifdef DUMP_DENSE_BLOCKS
                ,
                test_p2oPermutation, trial_p2oPermutation,
                testDofCenters, trialDofCenters
#endif // DUMP_DENSE_BLOCKS
                );

    std::auto_ptr<DiscreteBndOp> result;
    if (indexWithGlobalDofs)
        result = acaOp;
    else {
#ifdef WITH_TRILINOS
        // without Trilinos, this code will never be reached -- an exception
        // will be thrown earlier in this function
        typedef DiscreteBoundaryOperatorComposition<ResultType> DiscreteBndOpComp;
        shared_ptr<DiscreteBndOp> acaOpShared(acaOp.release());
        shared_ptr<DiscreteBndOp> trialGlobalToLocal =
                constructOperatorMappingGlobalToFlatLocalDofs<
                BasisFunctionType, ResultType>(trialSpace);
        shared_ptr<DiscreteBndOp> testLocalToGlobal =
                constructOperatorMappingFlatLocalToGlobalDofs<
                BasisFunctionType, ResultType>(testSpace);
        shared_ptr<DiscreteBndOp> tmp(
                    new DiscreteBndOpComp(acaOpShared, trialGlobalToLocal));
        result.reset(new DiscreteBndOpComp(testLocalToGlobal, tmp));
#endif // WITH_TRILINOS
    }
    return result;

#else // without Ahmed
    throw std::runtime_error("AcaGlobalAssembler::assembleDetachedWeakForm(): "
                             "To enable assembly in ACA mode, recompile BEM++ "
                             "with the symbol WITH_AHMED defined.");
#endif // WITH_AHMED
}

template <typename BasisFunctionType, typename ResultType>
std::auto_ptr<DiscreteBoundaryOperator<ResultType> >
AcaGlobalAssembler<BasisFunctionType, ResultType>::assembleDetachedWeakForm(
        const Space<BasisFunctionType>& testSpace,
        const Space<BasisFunctionType>& trialSpace,
        LocalAssemblerForBoundaryOperators& localAssembler,
        const AssemblyOptions& options,
        int symmetry)
{
    std::vector<LocalAssemblerForBoundaryOperators*> localAssemblers(
                1, &localAssembler);
    std::vector<const DiscreteBndOp*> sparseTermsToAdd;
    std::vector<ResultType> denseTermsMultipliers(1, 1.0);
    std::vector<ResultType> sparseTermsMultipliers;

    return assembleDetachedWeakForm(testSpace, trialSpace, localAssemblers,
                            sparseTermsToAdd,
                            denseTermsMultipliers,
                            sparseTermsMultipliers,
                            options, symmetry);
}

template <typename BasisFunctionType, typename ResultType>
std::auto_ptr<DiscreteBoundaryOperator<ResultType> >
AcaGlobalAssembler<BasisFunctionType, ResultType>::assemblePotentialOperator(
        const arma::Mat<CoordinateType>& points,
        const Space<BasisFunctionType>& trialSpace,
        const std::vector<LocalAssemblerForPotentialOperators*>& localAssemblers,
        const std::vector<ResultType>& termMultipliers,
        const EvaluationOptions& options)
{
    const int symmetric = false;

#ifdef WITH_AHMED
    typedef AhmedDofWrapper<CoordinateType> AhmedDofType;
    typedef ExtendedBemCluster<AhmedDofType> AhmedBemCluster;
    typedef bemblcluster<AhmedDofType, AhmedDofType> AhmedBemBlcluster;
    typedef DiscreteAcaBoundaryOperator<ResultType> DiscreteAcaLinOp;

    const AcaOptions& acaOptions = options.acaOptions();
    const bool indexWithGlobalDofs = acaOptions.globalAssemblyBeforeCompression;
    const bool verbosityAtLeastDefault =
            (options.verbosityLevel() >= VerbosityLevel::DEFAULT);
    const bool verbosityAtLeastHigh =
            (options.verbosityLevel() >= VerbosityLevel::HIGH);

#ifndef WITH_TRILINOS
    if (!indexWithGlobalDofs)
        throw std::runtime_error("AcaGlobalAssembler::assemblePotentialOperator(): "
                                 "ACA assembly with globalAssemblyBeforeCompression "
                                 "set to false requires BEM++ to be linked with "
                                 "Trilinos");
#endif // WITH_TRILINOS

    if (localAssemblers.empty())
        throw std::runtime_error("AcaGlobalAssembler::assemblePotentialOperator(): "
                                 "the 'localAssemblers' vector must not be empty");

    const size_t pointCount = points.n_cols;
    const int componentCount = localAssemblers[0]->resultDimension();
    const size_t testDofCount = pointCount * componentCount;
    const size_t trialDofCount = indexWithGlobalDofs ?
                trialSpace.globalDofCount() : trialSpace.flatLocalDofCount();

    // o2p: map of original indices to permuted indices
    // p2o: map of permuted indices to original indices
    typedef ClusterConstructionHelper<BasisFunctionType> CCH;
    shared_ptr<AhmedBemCluster> testClusterTree;
    shared_ptr<IndexPermutation> test_o2pPermutation, test_p2oPermutation;
    CCH::constructBemCluster(points, componentCount, acaOptions,
                             testClusterTree,
                             test_o2pPermutation, test_p2oPermutation);
    shared_ptr<AhmedBemCluster> trialClusterTree;
    shared_ptr<IndexPermutation> trial_o2pPermutation, trial_p2oPermutation;
    CCH::constructBemCluster(trialSpace, indexWithGlobalDofs, acaOptions,
                             trialClusterTree,
                             trial_o2pPermutation, trial_p2oPermutation);

    // Print the distribution of cluster ids
#ifdef DUMP_DENSE_BLOCKS
    std::vector<Point3D<CoordinateType> > testDofCenters, trialDofCenters;
    CCH::getComponentDofPositions(points, componentCount, testDofCenters);
    if (indexWithGlobalDofs)
        trialSpace.getGlobalDofPositions(trialDofCenters);
    else
        trialSpace.getFlatLocalDofPositions(trialDofCenters);
#endif // DUMP_DENSE_BLOCKS

    if (verbosityAtLeastHigh)
        std::cout << "Test cluster count: " << testClusterTree->getncl()
                  << "\nTrial cluster count: " << trialClusterTree->getncl()
                  << std::endl;

    unsigned int blockCount = 0;
    shared_ptr<AhmedBemBlcluster> bemBlclusterTree(
                CCH::constructBemBlockCluster(acaOptions, false /* symmetric */,
                                              *testClusterTree, *trialClusterTree,
                                              blockCount).release());

    if (verbosityAtLeastHigh)
        std::cout << "Mblock count: " << blockCount << std::endl;

    std::vector<unsigned int> p2oPoints =
        test_p2oPermutation->permutedIndices();
    std::vector<unsigned int> p2oTrialDofs =
        trial_p2oPermutation->permutedIndices();
    typedef PotentialOperatorAcaAssemblyHelper<BasisFunctionType, ResultType>
            AcaAssemblyHelper;
    AcaAssemblyHelper helper(points, trialSpace, p2oPoints, p2oTrialDofs,
                             localAssemblers, termMultipliers, options);

    std::auto_ptr<DiscreteAcaBoundaryOperator<ResultType> > acaOp =
    assembleAcaOperator<AcaAssemblyHelper, BasisFunctionType, ResultType>(
                helper, bemBlclusterTree,
                options.parallelizationOptions(), options.acaOptions(),
                verbosityAtLeastDefault, symmetric,
                test_o2pPermutation, trial_o2pPermutation
#ifdef DUMP_DENSE_BLOCKS
                ,
                test_p2oPermutation, trial_p2oPermutation,
                testDofCenters, trialDofCenters
#endif // DUMP_DENSE_BLOCKS
                );

    std::auto_ptr<DiscreteBndOp> result;
    if (indexWithGlobalDofs)
        result = acaOp;
    else {
#ifdef WITH_TRILINOS
        // without Trilinos, this code will never be reached -- an exception
        // will be thrown earlier in this function
        typedef DiscreteBoundaryOperatorComposition<ResultType> DiscreteBndOpComp;
        shared_ptr<DiscreteBndOp> acaOpShared(acaOp.release());
        shared_ptr<DiscreteBndOp> trialGlobalToLocal =
                constructOperatorMappingGlobalToFlatLocalDofs<
                BasisFunctionType, ResultType>(trialSpace);
        result.reset(new DiscreteBndOpComp(acaOpShared, trialGlobalToLocal));
#endif // WITH_TRILINOS
    }
    return result;

#else // without Ahmed
    throw std::runtime_error("AcaGlobalAssembler::assemblePotentialOperator(): "
                             "To enable assembly in ACA mode, recompile BEM++ "
                             "with the symbol WITH_AHMED defined.");
#endif // WITH_AHMED
}

template <typename BasisFunctionType, typename ResultType>
std::auto_ptr<DiscreteBoundaryOperator<ResultType> >
AcaGlobalAssembler<BasisFunctionType, ResultType>::assemblePotentialOperator(
        const arma::Mat<CoordinateType>& points,
        const Space<BasisFunctionType>& trialSpace,
        LocalAssemblerForPotentialOperators& localAssembler,
        const EvaluationOptions& options)
{
    std::vector<LocalAssemblerForPotentialOperators*> localAssemblers(
                1, &localAssembler);
    std::vector<ResultType> termMultipliers(1, 1.0);

    return assemblePotentialOperator(points, trialSpace, localAssemblers,
                                     termMultipliers, options);
}

FIBER_INSTANTIATE_CLASS_TEMPLATED_ON_BASIS_AND_RESULT(AcaGlobalAssembler);

} // namespace Bempp