// Copyright (c) fortiss GmbH
//
// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle, and
// Tobias Kessler
//
// This work is licensed under the terms of the MIT license.
// For a copy, see <https://opensource.org/licenses/MIT>.

#ifndef BARK_ML_OBSERVERS_NEAREST_OBSERVER_HPP_
#define BARK_ML_OBSERVERS_NEAREST_OBSERVER_HPP_

#include <memory>
#include <vector>
#include <tuple>
#include <map>
#include <functional>
#include <Eigen/Dense>

#include "bark/commons/params/params.hpp"
#include "bark/world/world.hpp"
#include "bark/world/observed_world.hpp"
#include "bark/world/goal_definition/goal_definition_state_limits_frenet.hpp"
#include "bark/models/dynamic/dynamic_model.hpp"
#include "bark_ml/commons/spaces.hpp"
#include "bark_ml/commons/commons.hpp"
#include "bark_ml/observers/base_observer.hpp"

namespace bark_ml {
namespace observers {
using bark::commons::ParamsPtr;
using spaces::Box;
using commons::Norm;
using spaces::Matrix_t;
using bark::world::AgentMap;
using bark::world::AgentPtr;
using bark::world::objects::AgentId;
using bark::world::WorldPtr;
using bark::world::goal_definition::GoalDefinitionStateLimitsFrenet;
using bark::world::map::LaneCorridorPtr;
using bark::world::ObservedWorldPtr;
using bark::geometry::Point2d;
using bark::commons::transformation::FrenetState;
using bark::commons::transformation::FrenetStateDifference;
using bark::geometry::B_2PI;
using bark::geometry::B_PI;
using bark::geometry::Line;
using bark::geometry::Distance;
using bark::geometry::NormToPI;
using bark::models::dynamic::StateDefinition::X_POSITION;
using bark::models::dynamic::StateDefinition::Y_POSITION;
using bark::models::dynamic::StateDefinition::THETA_POSITION;
using bark::models::dynamic::StateDefinition::VEL_POSITION;
using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
using bark::commons::transformation::FrenetPosition;
using State = Eigen::Matrix<double, Eigen::Dynamic, 1>;

/**
 * @brief  Concatenates the nearest agents state into a row-vecot.
 */
class NearestObserver : public BaseObserver {
 public:
  explicit NearestObserver(const ParamsPtr& params) :
    BaseObserver(params) {
      nearest_agent_num_ =
        params->GetInt(
          "ML::NearestObserver::NNearestAgents", "Nearest agents number", 4);
      min_vel_lon_ = params->GetReal("ML::NearestObserver::MinVelLon", "", -30.0);
      max_vel_lon_ = params->GetReal("ML::NearestObserver::MaxVelLon", "", 30.0);
      min_vel_lat_ = params->GetReal("ML::NearestObserver::MinVelLat", "", -10.0);
      max_vel_lat_ = params->GetReal("ML::NearestObserver::MaxVelLat", "", 10.0);
      max_dist_ = params->GetReal("ML::NearestObserver::MaxDist", "", 75.0);
      min_s_ = params->GetReal("ML::NearestObserver::MinS", "", -100.0);
      min_abs_s_ = params->GetReal("ML::NearestObserver::MinAbsS", "", 0.0);
      max_s_ = params->GetReal("ML::NearestObserver::MaxS", "", 100.0);
      min_d_ = params->GetReal("ML::NearestObserver::MinD", "", -20.0);
      max_d_ = params->GetReal("ML::NearestObserver::MaxD", "", 20.0);
      min_theta_ = params->GetReal("ML::NearestObserver::MinTheta", "", -B_PI);
      max_theta_ = params->GetReal("ML::NearestObserver::MaxTheta", "", B_PI);
      // Ego agent has s,d, vlat, vlon, theta, laneid other agents have del s, del d, del vlat, del vlon
      observation_len_ = nearest_agent_num_ * 4 + 6;
  }

  double Norm(const double val, const double mi, const double ma) const {
    LOG_IF_EVERY_N(WARNING, (val < mi), 100) << "Val=" << val << " < Lower Limit=" << mi;
    LOG_IF_EVERY_N(WARNING, (val > ma), 100) << "Val=" << val << " > Upper Limit=" << ma;
    // Normalize to be within -1 to 1
    return (val - (mi+ma)/2.0)/((ma-mi)/2.0);
  }

  ObservedState GetEgoState(const ObservedWorld& observed_world) const {
    const auto ego_state = observed_world.CurrentEgoState();
    const auto ego_pos = observed_world.CurrentEgoPosition();
    const auto corridor_and_idx = observed_world.GetEgoAgent()
      ->GetRoadCorridor()->GetNearestLaneCorridorAndIndex(ego_pos);
    const auto& ego_corridor = corridor_and_idx.first;
    const unsigned corr_idx = corridor_and_idx.second;
    FrenetState current_ego_frenet(ego_state, ego_corridor->GetCenterLine());
    ObservedState ego_nn_state(1, 6);
    const double normalized_angle = NormToPI(current_ego_frenet.angle);
    LOG_IF_EVERY_N(WARNING, (current_ego_frenet.lon < min_abs_s_), 100)
         << "Sego=" << current_ego_frenet.lon << " < Min abs s=" << min_abs_s_;
    ego_nn_state << Norm(current_ego_frenet.lon - min_abs_s_, min_s_, max_s_),
                 Norm(current_ego_frenet.lat, min_d_, max_d_),
                 Norm(normalized_angle, min_theta_, max_theta_),
                 Norm(current_ego_frenet.vlon, min_vel_lon_, max_vel_lon_),
                 Norm(current_ego_frenet.vlat, min_vel_lat_, max_vel_lat_),
                corr_idx;
    return ego_nn_state;
  }

  ObservedState GetOtherAgentState(const AgentId& agent_id, const ObservedWorld& observed_world) const {
    const auto ego_pos = observed_world.CurrentEgoPosition();
    const auto corridor_and_idx = observed_world.GetEgoAgent()
            ->GetRoadCorridor()->GetNearestLaneCorridorAndIndex(ego_pos);

    const auto ego_state = observed_world.CurrentEgoState();
    FrenetState current_ego_frenet(ego_state, corridor_and_idx.first->GetCenterLine());
    const auto ego_shape = observed_world.GetEgoAgent()->GetShape();
    const auto& other_state = observed_world.GetAgent(agent_id)->GetCurrentState();
    FrenetState other_frenet(other_state, corridor_and_idx.first->GetCenterLine());
    const auto other_shape = observed_world.GetAgent(agent_id)->GetShape();
    FrenetStateDifference state_difference(current_ego_frenet, ego_shape, other_frenet, other_shape);

    ObservedState other_nn_state(1, 4);
    other_nn_state << Norm(state_difference.lon_zeroed ? 0.0 : state_difference.lon, min_s_, max_s_),
                 Norm(state_difference.lat_zeroed ? 0.0 : state_difference.lat, min_d_, max_d_),
                 Norm(state_difference.vlon, min_vel_lon_, max_vel_lon_),
                 Norm(state_difference.vlat, min_vel_lat_, max_vel_lat_);
    return other_nn_state;
  }

  ObservedState Observe(const ObservedWorld& observed_world) const {
    // find near agents (n)
    AgentMap nearest_agents = observed_world.GetNearestAgents(
      observed_world.CurrentEgoPosition(), nearest_agent_num_ + 1);

    // Build state
    int state_start_idx = 0;
    ObservedState state(1, observation_len_);
    state.setZero();

    // add ego agent state
    ObservedState obs_ego_agent_state = GetEgoState(observed_world);
    state.block(0, state_start_idx, 1, obs_ego_agent_state.cols()) = obs_ego_agent_state;
    state_start_idx = obs_ego_agent_state.cols();

    // add other states
    for (const auto& agent : nearest_agents) {
      if(agent.first == observed_world.GetEgoAgentId()) continue;

      const auto& agent_state = agent.second->GetCurrentPosition();
      double distance = Distance(
        observed_world.CurrentEgoPosition(), agent_state);
      if (distance > max_dist_) continue;

      ObservedState other_agent_state = GetOtherAgentState(agent.second->GetAgentId(),
                      observed_world);
      state.block(0, state_start_idx, 1, other_agent_state.cols()) = other_agent_state;  // NOLINT
      state_start_idx += other_agent_state.cols();
    }
    return state;
  }

  WorldPtr Reset(const WorldPtr& world) {
    return world;
  }

  Box<double> ObservationSpace() const {
    Matrix_t<double> low(1, observation_len_);
    low.setZero();
    Matrix_t<double> high(1, observation_len_);
    high.setOnes();
    std::tuple<int> shape{observation_len_};
    return Box<double>(low, high, shape);
  }

 private:
  int nearest_agent_num_, observation_len_;
  double min_theta_, max_theta_, min_vel_lon_,
         max_vel_lon_, min_vel_lat_,
         max_vel_lat_, max_dist_,
         min_d_, max_d_, min_s_, max_s_, min_abs_s_;
};

}  // namespace observers
}  // namespace bark_ml

#endif  // BARK_ML_OBSERVERS_NEAREST_OBSERVER_HPP_


