//
// Created by cc on 2020/8/5.
//

#include <ros/ros.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>
#include <geometry_msgs/PoseStamped.h>
#include <math.h>
#include <Eigen/Eigen>
#include <mavros_msgs/State.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>

using namespace Eigen;

Vector3d current_p;
mavros_msgs::State current_state;
ros::Publisher pva_pub;

void positionCallback(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    /// ENU frame to NWU
    current_p << msg->pose.position.y, -msg->pose.position.x, msg->pose.position.z;

    //ROS_INFO("X:%f Y:%f  Z:%f",-current_p(0),current_p(1),current_p(2));
}

void stateCallback(const mavros_msgs::State::ConstPtr &msg)
{
    current_state = *msg;

}

void setPVA(Eigen::Vector3d p, Eigen::Vector3d v, Eigen::Vector3d a, double yaw=0.0)
{
    trajectory_msgs::JointTrajectoryPoint pva_setpoint;

    pva_setpoint.positions.push_back(p(0)); //x
    pva_setpoint.positions.push_back(p(1)); //y
    pva_setpoint.positions.push_back(p(2)); //z
    pva_setpoint.positions.push_back(yaw);

    pva_setpoint.velocities.push_back(v(0));
    pva_setpoint.velocities.push_back(v(1));
    pva_setpoint.velocities.push_back(v(2));

    pva_setpoint.accelerations.push_back(a(0));
    pva_setpoint.accelerations.push_back(a(1));
    pva_setpoint.accelerations.push_back(a(2));

    pva_pub.publish(pva_setpoint);

//    ROS_INFO_THROTTLE(1.0, "P x=%f, y=%f, z=%f", pva_setpoint.positions[0], pva_setpoint.positions[1], pva_setpoint.positions[2]);
//    ROS_INFO_THROTTLE(1.0, "V x=%f, y=%f, z=%f", pva_setpoint.velocities[0], pva_setpoint.velocities[1], pva_setpoint.velocities[2]);
//    ROS_INFO_THROTTLE(1.0, "A x=%f, y=%f, z=%f", pva_setpoint.accelerations[0], pva_setpoint.accelerations[1], pva_setpoint.accelerations[2]);
//
//    ROS_INFO("P x=%f, y=%f, z=%f", pva_setpoint.positions[0], pva_setpoint.positions[1], pva_setpoint.positions[2]);
//    ROS_INFO("V x=%f, y=%f, z=%f", pva_setpoint.velocities[0], pva_setpoint.velocities[1], pva_setpoint.velocities[2]);
//    ROS_INFO("A x=%f, y=%f, z=%f", pva_setpoint.accelerations[0], pva_setpoint.accelerations[1], pva_setpoint.accelerations[2]);
}

/** This function is to generate state to state trajectory **/
void motion_primitives(Eigen::Vector3d p0, Eigen::Vector3d v0, Eigen::Vector3d a0,
                       Eigen::Vector3d pf, Eigen::Vector3d vf, Eigen::Vector3d af, double v_max, double delt_t,
                       Eigen::MatrixXd &p, Eigen::MatrixXd &v, Eigen::MatrixXd &a, Eigen::VectorXd &t)
{
    // % Choose the time as running in average velocity
    // double decay_parameter = 0.5;
    // double T = 0.2;

    double j_limit = 5;
    double a_limit = 3;
    double v_limit = v_max;

//    double T1 = fabs(af(0)-a0(0))/j_limit > fabs(af(1)-a0(1))/j_limit ? fabs(af(0)-a0(0))/j_limit : fabs(af(1)-a0(1))/j_limit;
//    T1 = T1 > fabs(af(2)-a0(2))/j_limit ? T1 : fabs(af(2)-a0(2))/j_limit;
    double T2 = fabs(vf(0)-v0(0))/a_limit > fabs(vf(1)-v0(1))/a_limit ? fabs(vf(0)-v0(0))/a_limit : fabs(vf(1)-v0(1))/a_limit;
    T2 = T2 > fabs(vf(2)-v0(2))/a_limit ? T2 : fabs(vf(2)-v0(2))/a_limit;
    double T3 = fabs(pf(0)-p0(0))/v_limit > fabs(pf(1)-p0(1))/v_limit ? fabs(pf(0)-p0(0))/v_limit : fabs(pf(1)-p0(1))/v_limit;
    T3 = T3 > fabs(pf(2)-p0(2))/v_limit ? T3 : fabs(pf(2)-p0(2))/v_limit;

//    double T = T1 > T2 ? T1 : T2;
//    T = T > T3 ? T : T3;
    double T = T2;
    T = T > T3 ? T : T3;
    T = T < 0.3 ? 0.3 : T;

//    ROS_INFO_THROTTLE(2, "T=%lf", T);

    int times = T / delt_t;

    p = Eigen::MatrixXd::Zero(times, 3);
    v = Eigen::MatrixXd::Zero(times, 3);
    a = Eigen::MatrixXd::Zero(times, 3);
    t = Eigen::VectorXd::Zero(times);

    // % calculate optimal jerk controls by Mark W. Miller
    for(int ii=0; ii<3; ii++)
    {
        double delt_a = af(ii) - a0(ii);
        double delt_v = vf(ii) - v0(ii) - a0(ii)*T;
        double delt_p = pf(ii) - p0(ii) - v0(ii)*T - 0.5*a0(ii)*T*T;

        //%  if vf is not free
        double alpha = delt_a*60/pow(T,3) - delt_v*360/pow(T,4) + delt_p*720/pow(T,5);
        double beta = -delt_a*24/pow(T,2) + delt_v*168/pow(T,3) - delt_p*360/pow(T,4);
        double gamma = delt_a*3/T - delt_v*24/pow(T,2) + delt_p*60/pow(T,3);

        for(int jj=0; jj<times; jj++)
        {
            double tt = (jj + 1)*delt_t;
            t(jj) = tt;
            p(jj,ii) = alpha/120*pow(tt,5) + beta/24*pow(tt,4) + gamma/6*pow(tt,3) + a0(ii)/2*pow(tt,2) + v0(ii)*tt + p0(ii);
            v(jj,ii) = alpha/24*pow(tt,4) + beta/6*pow(tt,3) + gamma/2*pow(tt,2) + a0(ii)*tt + v0(ii);
            a(jj,ii) = alpha/6*pow(tt,3) + beta/2*pow(tt,2) + gamma*tt + a0(ii);
        }
    }
}

void compute_circular_traj(const double r, const double vel, const Eigen::Vector3d p0, const double t,
                           Eigen::Vector3d &p, Eigen::Vector3d &v, Eigen::Vector3d &a)
//@requires r > 0 && vel > 0 && t >= 0;
{
    const double theta = vel*t/r;

    p(0) = r*cos(theta) + p0(0) - r;
    p(1) = r*sin(theta) + p0(1);
    p(2) = p0(2);

    v(0) = -vel*sin(theta);
    v(1) = vel*cos(theta);
    v(2) = 0.0;

    a(0) = -vel*vel/r*cos(theta);
    a(1) = -vel*vel/r*sin(theta);
    a(2) = 0.0;
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "hover");
    ros::NodeHandle nh;

    ros::Subscriber state_sub = nh.subscribe<mavros_msgs::State>("/mavros/state", 1, stateCallback);
    ros::Subscriber pose_sub = nh.subscribe<geometry_msgs::PoseStamped>("/mavros/local_position/pose", 1, positionCallback);

    pva_pub = nh.advertise<trajectory_msgs::JointTrajectoryPoint>("/pva_setpoint", 1);

    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>
            ("mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>
            ("mavros/set_mode");

    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>("/mavros/setpoint_position/local", 10);

    const int LOOPRATE = 40;
    ros::Rate loop_rate(LOOPRATE);

    // wait for FCU connection
    while(ros::ok() && !current_state.connected){
        ros::spinOnce();
        loop_rate.sleep();
    }

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";

    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;
    ros::Time last_request = ros::Time::now();

    /// Take off with constant acceleration
    double take_off_height = 2.0;
    double take_off_acc = 1.0;

    double take_off_time_half = sqrt(take_off_height/take_off_acc);
    double delt_t = 1.0 / LOOPRATE;
    double take_off_send_times = take_off_time_half / delt_t * 2;
    int counter = 0;
    Vector3d recorded_takeoff_position(current_p(0), current_p(1), current_p(2));

    double yaw_set = 3.1415926;

    ROS_INFO("Arm and ta1111keoff");
    while(ros::ok()){
        if( current_state.mode != "OFFBOARD" &&
            (ros::Time::now() - last_request > ros::Duration(5.0))){
            if( set_mode_client.call(offb_set_mode) &&
                offb_set_mode.response.mode_sent){
                ROS_INFO("Offboard enabled");
            }
            last_request = ros::Time::now();
        } else {
            if( !current_state.armed &&
                (ros::Time::now() - last_request > ros::Duration(5.0))){
                if( arming_client.call(arm_cmd) &&
                    arm_cmd.response.success){
                    ROS_INFO("Vehicle armed");
                }
                last_request = ros::Time::now();
            }
        }

        trajectory_msgs::JointTrajectoryPoint pva_setpoint;

        static unsigned int count=0;
        count++;
        if(count>1000)
        {
            break;
        }
        double theta=180.0/180.0*3.14;
        geometry_msgs::PoseStamped pose;
        pose.pose.position.x = 0;
        pose.pose.position.y = 0;
        pose.pose.position.z =3;
        pose.pose.orientation.w=cos(theta/2);
//        ROS_INFO("theta:%f",theta);
//        ROS_INFO("W:%f",pose.pose.orientation.w);
        pose.pose.orientation.x=-0.0;
        pose.pose.orientation.y=0;
        pose.pose.orientation.z=sin(theta/2);
        local_pos_pub.publish(pose);

        ros::spinOnce();
        loop_rate.sleep();

    }

    /** Take off complete. Go to a point with minimum jerk trajectory **/

    /** Now draw circle **/

    double init_t = 0.0;
    int mode=1;
    ROS_INFO("HOVER!!");
    while(ros::ok())
    {
//        Eigen::Vector3d p, v ,a;
//        compute_circular_traj(circle_radius, circle_speed, circle_p0, init_t, p, v, a);
        static int count=0;
        count++;
        if(count==800)
        {
            mode=-mode;
            count=0;
        }
        if(mode==1)
        {
            Eigen::Vector3d p(0,0,3);
            Eigen::Vector3d v(0,0,0);
            Eigen::Vector3d a(0,0,0);
            yaw_set=180.0/180*3.1415926;
            setPVA(p, v, a, yaw_set);//a_t.row(last_index));

//            geometry_msgs::PoseStamped pose;
//            pose.pose.position.x = 3;
//            pose.pose.position.y = 0;
//            pose.pose.position.z =3;
//            double theta=180.0/180.0*3.14;
//            pose.pose.orientation.w=cos(theta/2);
////        ROS_INFO("theta:%f",theta);
////        ROS_INFO("W:%f",pose.pose.orientation.w);
//            pose.pose.orientation.x=-0.0;
//            pose.pose.orientation.y=0;
//            pose.pose.orientation.z=sin(theta/2);
//            local_pos_pub.publish(pose);

        }
        else {
            Eigen::Vector3d p(0, 0, 2);
            Eigen::Vector3d v(0, 0, 0);
            Eigen::Vector3d a(0, 0, 0);
            yaw_set = 180.0 / 180 * 3.1415926;
            setPVA(p, v, a, yaw_set);//a_t.row(last_index));

//            geometry_msgs::PoseStamped pose;
//            pose.pose.position.x = 6;
//            pose.pose.position.y = 0;
//            pose.pose.position.z =3;
//            double theta=180.0/180.0*3.14;
//            pose.pose.orientation.w=cos(theta/2);
////        ROS_INFO("theta:%f",theta);
////        ROS_INFO("W:%f",pose.pose.orientation.w);
//            pose.pose.orientation.x=-0.0;
//            pose.pose.orientation.y=0;
//            pose.pose.orientation.z=sin(theta/2);
//            local_pos_pub.publish(pose);
//        }
        }

        init_t += delt_t;
        loop_rate.sleep();
        ros::spinOnce();
    }

    return 0;
}


#define BOOST_TEST_MODULE "test_hard_core_excluded_volume_potential"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/forcefield/global/HardCoreExcludedVolumePotential.hpp>

BOOST_AUTO_TEST_CASE(HCEXV_double)
{
    mjolnir::LoggerManager::set_default_logger("test_hard_core_excluded_volume_potential.log");

    using traits_type = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type   = typename traits_type::real_type;
    using molecule_id_type = mjolnir::Topology::molecule_id_type;
    using group_id_type    = mjolnir::Topology::group_id_type;
    using potential_type   = mjolnir::HardCoreExcludedVolumePotential<traits_type>;
    using parameter_type   = potential_type::parameter_type;

    constexpr static std::size_t N = 10000;
    constexpr static real_type   h = 1e-6;

    const real_type epsilon = 1.0;
    const real_type hard_core_radius = 10.0;
    const real_type soft_shell_thickness = 3.0;
    const parameter_type param{hard_core_radius, soft_shell_thickness};
    potential_type hdexv{
      epsilon, potential_type::default_cutoff(), {{0, param}, {1, param}}, {},
      mjolnir::IgnoreMolecule<molecule_id_type>("Nothing"),
      mjolnir::IgnoreGroup<group_id_type>({})
    };

    const real_type x_min = 0.8 * soft_shell_thickness + hard_core_radius * 2;
    const real_type x_max = hdexv.cutoff_ratio() * soft_shell_thickness + hard_core_radius * 2;
    const real_type dx = (x_max - x_min) / N;

    for(std::size_t i=0; i<N; ++i)
    {
      const real_type x = x_min + i * dx;
      const real_type pot1 = hdexv.potential(0, 1, x + h);
      const real_type pot2 = hdexv.potential(0, 1, x - h);
      const real_type dpot = (pot1 - pot2) / (2 * h);
      const real_type deri = hdexv.derivative(0, 1, x);

      BOOST_TEST(dpot == deri, boost::test_tools::tolerance(h));
    }
}

BOOST_AUTO_TEST_CASE(HCEXV_float)
{
    using traits_type = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type   = typename traits_type::real_type;
    using molecule_id_type = mjolnir::Topology::molecule_id_type;
    using group_id_type    = mjolnir::Topology::group_id_type;
    using potential_type   = mjolnir::HardCoreExcludedVolumePotential<traits_type>;
    using parameter_type   = potential_type::parameter_type;

    constexpr static std::size_t N = 1000;
    constexpr static real_type   h = 0.002f;
    constexpr static real_type tol = 0.005f;

    const real_type epsilon = 1.0;
    const real_type hard_core_radius = 10.0;
    const real_type soft_shell_thickness = 3.0;
    const parameter_type param{hard_core_radius, soft_shell_thickness};
    potential_type hdexv{
      epsilon, potential_type::default_cutoff(), {{0, param}, {1, param}}, {},
      mjolnir::IgnoreMolecule<molecule_id_type>("Nothing"),
      mjolnir::IgnoreGroup<group_id_type>({})
    };

    const real_type x_min = 0.8 * soft_shell_thickness + hard_core_radius * 2;
    const real_type x_max = hdexv.cutoff_ratio() * soft_shell_thickness + hard_core_radius * 2;
    const real_type dx = (x_max - x_min) / N;

    for(std::size_t i=0; i<N; ++i)
    {
      const real_type x = x_min + i * dx;
      const real_type pot1 = hdexv.potential(0, 1, x + h);
      const real_type pot2 = hdexv.potential(0, 1, x - h);
      const real_type dpot = (pot1 - pot2) / (2 * h);
      const real_type deri = hdexv.derivative(0, 1, x);

      BOOST_TEST(dpot == deri, boost::test_tools::tolerance(tol));
    }
}


/**
 * @file init_rules_test.cpp
 * @author Marcus Edel
 *
 *  Tests for the various weight initialize methods.
 */
#include <mlpack/core.hpp>

#include <mlpack/methods/ann/init_rules/random_init.hpp>
#include <mlpack/methods/ann/init_rules/orthogonal_init.hpp>

#include <boost/test/unit_test.hpp>
#include "old_boost_test_definitions.hpp"

using namespace mlpack;
using namespace mlpack::ann;

BOOST_AUTO_TEST_SUITE(InitRulesTest);

// Be careful!  When writing new tests, always get the boolean value and store
// it in a temporary, because the Boost unit test macros do weird things and
// will cause bizarre problems.

// Test the RandomInitialization class with a constant value.
BOOST_AUTO_TEST_CASE(ConstantInitTest)
{
  arma::mat weights;
  RandomInitialization constantInit(1, 1);
  constantInit.Initialize(weights, 100, 100);

  bool b = arma::all(arma::vectorise(weights) == 1);
  BOOST_REQUIRE_EQUAL(b, 1);
}

// Test the OrthogonalInitialization class.
BOOST_AUTO_TEST_CASE(OrthogonalInitTest)
{
  arma::mat weights;
  OrthogonalInitialization orthogonalInit;
  orthogonalInit.Initialize(weights, 100, 200);

  arma::mat orthogonalWeights = arma::eye<arma::mat>(100, 100);
  weights *= weights.t();

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
        BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);

  orthogonalInit.Initialize(weights, 200, 100);
  weights = weights.t() * weights;

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
      BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);
}