/*
  Copyright (c) 2014-2016 DataStax

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#ifdef STAND_ALONE
#   define BOOST_TEST_MODULE cassandra
#endif

#include "cassandra.h"

#include <boost/test/unit_test.hpp>
#include <time.h>

#define CASS_DATE_EPOCH 2147483648U

BOOST_AUTO_TEST_SUITE(date_time)

BOOST_AUTO_TEST_CASE(simple)
{
  time_t now = ::time(NULL);
  cass_uint32_t date = cass_date_from_epoch(now);
  cass_int64_t time = cass_time_from_epoch(now);
  BOOST_CHECK(cass_date_time_to_epoch(date, time) == now);
}

BOOST_AUTO_TEST_CASE(date)
{
  BOOST_CHECK_EQUAL(cass_date_from_epoch(0), CASS_DATE_EPOCH);
  BOOST_CHECK_EQUAL(cass_date_from_epoch(24 * 3600), CASS_DATE_EPOCH + 1);
  BOOST_CHECK_EQUAL(cass_date_from_epoch(2 * 24 * 3600), CASS_DATE_EPOCH + 2);
}

BOOST_AUTO_TEST_CASE(time)
{
  time_t now = ::time(NULL);
  struct tm* tm = gmtime(&now);
  cass_int64_t actual = cass_time_from_epoch(now);
  cass_int64_t expected = 1000000000LL * (tm->tm_hour * 3600  + 60 * tm->tm_min + tm->tm_sec);
  BOOST_CHECK(actual == expected);
}

BOOST_AUTO_TEST_CASE(date_time_to_epoc)
{
  BOOST_CHECK_EQUAL(cass_date_time_to_epoch(CASS_DATE_EPOCH, 0), 0L); // Epoch
  BOOST_CHECK_EQUAL(cass_date_time_to_epoch(CASS_DATE_EPOCH - 1, 0), -24L * 3600L); // Epoch - 1 day
  BOOST_CHECK_EQUAL(cass_date_time_to_epoch(CASS_DATE_EPOCH + 1, 0), 24L * 3600L); // Epoch + 1 day
}

BOOST_AUTO_TEST_SUITE_END()


#include <boost/test/unit_test.hpp>

#include <crave/experimental/Coverage.hpp>
#include <crave/experimental/Variable.hpp>

#include <vector>

using namespace crave;

BOOST_FIXTURE_TEST_SUITE(CoverageSampling, Context_Fixture)

template<typename T>
void sample_from_list(crv_covergroup& cg, crv_variable<T>& v, std::vector<T>&& values) {
  for (T val : values) {
    v = val;
    cg.sample();
  }
}

struct my_covergroup : public crv_covergroup {
  crv_variable<unsigned> a;
  crv_variable<unsigned> b;

  crv_coverpoint cp1{"coverpoint1"};
  crv_coverpoint cp2{"coverpoint2"};

  my_covergroup(crv_object_name) {
    cp1.simple_bins() = {crv_bin(a() == 0, 10), crv_bin(a() == 1, 10), crv_bin(a() == 2, 5)};
    cp2.simple_bins() = {crv_bin(a() == 0, 10), crv_bin(a() == 1, 10), crv_bin(a() == 2, 5)};
  }
};

BOOST_AUTO_TEST_CASE(basic_test) {
  my_covergroup cg{"covergroup1"};
  for (int i = 0; i < 5; i++)
    for (int j = 0; j < 5; j++) {
      cg.a = i % 3;
      cg.b = j * 11 % 20;
      cg.sample();
    }

  BOOST_REQUIRE(cg.covered());
}

struct my_transitionCovergroup_consec : public crv_covergroup {
  crv_variable<unsigned> a;

  crv_coverpoint cp1{"coverpoint1"};

  my_transitionCovergroup_consec(crv_object_name) {
      //1
      cp1.start_with(a() == 1);
      cp1.last_transition_bin().hit_minimum(4);
      //1 => 2
      cp1.start_with(a() == 1)->next(a() == 2);
      cp1.last_transition_bin().hit_minimum(1);
      //0 => 1,2[*2] => 0
      cp1.start_with(a() == 0)->consecutive_next<2>(a() == 1 || a() == 2)->next(a() == 0);
      cp1.last_transition_bin().hit_minimum(1);
      //0 [*2:3] => 1
      cp1.start_with_consecutive<2,3>(a() == 0)->next(a() == 1);
      cp1.last_transition_bin().hit_minimum(2);
  }
};

BOOST_AUTO_TEST_CASE(transition_consec_test) {
  my_transitionCovergroup_consec cg{"covergroup2"};

  sample_from_list(cg, cg.a, {0, 1, 0, 1, 2, 0, 0, 1, 0, 0, 0, 1});

  BOOST_REQUIRE(cg.covered());
}

struct my_transitionCovergroup_goto : public crv_covergroup {
  crv_variable<unsigned> a;

  crv_coverpoint cp1{"coverpoint1"};

  my_transitionCovergroup_goto(crv_object_name) {
      //1
      cp1.start_with(a() == 1);
      cp1.last_transition_bin().hit_minimum(4);
      //1 [->3]
      cp1.start_with_goto<3>(a() == 1);
      cp1.last_transition_bin().hit_minimum(2);
      //0 [->2] => 1
      cp1.start_with_goto<2>(a() == 0)->next(a() == 1);
      cp1.last_transition_bin().hit_minimum(3);
      //0 => 1[->2:4] => 0
      cp1.start_with(a() == 0)->goto_next<2,4>(a() == 1)->next(a() == 0);
      cp1.last_transition_bin().hit_minimum(1);
  }
};

BOOST_AUTO_TEST_CASE(transition_goto_test) {
  my_transitionCovergroup_goto cg{"covergroup3"};

  sample_from_list(cg, cg.a, {0, 1, 0, 1, 2, 0, 0, 1, 0, 0, 0, 1});

  BOOST_REQUIRE(cg.covered());
}

struct my_transitionCovergroup_ncons : public crv_covergroup {
  crv_variable<unsigned> a;

  crv_coverpoint cp1{"coverpoint1"};

  my_transitionCovergroup_ncons(crv_object_name) {
      //1
      cp1.start_with(a() == 1);
      cp1.last_transition_bin().hit_minimum(4);
      //1 [= 3]
      cp1.start_with_nonconsecutive<3>(a() == 1);
      cp1.last_transition_bin().hit_minimum(2);
      //0 [= 2] => 2
      cp1.start_with_nonconsecutive<2>(a() == 0)->next(a() == 2);
      cp1.last_transition_bin().hit_minimum(1);
      //0 => 1[= 2:4] => 0
      cp1.start_with(a() == 0)->nonconsecutive_next<2,4>(a() == 1)->next(a() == 0);
      cp1.last_transition_bin().hit_minimum(2);
  }
};

BOOST_AUTO_TEST_CASE(transition_ncons_test) {
  my_transitionCovergroup_ncons cg{"covergroup4"};

  sample_from_list(cg, cg.a, {0, 1, 0, 1, 2, 0, 0, 1, 0, 0, 0, 1});

  BOOST_REQUIRE(cg.covered());
}

struct my_transitionCovergroup_overlap : public crv_covergroup {
  crv_variable<unsigned> a;

  crv_coverpoint cp1{"coverpoint1"};

  my_transitionCovergroup_overlap(crv_object_name) {
      //0 => 0
      cp1.start_with_consecutive<2>(a() == 0);
      cp1.last_transition_bin().hit_minimum(6);
      //2 => 0[1:2]
      cp1.start_with(a() == 2)->consecutive_next<1,2>(a()==0);
      cp1.last_transition_bin().hit_minimum(2);
      //0 => 1[->2]
      cp1.start_with(a() == 0)->goto_next<2>(a() == 1);
      cp1.last_transition_bin().hit_minimum(1);
      //0[*4] => 2 => 0[*2] => 1 => 0[*3] => 1
      cp1.start_with_consecutive<4>(a() == 0)->next(a() == 2)->consecutive_next<2>(a() == 0)->next(a()==1)->consecutive_next<3>(a() == 0)->next(a()==1);
      cp1.last_transition_bin().hit_minimum(1);
  }
};

BOOST_AUTO_TEST_CASE(transition_overlap_test) {
  my_transitionCovergroup_overlap cg{"covergroup5"};

  sample_from_list(cg, cg.a, {0, 0, 0, 0, 2, 0, 0, 1, 0, 0, 0, 1});

  BOOST_REQUIRE(cg.covered());
}

struct ieee1800_2012_p530_covergroup : public crv_covergroup {
  crv_variable<unsigned> a;

  crv_coverpoint cp1{"coverpoint1"};

  ieee1800_2012_p530_covergroup(crv_object_name) {
      // b2 = 2 [-> 3:5]
      // 1, 4, 3, [[2, 3, 3, 2, 2], 3, 2], 3, 1, 5, 5, 5, 5, 5, 5 (2 Hits)
      cp1.start_with_goto<3, 5>(a() == 2);
      cp1.last_transition_bin().hit_minimum(2);
      // b3 = 3 [-> 3:5]
      // 1, 4, [[[3, 2, 3, 3], 2, 2, 3], 2, 3], 1, 5, 5, 5, 5, 5, 5 (3 Hits)
      cp1.start_with_goto<3, 5>(a() == 3);
      cp1.last_transition_bin().hit_minimum(3);
      // b5 = 5 [* 3]
      // 1, 4, 3, 2, 3, 3, 2, 2, 3, 2, 3, 1, [5, (5, {5], [5), 5}, 5] (4 Hits)
      cp1.start_with_consecutive<3>(a() == 5);
      cp1.last_transition_bin().hit_minimum(4);
      // b6 = 1 => 3 [-> 4:6] => 1
      // [1, 4, 3, 2, 3, 3, 2, 2, 3, 2, 3, 1], 5, 5, 5, 5, 5, 5 (1 Hit)
      cp1.start_with(a() == 1)->goto_next<4, 6>(a() == 3)->next(a() == 1);
      cp1.last_transition_bin().hit_minimum(1);
      // b7 = 1 => 2 [= 3:6] => 5
      // [1, 4, 3, 2, 3, 3, 2, 2, 3, 2, 3, 1, 5], 5, 5, 5, 5, 5 (1 Hit)
      cp1.start_with(a() == 1)->nonconsecutive_next<3, 6>(a() == 2)->next(a() == 5);
      cp1.last_transition_bin().hit_minimum(1);
  }
};

BOOST_AUTO_TEST_CASE(ieee1800_2012_p530_covergroup_test) {
  ieee1800_2012_p530_covergroup cg{"covergroup5"};

  sample_from_list(cg, cg.a, {1, 4, 3, 2, 3, 3, 2, 2, 3, 2, 3, 1, 5, 5, 5, 5, 5, 5});

  cg.report();

  BOOST_REQUIRE(cg.covered());
}

BOOST_AUTO_TEST_SUITE_END()  // CoverageSampling

//  vim: ft=cpp:ts=2:sw=2:expandtab


// Copyright (c) 2018 Graphcore Ltd. All rights reserved.

//
// A simple addition that returns its result.
// Similar to the Python simple_addition example
// Recreated as a test to demonstrate a simple example.

#define BOOST_TEST_MODULE simple_addition_test

#include <boost/test/unit_test.hpp>
#include <iostream>
#include <popart/builder.hpp>
#include <popart/devicemanager.hpp>
#include <popart/logging.hpp>
#include <popart/ndarraywrapper.hpp>
#include <popart/op.hpp>
#include <popart/op/l1.hpp>
#include <popart/opmanager.hpp>
#include <popart/popx/opx.hpp>
#include <popart/popx/opxmanager.hpp>
#include <popart/session.hpp>
#include <popart/tensordata.hpp>
#include <popart/tensorinfo.hpp>
#include <popart/tensornames.hpp>

using namespace popart;

int runTest(float rawInputData[2]) {

  // Generate an ONNX inference model
  auto builder = Builder::create();

  std::cout << "Setting Opset9\n";
  auto aiOnnx = builder->aiOnnxOpset9();

  // Add input tensors
  TensorInfo inputInfo{"FLOAT", std::vector<int64_t>{2}};
  std::cout << "Adding input tensor a\n";
  auto a = builder->addInputTensor(inputInfo);
  std::cout << "Adding input tensor b\n";
  auto b = builder->addInputTensor(inputInfo);

  // Add operation
  std::cout << "Adding operation add(a,b)\n";
  auto o = aiOnnx.add({a, b});

  // Add output tensor
  std::cout << "Adding output tensor o\n";
  builder->addOutputTensor(o);

  std::cout << "Getting model proto\n";
  auto proto = builder->getModelProto();

  std::cout << "Constructing DataFlow\n";
  auto dataFlow = DataFlow(1, {{o, AnchorReturnType("ALL")}});

  std::map<std::string, std::string> deviceOpts{{"numIPUs", "1"}};
  auto ipuModelDevice =
      DeviceManager::createDeviceManager().createIpuModelDevice(deviceOpts);
  // or acquireAvailableDevice();

  std::cout << "Creating session from Onnx Model...\n";
  auto session =
      InferenceSession::createFromOnnxModel(proto, dataFlow, ipuModelDevice);
  std::cout << "Creating session from Onnx Model...done\n";

  NDArrayWrapper<float> inDataA(rawInputData, {2});
  NDArrayWrapper<float> inDataB(rawInputData, {2});
  std::map<TensorId, IArray &> inputs = {{a, inDataA}, {b, inDataB}};

  // Prepare output tensor
  float rawOutputData[2] = {0, 0};
  NDArrayWrapper<float> outData(rawOutputData, {2});
  std::map<TensorId, IArray &> anchors = {{o, outData}};

  std::cout << "Preparing session device...\n";
  session->prepareDevice();
  std::cout << "Preparing session device...done\n";

  StepIO stepio(inputs, anchors);

  std::cout << "Running..."
            << "\n";
  session->run(stepio);
  std::cout << "Running...done"
            << "\n";

  std::cout << "Input Data:  " << inDataA << "\n";
  std::cout << "Input Data:  " << inDataB << "\n";
  std::cout << "Output Data: " << outData << "\n";

  logging::ir::info("inputs : {}, {}", inDataA, inDataB);
  logging::ir::info("output : {}", outData);

  BOOST_CHECK(outData[0] == 4.0f);
  BOOST_CHECK(outData[1] == 8.0f);

  return 0;
}

BOOST_AUTO_TEST_CASE(simple_addition_test) {
  float rawInputData[2] = {2.0f, 4.0f};
  runTest(rawInputData);
}


/*!
@file
Defines `boost::hana::ext::std::IntegerSequence`.

@copyright Louis Dionne 2015
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#ifndef BOOST_HANA_EXT_STD_INTEGER_SEQUENCE_HPP
#define BOOST_HANA_EXT_STD_INTEGER_SEQUENCE_HPP

#include <boost/hana/bool.hpp>
#include <boost/hana/comparable.hpp>
#include <boost/hana/core/datatype.hpp>
#include <boost/hana/detail/std/forward.hpp>
#include <boost/hana/detail/std/is_same.hpp>
#include <boost/hana/ext/std/integral_constant.hpp>
#include <boost/hana/foldable.hpp>
#include <boost/hana/iterable.hpp>
#include <boost/hana/searchable.hpp>

#include <type_traits>
#include <utility>

namespace boost { namespace hana {
    namespace ext { namespace std { struct IntegerSequence; }}

    template <typename T, T ...v>
    struct datatype< ::std::integer_sequence<T, v...>> {
        using type = ext::std::IntegerSequence;
    };

    //////////////////////////////////////////////////////////////////////////
    // Comparable
    //////////////////////////////////////////////////////////////////////////
    template <>
    struct equal_impl<ext::std::IntegerSequence, ext::std::IntegerSequence> {
        template <typename X, X ...xs, typename Y, Y ...ys>
        static constexpr auto apply(
            ::std::integer_sequence<X, xs...>,
            ::std::integer_sequence<Y, ys...>,
            // this dummy parameter disables this specialization if
            // sizeof...(xs) != sizeof...(ys)
            char(*)[sizeof...(xs) == sizeof...(ys)] = 0)
        {
            return bool_<detail::std::is_same<
                ::std::integer_sequence<bool, (xs == ys)...>,
                ::std::integer_sequence<bool, (xs, true)...>
            >::value>;
        }

        template <typename Xs, typename Ys>
        static constexpr auto apply(Xs, Ys, ...)
        { return false_; }
    };

    //////////////////////////////////////////////////////////////////////////
    // Foldable
    //////////////////////////////////////////////////////////////////////////
    template <>
    struct unpack_impl<ext::std::IntegerSequence> {
        template <typename T, T ...v, typename F>
        static constexpr decltype(auto)
        apply(::std::integer_sequence<T, v...>, F&& f) {
            return detail::std::forward<F>(f)(
                std::integral_constant<T, v>{}...
            );
        }
    };

    //////////////////////////////////////////////////////////////////////////
    // Searchable
    //////////////////////////////////////////////////////////////////////////
    template <>
    struct find_if_impl<ext::std::IntegerSequence>
        : Iterable::find_if_impl<ext::std::IntegerSequence>
    { };

    template <>
    struct any_of_impl<ext::std::IntegerSequence>
        : Iterable::any_of_impl<ext::std::IntegerSequence>
    { };

    //////////////////////////////////////////////////////////////////////////
    // Iterable
    //////////////////////////////////////////////////////////////////////////
    template <>
    struct head_impl<ext::std::IntegerSequence> {
        template <typename T, T x, T ...xs>
        static constexpr auto apply(::std::integer_sequence<T, x, xs...>)
        { return ::std::integral_constant<T, x>{}; }
    };

    template <>
    struct tail_impl<ext::std::IntegerSequence> {
        template <typename T, T x, T ...xs>
        static constexpr auto apply(::std::integer_sequence<T, x, xs...>)
        { return ::std::integer_sequence<T, xs...>{}; }
    };

    template <>
    struct is_empty_impl<ext::std::IntegerSequence> {
        template <typename T, T ...xs>
        static constexpr auto apply(::std::integer_sequence<T, xs...>)
        { return bool_<sizeof...(xs) == 0>; }
    };
}} // end namespace boost::hana

#endif // !BOOST_HANA_EXT_STD_INTEGER_SEQUENCE_HPP


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

#include <queue>

#include "../parallel/tbb.h"

template<class T>
hera::bt::dnn::KDTree<T>::KDTree(const Traits& traits, HandleContainer&& handles):
    traits_(traits),
    tree_(std::move(handles)),
    delete_flags_(handles.size(), static_cast<char>(0) ),
    subtree_n_elems(handles.size(), static_cast<size_t>(0)),
    num_points_(handles.size())
{
    init();
}

template<class T>
template<class Range>
hera::bt::dnn::KDTree<T>::KDTree(const Traits& traits, const Range& range):
    traits_(traits)
{
    init(range);
}

template<class T>
template<class Range>
void hera::bt::dnn::KDTree<T>::init(const Range& range)
{
    size_t sz = std::distance(std::begin(range), std::end(range));
    subtree_n_elems = std::vector<int>(sz, 0);
    delete_flags_ = std::vector<char>(sz, 0);
    num_points_ = sz;
    tree_.reserve(sz);
    for (PointHandle h : range)
        tree_.push_back(h);
    parents_.resize(sz, -1);
    init();
}

template<class T>
void hera::bt::dnn::KDTree<T>::init()
{
    if (tree_.empty())
        return;

#if defined(TBB)
    task_group g;
    g.run(OrderTree(this, tree_.begin(), tree_.end(), -1, 0, traits()));
    g.wait();
#else
    OrderTree(this, tree_.begin(), tree_.end(), -1, 0, traits()).serial();
#endif

    for (size_t i = 0; i < tree_.size(); ++i)
        indices_[tree_[i]] = i;
    init_n_elems();
}

template<class T>
struct
hera::bt::dnn::KDTree<T>::OrderTree
{
                OrderTree(KDTree* tree_, HCIterator b_, HCIterator e_, ssize_t p_, size_t i_, const Traits& traits_):
                    tree(tree_), b(b_), e(e_), p(p_), i(i_), traits(traits_)            {}

    void        operator()() const
    {
        if (e - b < 1000)
        {
            serial();
            return;
        }

        HCIterator m  = b + (e - b)/2;
        ssize_t    im = m - tree->tree_.begin();
        tree->parents_[im]  = p;

        CoordinateComparison cmp(i, traits);
        std::nth_element(b,m,e, cmp);
        size_t next_i = (i + 1) % traits.dimension();

        task_group g;
        if (b < m - 1)  g.run(OrderTree(tree, b,   m, im, next_i, traits));
        if (e > m + 2)  g.run(OrderTree(tree, m+1, e, im, next_i, traits));
        g.wait();
    }

    void        serial() const
    {
        std::queue<KDTreeNode> q;
        q.push(KDTreeNode(b,e,p,i));
        while (!q.empty())
        {
            HCIterator b, e; ssize_t p; size_t i;
            std::tie(b,e,p,i) = q.front();
            q.pop();
            HCIterator m  = b + (e - b)/2;
            ssize_t    im = m - tree->tree_.begin();
            tree->parents_[im]  = p;

            CoordinateComparison cmp(i, traits);
            std::nth_element(b,m,e, cmp);
            size_t next_i = (i + 1) % traits.dimension();

            // Replace with a size condition instead?
            if (b < m - 1)
                q.push(KDTreeNode(b,   m, im, next_i));
            else if (b < m)
                tree->parents_[im - 1] = im;
            if (e > m + 2)
                q.push(KDTreeNode(m+1, e, im, next_i));
            else if (e > m + 1)
                tree->parents_[im + 1] = im;
        }
    }

    KDTree*         tree;
    HCIterator      b, e;
    ssize_t         p;
    size_t          i;
    const Traits&   traits;
};

template<class T>
void hera::bt::dnn::KDTree<T>::update_n_elems(ssize_t idx, const int delta)
// add delta to the number of points in node idx and update subtree_n_elems
// for all parents of the node idx
{
    //std::cout << "subtree_n_elems.size = " << subtree_n_elems.size() << std::endl;
    // update the node itself
    while (idx != -1)
    {
        //std::cout << idx << std::endl;
        subtree_n_elems[idx] += delta;
        idx = parents_[idx];
    }
}

template<class T>
void hera::bt::dnn::KDTree<T>::increase_n_elems(const ssize_t idx)
{
    update_n_elems(idx, static_cast<ssize_t>(1));
}

template<class T>
void hera::bt::dnn::KDTree<T>::decrease_n_elems(const ssize_t idx)
{
    update_n_elems(idx, static_cast<ssize_t>(-1));
}

template<class T>
void hera::bt::dnn::KDTree<T>::init_n_elems()
{
    for(size_t idx = 0; idx < tree_.size(); ++idx) {
        increase_n_elems(idx);
    }
}

template<class T>
template<class ResultsFunctor>
void hera::bt::dnn::KDTree<T>::search(PointHandle q, ResultsFunctor& rf) const
{
    typedef         typename HandleContainer::const_iterator        HCIterator;
    typedef         std::tuple<HCIterator, HCIterator, size_t>      KDTreeNode;

    if (tree_.empty())
        return;

    DistanceType    D  = std::numeric_limits<DistanceType>::infinity();

    // TODO: use tbb::scalable_allocator for the queue
    std::queue<KDTreeNode>  nodes;

    nodes.push(KDTreeNode(tree_.begin(), tree_.end(), 0));

    //std::cout << "started kdtree::search" << std::endl;

    while (!nodes.empty())
    {
        HCIterator b, e; size_t i;
        std::tie(b,e,i) = nodes.front();
        nodes.pop();

        CoordinateComparison cmp(i, traits());
        i = (i + 1) % traits().dimension();

        HCIterator   m    = b + (e - b)/2;
        size_t m_idx = m - tree_.begin();
        // ignore deleted points
        if ( delete_flags_[m_idx] == 0 ) {
            DistanceType dist = traits().distance(q, *m);
            // + weights_[m - tree_.begin()];
            //std::cout << "Supplied to functor: m : ";
            //std::cout << "(" << (*(*m))[0] << ", " << (*(*m))[1] << ")";
            //std::cout << " and q : ";
            //std::cout << "(" << (*q)[0] << ", " << (*q)[1] << ")" << std::endl;
            //std::cout << "dist^q + weight = " << dist << std::endl;
            //std::cout << "weight = " << weights_[m - tree_.begin()] << std::endl;
            //std::cout << "dist = " << traits().distance(q, *m) << std::endl;
            //std::cout << "dist^q = " << pow(traits().distance(q, *m), wassersteinPower) << std::endl;

            D = rf(*m, dist);
        }
        // we are really searching w.r.t L_\infty ball; could prune better with an L_2 ball
        Coordinate diff = cmp.diff(q, *m);     // diff returns signed distance
        DistanceType diffToWasserPower = (diff > 0 ? 1.0 : -1.0) * fabs(diff);

        size_t       lm   = m + 1 + (e - (m+1))/2 - tree_.begin();
        if ( subtree_n_elems[lm] > 0 ) {
            if (e > m + 1 && diffToWasserPower  >= -D) {
                nodes.push(KDTreeNode(m+1, e, i));
            }
        }

        size_t       rm   = b + (m - b) / 2 - tree_.begin();
        if ( subtree_n_elems[rm] > 0 ) {
            if (b < m && diffToWasserPower  <= D) {
                nodes.push(KDTreeNode(b,   m, i));
            }
        }
    }
    //std::cout << "exited kdtree::search" << std::endl;
}

template<class T>
typename hera::bt::dnn::KDTree<T>::HandleDistance hera::bt::dnn::KDTree<T>::find(PointHandle q) const
{
    hera::bt::dnn::NNRecord<HandleDistance> nn;
    search(q, nn);
    return nn.result;
}

template<class T>
typename hera::bt::dnn::KDTree<T>::Result hera::bt::dnn::KDTree<T>::findR(PointHandle q, DistanceType r) const
{
    hera::bt::dnn::rNNRecord<HandleDistance> rnn(r);
    search(q, rnn);
    //std::sort(rnn.result.begin(), rnn.result.end());
    return rnn.result;
}

template<class T>
typename hera::bt::dnn::KDTree<T>::Result hera::bt::dnn::KDTree<T>::findFirstR(PointHandle q, DistanceType r) const
{
    hera::bt::dnn::firstrNNRecord<HandleDistance> rnn(r);
    search(q, rnn);
    return rnn.result;
}

template<class T>
typename hera::bt::dnn::KDTree<T>::Result hera::bt::dnn::KDTree<T>::findK(PointHandle q, size_t k) const
{
    hera::bt::dnn::kNNRecord<HandleDistance> knn(k);
    search(q, knn);
    // do we need this???
    std::sort(knn.result.begin(), knn.result.end());
    return knn.result;
}

template<class T>
struct hera::bt::dnn::KDTree<T>::CoordinateComparison
{
                CoordinateComparison(size_t i, const Traits& traits):
                    i_(i), traits_(traits)                              {}

    bool        operator()(PointHandle p1, PointHandle p2) const        { return coordinate(p1) < coordinate(p2); }
    Coordinate  diff(PointHandle p1, PointHandle p2) const              { return coordinate(p1) - coordinate(p2); }

    Coordinate  coordinate(PointHandle p) const                         { return traits_.coordinate(p, i_); }
    size_t      axis() const                                            { return i_; }

    private:
        size_t          i_;
        const Traits&   traits_;
};

template<class T>
void hera::bt::dnn::KDTree<T>::delete_point(const size_t idx)
{
    // prevent double deletion
    assert(delete_flags_[idx] == 0);
    delete_flags_[idx] = 1;
    decrease_n_elems(idx);
    --num_points_;
}

template<class T>
void hera::bt::dnn::KDTree<T>::delete_point(PointHandle p)
{
    delete_point(indices_[p]);
}


#include <boost/fusion/tuple.hpp>
#include <boost/fusion/view.hpp>
#include <boost/fusion/algorithm.hpp>
#include <boost/type_traits.hpp>
#include <boost/mpl/arg.hpp>
#include <string>
#include <iostream>

using namespace boost::fusion;

struct print
{
  template <typename T>
  void operator()(const T &t) const
  {
    std::cout << std::boolalpha << t << '\n';
  }
};

int main()
{
  typedef tuple<int, std::string, bool, double> tuple_type;
  tuple_type t{10, "Boost", true, 3.14};
  filter_view<tuple_type, boost::is_integral<boost::mpl::arg<1>>> v{t};
  for_each(v, print{});
}


// Copyright (C) 2011-2012 The Trustees of Indiana University.

// Use, modification and distribution is subject to the Boost Software
// License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

//  Authors: Jeremiah Willcock
//           Andrew Lumsdaine

#ifndef BOOST_NWAY_BERNOULLI_DISTRIBUTION_HPP
#define BOOST_NWAY_BERNOULLI_DISTRIBUTION_HPP

#include <boost/random/uniform_01.hpp>
#include <limits>

namespace boost { namespace graph { namespace random {

  template <int NCases, typename RAIter>
  struct find_bucket_impl {
    typedef typename std::iterator_traits<RAIter>::value_type value_type;
    static int go(value_type x, RAIter start) {
      RAIter middle = start + NCases / 2;
      if (x < *middle) {
        return find_bucket_impl<NCases / 2, RAIter>::go(x, start);
      } else {
        return NCases / 2 + find_bucket_impl<NCases - NCases / 2, RAIter>::go(x, middle);
      }
    }
  };

  template <typename RAIter>
  struct find_bucket_impl<1, RAIter> {
    typedef typename std::iterator_traits<RAIter>::value_type value_type;
    static int go(value_type /* x */, RAIter /* start */) {
      return 0;
    }
  };

  template <typename RAIter>
  struct find_bucket_impl<0, RAIter> {};

  // Find the last bucket whose value is <= x
  template <int NCases, typename RAIter>
  int find_bucket(typename std::iterator_traits<RAIter>::value_type x,
                  RAIter start) {
    return find_bucket_impl<NCases, RAIter>::go(x, start);
  }

  template <int NCases, typename RealType = double, typename RNGResult = RealType>
  class nway_bernoulli_distribution {
    RNGResult ranges[NCases + 1];

    public:
    nway_bernoulli_distribution() {}

    // Range of generator must be [min_, max_)
    template <typename RandomAccessIterator>
    explicit nway_bernoulli_distribution(RandomAccessIterator it, RNGResult min_, RNGResult max_) {
      set_probabilities(it, min_, max_);
    }

    template <typename RandomAccessIterator>
    void set_probabilities(RandomAccessIterator it, RNGResult min_, RNGResult max_) {
      RealType total = 0;
      for (size_t i = 0; i < NCases; ++i) {
	ranges[i] = RNGResult(min_ + (max_ - min_) * total);
	total += *it++;
      }
      ranges[NCases] = max_;
    }

    int compute_from_random_number(RNGResult x) const {
      return find_bucket<NCases>(x, ranges);
    }

    template <typename Gen>
    int operator()(Gen& gen) const {
      return compute_from_random_number(gen());
    }
  };

  template <int NCases, typename RealType, typename RandomAccessIterator, typename Gen>
  nway_bernoulli_distribution<NCases, RealType, typename Gen::result_type>
  make_nway_bernoulli_distribution(RandomAccessIterator it, const Gen& gen) {
    return nway_bernoulli_distribution<NCases, RealType, typename Gen::result_type>
             (it,
              (gen.min)(),
              (gen.max)() + std::numeric_limits<typename Gen::result_type>::is_integer);
  }

} } } // end namespace boost::graph::random

#endif // BOOST_NWAY_BERNOULLI_DISTRIBUTION_HPP


/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <folly/experimental/QuotientMultiSet.h>

#include <boost/sort/spreadsort/integer_sort.hpp>
#include <folly/Benchmark.h>
#include <folly/Format.h>
#include <folly/Random.h>
#include <folly/String.h>
#include <folly/container/Enumerate.h>
#include <folly/container/F14Set.h>
#include <folly/container/Foreach.h>
#include <folly/experimental/EliasFanoCoding.h>
#include <folly/experimental/test/CodingTestUtils.h>
#include <folly/init/Init.h>

DEFINE_int64(
    key_bits,
    32,
    "The number of key bits used in quotient multiset. Should always <= 64");
DEFINE_uint64(
    num_elements,
    100000000,
    "The number of elements inserted into quotient multiset");
DEFINE_double(load_factor, 0.95, "Load factor of the multiset");

#if FOLLY_QUOTIENT_MULTI_SET_SUPPORTED

using std::string;

namespace {

static const unsigned int kRunsPerIteration = 5000000;

std::mt19937 rng;

// Uniformly distributed keys.
std::vector<uint64_t> uniform;
string qmsData;

uint64_t maxValue(uint32_t nbits) {
  return nbits == 64 ? std::numeric_limits<uint64_t>::max()
                     : (uint64_t(1) << nbits) - 1;
}

// Hash 64 bits into 1.
uint64_t mix64_to_1(uint64_t word) {
  // Constant from MurmurHash3 final mixer.
  return (word * UINT64_C(0xff51afd7ed558ccd)) >> 63;
}

void buildQuotientMultiSet(std::vector<uint64_t>& keys) {
  folly::QuotientMultiSetBuilder builder(
      FLAGS_key_bits, keys.size(), FLAGS_load_factor);
  folly::IOBufQueue buff;
  for (const auto& iter : folly::enumerate(keys)) {
    if (builder.insert(*iter)) {
      builder.setBlockPayload(iter.index);
    }
    if (builder.numReadyBlocks() >= 1) {
      builder.flush(buff);
    }
  }
  builder.close(buff);
  qmsData = buff.move()->to<string>();
}

std::vector<uint64_t> makeLookupKeys(size_t n, double hitRate) {
  folly::BenchmarkSuspender guard;
  std::vector<uint64_t> keys;
  keys.reserve(n);
  uint64_t maxKey = maxValue(FLAGS_key_bits);
  for (uint64_t idx = 0; idx < n; idx++) {
    keys.push_back(
        (folly::Random::randDouble01(rng) < hitRate)
            ? uniform[folly::Random::rand64(uniform.size(), rng)]
            : folly::Random::rand64(rng) & maxKey);
  }
  return keys;
}

const folly::F14FastSet<uint64_t>& getF14Baseline() {
  folly::BenchmarkSuspender guard;
  static const auto set = [] {
    folly::F14FastSet<uint64_t> ret(uniform.begin(), uniform.end());
    LOG(INFO) << folly::sformat(
        "Built F14FastSet, size: {}, space: {}",
        ret.size(),
        folly::prettyPrint(
            ret.getAllocatedMemorySize(), folly::PrettyType::PRETTY_BYTES_IEC));
    return ret;
  }();
  return set;
}

using EFEncoder =
    folly::compression::EliasFanoEncoderV2<uint64_t, uint64_t, 128, 128>;

const folly::compression::MutableEliasFanoCompressedList& getEFBaseline() {
  folly::BenchmarkSuspender guard;
  static auto list = [] {
    auto ret = EFEncoder::encode(uniform.begin(), uniform.end());
    LOG(INFO) << folly::sformat(
        "Built Elias-Fano list, space: {}",
        folly::prettyPrint(
            ret.data.size(), folly::PrettyType::PRETTY_BYTES_IEC));
    return ret;
  }();
  return list;
}

template <class Lookup>
size_t benchmarkHits(const Lookup& lookup) {
  auto keys = makeLookupKeys(kRunsPerIteration, /* hitRate */ 1);
  for (const auto& key : keys) {
    auto found = lookup(key);
    folly::doNotOptimizeAway(found);
    DCHECK(found);
  }
  return kRunsPerIteration;
}

template <class Lookup>
size_t benchmarkRandom(const Lookup& lookup) {
  auto keys = makeLookupKeys(kRunsPerIteration, /* hitRate */ 0);
  for (const auto& key : keys) {
    auto found = lookup(key);
    folly::doNotOptimizeAway(found);
  }
  return kRunsPerIteration;
}

template <class Lookup>
size_t benchmarkMixtureSerialized(const Lookup& lookup) {
  // Make the result unpredictable so we can use it to introduce a
  // serializing dependency in the loop.
  auto keys = makeLookupKeys(kRunsPerIteration * 2, /* hitRate */ 0.5);
  size_t unpredictableBit = 0;
  for (size_t i = 0; i < kRunsPerIteration * 2; i += 2) {
    unpredictableBit = lookup(keys[i + unpredictableBit]) ? 1 : 0;
    folly::doNotOptimizeAway(unpredictableBit);
  }
  return kRunsPerIteration;
}

template <class Lookup>
size_t benchmarkSerializedOnResultBits(const Lookup& lookup, double hitRate) {
  auto keys = makeLookupKeys(kRunsPerIteration * 2, hitRate);
  size_t unpredictableBit = 0;
  for (size_t i = 0; i < kRunsPerIteration * 2; i += 2) {
    // When all keys are hits we can use a hash of the location of the
    // found key to introduce a serializing dependency in the loop.
    unpredictableBit = mix64_to_1(lookup(keys[i + unpredictableBit]));
    folly::doNotOptimizeAway(unpredictableBit);
  }
  return kRunsPerIteration;
}

template <class Lookup>
size_t benchmarkHitsSerialized(const Lookup& lookup) {
  return benchmarkSerializedOnResultBits(lookup, /* hitRate */ 1);
}

template <class Lookup>
size_t benchmarkRandomSerialized(const Lookup& lookup) {
  return benchmarkSerializedOnResultBits(lookup, /* hitRate */ 0);
}

template <class Lookup>
size_t benchmarkHitsSmallWorkingSet(const Lookup& lookup) {
  // Loop over a small set of keys so that after the first iteration
  // all the relevant parts of the data structure are in LLC cache.
  constexpr size_t kLookupSetSize = 1 << 12; // Should be a power of 2.
  auto keys = makeLookupKeys(kLookupSetSize, /* hitRate */ 1);
  for (size_t i = 0; i < kRunsPerIteration; ++i) {
    auto found = lookup(keys[i % kLookupSetSize]);
    folly::doNotOptimizeAway(found);
    DCHECK(found);
  }
  return kRunsPerIteration;
}

} // namespace