// Test the OrthogonalInitialization class with a non default gain.
BOOST_AUTO_TEST_CASE(OrthogonalInitGainTest)
{
  arma::mat weights;

  const double gain = 2;
  OrthogonalInitialization orthogonalInit(gain);
  orthogonalInit.Initialize(weights, 100, 200);

  arma::mat orthogonalWeights = arma::eye<arma::mat>(100, 100);
  orthogonalWeights *= (gain * gain);
  weights *= weights.t();

  for (size_t i = 0; i < weights.n_rows; i++)
    for (size_t j = 0; j < weights.n_cols; j++)
        BOOST_REQUIRE_SMALL(weights.at(i, j) - orthogonalWeights.at(i, j), 1e-3);
}

BOOST_AUTO_TEST_SUITE_END();


/*=============================================================================
    Copyright (c) 2001-2011 Joel de Guzman
    Copyright (c) 2007 Dan Marsden

    Distributed under 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)
==============================================================================*/
#include <boost/detail/lightweight_test.hpp>
#include <boost/fusion/container/vector/vector.hpp>
#include <boost/fusion/adapted/mpl.hpp>
#include <boost/fusion/sequence/io/out.hpp>
#include <boost/fusion/container/generation/make_vector.hpp>
#include <boost/fusion/sequence/comparison/equal_to.hpp>
#include <boost/fusion/algorithm/transformation/transform.hpp>
#include <boost/type_traits/is_class.hpp>
#include <boost/type_traits/is_same.hpp>
#include <boost/mpl/range_c.hpp>
#include <boost/type_traits/is_reference.hpp>

struct square
{
    template<typename Sig>
    struct result;

    template <typename T>
    struct result<square(T)>
    {
        typedef int type;
    };

    template <typename T>
    int operator()(T x) const
    {
        return x * x;
    }
};

struct add
{
    template<typename Sig>
    struct result;

    template <typename A, typename B>
    struct result<add(A, B)>
    {
        typedef int type;
    };

    template <typename A, typename B>
    int operator()(A a, B b) const
    {
        return a + b;
    }
};

struct unary_lvalue_transform
{
    template<typename Sig>
    struct result;

    template<typename T>
    struct result<unary_lvalue_transform(T&)>
    {
        typedef T* type;
    };

    template<typename T>
    T* operator()(T& t) const
    {
        return &t;
    }
};

int twice(int v)
{
    return v*2;
}

struct binary_lvalue_transform
{
    template<typename Sig>
    struct result;

    template<typename T0, typename T1>
    struct result<binary_lvalue_transform(T0&,T1&)>
    {
        typedef T0* type;
    };

    template<typename T0, typename T1>
    T0* operator()(T0& t0, T1&) const
    {
        return &t0;
    }
};

int
main()
{
    using namespace boost::fusion;
    using boost::mpl::range_c;

    std::cout << tuple_open('[');
    std::cout << tuple_close(']');
    std::cout << tuple_delimiter(", ");

/// Testing the transform

    {
        typedef range_c<int, 5, 9> sequence_type;
        sequence_type sequence;
        std::cout << boost::fusion::transform(sequence, square()) << std::endl;
        BOOST_TEST((boost::fusion::transform(sequence, square()) == make_vector(25, 36, 49, 64)));
    }

    {
        typedef range_c<int, 5, 9> mpl_list1;
        std::cout << boost::fusion::transform(mpl_list1(), square()) << std::endl;
        BOOST_TEST((boost::fusion::transform(mpl_list1(), square()) == make_vector(25, 36, 49, 64)));
    }

    {
        vector<int, int, int> tup(1, 2, 3);
        std::cout << boost::fusion::transform(tup, square()) << std::endl;
        BOOST_TEST((boost::fusion::transform(tup, square()) == make_vector(1, 4, 9)));
    }

    {
        vector<int, int, int> tup1(1, 2, 3);
        vector<int, int, int> tup2(4, 5, 6);
        std::cout << boost::fusion::transform(tup1, tup2, add()) << std::endl;
        BOOST_TEST((boost::fusion::transform(tup1, tup2, add()) == make_vector(5, 7, 9)));
    }

    {
        // Unary transform that requires lvalues, just check compilation
        vector<int, int, int> tup1(1, 2, 3);
        BOOST_TEST(at_c<0>(boost::fusion::transform(tup1, unary_lvalue_transform())) == &at_c<0>(tup1));
        BOOST_TEST(*begin(boost::fusion::transform(tup1, unary_lvalue_transform())) == &at_c<0>(tup1));
    }

    {
        vector<int, int, int> tup1(1, 2, 3);
        vector<int, int, int> tup2(4, 5, 6);
        BOOST_TEST(at_c<0>(boost::fusion::transform(tup1, tup2, binary_lvalue_transform())) == &at_c<0>(tup1));
        BOOST_TEST(*begin(boost::fusion::transform(tup1, tup2, binary_lvalue_transform())) == &at_c<0>(tup1));
    }

    {
        vector<int, int, int> tup1(1, 2, 3);
        BOOST_TEST(boost::fusion::transform(tup1, twice) == make_vector(2,4,6));
    }

    return boost::report_errors();
}


#include "experiments/util/bitmap_db.hpp"
#include "experiments/util/threading.hpp"

#include <dtl/bitmap/dynamic_bitmap.hpp>
#include <dtl/bitmap/dynamic_roaring_bitmap.hpp>
#include <dtl/bitmap/dynamic_wah.hpp>
#include <dtl/bitmap/teb_wrapper.hpp>
#include <dtl/bitmap/util/convert.hpp>
#include <dtl/bitmap/util/random.hpp>
#include <dtl/dtl.hpp>
#include <dtl/env.hpp>

#include <boost/algorithm/string.hpp>
#include <boost/filesystem.hpp>