void benchmarkSetup() {
  rng.seed(UINT64_C(12345678));

  uint64_t maxKey = maxValue(FLAGS_key_bits);
  // Uniformly distributed keys.
  const uint64_t size = FLAGS_num_elements;
  for (uint64_t idx = 0; idx < size; idx++) {
    uint64_t key = folly::Random::rand64(rng) & maxKey;
    uniform.emplace_back(key);
  }
  boost::sort::spreadsort::integer_sort(uniform.begin(), uniform.end());
  buildQuotientMultiSet(uniform);

  LOG(INFO) << folly::sformat(
      "Built QuotientMultiSet, space: {}",
      folly::prettyPrint(qmsData.size(), folly::PrettyType::PRETTY_BYTES_IEC));
}

BENCHMARK_MULTI(QuotientMultiSetGetHits) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkHits([&](uint64_t key) { return reader.equalRange(key); });
  });
  return ret;
}

BENCHMARK_MULTI(QuotientMultiSetGetRandom) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkRandom([&](uint64_t key) { return reader.equalRange(key); });
  });
  return ret;
}

BENCHMARK_MULTI(QuotientMultiSetGetHitsSmallWorkingSet) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkHitsSmallWorkingSet(
        [&](uint64_t key) { return reader.equalRange(key); });
  });
  return ret;
}

BENCHMARK_MULTI(QuotientMultiSetGetMixtureSerialized) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkMixtureSerialized(
        [&](uint64_t key) { return reader.equalRange(key); });
  });
  return ret;
}

BENCHMARK_MULTI(QuotientMultiSetGetHitsSerialized) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkHitsSerialized(
        [&](uint64_t key) { return reader.equalRange(key).begin; });
  });
  return ret;
}

BENCHMARK_MULTI(QuotientMultiSetGetRandomSerialized) {
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = folly::QuotientMultiSet<decltype(instructions)>(qmsData);
    ret = benchmarkRandomSerialized([&](uint64_t key) {
      // Even without a match, begin will depend on the whole
      // computation.
      return reader.equalRange(key).begin;
    });
  });
  return ret;
}

BENCHMARK_DRAW_LINE();

BENCHMARK_MULTI(F14MapHits) {
  const auto& baseline = getF14Baseline();
  return benchmarkHits(
      [&](uint64_t key) { return baseline.find(key) != baseline.end(); });
}

BENCHMARK_MULTI(F14MapRandom) {
  const auto& baseline = getF14Baseline();
  return benchmarkRandom(
      [&](uint64_t key) { return baseline.find(key) != baseline.end(); });
}

BENCHMARK_MULTI(F14MapHitsSmallWorkingSet) {
  const auto& baseline = getF14Baseline();
  return benchmarkHitsSmallWorkingSet(
      [&](uint64_t key) { return baseline.find(key) != baseline.end(); });
}

BENCHMARK_MULTI(F14MapMixtureSerialized) {
  const auto& baseline = getF14Baseline();
  return benchmarkMixtureSerialized(
      [&](uint64_t key) { return baseline.find(key) != baseline.end(); });
}

BENCHMARK_MULTI(F14MapHitsSerialized) {
  const auto& baseline = getF14Baseline();
  return benchmarkHitsSerialized([&](uint64_t key) {
    return reinterpret_cast<uintptr_t>(&*baseline.find(key));
  });
}

// benchmarkRandomSerialized() cannot be used with F14.

BENCHMARK_DRAW_LINE();

using folly::compression::EliasFanoReader;

BENCHMARK_MULTI(EliasFanoGetHits) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkHits([&](uint64_t key) {
      return reader.jumpTo(key) && reader.value() == key;
    });
  });
  return ret;
}

BENCHMARK_MULTI(EliasFanoGetRandom) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkRandom([&](uint64_t key) {
      return reader.jumpTo(key) && reader.value() == key;
    });
  });
  return ret;
}

BENCHMARK_MULTI(EliasFanoGetHitsSmallWorkingSet) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkHitsSmallWorkingSet([&](uint64_t key) {
      return reader.jumpTo(key) && reader.value() == key;
    });
  });
  return ret;
}

BENCHMARK_MULTI(EliasFanoGetMixtureSerialized) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkMixtureSerialized([&](uint64_t key) {
      return reader.jumpTo(key) && reader.value() == key;
    });
  });
  return ret;
}

BENCHMARK_MULTI(EliasFanoGetHitsSerialized) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkHitsSerialized([&](uint64_t key) {
      reader.jumpTo(key);
      DCHECK(reader.valid());
      return reader.position();
    });
  });
  return ret;
}

BENCHMARK_MULTI(EliasFanoGetRandomSerialized) {
  auto list = getEFBaseline();
  size_t ret = 0;
  folly::compression::dispatchInstructions([&](auto instructions) {
    auto reader = EliasFanoReader<EFEncoder, decltype(instructions)>(list);
    ret = benchmarkRandomSerialized([&](uint64_t key) {
      reader.jumpTo(key);
      return reader.position();
    });
  });
  return ret;
}

#endif // FOLLY_QUOTIENT_MULTI_SET_SUPPORTED

int main(int argc, char** argv) {
  folly::init(&argc, &argv);

#if FOLLY_QUOTIENT_MULTI_SET_SUPPORTED
  benchmarkSetup();
  folly::runBenchmarks();
#endif // FOLLY_QUOTIENT_MULTI_SET_SUPPORTED
}

#if 0
// Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz, built with Clang.
$ buck run @mode/opt folly/experimental/test:quotient_multiset_benchmark -- --benchmark --bm_max_secs 3

I0917 10:28:54.497815 2874547 QuotientMultiSetBenchmark.cpp:195] Built QuotientMultiSet, space: 124.9 MiB
============================================================================
folly/experimental/test/QuotientMultiSetBenchmark.cpprelative  time/iter  iters/s
============================================================================
QuotientMultiSetGetHits                                    114.94ns    8.70M
QuotientMultiSetGetRandom                                   72.87ns   13.72M
QuotientMultiSetGetHitsSmallWorkingSet                      58.75ns   17.02M
QuotientMultiSetGetMixtureSerialized                       138.73ns    7.21M
QuotientMultiSetGetHitsSerialized                          138.56ns    7.22M
QuotientMultiSetGetRandomSerialized                        130.11ns    7.69M
----------------------------------------------------------------------------
I0917 10:29:33.808831 2874547 QuotientMultiSetBenchmark.cpp:83] Built F14FastSet, size: 98843868, space: 1 GiB
F14MapHits                                                  34.69ns   28.83M
F14MapRandom                                                47.22ns   21.18M
F14MapHitsSmallWorkingSet                                   14.42ns   69.34M
F14MapMixtureSerialized                                     94.64ns   10.57M
F14MapHitsSerialized                                       140.74ns    7.11M
----------------------------------------------------------------------------
I0917 10:29:52.722596 2874547 QuotientMultiSetBenchmark.cpp:100] Built Elias-Fano list, space: 101.5 MiB
EliasFanoGetHits                                           258.08ns    3.87M
EliasFanoGetRandom                                         247.18ns    4.05M
EliasFanoGetHitsSmallWorkingSet                             81.54ns   12.26M
EliasFanoGetMixtureSerialized                              253.63ns    3.94M
EliasFanoGetHitsSerialized                                 161.57ns    6.19M
EliasFanoGetRandomSerialized                               161.55ns    6.19M
============================================================================
#endif


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file    DoglegOptimizer.cpp
 * @brief
 * @author  Richard Roberts
 * @date 	Feb 26, 2012
 */

#include <gtsam/nonlinear/DoglegOptimizer.h>

#include <gtsam/inference/EliminationTree.h>
#include <gtsam/linear/GaussianJunctionTree.h>
#include <gtsam/nonlinear/DoglegOptimizerImpl.h>

#include <boost/algorithm/string.hpp>

using namespace std;

namespace gtsam {

/* ************************************************************************* */
DoglegParams::VerbosityDL DoglegParams::verbosityDLTranslator(const std::string &verbosityDL) const {
	std::string s = verbosityDL;  boost::algorithm::to_upper(s);
	if (s == "SILENT") return DoglegParams::SILENT;
	if (s == "VERBOSE") return DoglegParams::VERBOSE;

	/* default is silent */
	return DoglegParams::SILENT;
}

/* ************************************************************************* */
std::string DoglegParams::verbosityDLTranslator(VerbosityDL verbosityDL) const {
	std::string s;
	switch (verbosityDL) {
	case DoglegParams::SILENT:  s = "SILENT"; break;
	case DoglegParams::VERBOSE: s = "VERBOSE"; break;
	default:                    s = "UNDEFINED"; break;
	}
	return s;
}

/* ************************************************************************* */
void DoglegOptimizer::iterate(void) {

  // Linearize graph
  const Ordering& ordering = *params_.ordering;
  GaussianFactorGraph::shared_ptr linear = graph_.linearize(state_.values, ordering);

  // Pull out parameters we'll use
  const bool dlVerbose = (params_.verbosityDL > DoglegParams::SILENT);

  // Do Dogleg iteration with either Multifrontal or Sequential elimination
  DoglegOptimizerImpl::IterationResult result;

  if ( params_.isMultifrontal() ) {
    GaussianBayesTree bt;
    bt.insert(GaussianJunctionTree(*linear).eliminate(params_.getEliminationFunction()));
    result = DoglegOptimizerImpl::Iterate(state_.Delta, DoglegOptimizerImpl::ONE_STEP_PER_ITERATION, bt, graph_, state_.values, ordering, state_.error, dlVerbose);
  }
  else if ( params_.isSequential() ) {
    GaussianBayesNet::shared_ptr bn = EliminationTree<GaussianFactor>::Create(*linear)->eliminate(params_.getEliminationFunction());
    result = DoglegOptimizerImpl::Iterate(state_.Delta, DoglegOptimizerImpl::ONE_STEP_PER_ITERATION, *bn, graph_, state_.values, ordering, state_.error, dlVerbose);
  }
  else if ( params_.isCG() ) {
    throw runtime_error("todo: ");
  }
  else {
    throw runtime_error("Optimization parameter is invalid: DoglegParams::elimination");
  }

  // Maybe show output
  if(params_.verbosity >= NonlinearOptimizerParams::DELTA) result.dx_d.print("delta");

  // Create new state with new values and new error
  state_.values = state_.values.retract(result.dx_d, ordering);
  state_.error = result.f_error;
  state_.Delta = result.Delta;
  ++state_.iterations;
}

} /* namespace gtsam */


#ifndef FSCUTILS_HPP
#define FSCUTILS_HPP
#include <boost/graph/adjacency_list.hpp>
#include "FSC.hpp"
#include "NodeConfigurations.hpp"
#include "decpomdp/discrete/DiscreteElements.hpp"
#include "decpomdp/discrete/DiscreteSpaces.hpp"
namespace npgi {

template <typename Index, typename TimeStep>
NodeConfigurationSet<DiscreteLocalAction<Index>,
                        DiscreteLocalObservation<Index>,
                        fsc_node_t>
get_configurations(
    fsc_node_t q, const DiscreteLocalActionSpace<Index>& a,
    const DiscreteLocalObservationSpace<Index>& z,
    const fsc_graph_t<DiscreteLocalAction<Index>,
                      DiscreteLocalObservation<Index>, TimeStep>& g) {
  const int next_layer = g[q].time_ + 1;
  const auto next_nodes = get_layer(next_layer, g);
  auto edge_cfgs = make_edge_configuration_set(
      std::vector<DiscreteLocalObservation<Index>>(z.begin(), z.end()),
      next_nodes);
  return make_node_configuration_set(std::vector<DiscreteLocalAction<Index>>(a.begin(), a.end()), edge_cfgs);;
}

template <typename Index, typename TimeStep>
fsc_graph_t<DiscreteLocalAction<Index>, DiscreteLocalObservation<Index>, TimeStep>
get_fsc(const std::vector<std::size_t>& layer_width, const Agent<Index>& agent,
        const JointActionSpaceFlat<Index, TimeStep>& A,
        const JointObservationSpaceFlat<Index, TimeStep>& Z) {
  fsc_graph_t<DiscreteLocalAction<Index>, DiscreteLocalObservation<Index>, TimeStep> g;

  for (std::size_t l = 0; l < layer_width.size(); ++l)
    add_layer(g, l, layer_width[l], A.local_action_space(agent, l),
              Z.local_observation_space(agent, l + 1));

  // add end node
  const std::size_t last_layer = layer_width.size();
  add_layer(g, layer_width.size(), 1, A.local_action_space(agent, last_layer),
            Z.local_observation_space(agent, last_layer + 1));
  return g;
}

template <typename Graph, typename Index>
void add_layer(Graph& g, int layer_idx, std::size_t layer_width,
               const DiscreteLocalActionSpace<Index>& a,
               const DiscreteLocalObservationSpace<Index>& z) {
  std::vector<typename Graph::vertex_descriptor> added;
  for (std::size_t i = 0; i < layer_width; ++i) {
    added.push_back(boost::add_vertex({*a.begin(), layer_idx}, g));
  }

  if (layer_idx == 0) return;

  auto prev_layer = get_layer(layer_idx - 1, g);
  if (prev_layer.empty()) throw std::runtime_error("Previous layer empty!");

  const typename Graph::vertex_descriptor dst = added.front();
  for (const auto& src : prev_layer) {
    for (auto it = z.begin(); it != z.end(); ++it) {
      boost::add_edge(src, dst, *it, g);
    }
  }
}

template <typename Graph, typename Index, typename TimeStep,
          typename RandomNumberGenerator>
void randomize(Graph& g, const Agent<Index>& agent,
               const JointActionSpaceFlat<Index, TimeStep>& A,
               const JointObservationSpaceFlat<Index, TimeStep>& Z,
               RandomNumberGenerator& gen) {
  auto vit = boost::vertices(g);
  std::vector<fsc_node_t> nodes(vit.first, vit.second);
  // sort by layer, have to randomize starting from the end.
  std::sort(nodes.begin(), nodes.end(),
            [&g](const fsc_node_t& a, const fsc_node_t& b) {
              return g[a].time_ > g[b].time_;
            });
  for (auto q : nodes) {
  	auto layer = g[q].time_;
    randomize_node_configuration(q, g, A.local_action_space(agent, layer),
              Z.local_observation_space(agent, layer + 1), gen);
  }
}

template <typename Graph, typename Index, typename RandomNumberGenerator>
void randomize_node_configuration(fsc_node_t node, Graph& g,
                                  const DiscreteLocalActionSpace<Index>& a,
                                  const DiscreteLocalObservationSpace<Index>& z,
                                  RandomNumberGenerator& gen) {
  using local_action_type = typename DiscreteLocalActionSpace<Index>::local_element_type;
  using local_observation_type = typename DiscreteLocalObservationSpace<Index>::local_element_type;

  const int next_layer = g[node].time_ + 1;
  const auto next_nodes = get_layer(next_layer, g);

  auto edge_cfgs = make_edge_configuration_set(std::vector<local_observation_type>(z.begin(), z.end()), next_nodes);
  auto node_cfgs = make_node_configuration_set(std::vector<local_action_type>(a.begin(), a.end()), edge_cfgs);
  set_unique_random_configuration(node_cfgs, node, g, gen);
}

template <typename Graph, typename NodeConfigurationSet,
          typename RandomNumberGenerator>
void set_unique_random_configuration(const NodeConfigurationSet& confs,
                                     fsc_node_t node, Graph& g,
                                     RandomNumberGenerator& gen) {
  auto this_l = get_layer(g[node].time_, g);
  this_l.erase(std::remove(this_l.begin(), this_l.end(), node),
               this_l.end());  // without node
  do {
    set_configuration(*reservoir_sample(confs.begin(), confs.end(), gen), node,
                      g);
  } while (std::any_of(
      this_l.begin(), this_l.end(),
      [&node, &g](fsc_node_t other) { return same_policy(node, other, g); }));
}

}  // namespace npgi