#include <atomic>
#include <cstddef>
#include <fstream>
#include <iostream>
#include <limits>
#include <ostream>
#include <sstream>
#include <string>
#include <vector>
//===----------------------------------------------------------------------===//
namespace fs = boost::filesystem;
//===----------------------------------------------------------------------===//
void run(const std::string& dir, std::ostream& result_out) {
//  result_out
//      << "//===----------------------------------------------------------------------===//"
//      << "\ndata set: '" << dir << "'"
//      << std::endl;

  std::size_t min_val = std::numeric_limits<std::size_t>::max();
  std::size_t max_val = 0;
  std::size_t file_cnt = 0;

  std::vector<std::string> filenames;
  fs::directory_iterator end_it;
  for (fs::directory_iterator dir_it(dir); dir_it != end_it; ++dir_it) {
    const auto file = dir_it->path().string();
    filenames.push_back(file);
    ++file_cnt;
    std::cout << "reading file: " << file << std::endl;
    std::ifstream is(file);
    std::string token;
    while (std::getline(is, token, ',')) {
      const auto val = std::stoull(token);
      if (val < min_val) min_val = val;
      if (val > max_val) max_val = val;
    }
    is.close();
  }
  std::sort(filenames.begin(), filenames.end());

  const auto n = max_val + 1;
  const auto n_pow2 = dtl::next_power_of_two(n);
  const auto c = file_cnt;

//  result_out << "min_val: " << min_val
//             << " max_val: " << max_val
//             << " bitmap_cnt: " << file_cnt
//             << " n: " << n
//             << std::endl;

  $f64 min_d = std::numeric_limits<f64>::max();
  $f64 max_d = std::numeric_limits<f64>::min();
  $f64 sum_d = 0.0;

  $f64 min_f = std::numeric_limits<f64>::max();
  $f64 max_f = std::numeric_limits<f64>::min();
  $f64 sum_f = 0.0;

  std::atomic<std::size_t> bytes_roaring { 0 };
  std::atomic<std::size_t> bytes_wah { 0 };
  std::atomic<std::size_t> bytes_wah64 { 0 };
  std::atomic<std::size_t> bytes_teb { 0 };
  std::atomic<std::size_t> bytes_teb_rank { 0 };

  std::vector<dtl::bitmap> bitmaps;
  std::vector<dtl::bitmap> bitmaps_pow2;

  std::size_t total_bit_cnt = 0;
  for (auto& file : filenames) {
    std::ifstream is(file);
    std::string token;

    dtl::bitmap bm(n);
    dtl::bitmap bm_pow2(n_pow2);

    while (std::getline(is, token, ',')) {
      const auto val = std::stoull(token);
      bm[val] = true;
      bm_pow2[val] = true;
    }
    is.close();

    bitmaps.push_back(bm);
    bitmaps_pow2.push_back(bm_pow2);

    const auto d = dtl::determine_bit_density(bm);
    min_d = std::min(d, min_d);
    max_d = std::max(d, max_d);
    sum_d += d;

    const auto f = dtl::determine_clustering_factor(bm);
    min_f = std::min(f, min_f);
    max_f = std::max(f, max_f);
    sum_f += f;

    total_bit_cnt += bm.count();
  }
//  result_out << "cardinality: " << total_bit_cnt << std::endl;
//  result_out << "d (min/avg/max): "
//             << min_d << " / "
//             << (sum_d / file_cnt) << " / "
//             << max_d << std::endl;
//  result_out << "f (min/avg/max): "
//             << min_f << " / "
//             << (sum_f / file_cnt) << " / "
//             << max_f << std::endl;

  auto thread_fn = [&](const std::size_t bid, std::ostream& os) {
    auto& bm = bitmaps[bid];
    auto& bm_pow2 = bitmaps_pow2[bid];

    std::size_t r = 0;
    std::size_t t = 0;
    std::size_t t_rank = 0;
    std::size_t w = 0;
    std::size_t w64 = 0;
    {
      dtl::dynamic_roaring_bitmap roaring(bm);
      r = roaring.size_in_bytes();
    }
    {
      dtl::dynamic_wah32 wah(bm);
      w = wah.size_in_bytes();
    }
    {
      dtl::dynamic_wah64 wah(bm);
      w64 = wah.size_in_bytes();
    }
    {
      dtl::teb_wrapper teb(bm_pow2);
//      dtl::teb<> teb(bm_pow2);
      t = teb.size_in_bytes();
//      t_rank = 0; //teb.rank_.size_in_bytes();
      t_rank =
          teb.teb_->get_rank_word_cnt(teb.data_.data()) * sizeof(dtl::teb_word_type);
      const auto dec = dtl::to_bitmap_using_iterator(teb);
      if ((bm_pow2 & dec) != bm_pow2) {
        std::cerr << "Validation failed." << std::endl;
        std::exit(1);
      }
    }
    // clang-format off
// TEB (lossy compressed)
//    {
//      const auto fpr = 0.0001;
//      dtl::teb_wrapper teb(bm_pow2, fpr);
//      t = teb.size_in_bytes();
//      const auto dec = dtl::to_bitmap_using_iterator(teb);
//      if ((bm_pow2 & dec) != bm_pow2) {
//        std::cerr << "Validation failed." << std::endl;
//        std::exit(1);
//      }
//      u64 max_fp_cnt = static_cast<u64>(teb.size() * fpr);
//      u64 fp_cnt = (bm_pow2 ^ dec).count();
//      std::cout << "fp_cnt=" << fp_cnt << std::endl;
//      if (fp_cnt > max_fp_cnt) {
//        std::cerr << "Validation failed. Max FP count exceeded." << std::endl;
//        std::exit(1);
//      }
//      os << teb.info() << std::endl;
//    }
    // clang-format on

    bytes_roaring += r;
    bytes_wah += w;
    bytes_wah64 += w64;
    bytes_teb += t;
    bytes_teb_rank += t_rank;
  };

  dispatch(0, bitmaps.size(), thread_fn);

//  result_out << dir << std::endl;
//  result_out << "roaring: " << std::setw(15) << bytes_roaring << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_roaring * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
//  result_out << "teb:     " << std::setw(15) << bytes_teb     << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_teb     * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
//  result_out << "wah:     " << std::setw(15) << bytes_wah     << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_wah     * 8.0) / total_bit_cnt) << " bits/int" << std::endl;
//  result_out << "wah64:   " << std::setw(15) << bytes_wah64   << " bytes, " << std::setw(15) << std::setprecision(4) << ((bytes_wah64   * 8.0) / total_bit_cnt) << " bits/int" << std::endl;

//  result_out << "TEB details:" << std::endl;
  result_out
      << "\"" << dir << "\","
      << total_bit_cnt << ","
      << bytes_teb << "," << bytes_teb_rank << ","
      << std::setprecision(4) << ((bytes_teb * 8.0) / total_bit_cnt)
      << std::endl;

//    result_out
//      << "\"" << dir << "\","
//      << total_bit_cnt << ","
//      << bytes_roaring << ","
//      << std::setprecision(4) << ((bytes_roaring * 8.0) / total_bit_cnt)
//      << std::endl;
}
//===----------------------------------------------------------------------===//
$i32 main() {
  const auto basedir = dtl::env<std::string>::get("DIR",
      "./real-roaring-dataset");
  std::vector<std::string> dirs {
    basedir + "/census1881/",
    basedir + "/census1881_srt/",
    basedir + "/census-income/",
    basedir + "/census-income_srt/",
    basedir + "/weather_sept_85/",
    basedir + "/weather_sept_85_srt/",
    basedir + "/wikileaks-noquotes/",
    basedir + "/wikileaks-noquotes_srt/",
  };

  std::stringstream results;
  for (auto& dir : dirs) {
    run(dir, results);
  }
  std::cout << "Results" << std::endl;
  std::cout << results.str() << std::endl;
}
//===----------------------------------------------------------------------===//


#include "color.hpp"
#include "image.hpp"
#include "2d/cube.hpp"

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE c10t_tests
#include <boost/test/unit_test.hpp>

#include <iostream>

BOOST_AUTO_TEST_CASE( test_cube_projection_1 )
{
  // x, y, z
  Cube c(10, 10, 20);
  point p1(0, 0, 0), p2(0, 0, 20), p3(10, 0, 20), p4(10, 0, 0);
  int x, y;

  {
    c.project_top(p1, x, y);
    BOOST_REQUIRE(x == 0 && y == 0);
    c.project_top(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 20);
    c.project_top(p3, x, y);
    BOOST_REQUIRE(x == 10 && y == 20);
    c.project_top(p4, x, y);
    BOOST_REQUIRE(x == 10 && y == 0);
  }

  {
    c.project_oblique(p1, x, y);
    BOOST_REQUIRE(x == 0 && y == 10);
    c.project_oblique(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 30);
    c.project_oblique(p3, x, y);
    BOOST_REQUIRE(x == 10 && y == 30);
    c.project_oblique(p4, x, y);
    BOOST_REQUIRE(x == 10 && y == 10);
  }

  {
    c.project_obliqueangle(p1, x, y);
    BOOST_REQUIRE(x == 20 && y == 10);
    c.project_obliqueangle(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 30);
    c.project_obliqueangle(p3, x, y);
    BOOST_REQUIRE(x == 10 && y == 40);
    c.project_obliqueangle(p4, x, y);
    BOOST_REQUIRE(x == 30 && y == 20);
  }
}

BOOST_AUTO_TEST_CASE( test_cube_projection_2 )
{
  // x, y, z
  Cube c(20, 10, 10);
  point p1(0, 0, 0), p2(0, 0, 10), p3(20, 0, 10), p4(20, 0, 0);
  int x, y;

  {
    c.project_top(p1, x, y);
    BOOST_REQUIRE(x == 0 && y == 0);
    c.project_top(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 10);
    c.project_top(p3, x, y);
    BOOST_REQUIRE(x == 20 && y == 10);
    c.project_top(p4, x, y);
    BOOST_REQUIRE(x == 20 && y == 0);
  }

  {
    c.project_oblique(p1, x, y);
    BOOST_REQUIRE(x == 0 && y == 10);
    c.project_oblique(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 20);
    c.project_oblique(p3, x, y);
    BOOST_REQUIRE(x == 20 && y == 20);
    c.project_oblique(p4, x, y);
    BOOST_REQUIRE(x == 20 && y == 10);
  }

  {
    c.project_obliqueangle(p1, x, y);
    BOOST_REQUIRE(x == 10 && y == 10);
    c.project_obliqueangle(p2, x, y);
    BOOST_REQUIRE(x == 0 && y == 20);
    c.project_obliqueangle(p3, x, y);
    BOOST_REQUIRE(x == 20 && y == 40);
    c.project_obliqueangle(p4, x, y);
    BOOST_REQUIRE(x == 30 && y == 30);
  }
}


#include <string>
#include <math.h>

#include "global-loc/place-retrieval.h"
#include "global-loc/pca-reduction.h"

#include <glog/logging.h>
#include <opencv2/opencv.hpp>
#include <Eigen/Core>

#include <vi-map/vi-map.h>
#include <vi-map/unique-id.h>
#include <posegraph/pose-graph.h>
#include <posegraph/unique-id.h>
#include <aslam/common/pose-types.h>
#include <aslam/common/timer.h>
#include <maplab-common/aslam-id-proto.h>
#include <maplab-common/eigen-proto.h>
#include <maplab-common/progress-bar.h>
#include <kindr/minimal/quat-transformation.h>
#include <map-resources/resource-common.h>

DEFINE_uint64(
        subsampling, 1, "Interval at which the frames should be indexed.");
DEFINE_bool(
        index_pose, false,
        "Whether the frame pose (i.e. position vector and rotation matrix) should be"
        "added to the index.");