#endif  // FSCUTILS_HPP


#ifndef TVTML_FUNCTIONAL3D_HPP
#define TVTML_FUNCTIONAL3D_HPP

// Eigen includes
#include <Eigen/Core>
#include <Eigen/Sparse>

// video++ includes
#include <vpp/vpp.hh>

// system includes
#include <cmath>
#include <iostream>
#include <fstream>
#include <vector>

// own includes
#include "enumerators.hpp"
#include "manifold.hpp"
#include "data.hpp"

namespace tvmtl{

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA>
class Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>{

    public:
	// Manifold typedefs and constants
	static const MANIFOLD_TYPE mf_type = MANIFOLD::MyType;
	static const int value_dim = MANIFOLD::value_dim;
	static const int manifold_dim = MANIFOLD::manifold_dim;
	typedef typename MANIFOLD::scalar_type scalar_type;
	typedef typename MANIFOLD::value_type value_type;
	typedef typename MANIFOLD::ref_type ref_type;
	typedef typename MANIFOLD::cref_type cref_type;
	typedef typename MANIFOLD::deriv1_type deriv1_type;
	typedef typename MANIFOLD::deriv2_type deriv2_type;
	typedef typename MANIFOLD::restricted_deriv2_type restricted_deriv2_type;
	typedef typename MANIFOLD::tm_base_type tm_base_type;

	// Data typedef and constants
	static const int img_dim = DATA::img_dim;
	typedef typename DATA::storage_type img_type;
	typedef typename DATA::weights_type weights_type;
	typedef typename DATA::weights_mat weights_mat;
	typedef typename DATA::inp_mat inp_mat;

	// Functional parameters and return types
	static const FUNCTIONAL_DISC disc_type;
	typedef double param_type;
	typedef double result_type;

	// Tangent space transformation matrix types
	typedef vpp::imageNd<tm_base_type, img_dim> tm_base_mat_type;

	// Gradient and Hessian types
	typedef Eigen::Matrix<scalar_type, Eigen::Dynamic, 1> gradient_type;
	typedef vpp::imageNd<deriv2_type, img_dim> hessian_type;
	typedef Eigen::SparseMatrix<scalar_type> sparse_hessian_type;

	//Constructor
	Functional(param_type lambda, DATA& dat):
	    lambda_(lambda),
	    data_(dat)
	{
	    eps2_=1e-10;
	   static_assert(img_dim == 3, "Dimension of data and functional must match!");
	}

	void updateWeights();

	void updateTMBase();

	// Evaluation functions
	result_type evaluateJ();
	void  evaluateDJ();
	void  evaluateHJ();

	template <class IMG>
	void output_img(const IMG& img, const char* filename) const;
	template <class IMG>
	void output_matval_img(const IMG& img, const char* filename) const;

	// Getter and Setter
	inline param_type getlambda() const { return lambda_; }
	inline void setlambda(param_type lam) { lambda_=lam; }
	inline param_type geteps2() const { return eps2_; }
	inline void seteps2(param_type eps) { eps2_=eps; }

	inline const weights_mat& getweightsX() const { return weightsX_; }
	inline const weights_mat& getweightsY() const { return weightsY_; }
	inline const weights_mat& getweightsZ() const { return weightsZ_; }

	inline const gradient_type& getDJ() const { return DJ_; }
	inline const sparse_hessian_type& getHJ() const { return HJ_; }
	inline const tm_base_mat_type& getT() const { return T_; }

    private:
	DATA& data_;

	param_type lambda_, eps2_;
	weights_mat weightsX_, weightsY_, weightsZ_;