DEFINE_string(
        target_mission, "", "ID of the mission to be indexed.");
DEFINE_bool(
        use_pca, false,
        "Whether a PCA projection matrix should be computed from the descriptors.");
DEFINE_uint64(
        pca_dims, 40, "The number of dimensions of the PCA projection.");

PlaceRetrieval::PlaceRetrieval(const std::string model_path):
        network_(model_path, "image", "descriptor") {
    unsigned index_dims = FLAGS_use_pca ? FLAGS_pca_dims : network_.descriptor_size();
    index_.reset(new KDTreeIndex(index_dims));
}

void PlaceRetrieval::BuildIndexFromMap(
        const vi_map::VIMap& map,
        global_loc::proto::DescriptorIndex* proto_index) {
    CHECK_NOTNULL(proto_index);
    proto_index->set_descriptor_size(network_.descriptor_size());

    vi_map::MissionIdList mission_ids;
    if(!FLAGS_target_mission.empty()) {
        vi_map::MissionId target_mission;
        map.ensureMissionIdValid(FLAGS_target_mission, &target_mission);
        CHECK(target_mission.isValid());
        mission_ids.push_back(target_mission);
    } else {
        map.getAllMissionIdsSortedByTimestamp(&mission_ids);
        CHECK(!mission_ids.empty());
    }

    for (const vi_map::MissionId& mission_id : mission_ids) {
        pose_graph::VertexIdList vertex_ids;
        map.getAllVertexIdsInMissionAlongGraph(mission_id, &vertex_ids);
        CHECK(!vertex_ids.empty());

        unsigned num_indexed = 1 + (vertex_ids.size() - 1) / FLAGS_subsampling;
        LOG(INFO) << "Computing " << num_indexed
                  << " descriptors for mission " << mission_id.printString();
        common::ProgressBar progress_bar(num_indexed);

        for (int i = 0; i < vertex_ids.size(); i += FLAGS_subsampling) {
            progress_bar.increment();
            const pose_graph::VertexId& vertex_id = vertex_ids[i];
            const vi_map::Vertex& vertex = map.getVertex(vertex_id);
            if(!vertex.numFrames())
                continue;

            unsigned frame_index = 0;  // Add only the first frame

            global_loc::proto::DescriptorIndex::Frame* proto_frame =
                    proto_index->add_frames();
            vertex_id.serialize(proto_frame->mutable_vertex_id());
            proto_frame->set_frame_index(frame_index);
            backend::ResourceIdSet resource_ids;
            vertex.getFrameResourceIdsOfType(
                    frame_index, backend::ResourceType::kRawImage, &resource_ids);
            if(!resource_ids.size()) {
                LOG(WARNING)
                    << "Frame " << frame_index << " of vertex " << vertex_id
                    << " had no resource, skipping.";
                proto_index->mutable_frames()->RemoveLast();
                continue;
            }
            proto_frame->set_resource_name(resource_ids.begin()->hexString());

            cv::Mat image;
            CHECK(map.getRawImage(vertex, frame_index, &image))
                << "Vertex " << vertex_id << " does not have a raw image for frame "
                << frame_index;
            TensorflowNet::DescriptorType descriptor;
            descriptor.resize(network_.descriptor_size(), Eigen::NoChange);
            network_.PerformInference(image, &descriptor);
            Eigen::MatrixXf generic_descriptor = descriptor;  // copy constructor crashes
            common::eigen_proto::serialize(
                    generic_descriptor, proto_frame->mutable_global_descriptor());

            if(FLAGS_index_pose) {
                aslam::Transformation transf = map.getVertex_T_G_I(vertex_id);
                common::eigen_proto::serialize(
                        Eigen::MatrixXd(transf.getPosition()),
                        proto_frame->mutable_position_vector());
                common::eigen_proto::serialize(
                        Eigen::MatrixXd(transf.getRotationMatrix()),
                        proto_frame->mutable_rotation_matrix());
            }
        }
    }
}

void PlaceRetrieval::LoadIndex(
        const global_loc::proto::DescriptorIndex& proto_index) {
    CHECK_EQ(proto_index.descriptor_size(), network_.descriptor_size());

    KDTreeIndex::DescriptorMatrixType descriptors(
            network_.descriptor_size(), proto_index.frames().size());
    unsigned i = 0;

    LOG(INFO) << "Loading " << proto_index.frames_size()
              << " reference descriptors into index.";
    common::ProgressBar progress_bar(proto_index.frames_size());
    for (const global_loc::proto::DescriptorIndex::Frame& proto_frame :
            proto_index.frames()) {
        pose_graph::VertexId vertex_id;
        vertex_id.deserialize(proto_frame.vertex_id());
        size_t frame_index = proto_frame.frame_index(); // static_cast<size_t>()
        indexed_frame_identifiers_.emplace_back(vertex_id, frame_index);

        KDTreeIndex::DescriptorMatrixType descriptor;
        common::eigen_proto::deserialize(proto_frame.global_descriptor(), &descriptor);
        CHECK_EQ(descriptor.rows(), network_.descriptor_size());

        descriptors.col(i++) = descriptor;
        progress_bar.increment();
    }

    if(FLAGS_use_pca) {
        pca_reduction_.reset(new PcaReduction(descriptors));
        pca_reduction_->project(&descriptors, FLAGS_pca_dims);
    }
    index_->AddDescriptors(descriptors);
    index_->RefreshIndex();
}

void PlaceRetrieval::RetrieveNearestNeighbors(
        const cv::Mat& input_image, const unsigned num_neighbors,
        const float max_distance,
        vi_map::VisualFrameIdentifierList* retrieved_frame_identifiers) {
    TensorflowNet::DescriptorType descriptor;
    descriptor.resize(network_.descriptor_size(), Eigen::NoChange);
    network_mutex_.lock();
    timing::Timer timer_inference("Deep Relocalization: Compute descriptor");
    network_.PerformInference(input_image, &descriptor);
    timer_inference.Stop();
    network_mutex_.unlock();

    KDTreeIndex::DescriptorMatrixType descriptor_matrix(descriptor);
    if(FLAGS_use_pca) {
        CHECK_NOTNULL(pca_reduction_);
        pca_reduction_->project(&descriptor_matrix, FLAGS_pca_dims);
    }

    Eigen::MatrixXi indices;
    indices.resize(num_neighbors, 1);
    Eigen::MatrixXf distances;
    distances.resize(num_neighbors, 1);
    timing::Timer timer_get_nn("Deep Relocalization: Get neighbors");
    index_->GetNNearestNeighbors(
            descriptor_matrix, num_neighbors, &indices, &distances, max_distance);
    timer_get_nn.Stop();

    for (unsigned nn_search_idx = 0; nn_search_idx < num_neighbors; ++nn_search_idx) {
        const int nn_database_idx = indices(nn_search_idx, 0);
        const float nn_distance = distances(nn_search_idx, 0);
        if (nn_database_idx == -1 ||
                nn_distance == std::numeric_limits<float>::infinity()) {
            break;  // No more results
        }
        retrieved_frame_identifiers->push_back(
                indexed_frame_identifiers_[nn_database_idx]);
    }
}


#include "fbstab/components/full_feasibility.h"

#include <Eigen/Dense>
#include <algorithm>
#include <stdexcept>

#include "fbstab/components/abstract_components.h"
#include "fbstab/components/full_variable.h"

namespace fbstab {

FullFeasibility::FullFeasibility(int nz, int nl, int nv) {
  if (nz <= 0 || nl < 0 || nv <= 0) {
    throw std::runtime_error("Incorrect size inputs to FullFeasibility.");
  }
  nz_ = nz;
  nl_ = nl;
  nv_ = nv;

  tz_.resize(nz_);
  tl_.resize(nl_);
  tv_.resize(nv_);
}

FullFeasibility::FeasibilityStatus FullFeasibility::CheckFeasibility(
    const FullVariable& x, double tol) {
  NullDataCheck();

  // The conditions for dual-infeasibility are:
  // max(Az) <= 0 and f'*z < 0 and |Hz| <= tol * |z| and |Gz| <= tol*|z|

  // Compute d1 = max(Az).
  data_->gemvA(x.z(), 1.0, 0.0, &tv_);
  const double d1 = tv_.maxCoeff();

  // Compute d2 = infnorm(Gz).
  data_->gemvG(x.z(), 1.0, 0.0, &tl_);
  const double d2 = tl_.lpNorm<Eigen::Infinity>();

  // Compute d3 = infnorm(Hz).
  data_->gemvH(x.z(), 1.0, 0.0, &tz_);
  const double d3 = tz_.lpNorm<Eigen::Infinity>();

  // Compute d4 = f'*z
  tz_.setConstant(0.0);
  data_->axpyf(1.0, &tz_);
  const double d4 = tz_.dot(x.z());

  double w = x.z().lpNorm<Eigen::Infinity>();
  bool dual_feasible = true;
  if ((d1 <= w * tol) && (d2 <= tol * w) && (d3 <= tol * w) && (d4 < 0) &&
      (w > 1e-14)) {
    dual_feasible = false;
  }

  // The conditions for primal infeasibility are:
  // v'*b + l'*h < 0 and |A'*v + G'*l| \leq tol * |(v,l)|

  // Compute p1 = infnorm(G'*l + A'*v).
  tz_.fill(0.0);
  data_->gemvAT(x.v(), 1.0, 1.0, &tz_);
  data_->gemvGT(x.l(), 1.0, 1.0, &tz_);
  const double p1 = tz_.lpNorm<Eigen::Infinity>();

  // Compute p2 = v'*b + l'*h.
  tv_.setConstant(0.0);
  data_->axpyb(1.0, &tv_);
  tl_.setConstant(0.0);
  data_->axpyh(1.0, &tl_);
  const double p2 = tl_.dot(x.l()) + tv_.dot(x.v());

  const double u = std::max(x.v().lpNorm<Eigen::Infinity>(),
                            x.l().lpNorm<Eigen::Infinity>());
  bool primal_feasible = true;
  if ((p1 <= tol * u) && (p2 < 0)) {
    primal_feasible = false;
  }

  if (primal_feasible && dual_feasible) {
    return FeasibilityStatus::FEASIBLE;
  } else if (primal_feasible && !dual_feasible) {
    return FeasibilityStatus::DUAL_INFEASIBLE;
  } else if (!primal_feasible && dual_feasible) {
    return FeasibilityStatus::PRIMAL_INFEASIBLE;
  } else {
    return FeasibilityStatus::BOTH;
  }
}

void FullFeasibility::NullDataCheck() const {
  if (data_ == nullptr) {
    throw std::runtime_error(
        "FullFeasibility tried to access problem data before it's "
        "linked.");
  }
}

}  // namespace fbstab


#include "ros/ros.h"
#include "geometry_msgs/Vector3.h"

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

using namespace Eigen;
Vector3d position(-5,-5,-5);

void position_callback(const geometry_msgs::Vector3::ConstPtr& msg)
{
	position(0) = msg->x;
	position(1) = msg->y;
	position(2) = msg->z;
}

int main(int argc, char **argv)
{
	// ROS Initalization
	ros::init(argc, argv, "control");

	ros::NodeHandle n;
	ros::NodeHandle private_n("~");

	ros::Subscriber position_sub = n.subscribe("position", 10, position_callback);
	ros::Publisher force_pub = n.advertise<geometry_msgs::Vector3>("force", 10);
	// ROS loop
	ros::Rate loop_rate(1000); // 1 kHz
	bool uninitialized = true;
	while (ros::ok())
	{
		ros::spinOnce();

		if (uninitialized)
		{
			if (position(0) < 0.015 && position(0) > -0.015 && position(1) < 0.015 && position(1) > -0.015)
			{
				uninitialized = false;
				//std::cout << "initialized" << std::endl;
			}

		} else {
			double dist = 10000, m_stiffness = 1000;
			int closest = -1, outside=2, axis;

			Vector3d force;
			for (axis=0; axis<2; axis++)
			{
				force(axis) = 0;
				dist = 10000;
				if (position(axis) > 0.015 || position(axis) < -0.015)
				{
					double dA = position(axis)-(0.015);
					double dB = position(axis)-(-0.015);
					if (fabs(dA) > fabs(dB)) { dist = dB; } else { dist = dA; }
					//if (fabs(dA) < fabs(dist)) { dist = dA; closest = axis; }
					//if (fabs(dB) < fabs(dist)) { dist = dB; closest = axis; }
					//outside--;
					force(axis) = -m_stiffness*dist;
				}
			}
			force(2) = 0;

			geometry_msgs::Vector3 msg;
			msg.x = force(0);
			msg.y = force(1);
			msg.z = force(2);
			force_pub.publish(msg);
		}

		loop_rate.sleep();
	}

	return 0;
}


// Copyright (c) 2019 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE sgd_mixed_mode_test_0

#include <algorithm>
#include <boost/test/unit_test.hpp>
#include <map>
#include <random>
#include <tuple>
#include <vector>

#define protected public
#include <filereader.hpp>
#include <popart/builder.hpp>
#include <popart/dataflow.hpp>
#include <popart/devicemanager.hpp>
#include <popart/inputshapeinfo.hpp>
#include <popart/ndarraywrapper.hpp>
#include <popart/op/identity.hpp>
#include <popart/op/ipucopy.hpp>
#include <popart/op/l1.hpp>
#include <popart/op/restore.hpp>
#include <popart/op/sgd0varupdate.hpp>
#include <popart/op/stash.hpp>
#include <popart/optimizer.hpp>
#include <popart/sgd.hpp>
#include <popart/tensorinfo.hpp>
#include <popart/tensornames.hpp>
#include <popart/testdevice.hpp>
#undef protected

#include <popart/session.hpp>