	tm_base_mat_type T_;
	gradient_type DJ_;
	sparse_hessian_type HJ_;
};

//--------Implementation FIRSTORDER-----/

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
const FUNCTIONAL_DISC Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>::disc_type = disc;

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA>
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>::updateWeights(){

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t Update Weights..." << std::endl;
    #endif

    // Neighbourhood box
    int ns = data_.img_.nslices();  // z
    int nr = data_.img_.nrows();    // y
    int nc = data_.img_.ncols();    // x

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    weightsX_ = weights_mat(data_.img_.domain());
    weightsY_ = weights_mat(data_.img_.domain());
    weightsZ_ = weights_mat(data_.img_.domain());

    auto calc_dist = [&] (weights_type& w, const value_type i, const value_type n) {
	w = MANIFOLD::dist_squared(i, n);
    };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...X neighbours " << std::endl;
    #endif

    // X Neighbours
    fill3d(weightsX_, 0.0);
    pixel_wise3d(calc_dist, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
	std::cout << "\t\t...Y neighbours" << std::endl;
    #endif

    // Y Neighbours
    fill3d(weightsY_, 0.0);
    pixel_wise3d(calc_dist, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
	std::cout << "\t\t...Z neighbours" << std::endl;
    #endif

    // Z Neighbours
    fill3d(weightsZ_, 0.0);
    pixel_wise3d(calc_dist, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
    #endif

    if(disc==ISO){
	auto g =  [&] (const weights_type& ew, weights_type& x, weights_type& y, weights_type& z) { x = ew / std::sqrt(x + y + z + eps2_); y = x; z = x; };
	pixel_wise3d(g, data_.edge_weights_, weightsX_, weightsY_, weightsZ_);
    }
    else{
	auto g =  [&] (const weights_type& ew, weights_type& w) { w = ew / std::sqrt(w+eps2_); };
	pixel_wise3d(g, data_.edge_weights_, weightsX_);
	pixel_wise3d(g, data_.edge_weights_, weightsY_);
	pixel_wise3d(g, data_.edge_weights_, weightsZ_);
    }
}

// Update the Tangent space ONB
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::updateTMBase(){

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tUpdate tangent space basis..." << std::endl;
    #endif

   tm_base_mat_type T(data_.img_.domain());
   pixel_wise3d([&] (tm_base_type& t, const value_type& i) { MANIFOLD::tangent_plane_base(i,t); }, T, data_.img_);
   T_=T;

}

// Evaluation of J
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
typename Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::result_type Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateJ(){

    // sum d^2(img, img_noise)
    result_type J1, J2;
    J1 = J2 = 0.0;

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tFunctional evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif

    if(data_.doInpaint()){
	auto f = [&] (const value_type& i, const value_type& n, const bool inp ) { J1 += MANIFOLD::dist_squared(i,n)*(1-inp); };
	pixel_wise3d_nothreads(f, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [&] (const value_type& i, const value_type& n) { J1 += MANIFOLD::dist_squared(i,n); };
	pixel_wise3d_nothreads(f, data_.img_, data_.noise_img_);
    }

	updateWeights();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...TV part." << std::endl;
    #endif

    if(disc==ISO)
	pixel_wise3d_nothreads([&] (const weights_type& w) { J2 += 1.0/w;} ,weightsX_);
    else
	pixel_wise3d_nothreads([&] (const weights_type& wx, const weights_type& wy, const weights_type& wz) { J2 += 1.0 / wx + 1.0 / wy+ 1.0 / wz;}, weightsX_, weightsY_, weightsZ_);
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "J1: " << J1 << std::endl;
	std::cout << "J2: " << J2 << std::endl;
    #endif

    return 0.5 * J1 + lambda_* J2;
}

// Evaluation of J'
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateDJ(){

    img_type grad = img_type(data_.img_.domain());
    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tGradient evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif
    //GRADIENT OF FIDELITY TERM
    if(data_.doInpaint()){
	auto f = [] (value_type& g, const value_type& i, const value_type& n, const bool inp ) { MANIFOLD::deriv1x_dist_squared(i,n,g); g*=(1-inp); };
	pixel_wise3d(f, grad, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [] (value_type& g, const value_type& i, const value_type& n) { MANIFOLD::deriv1x_dist_squared(i,n,g); };
	pixel_wise3d(f, grad, data_.img_, data_.noise_img_);
    }

    //GRADIENT OF TV TERM

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    auto calc_first_arg_deriv = [&] (value_type& x, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv1x_dist_squared(i, n, x); x *= w; };
    auto calc_second_arg_deriv = [&] (value_type& y, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv1y_dist_squared(i, n, y); y *= w; };
    auto add_to_gradient = [&] (value_type& g, const value_type& d) { g+=d*lambda_; };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tGradient evaluation..." << std::endl;
	std::cout << "\t\t...TV part" << std::endl;
	std::cout << "\t\t...-> XD1" << std::endl;
    #endif
    // X neighbors and reweighting
    // ... w.r.t. to first argument
    { // Temporary image XD1 is deallocated after this scope
	img_type XD1 = img_type(without_last_x);
	pixel_wise3d(calc_first_arg_deriv, XD1, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_gradient, grad | without_last_x, XD1);
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> XD2" << std::endl;
    #endif
    // ... w.r.t. second argument
    {
	img_type XD2 = img_type(without_last_x);
	pixel_wise3d(calc_second_arg_deriv, XD2, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_gradient, grad | without_first_x,  XD2);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> YD1" << std::endl;
    #endif
    // Vertical derivatives and weighting
    // ... w.r.t. first argument
    {
	img_type YD1 = img_type(without_last_y);
	pixel_wise3d(calc_first_arg_deriv, YD1, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_gradient, grad | without_first_y, YD1);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> YD2" << std::endl;
    #endif
    // ... w.r.t second argument
    {
	img_type YD2 = img_type(without_last_y);
	pixel_wise3d(calc_second_arg_deriv, YD2, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
        pixel_wise3d(add_to_gradient, grad | without_first_y, YD2);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> ZD1" << std::endl;
    #endif
    // Vertical derivatives and weighting
    // ... w.r.t. first argument
    {
	img_type ZD1 = img_type(without_last_z);
	pixel_wise3d(calc_first_arg_deriv, ZD1, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_gradient, grad | without_first_z, ZD1);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> ZD2" << std::endl;
    #endif
    // ... w.r.t second argument
    {
	img_type ZD2 = img_type(without_last_z);
	pixel_wise3d(calc_second_arg_deriv, ZD2, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
        pixel_wise3d(add_to_gradient, grad | without_first_z, ZD2);
    }

    #ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(grad,"3dgrad.csv");
    #endif

    DJ_ = gradient_type::Zero(ns*nr*nc*manifold_dim);

    updateTMBase();
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif

   for(int s = 0; s < ns; ++s){
	//#pragma omp parallel for
	for(int r = 0; r < nr; ++r){
	    // Start of row pointers
	    value_type* p = &grad(s, r, 0);
	    tm_base_type* t = &T_(s, r, 0);
	    for(int c = 0; c < nc; ++c)
		DJ_.segment(manifold_dim * (nr * nc * s + nc * r + c), manifold_dim) = t[c].transpose() * Eigen::Map<const Eigen::VectorXd>(p[c].data(), p[c].size());
	//	DJ_.segment(manifold_dim * (s + ns * r + ns * nr * c), manifold_dim) = t[c].transpose() * Eigen::Map<const Eigen::VectorXd>(p[c].data(), p[c].size());
	}
    }

    #ifdef TV_FUNC_DEBUG
	std::fstream f;
	f.open("3dDJ.csv",std::fstream::out);
	f << DJ_;
	f.close();
    #endif
}

// Evaluation of Hessian J
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateHJ(){
    hessian_type hessian(data_.img_.domain());

    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();
    int sparsedim = ns*nr*nc*manifold_dim;

    //HESSIAN OF FIDELITY TERM
     #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tHessian evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif

    sparse_hessian_type HF(sparsedim,sparsedim);

    //HF.reserve(Eigen::VectorXi::Constant(nc,manifold_dim));
    typedef Eigen::Triplet<double> Trip;
    std::vector<Trip> triplist;
    triplist.reserve(sparsedim*manifold_dim);

    if(data_.doInpaint()){
	auto f = [] (deriv2_type& h, const value_type& i, const value_type& n, const bool inp ) { MANIFOLD::deriv2xx_dist_squared(i,n,h); h*=(1-inp); };
	pixel_wise3d(f, hessian, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [] (deriv2_type& h, const value_type& i, const value_type& n) { MANIFOLD::deriv2xx_dist_squared(i,n,h); };
	pixel_wise3d(f, hessian, data_.img_, data_.noise_img_);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif
   for(int s = 0; s < ns; ++s){
    for(int r = 0; r < nr; ++r){
	    // Start of row pointers
	    deriv2_type* h = &hessian(s, r, 0);
	    tm_base_type* t = &T_(s, r, 0);
	    for(int c = 0; c < nc; ++c){
		int pos = manifold_dim * (nr * nc * s + nc * r + c); // rowwise flattening
		//int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		restricted_deriv2_type ht=t[c].transpose()*h[c]*t[c];

		for(int local_row = 0; local_row<ht.rows(); local_row++)
		    for(int local_col = local_row; local_col < ht.cols(); local_col++){
			scalar_type e = ht(local_row, local_col);
			    if(e!=0){
				int global_row = pos + local_row;
				int global_col = pos + local_col;
				triplist.push_back(Trip(global_row,global_col,e));
				if(global_row != global_col)
				    triplist.push_back(Trip(global_col,global_row,e));
			    }
		    }

	    }
	}
    }

     #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Triplet list created" << std::endl;
    #endif
    HF.setFromTriplets(triplist.begin(),triplist.end());
    HF.makeCompressed();

    //HESSIAN OF TV TERM
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tHessian evaluation..." << std::endl;
	std::cout << "\t\t...TV part" << std::endl;
    #endif

    sparse_hessian_type HTV(sparsedim,sparsedim);

    //HTV.reserve(Eigen::VectorXi::Constant(nc,5*manifold_dim));
    triplist.clear();
    triplist.reserve(7 * sparsedim*manifold_dim);

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    auto calc_xx_der = [&] (deriv2_type& x, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2xx_dist_squared(i, n, x); x*=w; };
    auto calc_xy_der = [&] (deriv2_type& xy, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2xy_dist_squared(i, n, xy); xy*=w; };
    auto calc_yy_der = [&] (deriv2_type& y, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2yy_dist_squared(i, n, y); y*=w; };
    auto add_to_hessian =  [&] (deriv2_type& h, const deriv2_type& d) { h += d; };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD11" << std::endl;
    #endif

    // Horizontal Second Derivatives and weighting
    // ... w.r.t. first arguments
    { // Temporary image XD11 is deallocated after this scope
	hessian_type XD11(without_last_x);
        pixel_wise3d(calc_xx_der, XD11, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d([&] (deriv2_type& h, const deriv2_type& d) { h=d; }, hessian | without_last_x, XD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD11,"3dXD11.csv");
	#endif
    }

    for(int s=0; s < ns; ++s)
	for(int r=0; r<nr; ++r)
	    hessian(s,r,nc-1)=deriv2_type::Zero();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type XD22(without_last_x);
        pixel_wise3d(calc_yy_der, XD22, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_hessian, hessian | without_first_x, XD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD22,"3dXD22.csv");
	#endif
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD11" << std::endl;
    #endif
    // Vertical Second Derivatives weighting
    //... w.r.t. first arguments
    {
	hessian_type YD11(without_last_y);
        pixel_wise3d(calc_xx_der, YD11, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_hessian, hessian | without_last_y, YD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD11,"3dYD11.csv");
	#endif
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type YD22(without_last_y);
        pixel_wise3d(calc_yy_der, YD22, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_hessian, hessian | without_first_y, YD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD22,"3dYD22.csv");
	#endif
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD11" << std::endl;
    #endif
    // Vertical Second Derivatives weighting
    //... w.r.t. first arguments
    {
	hessian_type ZD11(without_last_z);
        pixel_wise3d(calc_xx_der, ZD11, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_hessian, hessian | without_last_z, ZD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD11,"3dZD11.csv");
	#endif
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type ZD22(without_last_z);
        pixel_wise3d(calc_yy_der, ZD22, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_hessian, hessian | without_first_z, ZD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD22,"3dZD22.csv");
	#endif
    }
    	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(hessian,"3dhessian.csv");
	#endif
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif
    // Insert elementwise into sparse Hessian
    // NOTE: Eventually make single version for both cases, including an offset
    // --> additional parameters sparse_mat, offset
    auto HTV_insert = [&]( hessian_type H, tm_base_mat_type T1, tm_base_mat_type T2, int row_offset, int col_offset){
	int Hns = H.nslices();
	int Hnr = H.nrows();
	int Hnc = H.ncols();

	for(int s = 0; s < Hns; ++s){
	    for(int r = 0; r < Hnr; ++r){
		// Start of row pointers
		deriv2_type* h = &H(s, r, 0);
		tm_base_type* t1 = &T1(s, r, 0);
		tm_base_type* t2 = &T2(s, r, 0);
		for(int c = 0; c < Hnc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
		   // int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    restricted_deriv2_type ht = t1[c].transpose()*h[c]*t2[c];

		    for(int local_row = 0; local_row<ht.rows(); local_row++)
			for(int local_col = 0; local_col < ht.cols(); local_col++){
			    scalar_type e = ht(local_row, local_col);
				if(e!=0){
				    int global_row = pos + row_offset + local_row;
				    int global_col = pos + col_offset + local_col;
				    triplist.push_back(Trip(global_row,global_col,e));
				    if(row_offset > 0 || col_offset > 0)
					triplist.push_back(Trip(global_col,global_row,e));
				}
			}

		}
	    }
	}
    };

    auto HTV_single_insert = [&] (deriv2_type& h, tm_base_type& t1, tm_base_type& t2, int pos, int row_offset, int col_offset){
   		    restricted_deriv2_type ht = t1.transpose()*h*t2;
		    for(int local_row = 0; local_row<ht.rows(); local_row++)
			for(int local_col = 0; local_col < ht.cols(); local_col++){
			    scalar_type e = ht(local_row, local_col);
				if(e!=0){
				    int global_row = pos + row_offset + local_row;
				    int global_col = pos + col_offset + local_col;
				    triplist.push_back(Trip(global_row,global_col,e));
				    if(row_offset > 0 || col_offset > 0)
					triplist.push_back(Trip(global_col,global_row,e));
				}
			}
    };

    HTV_insert(hessian, T_, T_, 0, 0);

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD12" << std::endl;
    #endif
    // Z Neighbors Second Derivatives and weighting
    // ... w.r.t. first and second arguments
    {
	hessian_type ZD12(without_last_z);
	pixel_wise3d(calc_xy_der, ZD12, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z );
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD12,"3dZD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif
	// Offsets for upper nx*ny-th subdiagonal
	int offset =  manifold_dim * nr * nc;
	HTV_insert(ZD12, T_ | without_last_z, T_ | without_first_z, 0, offset);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD12" << std::endl;
    #endif
    // Y Neighbors Second Derivatives and weighting
    //... w.r.t. second arguments
    {
	hessian_type YD12(data_.img_.domain());
	pixel_wise3d(calc_xy_der, YD12 | without_last_y, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y );

	//Set last y-slice to zero
	#pragma omp parallel for
	for(int s = 0; s < ns ; ++s){
	    deriv2_type* row = &YD12(s, nr-1 ,0);
	    for(int c = 0; c < nc; ++c)
		row[c] = deriv2_type::Zero();
	}
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD12,"3dYD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif
	// Offsets for first nxth subdiagonal
	int offset =  manifold_dim * nc;

	// Insert YD12 except last ns entries
	auto t1_it = T_.begin();
	auto t2_it = typename tm_base_mat_type::iterator(vpp::vint3(0,1,0), T_); // Iterator at beginning of second row
	bool do_break = false;
	int k = 0;
	int max_entry = ns * nr * nc - nc;

	for(int s=0; s < ns; ++s){
	    for(int r=0; r < nr; ++r){
		 deriv2_type* row = &YD12(s, r, 0);
		 for(int c=0; c < nc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
	//	    int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    HTV_single_insert(row[c], *t1_it, *t2_it, pos, 0, offset);
		    t1_it.next();
		    t2_it.next();
		    #ifdef TV_FUNC_DEBUG_VERBOSE2
			std::cout << "\n\t\t\tPosition: (" << s << ","<< c << "," << r << ")" << std::endl;
			std::cout << "\t\t\tValue(YD12): " << row[c] << std::endl;
			std::cout << "\t\t\tEntry number: " << k << std::endl;
			std::cout << "\t\t\tPos: " << pos << std::endl;
		    #endif
		    ++k;
		    if(k >= max_entry) {do_break = true; break;}
		 }
		 if(do_break) break;
	    }
	    if(do_break) break;
	}

    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD12" << std::endl;
    #endif
    // Z neightbors Second Derivatives and weighting
    //... w.r.t. second arguments
    {
	hessian_type XD12(data_.img_.domain());
	pixel_wise3d(calc_xy_der, XD12 | without_last_x, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);

	//Set last slice to zero, - actually not necessary but safer than keeping uninitialized data
	#pragma omp parallel for
	for(int s = 0; s < ns; ++s){
	    for(int r = 0; r < nr; ++r)
		XD12(s, r, nc-1) = deriv2_type::Zero();
	}

	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD12,"3dXD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif

	// Insert into first upper subdiagonal
	auto t1_it = T_.begin();
	auto t2_it = T_.begin(); t2_it.next();
	for(int s=0; s < ns; ++s){
	    for(int r=0; r < nr; ++r){
		 deriv2_type* row = &XD12(s, r, 0);
		 for(int c=0; c < nc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
		    //int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    HTV_single_insert(row[c], *t1_it, *t2_it, pos, 0, manifold_dim);
		    t1_it.next();
		    t2_it.next();
		    if(s==ns-1 && r == nc-1 && c >= nc-2 ) break;
		 }
	    }
	}

	HTV.setFromTriplets(triplist.begin(),triplist.end());
	HTV.makeCompressed();
	triplist.clear();
    }
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Combine Fidelity and TV parts:" << std::endl;
	#endif

	HJ_= HF + lambda_*HTV;

	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Output Hessian (stats):" << std::endl;
	#endif
    #ifdef TV_FUNC_DEBUG
	if (sparsedim<200){
	    if(sparsedim<80){
		std::cout << "\nFidelity\n" << HF << std::endl;
		std::cout << "\nTV\n" << HTV << std::endl;
		std::cout << "\nHessian\n" << HJ_ << std::endl;
	    }

	    std::fstream f;
	    f.open("H.csv",std::fstream::out);
	    Eigen::IOFormat CommaInitFmt(Eigen::StreamPrecision, Eigen::DontAlignCols, ", ", ", ", "", "", "", "");
	    //Eigen::MatrixXd M(HJ_);
	    //f << M.format(CommaInitFmt);
	    f << HJ_;
	    f.close();

	}
	else{
	    std::cout << "\nFidelity Non-Zeros: " << HJ_.nonZeros() << std::endl;
	    std::cout << "\nTV Non-Zeros: " << HJ_.nonZeros() << std::endl;
	    std::cout << "\nHessian Non-Zeros: " << HJ_.nonZeros() << std::endl;
	}
    // Test Solver:
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Test Solve" << std::endl;
	#endif
	gradient_type x;

	Eigen::SparseLU<sparse_hessian_type> solver;
	solver.analyzePattern(HJ_);
	solver.factorize(HJ_);
	x = solver.solve(DJ_);

	std::fstream f;
	f.open("Sol.csv",std::fstream::out);
	f << x;
	f.close();
    #endif
}

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
template < class IMG >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::output_img(const IMG& img, const char* filename) const{
    //TODO
}

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
template < class IMG >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::output_matval_img(const IMG& img, const char* filename) const{
    //TODO
}

}// end namespace tvtml

#endif


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_REDUCTION_FUNCTIONS_SIMD_SSE_AVX_SUM_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_REDUCTION_FUNCTIONS_SIMD_SSE_AVX_SUM_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_AVX_SUPPORT

#include <boost/simd/toolbox/reduction/functions/sum.hpp>
#include <boost/simd/include/constants/one.hpp>

// /////////////////////////////////////////////////////////////////////////////
// // Implementation when type  is arithmetic_
// /////////////////////////////////////////////////////////////////////////////
namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::sum_, boost::simd::tag::avx_,
                                    (A0),
                                    ((simd_<double_<A0>,boost::simd::tag::sse_>))
                                    ((simd_<double_<A0>,boost::simd::tag::sse_>))
                                   )
  {
    typedef typename meta::scalar_of<A0>::type result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      A0 r = _mm256_dp_pd(a0, One<A0>(), 0x3);
      return r[0];
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::sum_, boost::simd::tag::avx_,
                                    (A0),
                                    ((simd_<single_<A0>,boost::simd::tag::avx_>))
                                    ((simd_<single_<A0>,boost::simd::tag::avx_>))
                                   )
  {
    typedef typename meta::scalar_of<A0>::type result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      A0 r = _mm256_dp_ps(a0, One<A0>(), 0xFF);
      return r[0]+r[4];
      // this in another solution for floats,  perhaps speedier ?
      // x3 = _mm256_add_ps(x0, _mm256_movehdup_ps(x0));
      // x4 = _mm256_unpackhi_ps(x3, x3)  ;
      // x4 = _mm256_add_ps(x3, x4)  ;
      // x5 = _mm256_permute2f128_ps(x4, x4, 0x01)  ;
      // x5 = _mm256_add_ps(x5, x4)  ;
    }
  };

} } }
#endif
#endif


//
// $Id$
//
//
// Original author: Jarrett Egertson <jegertso .a.t uw.edu>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "pwiz/data/msdata/MSDataFile.hpp"
#include "pwiz/data/msdata/Serializer_mzML.hpp"
#include "pwiz/data/msdata/Diff.hpp"
#include "pwiz/utility/misc/unit.hpp"
#include "pwiz/utility/misc/Std.hpp"
#include "pwiz_tools/common/FullReaderList.hpp"
#include <boost/program_options/options_description.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>
#include <pwiz/analysis/demux/EnumConstantNotPresentException.hpp>
#include <pwiz/analysis/spectrum_processing/SpectrumList_Demux.hpp>
#include <pwiz/analysis/spectrum_processing/SpectrumList_PeakPicker.hpp>
#include <pwiz/analysis/spectrum_processing/SpectrumListFactory.hpp>
#include <boost/chrono/chrono.hpp>
#include <src/Core/products/Parallelizer.h>

#define PRISM_VERSION "0.3"

using namespace pwiz::data;
using namespace pwiz::msdata;
using namespace pwiz::util;
using namespace pwiz::analysis;

namespace po = boost::program_options;

struct Optimization
{
    Optimization(const string& s) : value(s) {}
    string value;
};

// This overload is needed for using default_value when building a variable_map. The default_value function relies on boost::lexical_cast, which relies on this operator.
ostream& operator <<(ostream& stream, const Optimization& opt) { return stream << opt.value; }

/// Overload the 'validate' function for the user-defined Optimization option class.
void validate(boost::any& v,
    const vector<string>& values,
    Optimization*,
    int)
{
    // Make sure no previous assignment to 'a' was made.
    po::validators::check_first_occurrence(v);
    // Extract the first string from 'values'. If there is more than
    // one string, it's an error, and exception will be thrown.
    const string& s = po::validators::get_single_string(values);

    // Try conversion to desired enum type
    try
    {
        SpectrumList_Demux::Params::stringToOptimization(s);
        v = boost::any(Optimization(s));
    }
    catch (EnumConstantNotPresentException&)
    {
        throw po::invalid_option_value(s);
    }
}

class IterationListenerStream : public IterationListener
{
public:
    explicit IterationListenerStream(ostream &os = std::cerr) : os_(os), entfmt_("%5.1f"){}

    Status update(const UpdateMessage& updateMessage) override
    {
        auto percentProgress = 100.0 * float(updateMessage.iterationIndex) / float(updateMessage.iterationCount);
        os_ << "\rDemultiplexing: " << entfmt_ % percentProgress << "% (Spectrum "
            << updateMessage.iterationIndex << " of " << updateMessage.iterationCount << ")       ";
        return Status_Ok;
    }
private:
    ostream &os_;
    boost::format entfmt_;
};

po::variables_map parseCommandLine(int argc, char **argv);

void run(int argc, char ** argv)
{
    // The Eigen library must be initialized to support multithreaded operations before creating the threads
    Eigen::initParallel();
    auto start = boost::chrono::system_clock::now();
    std::cerr << "Welcome to Prism v" << PRISM_VERSION << endl;
    auto vm = parseCommandLine(argc, argv);

    SpectrumList_Demux::Params demuxParams;
    demuxParams.massError = vm["massError"].as<pwiz::chemistry::MZTolerance>();
    demuxParams.nnlsMaxIter = vm["nnlsMaxIter"].as<unsigned int>();
    demuxParams.nnlsEps = vm["nnlsEps"].as<double>();
    demuxParams.applyWeighting = !vm["noWeighting"].as<bool>();
    demuxParams.demuxBlockExtra = vm["demuxBlockExtra"].as<double>();
    demuxParams.variableFill = vm["variableFill"].as<bool>();
    demuxParams.regularizeSums = !vm["noSumNormalize"].as<bool>();
    demuxParams.optimization = SpectrumList_Demux::Params::stringToOptimization(vm["optimization"].as<Optimization>().value);
    demuxParams.padScanTimes = vm["padScanTimes"].as<bool>();
    demuxParams.interpolateRetentionTime = vm["interpolateRT"].as<bool>();
    bool skipCentroiding = vm["skipCentroiding"].as<bool>();

    FullReaderList readers;
    MSDataFile msd(vm["inputfile"].as<string>(), &readers);
    IntegerSet levelsToCentroid(1,2);

    // Bypass centroiding
    SpectrumListPtr toDemuxPtr = msd.run.spectrumListPtr;
    if (!skipCentroiding)
    {
        SpectrumListPtr tempPtr = toDemuxPtr;
        toDemuxPtr.reset(new SpectrumList_PeakPicker(tempPtr,
                PeakDetectorPtr(boost::make_shared<LocalMaximumPeakDetector>(3)),
                true,
                levelsToCentroid));
        msd.filterApplied();
    }

    SpectrumListPtr demux_list(new SpectrumList_Demux(toDemuxPtr, demuxParams));
    msd.filterApplied();
    msd.run.spectrumListPtr = demux_list;
    // set up progress indicator
    IterationListenerPtr listenerPtr(new IterationListenerStream(std::cerr));
    IterationListenerRegistry listenerRegistry;
    listenerRegistry.addListener(listenerPtr, 10);

    MSDataFile::WriteConfig writeConfig;
    MSDataFile::write(msd, vm["outputfile"].as<string>(), writeConfig, &listenerRegistry);
    std::cerr << endl;
    boost::chrono::duration<double> sec = boost::chrono::system_clock::now() - start;
    std::cout << "Demultiplexing took " << sec.count() / 60.0 << " minutes." << std::endl;
}

po::variables_map parseCommandLine(int argc, char **argv)
{
    po::options_description desc("PRISM arguments");
    // Get default values
    SpectrumList_Demux::Params demuxParams;
    Optimization defaultOptimization(SpectrumList_Demux::Params::optimizationToString(demuxParams.optimization));
    bool skipCentroiding = false;

    desc.add_options()
        ("help,h", "Produce this help message.")
        ("massError", po::value<pwiz::chemistry::MZTolerance>()->default_value(demuxParams.massError),
        "MS/MS error between spectra. Can specify number and units of either ppm or da, e.g. 10ppm or 0.4da. Default 10ppm.")
        ("variableFill", po::bool_switch()->default_value(demuxParams.variableFill),
        "MSX data was acquired with variable fill times.")
        ("demuxBlockExtra", po::value<double>()->default_value(demuxParams.demuxBlockExtra),
        "Extra slop to include in the set of spectra used for demux. Number of spectra added is input * num spectra in one cycle")
        ("nnlsMaxIter", po::value<unsigned int>()->default_value(demuxParams.nnlsMaxIter),
        "Maximum number of iterations for demux solve.")
        ("nnlsEps", po::value<double>()->default_value(demuxParams.nnlsEps),
        "Epsilon value for testing for convergence of demux solve.")
        ("noSumNormalize", po::bool_switch()->default_value(!demuxParams.regularizeSums),
        "Do not normalize the raw demux output for a spectrum so it sums up to the multiplexed peak intensity")
        ("noWeighting", po::bool_switch()->default_value(!demuxParams.applyWeighting),
        "Use non-weighted non negative least squares")
        ("optimization", po::value<Optimization>()->default_value(defaultOptimization),
        "Use optimizations. Available optimizations are \"none\", and \"overlap_only\"" )
        ("padScanTimes", po::bool_switch()->default_value(demuxParams.padScanTimes),
        "Pad scan times with small increment to prevent simultaneous spectra (which some software cannot handle)")
        ("interpolateRT", po::bool_switch()->default_value(demuxParams.interpolateRetentionTime),
        "Interpolate retention times. This only applies when overlap demultiplexing with no MSX (overlap only)")
        ("skipCentroiding", po::bool_switch()->default_value(skipCentroiding),
        "Set to true to use profile data if available")
        ("inputfile", po::value<string>()->required(),
        "Input spectra file")
        ("outputfile", po::value<string>()->required(),
        "Output spectra file");

    po::positional_options_description p;
    p.add("inputfile", 1);
    p.add("outputfile", 1);
    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
    if (vm.count("help"))
    {
        std::cout << desc << endl;
    }
    notify(vm);
    return vm;
}

int main(int argc, char **argv)
{
    try
    {
        run(argc, argv);

        return 0;
    }
    catch (std::exception& e)
    {
        std::cerr << e.what() << endl;
    }
    catch (...)
    {
        std::cerr << "Caught unknown exception.\n";
    }

    return 1;
}