BOOST_AUTO_TEST_CASE(SgdMixedModeTest0) {

  using namespace popart;

  // names of 3 weights used in model
  std::string w0name = "__w0__";
  std::string w1name = "__w1__";
  std::string w2name = "__w2__";

  class Expectation {
  public:
    // after one run with SGD opt0, and one run with SGD opt1, what do we expect
    // the final value of the weight to be?
    float finalValue;

    // Given opt0 (and opt1) do we expect the computation to be done with
    // constant optimizer values?
    bool scaledLearningRateIsConst;
    bool weightDecayScaleFactorIsConst;
  };

  auto test = [w0name, w1name, w2name](
                  const SGD &opt0,       // initial Optimizer
                  const SGD &opt1,       // Optimizer to switch to
                  const Expectation &e0, // expectation for weight 0
                  const Expectation &e1, // expectation for weight 1
                  const Expectation &e2  // expectation for weight 2
              ) {
    // Model
    // ----
    //
    // loss = l1_loss((((input + w0) + w1) + w2))
    //
    // where input is small positive and w0, w1, w2 are large positive
    //

    auto builder = Builder::create();
    auto aiOnnx  = builder->aiOnnxOpset9();

    int64_t sampleDim = 1;
    int64_t batchSize = 1;
    int64_t stepSize  = 1;

    std::vector<int64_t> batchShape{batchSize, sampleDim};
    TensorInfo batchInfo{"FLOAT", batchShape};

    std::vector<int64_t> sampleShape{sampleDim};
    TensorInfo sampleInfo{"FLOAT", sampleShape};

    std::vector<int64_t> stepDataShape{stepSize, batchSize, sampleDim};
    TensorInfo stepDataInfo{"FLOAT", stepDataShape};

    auto input0 = builder->addInputTensor(batchInfo, "0tupni");

    WeightsIO weightsRead;

    std::vector<float> weight0(sampleDim, 100.0f);
    // read-back buffer for weight0
    std::vector<float> rb0(sampleDim, -777.0f);
    ConstVoidData cvd0({weight0.data(), sampleInfo});

    auto w0Id = builder->addInitializedInputTensor(cvd0, w0name);
    weightsRead.insert(w0Id, {rb0.data(), sampleInfo});

    std::vector<float> weight1(sampleDim, 200.0f);
    std::vector<float> rb1(sampleDim, -777.0f);
    ConstVoidData cvd1({weight1.data(), sampleInfo});
    auto w1Id = builder->addInitializedInputTensor(cvd1, w1name);
    weightsRead.insert(w1Id, {rb1.data(), sampleInfo});

    std::vector<float> weight2(sampleDim, 300.0f);
    std::vector<float> rb2(sampleDim, -777.0f);
    ConstVoidData cvd2({weight2.data(), sampleInfo});
    auto w2Id = builder->addInitializedInputTensor(cvd2, w2name);
    weightsRead.insert(w2Id, {rb2.data(), sampleInfo});

    auto add0 = aiOnnx.add({w0Id, input0});
    auto add1 = aiOnnx.add({w1Id, add0});
    auto add2 = aiOnnx.add({w2Id, add1});
    auto l1   = builder->aiGraphcoreOpset1().l1loss({add2}, 1.0);

    builder->addOutputTensor(add2);
    auto proto    = builder->getModelProto();
    auto dataFlow = DataFlow(stepSize);

    SessionOptions userOptions;

    auto device = createTestDevice(TEST_TARGET, 1);

    auto session = popart::TrainingSession::createFromOnnxModel(
        proto,
        dataFlow,
        l1,
        opt0, // construct with opt0, will switch to opt1 later
        device,
        InputShapeInfo(),
        SessionOptions(),
        popart::Patterns(PatternsLevel::Default));

    session->prepareDevice();
    std::vector<float> v_input_x(stepDataInfo.nelms(), 3.1415);

    popart::NDArrayWrapper<float> input_x_wrapper(v_input_x.data(),
                                                  stepDataInfo);

    std::map<popart::TensorId, popart::IArray &> inputs = {
        {input0, input_x_wrapper}};

    popart::StepIO stepio(inputs, {});

    session->weightsFromHost();

    // run 1 with opt0
    session->run(stepio);

    // run 2 with opt1
    session->updateOptimizerFromHost(&opt1);

    session->run(stepio);

    // read final weights back
    session->weightsToHost();
    session->readWeights(weightsRead);

    std::cout << "Values read-back from device for weights 0,1,2:" << std::endl;
    std::cout << rb0[0] << "  " << rb1[0] << "  " << rb2[0] << "." << std::endl;

    // All the SGD0VarUpdateOps, in no particular order
    std::vector<SGD0VarUpdateOp *> sgdOpsOOO;
    for (auto op :
         session->getIr().getOpSchedule({}, RequireOptimalSchedule::No)) {
      auto asSgd = dynamic_cast<SGD0VarUpdateOp *>(op);
      if (asSgd) {
        sgdOpsOOO.push_back(asSgd);
      }
    }
    BOOST_CHECK(sgdOpsOOO.size() == 3);

    std::vector<SGD0VarUpdateOp *> sgdOps = sgdOpsOOO;
    for (auto op : sgdOpsOOO) {
      auto id = op->inId(op->getVarToUpdateInIndex());
      if (id == w0name) {
        sgdOps[0] = op;
      } else if (id == w1name) {
        sgdOps[1] = op;
      } else if (id == w2name) {
        sgdOps[2] = op;
      } else {
        throw error("failed to determine input id in sgd test");
      }
    }

    // check that expected values are the same as read-back values
    BOOST_CHECK(e0.finalValue == rb0[0]);
    BOOST_CHECK(e1.finalValue == rb1[0]);
    BOOST_CHECK(e2.finalValue == rb2[0]);

    BOOST_CHECK(e0.scaledLearningRateIsConst == sgdOps[0]->initSlr0.isConst());
    BOOST_CHECK(e1.scaledLearningRateIsConst == sgdOps[1]->initSlr0.isConst());
    BOOST_CHECK(e2.scaledLearningRateIsConst == sgdOps[2]->initSlr0.isConst());

    BOOST_CHECK(e0.weightDecayScaleFactorIsConst ==
                sgdOps[0]->initWdsf0.isConst());
    BOOST_CHECK(e1.weightDecayScaleFactorIsConst ==
                sgdOps[1]->initWdsf0.isConst());
    BOOST_CHECK(e2.weightDecayScaleFactorIsConst ==
                sgdOps[2]->initWdsf0.isConst());
  };

  // Test case 1
  // ------------
  std::cout << "TEST CASE 1" << std::endl;
  OptimizerValue defaultWeightDecay{0, false};
  OptimizerValue defaultLearningRate{0.1, true};
  OptimizerValue lossScaling{10, true};

  std::map<std::string, std::pair<float, bool>> optParams;
  optParams.insert({"defaultWeightDecay", {0.0f, false}});
  optParams.insert({"defaultLearningRate", {0.1f, true}});
  optParams.insert({"lossScaling", {10.0f, true}});

  auto opt0 = SGD(optParams);
  opt0.insertSpecific(
      w1name, {{"learningRate", {0.2, false}}, {"weightDecay", {0.0f, true}}});

  // same as opt1, but increased learning rate for weight 1
  auto opt1 = SGD(optParams);

  opt1.insertSpecific(
      w1name, {{"learningRate", {0.5, false}}, {"weightDecay", {0, true}}});

  // weight 0 uses the global optimizer values
  Expectation e0({
      100 - 0.1 - 0.1, // initial value - (lr run 1) - (lr run 2)
      true,            // constant scaled learning rate
      false            // non-constant weight decay scale factor
  });

  // weight 1 uses local optimizer values
  Expectation e1({
      200 - 0.2 - 0.5, // initial - (lr 1) - (lr 2)
      false,           // non-constant scaled learning rate
      false // non-constant scaled weight decay - this is because of the
            // dependance on the learning rate, which is non-constant
  });

  Expectation e2({300 - 0.1 - 0.1, true, false});

  test(opt0, opt1, e0, e1, e2);

  // Test case 2
  // -----------
  // We confirm that a non-const lossScaling results in all learning rates being
  // non-const
  std::cout << "TEST CASE 2" << std::endl;

  opt0 = SGD({{"defaultLearningRate", {2, true}},
              {"defaultWeightDecay", {0, true}},
              {"lossScaling", {20, false}}});

  opt0.insertSpecific(
      w1name, {{"learningRate", {4.0, true}}, {"weightDecay", {0.125, true}}});

  opt1 = opt0;

  // weight 0 uses the global optimizer values
  e0 = {
      100 - 2.0 - 2.0,
      false, // scaled learning rate is not constant, becase loss scaling isn't
      true   // weight decay scaling factor is constant
  };

  // 0.125 * 4.0 = 0.5 (half-life of 1)
  // 200 * (1 - 0.5) - 4 = 96
  // 96  * (1 - 0.5) - 4 = 44
  //
  // weight 1 uses local optimizer values
  e1 = {44, false, true};

  e2 = {300 - 2.0 - 2.0, false, true};

  test(opt0, opt1, e0, e1, e2);
}


#include <boost/accumulators/statistics/extended_p_square_quantile.hpp>


#include "lsgc.h"

#include <deal.II/base/quadrature_lib.h>

template <int dim>
LSGC<dim>::LSGC (const dealii::ParameterHandler &prm)
    :
    AQBase<dim> (prm) {}

template <int dim>
LSGC<dim>::~LSGC () {}

template <int dim>
void LSGC<dim>::ProduceAQ () {
  AssertThrow (this->n_azi_%2==0,
               dealii::ExcMessage("SN order must be even numbers"));
  AssertThrow (dim>1,
               dealii::ExcMessage("LSGC only exists in multi-D"));

  dealii::QGauss<1> mu_quad (this->n_azi_);

  switch (dim) {
    case 2: {
      this->n_dir_ = this->n_azi_ * (this->n_azi_ + 2) /
          (this->transport_model_name_=="ep" ? 4:2);
      this->total_angle_ = 8.0 * this->k_pi;
      break;
    }
    case 3: {
      this->n_dir_ = this->n_azi_ * (this->n_azi_ + 2) /
          (this->transport_model_name_=="ep" ? 2:1);
      this->total_angle_ = 4.0 * this->k_pi *
          (this->transport_model_name_=="ep" ? 2.0 : 1.0);
      break;
    }
    default:
      break;
  }

  int n_total_azi = ((dim==3 && this->transport_model_name_!="ep")?
                              this->n_azi_ : this->n_azi_ / 2);
  for (int i=0; i<n_total_azi; ++i) {
    double mu = mu_quad.point(i)[0] * 2.0 - 1.0;
    int n_level = ((i<this->n_azi_/2?4*(i+1):4*(this->n_azi_-i))/
                            ((dim==2&&this->transport_model_name_=="ep")?2:1));
    double dphi = 2.0 * this->k_pi /
        (n_level*(this->transport_model_name_=="ep"?2.0:1.0));
    double w_pt = mu_quad.weight(i) * this->total_angle_ / n_level;
    for (int j=0; j<n_level; ++j) {
      dealii::Tensor<1, dim> omega;
      double phi = (j + 0.5) * dphi;
      omega[0] = std::sqrt (1.0 - mu * mu) * cos (phi);
      omega[1] = std::sqrt (1.0 - mu * mu) * sin (phi);
      if (dim==3)
        omega[2] = mu;
      this->wi_.push_back (w_pt);
      this->omega_i_.push_back (omega);
    }
  }

  AssertThrow (this->n_dir_==static_cast<int>(this->wi_.size()),
               dealii::ExcMessage("calculated number of angles should be the same as number of angular weights"));
  AssertThrow (this->n_dir_==static_cast<int>(this->omega_i_.size()),
               dealii::ExcMessage("calculated number of angles should be the same as number of angles"));

  this->n_total_ho_vars_ = this->n_dir_ * this->n_group_;
}

template class LSGC<2>;
template class LSGC<3>;


#include <boost/graph/bipartite.hpp>


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

    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_ICEIL_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_ICEIL_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

  /*!

    @ingroup group-arithmetic
    This function object computes the integer conversion of the @ref ceil of its parameter.

    @par Header <boost/simd/function/iceil.hpp>

    @par Note:
    This operation is properly saturated

    @see ceil, ifloor, iround, ifix, inearbyint

    @par Example:

      @snippet iceil.cpp iceil

    @par Possible output:

      @snippet iceil.txt iceil

  **/
  as_integer_t<Value> iceil(Value const& x);
} }
#endif

#include <boost/simd/function/scalar/iceil.hpp>
#include <boost/simd/function/simd/iceil.hpp>

#endif


//--------------------------- Warp the input scalar 3D image with the given input displacement field-----------//
/* Options include :
--inputDisplacement inDisplacement   :
--inputImage        inImage          :
--outputImage       outImage
--scale             scaleValue       : Value with which the displacement field will be scaled before using it to warp the image. Default: 1
--invert            true/false       : Whether the field must be inverted or not.
--modulate          true/false       : Whether it should be modulated with the jacobian determinant or not.
--interpolator      l/b/n            : linear/bspline/nearestNeighbor
--order             bspline order    : if bspline interpolator selected

Author:      Bishesh Khanal
Asclepios, INRIA Sophia Antipolis
*/

#include <iostream>
#include <sstream>
#include <string>

#include <itkImage.h>
#include <itkImageAdaptor.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

#include "InverseDisplacementImageFilter.h"
#include <itkWarpImageFilter.h>

#include <itkNearestNeighborInterpolateImageFunction.h>
#include <itkBSplineInterpolateImageFunction.h>
#include <itkDivideImageFilter.h>

#include <itkDisplacementFieldJacobianDeterminantFilter.h>

#include <boost/program_options.hpp>

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc, char **argv)
{
    std::string inImgFile;
    std::string outImgFile;
    std::string displacementFile;
    std::string interpolator;

    double scale;
    bool invertField;
    bool modulate;
    int bsplineOrder;

    //------------------- Set up the command line options database and parse them -----------------------------//
    boost::program_options::options_description optionsDescription("Possible options");
    optionsDescription.add_options()
            ("help,h", "display help message")
            ("inImage", boost::program_options::value< std::string >(&inImgFile),
             "Filename of the input image to be warped")
            ("outImage", boost::program_options::value< std::string >(&outImgFile),
             "Filename of the output warped image")
            ("displacementImage", boost::program_options::value< std::string >(&displacementFile),
             "Filename of the displacement field image")
            ("interpolator", boost::program_options::value< std::string >(&interpolator)->default_value("linear"),
             "linear/bspline/nearestneighbor interpolator.")
            ("scaleValue",boost::program_options::value< double >(&scale)->default_value(1.),
             "scale the displacement field by s.")
            ("modulate",boost::program_options::value< bool >(&modulate)->default_value(false),
             "true or false: modulate with Jacobian determinant of the field.")
            ("invert",boost::program_options::value< bool >(&invertField)->default_value(false),
             "true/false invert displacement field before warping.")
            ("order",boost::program_options::value< int >(&bsplineOrder),
             "Bspline order, use only if bspline interpolator used")
            ;

    boost::program_options::variables_map options;
    boost::program_options::store(boost::program_options::parse_command_line(argc,argv,optionsDescription),options);
    boost::program_options::notify(options);

    //If help is asked!
    if(options.count("help")) {
        std::cout<<optionsDescription<<std::endl;
        return EXIT_SUCCESS;
    }

    //Confirm all the required options are given.
    if(!options.count("inImage") || !options.count("outImage") || !options.count("displacementImage")
            ) {
        std::cerr<<"invalid options! run with --help or -h to see the proper options."<<std::endl;
        return EXIT_FAILURE;
    }

    //---------------------  Read the displacement field ----------------------//
    typedef itk::Vector<double, 3>                      DisplacementPixelType;
    typedef itk::Image<DisplacementPixelType, 3>        DisplacementImageType;
    typedef itk::ImageFileReader<DisplacementImageType> DisplacementReaderType;
    DisplacementImageType::Pointer warperField;

    DisplacementReaderType::Pointer displacementReader = DisplacementReaderType::New();
    displacementReader->SetFileName(displacementFile);
    displacementReader->Update();
    warperField = displacementReader->GetOutput();

    //---------------------- Scale the displacement field --------------------//
    double eps = 0.001;
    if( !((scale>(1-eps)) && (scale < (1+eps))) ){
        itk::ImageRegionIterator<DisplacementImageType> iterator(warperField,warperField->GetLargestPossibleRegion());
        iterator.GoToBegin();
        DisplacementImageType::PixelType pixelValue;
        while(!iterator.IsAtEnd()) {
            pixelValue = iterator.Get();
            pixelValue[0] = pixelValue[0]*scale;
            pixelValue[1] = pixelValue[1]*scale;
            pixelValue[2] = pixelValue[2]*scale;
            iterator.Set(pixelValue);
            ++iterator;
        }
    }

    //----------------------- Invert the displacement field ----------------------//
    if(invertField) {
        typedef InverseDisplacementImageFilter<DisplacementImageType> FPInverseType;
        FPInverseType::Pointer inverter1 = FPInverseType::New();
        inverter1->SetInput(warperField);
        inverter1->SetErrorTolerance(1e-1);
        inverter1->SetMaximumNumberOfIterations(50);
        inverter1->Update();
        std::cout<<"tolerance not reached for "<<inverter1->GetNumberOfErrorToleranceFailures()<<" pixels"<<std::endl;
        warperField = inverter1->GetOutput();

        typedef itk::ImageFileWriter<DisplacementImageType> WriterType;
        WriterType::Pointer writer = WriterType::New();
        writer->SetFileName("displacementFieldInverted.nii.gz");
        writer->SetInput(warperField);
        writer->Update();
    }

    //--------------------  Read and warp the input image -------------------------//
    typedef double                                  ImagePixelType;
    typedef itk::Image<ImagePixelType, 3>           ImageType;
    typedef itk::ImageFileReader<ImageType>         ImageReaderType;
    ImageType::Pointer   inputImage;
    ImageType::Pointer  outputImage;
    {
        ImageReaderType::Pointer imageReader = ImageReaderType::New();
        imageReader->SetFileName(inImgFile);
        imageReader->Update();
        inputImage = imageReader->GetOutput();

        typedef itk::WarpImageFilter<ImageType,ImageType,DisplacementImageType> WarpFilterType;
        WarpFilterType::Pointer warper = WarpFilterType::New();
        warper->SetDisplacementField(warperField);
        warper->SetInput(inputImage);
        warper->SetOutputSpacing(inputImage->GetSpacing());
        warper->SetOutputOrigin(inputImage->GetOrigin());
        warper->SetOutputDirection(inputImage->GetDirection());

        if(interpolator.compare("bspline")==0) {
            typedef itk::BSplineInterpolateImageFunction<ImageType> InterpolatorFilterType;
            InterpolatorFilterType::Pointer interpolatorFilter = InterpolatorFilterType::New();
            interpolatorFilter->SetSplineOrder(bsplineOrder);
            warper->SetInterpolator(interpolatorFilter);
            //            warper->Update();
        } else if (interpolator.compare("nearestneighbor")==0) {
            typedef itk::NearestNeighborInterpolateImageFunction<ImageType> InterpolatorFilterType;
            InterpolatorFilterType::Pointer interpolatorFilter = InterpolatorFilterType::New();
            warper->SetInterpolator(interpolatorFilter);
            //            warper->Update();
        } else {
            //            warper->Update();
        }
        warper->Update();
        outputImage = warper->GetOutput();

        //-------------------- Jacobian determinant of the input field ----------------------------------//
        if(modulate) {
            typedef itk::DisplacementFieldJacobianDeterminantFilter< DisplacementImageType, double > JacobianFilterType;
            JacobianFilterType::Pointer jacobianFilter = JacobianFilterType::New();
            jacobianFilter->SetUseImageSpacingOff();
            jacobianFilter->SetInput(warperField);
            jacobianFilter->Update();

            typedef itk::ImageFileWriter<ImageType> WriterType;
            WriterType::Pointer writer1 = WriterType::New();
            writer1->SetFileName("jacobian.mha");
            writer1->SetInput(jacobianFilter->GetOutput());
            writer1->Update();

            //----------------------- modulate the warped atrophy map with the jacobian --------------------//
            typedef itk::DivideImageFilter< ImageType, ImageType, ImageType > ModulatorType;
            ModulatorType::Pointer modulator = ModulatorType::New();
            modulator->SetInput1(outputImage);
            modulator->SetInput2(jacobianFilter->GetOutput());
            modulator->Update();
            outputImage = modulator->GetOutput();
        }

        //---------------------- write the warped with modulation atrophy map -----------------------//
        typedef itk::ImageFileWriter<ImageType> WriterType;
        WriterType::Pointer writer = WriterType::New();
        writer->SetFileName(outImgFile);
        writer->SetInput(outputImage);
        writer->Update();
    }

    //use it to warp atrophy.
    return(EXIT_SUCCESS);

}