#include <algorithm>
#include <cassert>
#include <iostream>
#include <iterator>
#include <unordered_map>
#include <queue>
#include <deque>
#include <random>
#include <boost/bind.hpp>
#include <boost/function.hpp>
#include <boost/noncopyable.hpp>
#include "slide_algorithm/truncater.hpp"

typedef std::vector<std::vector<point_type>> matrix_type;
struct evaluate_set_type
{
    matrix_type matrix;
    point_type position;//場所
    point_type content;
    int score;
    std::string direct;
};

// class definition {{{1
class truncater::impl : boost::noncopyable
{
public:
    typedef truncater::return_type return_type;

    impl() = default;
    virtual ~impl() = default;

    auto get() -> boost::optional<return_type>;
    void reset(question_data const& data);

private:
    auto process() -> boost::optional<return_type>;

    boost::optional<question_data> data_;

    void ymove();
    int form_evaluate(matrix_type const& mat);
    point_type get_start_point(matrix_type const& mat);
    int eval_two_piece(evaluate_set_type const& eval_set, point_type const& new_position);
    evaluate_set_type try_u(evaluate_set_type return_set);
    evaluate_set_type try_r(evaluate_set_type return_set);
    evaluate_set_type try_d(evaluate_set_type return_set);
    evaluate_set_type try_l(evaluate_set_type return_set);

    matrix_type matrix;
    answer_type answer;
    int width;
    int height;
    int selectable;
};

// interfaces for pointer to implementation idiom {{{1
truncater::truncater()
    : pimpl_(new impl())
{
}

truncater::~truncater()
{
    delete pimpl_;
}

auto truncater::get() -> boost::optional<return_type>
{
    return pimpl_->get();
}

void truncater::reset(question_data const& data)
{
    pimpl_->reset(data);
}

auto truncater::impl::get() -> boost::optional<return_type>
{
    return this->process();
}

void truncater::impl::reset(question_data const& data)
{
    data_ = data.clone();
}

// yoshikawa {{{1
int truncater::impl::form_evaluate(matrix_type const& mat)
{
    int s = 0;
    for (int i = 0; i < mat.size(); ++i)
    {
        for (int j = 0; j < mat.at(0).size(); ++j)
        {
            s += mat[i][j].manhattan_pow({ j, i });
        }
    }
    return s;
}

point_type truncater::impl::get_start_point(matrix_type const& mat)
{
    int max_val = 0;
    point_type max_point;//実態
    point_type position;

    for (int i = 0; i < mat.size(); ++i)
    {
        for (int j = 0; j < mat.at(0).size(); ++j)
        {
            int temp = mat[i][j].manhattan_pow({ j, i });
            if (temp > max_val)
            {
                max_val = temp;
                max_point = point_type{ j, i };
                position = point_type{ mat[i][j].x, mat[i][j].y };
            }
        }
    }
    return position;
}

int truncater::impl::eval_two_piece(evaluate_set_type const& eval_set, point_type const& new_position)
{
    int s = 0;
    point_type const& content_a = eval_set.matrix[eval_set.position.y][eval_set.position.x];
    //aの新しい場所からの距離
    s += content_a.manhattan_pow(new_position);
    //aの古い場所からの距離
    s -= content_a.manhattan_pow(eval_set.position);
    point_type const& content_b = eval_set.matrix[new_position.y][new_position.x];
    //bの新しい場所からの距離
    s += content_b.manhattan_pow(eval_set.position);
    //bの古い場所からの距離
    s -= content_b.manhattan_pow(new_position);

    //assert(s == -2 || s == 0 || s == 2);
    return s;
}

evaluate_set_type truncater::impl::try_u(evaluate_set_type return_set)
{
    point_type const& new_position = { return_set.position.x, return_set.position.y - 1 };
    return_set.score += eval_two_piece(return_set, new_position);
    std::swap(return_set.matrix[return_set.position.y][return_set.position.x], return_set.matrix[new_position.y][new_position.x]);
    return_set.direct += "U";
    return_set.position = new_position;
    return std::move(return_set);
}

evaluate_set_type truncater::impl::try_r(evaluate_set_type return_set)
{
    point_type const& new_position = { return_set.position.x + 1, return_set.position.y };
    return_set.score += eval_two_piece(return_set, new_position);
    std::swap(return_set.matrix[return_set.position.y][return_set.position.x], return_set.matrix[new_position.y][new_position.x]);
    return_set.direct += "R";
    return_set.position = new_position;
    return std::move(return_set);
}

evaluate_set_type truncater::impl::try_d(evaluate_set_type return_set)
{
    point_type const& new_position = { return_set.position.x, return_set.position.y + 1 };
    return_set.score += eval_two_piece(return_set, new_position);
    std::swap(return_set.matrix[return_set.position.y][return_set.position.x], return_set.matrix[new_position.y][new_position.x]);
    return_set.direct += "D";
    return_set.position = new_position;
    return std::move(return_set);
}

evaluate_set_type truncater::impl::try_l(evaluate_set_type return_set)
{
    point_type const& new_position = { return_set.position.x - 1, return_set.position.y };
    return_set.score += eval_two_piece(return_set, new_position);
    std::swap(return_set.matrix[return_set.position.y][return_set.position.x], return_set.matrix[new_position.y][new_position.x]);
    return_set.direct += "L";
    return_set.position = new_position;
    return std::move(return_set);
}

void truncater::impl::ymove()
{
    std::cout << "ymove start" << std::endl;

    int const width = matrix.at(0).size();
    int const height = matrix.size();

    std::queue<evaluate_set_type> que;
    std::vector<evaluate_set_type> children;
    children.reserve(900);

    point_type const start_position = get_start_point(matrix);
    evaluate_set_type start = evaluate_set_type{ matrix, start_position, matrix[start_position.y][start_position.x], form_evaluate(start.matrix), " " };
    evaluate_set_type best = start;

    que.push(start);

    //std::cout << "select piece position = " << start.position << std::endl;
    //std::cout << "select piece content = " << start.content << std::endl;

    auto P = [](int e, int ne, double T) {
        if( ne < e ) {
            return 1.0;
        } else {
            return std::exp((e - ne)/T);
        }
    };
    double T = 1000.0;

    std::random_device sec_rng;
    std::mt19937 mt(sec_rng());
    std::uniform_real_distribution<double> prob_dist(0.0, 1.0);
    auto prob = [&]{return prob_dist(mt);};

    auto is_transition = [&](int s, int ns) {
        return prob() < P(s, ns, T);
    };

    while (que.size() > 0 && T > 1.0)
    {
        //std::cout << "T = " << T << ", score = " << best.score << " que size = " << que.size() << std::endl;
        while (que.size() > 0)
        {
            auto const node = que.front();
            que.pop();

            if (best.score > node.score) best = node;

            if (node.position.y != 0 && node.direct.back() != 'D')
            {
                auto child = try_u(node);
                //if (child.score <= node.score)
                if (is_transition(node.score, child.score))
                {
                    children.push_back(std::move(child));
                }
            }
            if (node.position.x != width - 1 && node.direct.back() != 'L')
            {
                auto child = try_r(node);
                //if (child.score <= node.score)
                if (is_transition(node.score, child.score))
                {
                    children.push_back(std::move(child));
                }
            }
            if (node.position.y != height - 1 && node.direct.back() != 'U')
            {
                auto child = try_d(node);
                //if (child.score <= node.score)
                if (is_transition(node.score, child.score))
                {
                    children.push_back(std::move(child));
                }
            }
            if (node.position.x != 0 && node.direct.back() != 'R')
            {
                auto child = try_l(node);
                //if (child.score <= node.score)
                if (is_transition(node.score, child.score))
                {
                    children.push_back(std::move(child));
                }
            }
        }

        std::sort(children.begin(), children.end(), [](evaluate_set_type a, evaluate_set_type b){return a.score < b.score; });

        for (int i = 0; i < children.size() && i < 100; ++i)
        {
            que.push(children.at(i));
        }

        children.clear();

        T *= 0.95;
    }

    matrix = best.matrix;
    answer.list.push_back({ start_position, best.direct.substr(1) });

    std::cout << "ymove done." << std::endl;
}

// process {{{1
auto truncater::impl::process() -> boost::optional<return_type>
{
    algorithm algo;
    matrix = data_->block;
    width = data_->size.first;
    height = data_->size.second;
    selectable = data_->selectable;

    if ((width > 3 || height > 3) && selectable >= 3) ymove();

    auto qdata = question_data(
        data_->problem_id,
        data_->player_id,
        data_->size,
        data_->selectable,
        data_->cost_select,
        data_->cost_change,
        std::move(matrix)
    );
    algo.reset(std::move(qdata));
    answer_type subsequent_answer = *algo.get();
    std::copy(subsequent_answer.list.begin(), subsequent_answer.list.end(), std::back_inserter(answer.list));
    return std::move(answer);
}


#include <actionlib/server/simple_action_server.h>
#include <kr_tracker_msgs/LineTrackerAction.h>
#include <kr_tracker_msgs/TrackerStatus.h>
#include <kr_trackers/initial_conditions.h>
#include <kr_trackers_manager/Tracker.h>
#include <ros/ros.h>

#include <Eigen/Core>
#include <memory>

class SmoothVelTracker : public kr_trackers_manager::Tracker
{
 public:
  SmoothVelTracker(void);

  void Initialize(const ros::NodeHandle &nh);
  bool Activate(const kr_mav_msgs::PositionCommand::ConstPtr &cmd);
  void Deactivate(void);

  kr_mav_msgs::PositionCommand::ConstPtr update(const nav_msgs::Odometry::ConstPtr &msg);
  uint8_t status() const;

 private:
  void goal_callback();
  void preempt_callback();

  using ServerType = actionlib::SimpleActionServer<kr_tracker_msgs::LineTrackerAction>;

  // Action server that takes a goal.
  // Must be a pointer, because plugin does not support a constructor
  // with inputs, but an action server must be initialized with a Nodehandle.
  std::unique_ptr<ServerType> tracker_server_;

  bool goal_set_, goal_reached_;
  float target_speed_;
  float ramp_dist_, total_dist_;
  float ramp_time_, total_time_;
  float start_yaw_, goal_yaw_;
  float yaw_dot_max_;
  bool active_;
  InitialConditions ICs_;
  ros::Time start_time_;
  Eigen::Vector3f start_pos_, dir_;
  Eigen::Matrix<float, 7, 1> vel_coeffs_;
  float current_traj_length_;
  Eigen::Vector3f prev_pos_;
};

SmoothVelTracker::SmoothVelTracker(void) : goal_set_(false), goal_reached_(true), active_(false) {}

void SmoothVelTracker::Initialize(const ros::NodeHandle &nh)
{
  ros::NodeHandle priv_nh(nh, "smooth_vel_tracker");

  // Set up the action server.
  tracker_server_ = std::unique_ptr<ServerType>(new ServerType(priv_nh, "SmoothVelTracker", false));
  tracker_server_->registerGoalCallback(boost::bind(&SmoothVelTracker::goal_callback, this));
  tracker_server_->registerPreemptCallback(boost::bind(&SmoothVelTracker::preempt_callback, this));

  tracker_server_->start();
}

bool SmoothVelTracker::Activate(const kr_mav_msgs::PositionCommand::ConstPtr &cmd)
{
  // Only allow activation if a goal has been set
  if(goal_set_)
    active_ = true;
  return active_;
}

void SmoothVelTracker::Deactivate(void)
{
  if(tracker_server_->isActive())
  {
    ROS_WARN("SmoothVelTracker::Deactivate: deactivated tracker while "
             "still tracking the goal.");
    tracker_server_->setAborted();
  }

  ICs_.reset();
  goal_set_ = false;
  active_ = false;
}

kr_mav_msgs::PositionCommand::ConstPtr SmoothVelTracker::update(const nav_msgs::Odometry::ConstPtr &msg)
{
  if(!active_)
  {
    ICs_.set_from_odom(msg);
    return kr_mav_msgs::PositionCommand::Ptr();
  }

  // Record distance between last position and current.
  const Eigen::Vector3f current_pos(msg->pose.pose.position.x, msg->pose.pose.position.y, msg->pose.pose.position.z);

  current_traj_length_ += (current_pos - prev_pos_).norm();

  prev_pos_ = current_pos;

  const ros::Time t_now = ros::Time::now();

  if(goal_set_)
  {
    start_time_ = t_now;
    goal_set_ = false;
  }

  // Set gains
  kr_mav_msgs::PositionCommand::Ptr cmd(new kr_mav_msgs::PositionCommand);
  cmd->header.stamp = t_now;
  cmd->header.frame_id = msg->header.frame_id;

  // Get elapsed time
  const ros::Duration elapsed_time = t_now - start_time_;
  const float t = elapsed_time.toSec();
  const float ts = t / ramp_time_;  // scaled time
  const float ts2 = ts * ts;
  const float ts3 = ts2 * ts;
  const float ts4 = ts3 * ts;
  const float ts5 = ts4 * ts;
  const float ts6 = ts5 * ts;
  const float ts7 = ts6 * ts;
  const float ts8 = ts7 * ts;

  // Test each case to generate trajectory
  Eigen::Vector3f pos, vel, acc, jrk;
  if(t > total_time_)
  {
    pos = start_pos_ + total_dist_ * dir_;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = 0, cmd->velocity.y = 0, cmd->velocity.z = 0;
    cmd->acceleration.x = 0, cmd->acceleration.y = 0, cmd->acceleration.z = 0;
    cmd->jerk.x = 0, cmd->jerk.y = 0, cmd->jerk.z = 0;
    cmd->yaw = goal_yaw_;
    cmd->yaw_dot = 0;
    goal_reached_ = true;
    if(tracker_server_->isActive())
    {
      // Send a success message and reset the length and duration variables.
      kr_tracker_msgs::LineTrackerResult result;
      result.duration = total_time_;
      result.length = current_traj_length_;
      result.x = pos(0);
      result.y = pos(1);
      result.z = pos(2);
      result.yaw = goal_yaw_;

      tracker_server_->setSucceeded(result);
      current_traj_length_ = 0.0;
    }
  }
  else if(t < ramp_time_)
  {
    float dist = ramp_time_ * (vel_coeffs_(0) / 2 * ts2 + vel_coeffs_(1) / 3 * ts3 + vel_coeffs_(2) / 4 * ts4 +
                               vel_coeffs_(3) / 5 * ts5 + vel_coeffs_(4) / 6 * ts6 + vel_coeffs_(5) / 7 * ts7 +
                               vel_coeffs_(6) / 8 * ts8);
    float speed = vel_coeffs_(0) * ts + vel_coeffs_(1) * ts2 + vel_coeffs_(2) * ts3 + vel_coeffs_(3) * ts4 +
                  vel_coeffs_(4) * ts5 + vel_coeffs_(5) * ts6 + vel_coeffs_(6) * ts7;
    float accel = (vel_coeffs_(0) + 2 * vel_coeffs_(1) * ts + 3 * vel_coeffs_(2) * ts2 + 4 * vel_coeffs_(3) * ts3 +
                   5 * vel_coeffs_(4) * ts4 + 6 * vel_coeffs_(5) * ts5 + 7 * vel_coeffs_(6) * ts6) /
                  ramp_time_;
    float jerk = (2 * vel_coeffs_(1) + 6 * vel_coeffs_(2) * ts + 12 * vel_coeffs_(3) * ts2 + 20 * vel_coeffs_(4) * ts3 +
                  30 * vel_coeffs_(5) * ts4 + 42 * vel_coeffs_(6) * ts5) /
                 (ramp_time_ * ramp_time_);
    pos = start_pos_ + dist * dir_;
    vel = speed * dir_;
    acc = accel * dir_;
    jrk = jerk * dir_;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = vel(0), cmd->velocity.y = vel(1), cmd->velocity.z = vel(2);
    cmd->acceleration.x = acc(0), cmd->acceleration.y = acc(1), cmd->acceleration.z = acc(2);
    cmd->jerk.x = jrk(0), cmd->jerk.y = jrk(1), cmd->jerk.z = jrk(2);
    cmd->yaw = start_yaw_ + yaw_dot_max_ / (2 * ramp_time_) * t * t;
    cmd->yaw_dot = yaw_dot_max_ / ramp_time_ * t;
  }
  else if(t < total_time_ - ramp_time_)
  {
    float dist = ramp_dist_ + target_speed_ * (t - ramp_time_);
    pos = start_pos_ + dist * dir_;
    vel = target_speed_ * dir_;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = vel(0), cmd->velocity.y = vel(1), cmd->velocity.z = vel(2);
    cmd->acceleration.x = 0, cmd->acceleration.y = 0, cmd->acceleration.z = 0;
    cmd->jerk.x = 0, cmd->jerk.y = 0, cmd->jerk.z = 0;
    cmd->yaw = start_yaw_ + yaw_dot_max_ * (t - ramp_time_ / 2);
    cmd->yaw_dot = yaw_dot_max_;
  }
  else
  {
    const float te = total_time_ - elapsed_time.toSec();  // time from end
    const float tes = te / ramp_time_;                    // scaled time from end
    const float tes2 = tes * tes;
    const float tes3 = tes2 * tes;
    const float tes4 = tes3 * tes;
    const float tes5 = tes4 * tes;
    const float tes6 = tes5 * tes;
    const float tes7 = tes6 * tes;
    const float tes8 = tes7 * tes;
    const float dist_from_end =
        ramp_time_ *
        (vel_coeffs_(0) / 2 * tes2 + vel_coeffs_(1) / 3 * tes3 + vel_coeffs_(2) / 4 * tes4 + vel_coeffs_(3) / 5 * tes5 +
         vel_coeffs_(4) / 6 * tes6 + vel_coeffs_(5) / 7 * tes7 + vel_coeffs_(6) / 8 * tes8);
    const float dist = total_dist_ - dist_from_end;
    const float speed = vel_coeffs_(0) * tes + vel_coeffs_(1) * tes2 + vel_coeffs_(2) * tes3 + vel_coeffs_(3) * tes4 +
                        vel_coeffs_(4) * tes5 + vel_coeffs_(5) * tes6 + vel_coeffs_(6) * tes7;
    const float accel =
        (vel_coeffs_(0) + 2 * vel_coeffs_(1) * tes + 3 * vel_coeffs_(2) * tes2 + 4 * vel_coeffs_(3) * tes3 +
         5 * vel_coeffs_(4) * tes4 + 6 * vel_coeffs_(5) * tes5 + 7 * vel_coeffs_(6) * tes6) /
        ramp_time_;
    const float jerk = (2 * vel_coeffs_(1) + 6 * vel_coeffs_(2) * tes + 12 * vel_coeffs_(3) * tes2 +
                        20 * vel_coeffs_(4) * tes3 + 30 * vel_coeffs_(5) * tes4 + 42 * vel_coeffs_(6) * tes5) /
                       (ramp_time_ * ramp_time_);
    pos = start_pos_ + dist * dir_;
    vel = speed * dir_;
    acc = -accel * dir_;
    jrk = jerk * dir_;
    cmd->position.x = pos(0), cmd->position.y = pos(1), cmd->position.z = pos(2);
    cmd->velocity.x = vel(0), cmd->velocity.y = vel(1), cmd->velocity.z = vel(2);
    cmd->acceleration.x = acc(0), cmd->acceleration.y = acc(1), cmd->acceleration.z = acc(2);
    cmd->jerk.x = jrk(0), cmd->jerk.y = jrk(1), cmd->jerk.z = jrk(2);
    cmd->yaw = goal_yaw_ - yaw_dot_max_ / (2 * ramp_time_) * (total_time_ - t) * (total_time_ - t);
    cmd->yaw_dot = yaw_dot_max_ - yaw_dot_max_ / ramp_time_ * (t - total_time_ + ramp_time_);
  }
  ICs_.set_from_cmd(cmd);

  if(!goal_reached_)
  {
    kr_tracker_msgs::LineTrackerFeedback feedback;
    Eigen::Vector3f goal = start_pos_ + total_dist_ * dir_;
    feedback.distance_from_goal = (current_pos - goal).norm();
    tracker_server_->publishFeedback(feedback);
  }

  return cmd;
}

void SmoothVelTracker::preempt_callback()
{
  ROS_INFO("SmoothVelTracker goal preempted.");
  tracker_server_->setPreempted();

  // TODO: How much overshoot will this cause at high velocities?
  total_time_ = (ros::Time::now() - start_time_).toSec();
  total_dist_ = (ICs_.pos() - start_pos_).norm();
  dir_ = (ICs_.pos() - start_pos_).normalized();
  goal_set_ = false;
  goal_reached_ = true;
}

void SmoothVelTracker::goal_callback()
{
  // Pointer to the goal recieved.
  const auto msg = tracker_server_->acceptNewGoal();

  // If preempt has been requested, then set this goal to preempted
  // and make no changes to the tracker state.
  if(tracker_server_->isPreemptRequested())
  {
    ROS_INFO("SmoothVelTracker going to goal (%f, %f, %f, %f) preempted.", msg->x, msg->y, msg->z, msg->yaw);
    tracker_server_->setPreempted();
    return;
  }

  current_traj_length_ = 0.0;

  // Make sure user specifies desired velocity
  if(msg->v_des > 0 && msg->a_des > 0)
  {
    // Set the start position
    Eigen::Vector3f start_pos = ICs_.pos();

    // Get target speed and acceleration
    target_speed_ = msg->v_des;
    float target_accel = msg->a_des;
    ramp_time_ = target_speed_ / target_accel;

    // Find goal position
    Eigen::Vector3f goal_pos;
    goal_pos(0) = msg->x, goal_pos(1) = msg->y, goal_pos(2) = msg->z;
    if(msg->relative)
      goal_pos += start_pos;

    // Find distance and direction to goal
    const float total_dist = (goal_pos - start_pos).norm();
    Eigen::Vector3f dir = (goal_pos - start_pos).normalized();

    // Compute the coefficients
    Eigen::Matrix<float, 7, 7> Ainv;
    Ainv << 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 1.0 / 6.0, 0, 35, -20, -15, -5, 2.5, -2.0 / 3.0,
        -1.0 / 6.0, -84, 45, 39, 10, -7, 1, 0.5, 70, -36, -34, -7.5, 6.5, -2.0 / 3.0, -0.5, -20, 10, 10, 2, -2,
        1.0 / 6.0, 1.0 / 6.0;

    Eigen::Matrix<float, 7, 1> b;
    b << target_speed_, 0, 0, 0, 0, 0, 0;
    vel_coeffs_ = Ainv * b;

    // Compute the ramp distance
    ramp_dist_ = (vel_coeffs_(0) / 2 + vel_coeffs_(1) / 3 + vel_coeffs_(2) / 4 + vel_coeffs_(3) / 5 +
                  vel_coeffs_(4) / 6 + vel_coeffs_(5) / 7 + vel_coeffs_(6) / 8) *
                 ramp_time_;

    // Check to make sure that twice the ramp distance is less than the entire distance
    if(2 * ramp_dist_ < total_dist)
    {
      // Set the parameters for the trajectory
      start_pos_ = start_pos;
      total_dist_ = total_dist;
      dir_ = dir;

      // Compute total time of trajectory
      total_time_ = (total_dist_ - 2 * ramp_dist_) / target_speed_ + 2 * ramp_time_;

      // Set the target yaw and max yaw dot
      start_yaw_ = ICs_.yaw();
      goal_yaw_ = std::fmod(msg->yaw, 2 * M_PI);
      yaw_dot_max_ = (goal_yaw_ - start_yaw_) / (total_time_ - ramp_time_);

      // Set goal_set to true and goal_reached to false
      goal_set_ = true;
      goal_reached_ = false;
    }
    else
    {
      ROS_ERROR("increase the ramp acceleration");
      tracker_server_->setAborted();
    }
  }
  else
  {
    ROS_ERROR("v_des and a_des must be nonzero!");
    tracker_server_->setAborted();
  }
}

uint8_t SmoothVelTracker::status() const
{
  return goal_reached_ ? static_cast<uint8_t>(kr_tracker_msgs::TrackerStatus::SUCCEEDED) :
                         static_cast<uint8_t>(kr_tracker_msgs::TrackerStatus::ACTIVE);
}

#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(SmoothVelTracker, kr_trackers_manager::Tracker);


/*
 * Copyright 2000-2014 NeuStar, Inc. All rights reserved.
 * NeuStar, the Neustar logo and related names and logos are registered
 * trademarks, service marks or tradenames of NeuStar, Inc. All other
 * product names, company names, marks, logos and symbols may be trademarks
 * of their respective owners.
 */

/*
 * exp_decay_sample.cpp
 *
 *  Created on: Jun 5, 2014
 *      Author: vpoliboy
 */

#include <boost/foreach.hpp>
#include "cppmetrics/core/exp_decay_sample.h"
#include "cppmetrics/core/utils.h"

namespace cppmetrics {
namespace core {

const double ExpDecaySample::DEFAULT_ALPHA = 0.015;
const Clock::duration ExpDecaySample::RESCALE_THRESHOLD(
        boost::chrono::hours(1));

ExpDecaySample::ExpDecaySample(boost::uint32_t reservoir_size, double alpha) :
        alpha_(alpha), reservoir_size_(reservoir_size), count_(0) {
    clear();
    rng_.seed(get_millis_from_epoch());
}

ExpDecaySample::~ExpDecaySample() {
}

void ExpDecaySample::clear() {
    boost::lock_guard<boost::mutex> lock(mutex_);
    values_.clear();
    count_ = 0;
    start_time_ = Clock::now();
    next_scale_time_ = start_time_ + RESCALE_THRESHOLD;
}

boost::uint64_t ExpDecaySample::size() const {
    return std::min(reservoir_size_, count_.load());
}

void ExpDecaySample::update(boost::int64_t value) {
    update(value, Clock::now());
}

void ExpDecaySample::update(boost::int64_t value,
        const Clock::time_point& timestamp) {
    boost::lock_guard<boost::mutex> rlock(mutex_);
    rescaleIfNeeded(timestamp);
    boost::random::uniform_real_distribution<> dist(0, 1);
    boost::chrono::seconds dur = boost::chrono::duration_cast<
            boost::chrono::seconds>(timestamp - start_time_);
    double priority = 0.0;
    do {
        priority = std::exp(alpha_ * dur.count()) / dist(rng_);
    } while (std::isnan(priority));

    boost::uint64_t count = ++count_;
    if (count <= reservoir_size_) {
        values_[priority] = value;
    } else {
        Double2Int64Map::iterator first_itt(values_.begin());
        double first = first_itt->first;
        if (first < priority
                && values_.insert(std::make_pair(priority, value)).second) {
            values_.erase(first_itt);
        }
    }
}

void ExpDecaySample::rescaleIfNeeded(const Clock::time_point& now) {
    if (next_scale_time_ < now) {
        Clock::time_point prevStartTime = start_time_;
        next_scale_time_ = now + RESCALE_THRESHOLD;
        prevStartTime = start_time_;
        start_time_ = now;
        rescale(prevStartTime);
    }
}

void ExpDecaySample::rescale(const Clock::time_point& prevStartTime) {
    Double2Int64Map old_values;
    std::swap(values_, old_values);
    BOOST_FOREACH (const Double2Int64Map::value_type& kv, old_values) {
        boost::chrono::seconds dur = boost::chrono::duration_cast<
                boost::chrono::seconds>(start_time_ - prevStartTime);
        double key = kv.first * std::exp(-alpha_ * dur.count());
        values_[key] = kv.second;
    }
    count_ = values_.size();
}

SnapshotPtr ExpDecaySample::getSnapshot() const {
    ValueVector vals;
    vals.reserve(values_.size());
    boost::lock_guard<boost::mutex> rlock(mutex_);
    BOOST_FOREACH (const Double2Int64Map::value_type& kv, values_) {
        vals.push_back(kv.second);
    }
    return SnapshotPtr(new Snapshot(vals));
}

} /* namespace core */
} /* namespace cppmetrics */


/*
 * Copyright (c) 2019, The Regents of the University of California
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the University nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include <boost/geometry.hpp>

#include "frProfileTask.h"
#include "gc/FlexGC_impl.h"
#include "odb/db.h"

using namespace fr;
typedef odb::dbTechLayerCutSpacingTableDefRule::LOOKUP_STRATEGY LOOKUP_STRATEGY;

inline frSquaredDistance getC2CDistSquare(
    const gtl::rectangle_data<frCoord>& rect1,
    const gtl::rectangle_data<frCoord>& rect2)
{
  gtl::point_data<frCoord> centerPt1, centerPt2;
  gtl::center(centerPt1, rect1);
  gtl::center(centerPt2, rect2);
  return gtl::distance_squared(centerPt1, centerPt2);
}

bool FlexGCWorker::Impl::checkLef58CutSpacingTbl_prlValid(
    const gtl::rectangle_data<frCoord>& viaRect1,
    const gtl::rectangle_data<frCoord>& viaRect2,
    const gtl::rectangle_data<frCoord>& markerRect,
    std::string cutClass1,
    std::string cutClass2,
    frCoord& prl,
    odb::dbTechLayerCutSpacingTableDefRule* dbRule)
{
  auto reqPrl = dbRule->getPrlEntry(cutClass1, cutClass2);

  // check for prl
  auto distX = gtl::euclidean_distance(viaRect1, viaRect2, gtl::HORIZONTAL);
  auto distY = gtl::euclidean_distance(viaRect1, viaRect2, gtl::VERTICAL);

  auto prlX = gtl::delta(markerRect, gtl::HORIZONTAL);
  auto prlY = gtl::delta(markerRect, gtl::VERTICAL);

  if (distX) {
    prlX = -prlX;
  }
  if (distY) {
    prlY = -prlY;
  }
  if (prlX > reqPrl || prlY > reqPrl) {
    prl = std::max(prlX, prlY);
    return true;
  } else
    return false;
}

bool FlexGCWorker::Impl::checkLef58CutSpacingTbl_helper(
    gcRect* viaRect1,
    gcRect* viaRect2,
    frString class1,
    frString class2,
    const frDirEnum dir,
    frSquaredDistance distSquare,
    frSquaredDistance c2cSquare,
    bool prlValid,
    frCoord prl,
    odb::dbTechLayerCutSpacingTableDefRule* dbRule)
{
  bool isSide1, isSide2;
  auto deltaH1 = gtl::delta(*viaRect1, gtl::HORIZONTAL);
  auto deltaV1 = gtl::delta(*viaRect1, gtl::VERTICAL);
  auto deltaH2 = gtl::delta(*viaRect2, gtl::HORIZONTAL);
  auto deltaV2 = gtl::delta(*viaRect2, gtl::VERTICAL);
  frSquaredDistance reqSpcSqr;

  if (dir == frDirEnum::N || dir == frDirEnum::S) {
    isSide1 = deltaH1 > deltaV1;
    isSide2 = deltaH2 > deltaV2;
  } else {
    isSide1 = deltaV1 > deltaH1;
    isSide2 = deltaV2 > deltaH2;
  }

  if (dbRule->isNoPrl() && dbRule->isCenterAndEdge(class1, class2)) {
    reqSpcSqr = dbRule->getSpacing(class1,
                                   isSide1,
                                   class2,
                                   isSide2,
                                   odb::dbTechLayerCutSpacingTableDefRule::MAX);
    reqSpcSqr *= reqSpcSqr;
    if (c2cSquare < reqSpcSqr) {
      return true;
    } else {
      reqSpcSqr
          = dbRule->getSpacing(class1,
                               isSide1,
                               class2,
                               isSide2,
                               odb::dbTechLayerCutSpacingTableDefRule::MIN);
      reqSpcSqr *= reqSpcSqr;
      if (distSquare < reqSpcSqr)
        return true;
    }
    return false;
  }
  if (class1 == class2 && !dbRule->isLayerValid()) {
    bool exactlyAligned = false;
    if (dir == frDirEnum::N || dir == frDirEnum::S)
      exactlyAligned = (prl == deltaH1) && !dbRule->isHorizontal();
    else
      exactlyAligned = (prl == deltaV1) && !dbRule->isVertical();

    auto exAlSpc = dbRule->getExactAlignedSpacing(class1);
    if (exactlyAligned && exAlSpc != -1) {
      reqSpcSqr = (frSquaredDistance) exAlSpc * exAlSpc;
      return distSquare < reqSpcSqr;
    }
  }
  LOOKUP_STRATEGY spcIdx
      = prlValid ? LOOKUP_STRATEGY::SECOND : LOOKUP_STRATEGY::FIRST;
  // check PRLALIGNED
  if (prlValid && dbRule->isPrlForAlignedCutClasses(class1, class2)) {
    gtl::rectangle_data<frCoord> edgeRect2;
    switch (dir) {
      case frDirEnum::S:
        gtl::set_points(edgeRect2, gtl::ul(*viaRect2), gtl::ur(*viaRect2));
        break;
      case frDirEnum::N:
        gtl::set_points(edgeRect2, gtl::ll(*viaRect2), gtl::lr(*viaRect2));
        break;
      case frDirEnum::E:
        gtl::set_points(edgeRect2, gtl::ll(*viaRect2), gtl::ul(*viaRect2));
        break;
      default:
        gtl::set_points(edgeRect2, gtl::lr(*viaRect2), gtl::ur(*viaRect2));
        break;
    }
    box_t qb(point_t(gtl::xl(edgeRect2), gtl::yl(edgeRect2)),
             point_t(gtl::xh(edgeRect2), gtl::yh(edgeRect2)));
    vector<pair<segment_t, gcSegment*>> results;
    auto& workerRegionQuery = getWorkerRegionQuery();
    workerRegionQuery.queryPolygonEdge(
        qb, viaRect2->getLayerNum() + 1, results);
    if (results.size() == 0)
      spcIdx = LOOKUP_STRATEGY::FIRST;
  }

  frCoord reqSpc = dbRule->getSpacing(class1, isSide1, class2, isSide2, spcIdx);
  reqSpcSqr = (frSquaredDistance) reqSpc * reqSpc;

  bool useCenter
      = dbRule->isCenterToCenter(class1, class2)
        || (dbRule->isCenterAndEdge(class1, class2)
            && reqSpc
                   == dbRule->getSpacing(
                       class1, isSide1, class2, isSide2, LOOKUP_STRATEGY::MAX));
  if (useCenter) {
    if (c2cSquare < reqSpcSqr)
      return true;
  } else if (distSquare < reqSpcSqr) {
    return true;
  }
  return false;
}

bool FlexGCWorker::Impl::checkLef58CutSpacingTbl_sameMetal(gcRect* viaRect1,
                                                           gcRect* viaRect2)
{
  box_t qb(point_t(gtl::xl(*viaRect1), gtl::yl(*viaRect1)),
           point_t(gtl::xh(*viaRect1), gtl::yh(*viaRect1)));
  vector<rq_box_value_t<gcRect*>> results;
  auto& workerRegionQuery = getWorkerRegionQuery();
  workerRegionQuery.queryMaxRectangle(qb, viaRect1->getLayerNum() - 1, results);
  for (auto res : results) {
    auto metalRect = res.second;
    if (gtl::intersects(*metalRect, *viaRect1, false))
      if (gtl::intersects(*metalRect, *viaRect2, false))
        return true;
  }
  return false;
}

bool FlexGCWorker::Impl::checkLef58CutSpacingTbl_stacked(gcRect* viaRect1,
                                                         gcRect* viaRect2)
{
  if (*viaRect1 == *viaRect2)
    return true;
  return gtl::contains(*viaRect1, *viaRect2)
         || gtl::contains(*viaRect2, *viaRect1);
}

void FlexGCWorker::Impl::checkLef58CutSpacingTbl_main(
    gcRect* viaRect1,
    gcRect* viaRect2,
    frLef58CutSpacingTableConstraint* con)
{
  auto dbRule = con->getODBRule();
  auto layerNum1 = viaRect1->getLayerNum();
  auto layer1 = getTech()->getLayer(layerNum1);
  auto layerNum2 = viaRect2->getLayerNum();
  auto layer2 = getTech()->getLayer(layerNum2);
  bool viol = false;
  if (dbRule->isLayerValid()) {
    if (dbRule->isSameNet()) {
      if (viaRect1->getNet() != viaRect2->getNet())
        return;
      if (layer1->hasLef58SameMetalInterCutSpcTblConstraint()
          && checkLef58CutSpacingTbl_sameMetal(viaRect1, viaRect2))
        return;
    } else if (dbRule->isSameMetal()) {
      if (viaRect1->getNet() != viaRect2->getNet())
        return;
      if (!checkLef58CutSpacingTbl_sameMetal(viaRect1, viaRect2))
        return;
    } else {
      if (viaRect1->getNet() == viaRect2->getNet()) {
        return;
      }
    }
  } else {
    if (dbRule->isSameNet()) {
      if (viaRect1->getNet() != viaRect2->getNet())
        return;
    } else if (dbRule->isSameMetal()) {
      if (viaRect1->getNet() != viaRect2->getNet())
        return;
      if (!checkLef58CutSpacingTbl_sameMetal(viaRect1, viaRect2))
        return;
    } else {
      if (viaRect1->getNet() == viaRect2->getNet()) {
        if (layer1->hasLef58SameNetCutSpcTblConstraint())
          return;
        else if (checkLef58CutSpacingTbl_sameMetal(viaRect1, viaRect2)
                 && layer1->hasLef58SameMetalCutSpcTblConstraint())
          return;
      }
    }
  }
  auto cutClassIdx1
      = layer1->getCutClassIdx(viaRect1->width(), viaRect1->length());
  auto cutClassIdx2
      = layer2->getCutClassIdx(viaRect2->width(), viaRect2->length());
  frString class1 = "";
  frString class2 = "";
  if (cutClassIdx1 != -1)
    class1 = layer1->getCutClass(cutClassIdx1)->getName();
  if (cutClassIdx2 != -1)
    class2 = layer2->getCutClass(cutClassIdx2)->getName();

  gtl::rectangle_data<frCoord> markerRect(*viaRect1);
  gtl::generalized_intersect(markerRect, *viaRect2);
  frSquaredDistance distSquare
      = gtl::square_euclidean_distance(*viaRect1, *viaRect2);
  if (distSquare == 0) {
    if (dbRule->getMaxSpacing(class1, class2) == 0)
      return;
    if (!dbRule->isLayerValid())
      checkCutSpacing_short(viaRect1, viaRect2, markerRect);
    else if (dbRule->isNoStack())
      viol = true;
    else
      viol = !checkLef58CutSpacingTbl_stacked(viaRect1, viaRect2);
  }
  frSquaredDistance c2cSquare = getC2CDistSquare(*viaRect1, *viaRect2);

  bool isRight = gtl::xl(*viaRect2) > gtl::xh(*viaRect1);
  bool isLeft = gtl::xh(*viaRect2) < gtl::xl(*viaRect1);
  bool isUp = gtl::yl(*viaRect2) > gtl::yh(*viaRect1);
  bool isDown = gtl::yh(*viaRect2) < gtl::yl(*viaRect1);

  frCoord prl = -1;
  auto prlValid = checkLef58CutSpacingTbl_prlValid(
      *viaRect1, *viaRect2, markerRect, class1, class2, prl, dbRule);

  if (!viol && (isUp || isDown)) {
    frDirEnum secondToFirstDir = isUp ? frDirEnum::N : frDirEnum::S;
    viol = checkLef58CutSpacingTbl_helper(viaRect1,
                                          viaRect2,
                                          class1,
                                          class2,
                                          secondToFirstDir,
                                          distSquare,
                                          c2cSquare,
                                          prlValid,
                                          prl,
                                          dbRule);
  }
  if (!viol && (isRight || isLeft)) {
    frDirEnum secondToFirstDir = isRight ? frDirEnum::E : frDirEnum::W;
    viol = checkLef58CutSpacingTbl_helper(viaRect1,
                                          viaRect2,
                                          class1,
                                          class2,
                                          secondToFirstDir,
                                          distSquare,
                                          c2cSquare,
                                          prlValid,
                                          prl,
                                          dbRule);
  }

  if (viol) {
    // violation
    auto net1 = viaRect1->getNet();
    auto net2 = viaRect2->getNet();
    auto marker = make_unique<frMarker>();
    frBox box(gtl::xl(markerRect),
              gtl::yl(markerRect),
              gtl::xh(markerRect),
              gtl::yh(markerRect));
    marker->setBBox(box);
    marker->setLayerNum(layerNum1);
    marker->setConstraint(con);
    marker->addSrc(net1->getOwner());
    frCoord llx = gtl::xl(*viaRect1);
    frCoord lly = gtl::yl(*viaRect1);
    frCoord urx = gtl::xh(*viaRect1);
    frCoord ury = gtl::xh(*viaRect1);

    marker->addVictim(
        net1->getOwner(),
        make_tuple(layerNum1, frBox(llx, lly, urx, ury), viaRect1->isFixed()));
    marker->addSrc(net2->getOwner());
    llx = gtl::xl(*viaRect2);
    lly = gtl::yl(*viaRect2);
    urx = gtl::xh(*viaRect2);
    ury = gtl::xh(*viaRect2);
    marker->addAggressor(
        net2->getOwner(),
        make_tuple(layerNum2, frBox(llx, lly, urx, ury), viaRect2->isFixed()));
    addMarker(std::move(marker));
  }
}

void FlexGCWorker::Impl::checkLef58CutSpacingTbl(
    gcRect* viaRect,
    frLef58CutSpacingTableConstraint* con)
{
  auto layerNum1 = viaRect->getLayerNum();
  auto layer1 = getTech()->getLayer(layerNum1);
  auto width = viaRect->width();
  auto length = viaRect->length();
  auto cutClassIdx = layer1->getCutClassIdx(width, length);
  frString cutClass = "";
  if (cutClassIdx != -1)
    cutClass = layer1->getCutClass(cutClassIdx)->getName();

  auto dbRule = con->getODBRule();
  frLayerNum queryLayerNum;
  if (dbRule->isLayerValid())
    queryLayerNum = getTech()
                        ->getLayer(dbRule->getSecondLayer()->getName())
                        ->getLayerNum();
  else
    queryLayerNum = layerNum1;
  frCoord maxSpc;
  if (width == length) {
    maxSpc = dbRule->getMaxSpacing(cutClass, false);
  } else {
    maxSpc = std::max(dbRule->getMaxSpacing(cutClass, true),
                      dbRule->getMaxSpacing(cutClass, false));
  }

  box_t queryBox;
  myBloat(*viaRect, maxSpc, queryBox);
  vector<rq_box_value_t<gcRect*>> results;
  auto& workerRegionQuery = getWorkerRegionQuery();
  workerRegionQuery.queryMaxRectangle(queryBox, queryLayerNum, results);
  for (auto& [box, ptr] : results) {
    if (ptr->isFixed() && viaRect->isFixed())
      continue;
    if (ptr->getPin() == viaRect->getPin())
      continue;
    checkLef58CutSpacingTbl_main(viaRect, ptr, con);
  }
}


#include "ahrs.hpp"

/* C/C++ Includes */
#include <stdint.h>
#include <stdlib.h>
#include <math.h>

/* Thor Includes */
#include "Thor/include/thor.h"
#include "Thor/include/spi.h"
#include "Thor/include/gpio.h"
#include "Thor/include/exceptions.h"

/* Boost Includes */
#include <boost/smart_ptr.hpp>
#include <boost/make_shared.hpp>

/* FreeRTOS Includes */
#include "FreeRTOS.h"
#include "task.h"
#include "queue.h"

/* Project Includes */
#include "ahrs.hpp"
#include "LSM9DS1.hpp"
#include "config.hpp"
#include "dataTypes.hpp"
#include "threading.hpp"

/* Madgwick Filter */
#include "madgwick.hpp"

/* Kalman Filter */
#include "kalman/SquareRootUnscentedKalmanFilter.hpp"
#include "IMUModel.hpp"

namespace SOAR_AHRS
{
	typedef float T;
	typedef IMU::State<T> State;
	typedef IMU::Control<T> Control;
	typedef IMU::Measurement<T> Measurement;
	typedef IMU::SystemModel<T> SystemModel;
	typedef IMU::MeasurementModel<T> MeasurementModel;

	const int updateRate_mS = (1.0 / SENSOR_UPDATE_FREQ_HZ) * 1000.0;
	const int magMaxUpdateRate_mS = (1.0 / LSM9DS1_M_MAX_BW) * 1000.0;

	const float accelUncertainty = 0.8f;
	const float gyroUncertainty = 1.05f;

	void ahrsTask(void* argument)
	{
		#ifdef DEBUG
		volatile float pitch;
		volatile float roll;
		volatile float yaw;
		volatile float ax;
		volatile float ay;
		volatile float az;
		volatile float gx;
		volatile float gy;
		volatile float gz;
		volatile float mx;
		volatile float my;
		volatile float mz;
		volatile UBaseType_t stackHighWaterMark_AHRS = 0;
		volatile size_t bytesRemaining = xPortGetFreeHeapSize();
		#endif

		AHRSData_t ahrsData;					/* Output struct to push to the motor controller thread */

		/*----------------------------------
		* Initialize the UKF
		*----------------------------------*/
		State x, x_ukf;
		Control u;
		SystemModel sys;
		MeasurementModel om;
		Measurement meas;

		Eigen::Matrix<T, 6, 6> R;
		Eigen::Matrix<T, 6, 6> processNoise;

		Kalman::SquareRootUnscentedKalmanFilter<State> ukf(0.5f, 3.0f, 0.0f);

		x.setZero();
		u.setZero();
		R.setZero();

		T cnst = 5.00e-4f;
		T dnst = 3.33e-4f;
		T enst = 5.00e-4f;

		processNoise <<
			cnst, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f,
			0.0f, dnst, 0.0f, 0.0f, 0.0f, 0.0f,
			0.0f, 0.0f, enst, 0.0f, 0.0f, 0.0f,
			0.0f, 0.0f, 0.0f, cnst, 0.0f, 0.0f,
			0.0f, 0.0f, 0.0f, 0.0f, dnst, 0.0f,
			0.0f, 0.0f, 0.0f, 0.0f, 0.0f, enst;

		sys.setCovariance(processNoise);

		R(0, 0) = accelUncertainty * accelUncertainty;
		R(1, 1) = accelUncertainty * accelUncertainty;
		R(2, 2) = accelUncertainty * accelUncertainty;
		R(3, 3) = gyroUncertainty * gyroUncertainty;
		R(4, 4) = gyroUncertainty * gyroUncertainty;
		R(5, 5) = gyroUncertainty * gyroUncertainty;

		om.setCovariance(R);

		ukf.init(x);

		/*----------------------------------
		* Initialize the IMU
		*----------------------------------*/
		GPIOClass_sPtr lsm_ss_xg = boost::make_shared<GPIOClass>(GPIOC, PIN_4, ULTRA_SPD, NOALTERNATE);
		GPIOClass_sPtr lsm_ss_m = boost::make_shared<GPIOClass>(GPIOC, PIN_3, ULTRA_SPD, NOALTERNATE);
		SPIClass_sPtr lsm_spi = spi2;

		LSM9DS1 imu(lsm_spi, lsm_ss_xg, lsm_ss_m);

		/* Force halt of the device if the IMU cannot be reached */
		if (imu.begin() == 0)
			BasicErrorHandler("The IMU WHO_AM_I registers did not return valid readings");

		//TODO: Switch this out to print over the serial port

		imu.calibrate(true); /* "true" forces an automatic software subtraction of the calculated bias from all further data */
		imu.calibrateMag(true); /* "true" writes the offest into the mag sensor hardware for automatic subtraction in results */

		int count = 0;

		/*----------------------------------
		* Initialize the Madgwick Filter
		*----------------------------------*/
		float beta = 10.0;
		Eigen::Vector3f accel_raw, gyro_raw, mag_raw;
		Eigen::Vector3f accel_filtered, gyro_filtered, mag_filtered, eulerDeg;

		MadgwickFilter ahrs((AHRS_UPDATE_RATE_MULTIPLIER * SENSOR_UPDATE_FREQ_HZ), beta);

		float dt = (1.0f / SENSOR_UPDATE_FREQ_HZ);
		float tau_accel = 0.01f;
		float alpha_lp_accel = dt / tau_accel;

		/* Tell init task that this thread's initialization is done and ok to run.
		* Wait for init task to resume operation. */
		xTaskSendMessage(INIT_TASK, 1u);
		vTaskSuspend(NULL);
		taskYIELD();

		TickType_t lastTimeWoken = xTaskGetTickCount();
		for (;;)
		{
			#ifdef DEBUG
			stackHighWaterMark_AHRS = uxTaskGetStackHighWaterMark(NULL);
			bytesRemaining = xPortGetFreeHeapSize();
			#endif

			/*----------------------------
			* Sensor Reading
			*---------------------------*/
			/* Update Accel & Gyro Data at whatever frequency set by user. Max bandwidth on
			* chip is 952Hz which will saturate FreeRTOS if sampled that often.*/
			//taskENTER_CRITICAL();
			imu.readAccel();
			imu.readGyro();

			/* Update Mag Data at a max frequency set by LSM9DS1_M_MAX_BW (~75Hz) */
			#if (SENSOR_UPDATE_FREQ_HZ > LSM9DS1_M_MAX_BW)
			if (count > magMaxUpdateRate_mS)
			{
				imu.readMag();
				imu.calcMag();
				count = 0;
			}
			else
				count += updateRate_mS;
			#else
			imu.readMag();
			imu.calcMag();
			#endif
			//taskEXIT_CRITICAL();

			/* Convert raw data from chip into meaningful data */
			imu.calcAccel(); imu.calcGyro();

// 			#ifdef DEBUG
// 			ax = imu.aRaw[0];
// 			ay = imu.aRaw[1];
// 			az = imu.aRaw[2];
//
// 			gx = imu.gRaw[0];
// 			gy = imu.gRaw[1];
// 			gz = imu.gRaw[2];
//
// 			mx = imu.mRaw[0];
// 			my = imu.mRaw[1];
// 			mz = imu.mRaw[2];
// 			#endif

			accel_raw << imu.aRaw[0], imu.aRaw[1], imu.aRaw[2];
			gyro_raw << imu.gRaw[0], imu.gRaw[1], imu.gRaw[2];
			mag_raw <<
				-imu.mRaw[0],  //Align mag x with accel x
				-imu.mRaw[1],  //Align mag y with accel y
				-imu.mRaw[2];  //from LSM9DS1, mag z is opposite direction of accel z

			/*----------------------------
			* UKF Algorithm
			*---------------------------*/
			//Simulate the system
			x = sys.f(x, u);

			//Predict state for current time step
			x_ukf = ukf.predict(sys);

			//Take a measurement given system state
			meas << accel_raw, gyro_raw;

			//Update the state equation given measurement
			x_ukf = ukf.update(om, meas);

			accel_filtered << x_ukf.ax(), x_ukf.ay(), x_ukf.az();
			gyro_filtered << x_ukf.gx(), x_ukf.gy(), x_ukf.gz();
			mag_filtered = mag_raw;

			/*----------------------------
			* AHRS Algorithm
			*---------------------------*/
			/* The Madgwick filter needs to run between 3-5 times as fast IMU measurements
			* to achieve decent convergence to a stable value. This thread only runs when
			* new data has arrived from the IMU, so frequency multiplication is as simple
			* as looping 3-5 times here. */
			for (int i = 0; i < AHRS_UPDATE_RATE_MULTIPLIER; i++)
				ahrs.update(accel_filtered, gyro_filtered, mag_filtered);

			ahrs.getEulerDeg(eulerDeg);
			ahrsData(eulerDeg, accel_filtered, gyro_filtered, mag_filtered);

			#ifdef DEBUG
			pitch = eulerDeg(0);
			roll = eulerDeg(1);
			yaw = eulerDeg(2);

			ax = accel_filtered(0);
			ay = accel_filtered(1);
			az = accel_filtered(2);

			gx = gyro_filtered(0);
			gy = gyro_filtered(1);
			gz = gyro_filtered(2);

			mx = mag_filtered(0);
			my = mag_filtered(1);
			mz = mag_filtered(2);
			#endif

			/* Send data over to the Serial thread*/
			if (xSemaphoreTake(ahrsBufferMutex, 0) == pdPASS)
			{
				xQueueOverwrite(qAHRS, &ahrsData);
				xSemaphoreGive(ahrsBufferMutex);
			}

			vTaskDelayUntil(&lastTimeWoken, pdMS_TO_TICKS(updateRate_mS));
		}
	}
}


//
// Created by dchansen on 10/12/18.
//
#include "ConvolutionMatrix.h"

#include <GadgetronTimer.h>
#include <numeric>
#include "vector_td_utilities.h"
#include <boost/range/irange.hpp>

namespace
{
    using namespace Gadgetron;

    template<int N>
    struct iteration_counter { };

    template<class REAL, unsigned int D, template<class, unsigned int> class K>
    void iterate_body(
        const vector_td<REAL, D> &point,
        const vector_td<size_t, D> &matrix_size,
        std::vector<size_t> &indices,
        std::vector<REAL> &weights,
        vector_td<REAL, D> &image_point,
        size_t index,
        const ConvolutionKernel<REAL, D, K>& kernel,
        iteration_counter<-1>)
    {
        auto delta = abs(image_point - point);
        indices.push_back(index);
        weights.push_back(kernel.get(delta));
    }

    template<class REAL, unsigned int D, template<class, unsigned int> class K, int N>
    void iterate_body(
        const vector_td<REAL, D> &point,
        const vector_td<size_t, D> &matrix_size,
        std::vector<size_t> &indices,
        std::vector<REAL> &weights,
        vector_td<REAL, D> &image_point,
        size_t index,
        const ConvolutionKernel<REAL, D, K>& kernel,
        iteration_counter<N>)
    {
        size_t frame_offset = std::accumulate(&matrix_size[0], &matrix_size[N], 1, std::multiplies<size_t>());

        for (int i = std::ceil(point[N] - kernel.get_radius());
             i <= std::floor(point[N] + kernel.get_radius());
             i++)
        {
            auto wrapped_i = (i + matrix_size[N]) % matrix_size[N];
            size_t index2 = index + frame_offset * wrapped_i;
            image_point[N] = i;
            iterate_body(point, matrix_size, indices, weights, image_point,
                         index2, kernel, iteration_counter<N - 1>());
        }
    }

    template<class REAL, unsigned int D, template<class, unsigned int> class K>
    std::tuple<std::vector<size_t>, std::vector<REAL>> get_indices(
        const vector_td<REAL, D> &point,
        const vector_td<size_t, D> &matrix_size,
        const ConvolutionKernel<REAL, D, K>& kernel)
    {
        std::vector<size_t> indices;
        indices.reserve(size_t(std::pow(std::ceil(kernel.get_width()), D)));

        std::vector<REAL> weights;
        weights.reserve(size_t(std::pow(std::ceil(kernel.get_width()), D)));

        vector_td<REAL, D> image_point;
        size_t index = 0;
        iterate_body(point, matrix_size, indices, weights, image_point, index,
                     kernel, iteration_counter<D - 1>());

        return std::make_tuple(std::move(indices), std::move(weights));
    }
}

template<class REAL, unsigned int D, template<class, unsigned int> class K>
Gadgetron::ConvInternal::ConvolutionMatrix<REAL>
Gadgetron::ConvInternal::make_conv_matrix(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<REAL, D>> trajectory,
    const Gadgetron::vector_td<size_t, D> &matrix_size,
    const ConvolutionKernel<REAL, D, K>& kernel)
{
    ConvolutionMatrix<REAL> matrix(trajectory.get_number_of_elements(),
                                   prod(matrix_size));

    #pragma omp parallel for
    for (int i = 0; i < (int)trajectory.get_number_of_elements(); i++)
    {
        std::tie(matrix.indices[i], matrix.weights[i]) = get_indices(
            trajectory[i], matrix_size, kernel);
    }

    return matrix;
}

template<class REAL>
Gadgetron::ConvInternal::ConvolutionMatrix<REAL>
Gadgetron::ConvInternal::transpose(const Gadgetron::ConvInternal::ConvolutionMatrix<REAL> &matrix) {

    auto work_index = boost::irange(size_t(0),matrix.n_cols);

    ConvolutionMatrix<REAL> transposed(matrix.n_rows, matrix.n_cols);

    std::vector<int> counts(matrix.n_rows, 0);

    for (size_t i : work_index)
    {
        auto &rows = matrix.indices[i];
        for (auto &row : rows)
            counts[row]++;
    }

    for (size_t i = 0; i < counts.size(); i++)
    {
        transposed.indices[i].reserve(counts[i]);
        transposed.weights[i].reserve(counts[i]);
    }

    for (size_t i : work_index )
    {
        auto &rows = matrix.indices[i];
        auto &weights = matrix.weights[i];
        for (size_t n = 0; n < rows.size(); n++)
        {
            transposed.indices[rows[n]].push_back(i);
            transposed.weights[rows[n]].push_back(weights[n]);
        }
    }

    return transposed;
}

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 1, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 1>> trajectory,
    const Gadgetron::vector_td<size_t, 1> &matrix_size,
    const ConvolutionKernel<float, 1, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 2, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 2>> trajectory,
    const Gadgetron::vector_td<size_t, 2> &matrix_size,
    const ConvolutionKernel<float, 2, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 3, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 3>> trajectory,
    const Gadgetron::vector_td<size_t, 3> &matrix_size,
    const ConvolutionKernel<float, 3, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 4, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 4>> trajectory,
    const Gadgetron::vector_td<size_t, 4> &matrix_size,
    const ConvolutionKernel<float, 4, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 1, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 1>> trajectory,
    const Gadgetron::vector_td<size_t, 1> &matrix_size,
    const ConvolutionKernel<double, 1, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 2, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 2>> trajectory,
    const Gadgetron::vector_td<size_t, 2> &matrix_size,
    const ConvolutionKernel<double, 2, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 3, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 3>> trajectory,
    const Gadgetron::vector_td<size_t, 3> &matrix_size,
    const ConvolutionKernel<double, 3, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 4, Gadgetron::KaiserKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 4>> trajectory,
    const Gadgetron::vector_td<size_t, 4> &matrix_size,
    const ConvolutionKernel<double, 4, Gadgetron::KaiserKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 1, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 1>> trajectory,
    const Gadgetron::vector_td<size_t, 1> &matrix_size,
    const ConvolutionKernel<float, 1, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 2, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 2>> trajectory,
    const Gadgetron::vector_td<size_t, 2> &matrix_size,
    const ConvolutionKernel<float, 2, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 3, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 3>> trajectory,
    const Gadgetron::vector_td<size_t, 3> &matrix_size,
    const ConvolutionKernel<float, 3, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::make_conv_matrix<float, 4, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<float, 4>> trajectory,
    const Gadgetron::vector_td<size_t, 4> &matrix_size,
    const ConvolutionKernel<float, 4, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 1, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 1>> trajectory,
    const Gadgetron::vector_td<size_t, 1> &matrix_size,
    const ConvolutionKernel<double, 1, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 2, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 2>> trajectory,
    const Gadgetron::vector_td<size_t, 2> &matrix_size,
    const ConvolutionKernel<double, 2, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 3, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 3>> trajectory,
    const Gadgetron::vector_td<size_t, 3> &matrix_size,
    const ConvolutionKernel<double, 3, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::make_conv_matrix<double, 4, Gadgetron::JincKernel>(
    const Gadgetron::hoNDArray<Gadgetron::vector_td<double, 4>> trajectory,
    const Gadgetron::vector_td<size_t, 4> &matrix_size,
    const ConvolutionKernel<double, 4, Gadgetron::JincKernel>& kernel);

template Gadgetron::ConvInternal::ConvolutionMatrix<float>
Gadgetron::ConvInternal::transpose(
    const Gadgetron::ConvInternal::ConvolutionMatrix<float> &matrix);

template Gadgetron::ConvInternal::ConvolutionMatrix<double>
Gadgetron::ConvInternal::transpose(
        const Gadgetron::ConvInternal::ConvolutionMatrix<double> &matrix);


/**
 * @file tests/q_learning_test.cpp
 * @author Shangtong Zhang
 * @author Rohan Raj
 *
 * Test for Q-Learning implementation
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */

#include <mlpack/core.hpp>

#include <mlpack/methods/ann/ffn.hpp>
#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
#include <mlpack/methods/ann/layer/layer.hpp>
#include <mlpack/methods/ann/loss_functions/mean_squared_error.hpp>
#include <mlpack/methods/ann/loss_functions/empty_loss.hpp>
#include <mlpack/methods/reinforcement_learning/q_learning.hpp>
#include <mlpack/methods/reinforcement_learning/sac.hpp>
#include <mlpack/methods/reinforcement_learning/q_networks/simple_dqn.hpp>
#include <mlpack/methods/reinforcement_learning/q_networks/dueling_dqn.hpp>
#include <mlpack/methods/reinforcement_learning/environment/pendulum.hpp>
#include <mlpack/methods/reinforcement_learning/q_networks/categorical_dqn.hpp>
#include <mlpack/methods/reinforcement_learning/environment/mountain_car.hpp>
#include <mlpack/methods/reinforcement_learning/environment/acrobot.hpp>
#include <mlpack/methods/reinforcement_learning/environment/cart_pole.hpp>
#include <mlpack/methods/reinforcement_learning/environment/double_pole_cart.hpp>
#include <mlpack/methods/reinforcement_learning/policy/greedy_policy.hpp>
#include <mlpack/methods/reinforcement_learning/training_config.hpp>

#include <ensmallen.hpp>
#include <numeric>

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

using namespace mlpack;
using namespace mlpack::ann;
using namespace ens;
using namespace mlpack::rl;

BOOST_AUTO_TEST_SUITE(QLearningTest);

template<typename AgentType>
bool testAgent(AgentType& agent,
               const double rewardThreshold,
               const size_t noOfEpisodes,
               const size_t consecutiveEpisodesTest = 50)
{
  bool converged = false;
  std::vector<double> returnList;
  size_t episodes = 0;
  while (true)
  {
    double episodeReturn = agent.Episode();
    episodes += 1;
    returnList.push_back(episodeReturn);

    if (returnList.size() > consecutiveEpisodesTest)
      returnList.erase(returnList.begin());

    double averageReturn = std::accumulate(returnList.begin(),
        returnList.end(), 0.0) / returnList.size();

    Log::Debug << "Average return in last " << returnList.size()
        << " consecutive episodes: " << averageReturn
        << " Episode return: " << episodeReturn << std::endl;

    // For the speed of the test case, a high criterion should not be set
    // for the rewardThreshold.
    if (averageReturn > rewardThreshold &&
        returnList.size() >= consecutiveEpisodesTest)
    {
      converged = true;
      agent.Deterministic() = true;
      arma::running_stat<double> testReturn;
      for (size_t i = 0; i < 10; ++i)
        testReturn(agent.Episode());

      Log::Debug << "Average return in deterministic test: "
          << testReturn.mean() << std::endl;
      break;
    }

    if (episodes > noOfEpisodes)
    {
      Log::Debug << "Agent failed." << std::endl;
      break;
    }
  }

  return converged;
}

//! Test DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithDQN)
{
  // Set up the network.
  SimpleDQN<> network(4, 128, 128, 2);

  // Set up the policy and replay method.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
  RandomReplay<CartPole> replayMethod(10, 10000);

  // Setting all training hyperparameters.
  TrainingConfig config;
  config.StepSize() = 0.01;
  config.Discount() = 0.9;
  config.TargetNetworkSyncInterval() = 100;
  config.ExplorationSteps() = 100;
  config.DoubleQLearning() = false;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 40, 1000);

  // To check if the action returned by the agent is not nan and is finite.
  BOOST_REQUIRE(std::isfinite(double(agent.Action().action)));
  BOOST_REQUIRE(converged);
}

//! Test DQN in Cart Pole task with Prioritized Replay.
BOOST_AUTO_TEST_CASE(CartPoleWithDQNPrioritizedReplay)
{
  // Set up the network.
  SimpleDQN<> network(4, 128, 128, 2);

  // Set up the policy and replay method.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1);
  PrioritizedReplay<CartPole> replayMethod(10, 10000, 0.6);

  TrainingConfig config;
  config.ExplorationSteps() = 100;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy),
      decltype(replayMethod)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 45, 1000);
  BOOST_REQUIRE(converged);
}

//! Test Double DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithDoubleDQN)
{
  // It isn't guaranteed that the network will converge in the specified number
  // of iterations using random weights. If this works 1 of 4 times, I'm fine
  // with that.
  bool converged = false;
  for (size_t trial = 0; trial < 4; ++trial)
  {
    // Set up the network.
    SimpleDQN<> network(4, 20, 20, 2);

    // Set up the policy and replay method.
    GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<CartPole> replayMethod(10, 10000);

    TrainingConfig config;
    config.ExplorationSteps() = 100;
    config.DoubleQLearning() = true;
    config.StepLimit() = 200;

    // Set up the DQN agent.
    QLearning<CartPole, decltype(network), RMSPropUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, 45, 1000);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test DQN in Acrobot task.
BOOST_AUTO_TEST_CASE(AcrobotWithDQN)
{
  // We will allow three trials, although it would be very uncommon for the test
  // to use more than one.
  bool converged = false;
  for (size_t trial = 0; trial < 3; ++trial)
  {
    // Set up the network.
    SimpleDQN<> network(4, 64, 32, 3);

    // Set up the policy and replay method.
    GreedyPolicy<Acrobot> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<Acrobot> replayMethod(20, 10000);

    TrainingConfig config;
    config.ExplorationSteps() = 100;
    config.StepLimit() = 400;

    // Set up DQN agent.
    QLearning<Acrobot, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, -380, 1000);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test DQN in Mountain Car task.
BOOST_AUTO_TEST_CASE(MountainCarWithDQN)
{
  // We will allow five trials total.
  bool converged = false;
  for (size_t trial = 0; trial < 5; trial++)
  {
    // Set up the network.
    SimpleDQN<> network(2, 64, 32, 3);

    // Set up the policy and replay method.
    GreedyPolicy<MountainCar> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<MountainCar> replayMethod(20, 10000);

    TrainingConfig config;
    config.StepSize() = 0.0001;
    config.ExplorationSteps() = 100;
    config.StepLimit() = 400;

    // Set up DQN agent.
    QLearning<MountainCar, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, -380, 1000);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test DQN in DoublePoleCart task.
BOOST_AUTO_TEST_CASE(DoublePoleCartWithDQN)
{
  bool converged = false;
  // We will allow four trials total.
  for (size_t trial = 0; trial < 4; trial++)
  {
    // Set up the module. Note that we use a custom network here.
    FFN<MeanSquaredError<>, GaussianInitialization> module(
        MeanSquaredError<>(), GaussianInitialization(0, 0.001));
    module.Add<Linear<>>(6, 256);
    module.Add<ReLULayer<>>();
    module.Add<Linear<>>(256, 3);

    // Adding the module to the SimpleDQN network containing required functions.
    SimpleDQN<> network(module);

    // Set up the policy and replay method.
    GreedyPolicy<DoublePoleCart> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<DoublePoleCart> replayMethod(20, 10000);

    TrainingConfig config;
    config.ExplorationSteps() = 100;
    config.StepLimit() = 600;

    // Set up DQN agent.
    QLearning<DoublePoleCart, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    size_t episodes = 0;
    size_t episodeSuccesses = 0;
    while (true)
    {
      double episodeReturn = agent.Episode();
      episodes += 1;

      if (episodeReturn >= 280)
        episodeSuccesses++;

      if (episodes > 2000)
      {
        Log::Debug << "Agent failed." << std::endl;
        break;
      }

      // If the network can solve the environment in two trials this is fine for
      // a simple test.
      Log::Debug << " Episode return: " << episodeReturn << std::endl;
      if (episodeSuccesses >= 2)
      {
        converged = true;
        Log::Debug << "QLearning has succeeded in the multiple pole cart" <<
            " environment." << std::endl;
        break;
      }
    }
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test Dueling DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDQN)
{
  // Set up the network.
  DuelingDQN<> network(4, 128, 64, 2);

  // Set up the policy and replay method.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
  RandomReplay<CartPole> replayMethod(32, 2000);

  TrainingConfig config;
  config.ExplorationSteps() = 50;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 45, 2000);
  BOOST_REQUIRE(converged);
}

//! Test Dueling DQN in Cart Pole task with Prioritized Replay.
BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDQNPrioritizedReplay)
{
  // Set up the network.
  DuelingDQN<> network(4, 128, 64, 2);

  // Set up the policy and replay method.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1);
  PrioritizedReplay<CartPole> replayMethod(32, 2000, 0.6);

  TrainingConfig config;
  config.ExplorationSteps() = 50;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy),
      decltype(replayMethod)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 50, 2000);
  BOOST_REQUIRE(converged);
}

//! Test Noisy DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithNoisyDQN)
{
  // It isn't guaranteed that the network will converge in the specified number
  // of iterations using random weights.
  bool converged = false;
  for (size_t trial = 0; trial < 3; ++trial)
  {
    Log::Debug << "Trial number: " << trial << std::endl;

    // Set up the policy and replay method.
    GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<CartPole> replayMethod(32, 2000);

    TrainingConfig config;
    config.StepLimit() = 200;
    config.NoisyQLearning() = true;

    // Set up the network with a flag to enable noisy layers.
    SimpleDQN<> network(4, 64, 32, 2, config.NoisyQLearning());

    // Set up DQN agent.
    QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, 45, 500, 30);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test Dueling-Double-Noisy DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithDuelingDoubleNoisyDQN)
{
  // It isn't guaranteed that the network will converge in the specified number
  // of iterations using random weights.
  bool converged = false;
  for (size_t trial = 0; trial < 5; ++trial)
  {
    Log::Debug << "Trial number: " << trial << std::endl;

    // Set up the policy and replay method.
    GreedyPolicy<CartPole> policy(1.0, 2000, 0.1, 0.99);
    RandomReplay<CartPole> replayMethod(32, 4000);

    TrainingConfig config;
    config.DoubleQLearning() = true;
    config.StepLimit() = 200;
    config.NoisyQLearning() = true;

    // Set up the network with a flag to enable noisy layers.
    DuelingDQN<> network(4, 64, 64, 2, config.NoisyQLearning());

    // Set up DQN agent.
    QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, 45, 500, 30);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test N-step DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithNStepDQN)
{
  // Set up the network.
  SimpleDQN<> network(4, 128, 128, 2);

  // Set up the policy.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
  /**
   * For N-step learning, we need to specify n as the last parameter in
   * the replay method. Here we use n = 3.
   */
  RandomReplay<CartPole> replayMethod(10, 10000, 3);

  // Setting all training hyperparameters.
  TrainingConfig config;
  config.ExplorationSteps() = 50;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 50, 1000);
  BOOST_REQUIRE(converged);
}

//! Test N-step Prioritized DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithNStepPrioritizedDQN)
{
  // Set up the network.
  SimpleDQN<> network(4, 128, 128, 2);

  // Set up the policy.
  GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
  /**
   * For N-step learning, we need to specify n as the last parameter in
   * the replay method. Here we use n = 3.
   */
  PrioritizedReplay<CartPole> replayMethod(10, 10000, 0.6, 3);

  // Setting all training hyperparameters.
  TrainingConfig config;
  config.ExplorationSteps() = 50;
  config.StepLimit() = 200;

  // Set up DQN agent.
  QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy),
      decltype(replayMethod)>
      agent(config, network, policy, replayMethod);

  bool converged = testAgent<decltype(agent)>(agent, 50, 1000);
  BOOST_REQUIRE(converged);
}

//! Test Categorical DQN in Cart Pole task.
BOOST_AUTO_TEST_CASE(CartPoleWithCategoricalDQN)
{
  // It isn't guaranteed that the network will converge in the specified number
  // of iterations.
  bool converged = false;
  for (size_t trial = 0; trial < 3; ++trial)
  {
    Log::Debug << "Trial number: " << trial << std::endl;

    // Set up the policy and replay method.
    GreedyPolicy<CartPole> policy(1.0, 1000, 0.1, 0.99);
    RandomReplay<CartPole> replayMethod(32, 4000);

    TrainingConfig config;
    config.IsCategorical() = true;
    config.ExplorationSteps() = 32;

    // Set up the module. Note that we use a custom network here.
    FFN<EmptyLoss<>, GaussianInitialization> module(
        EmptyLoss<>(), GaussianInitialization(0, 0.1));
    module.Add<Linear<>>(4, 128);
    module.Add<ReLULayer<>>();
    module.Add<Linear<>>(128, 2 * config.AtomSize());

    // Adding the module to the CategoricalDQN network.
    CategoricalDQN<> network(module, config);

    // Set up DQN agent.
    QLearning<CartPole, decltype(network), AdamUpdate, decltype(policy)>
        agent(config, network, policy, replayMethod);

    converged = testAgent<decltype(agent)>(agent, 40, 1000, 20);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

//! Test SAC on Pendulum task.
BOOST_AUTO_TEST_CASE(PendulumWithSAC)
{
  // It isn't guaranteed that the network will converge in the specified number
  // of iterations using random weights.
  bool converged = false;
  for (size_t trial = 0; trial < 3; ++trial)
  {
    Log::Debug << "Trial number: " << trial << std::endl;
    // Set up the policy and replay method.
    RandomReplay<Pendulum> replayMethod(32, 10000);

    TrainingConfig config;
    config.StepSize() = 0.001;
    config.TargetNetworkSyncInterval() = 1;
    config.UpdateInterval() = 3;

    FFN<EmptyLoss<>, GaussianInitialization>
        policyNetwork(EmptyLoss<>(), GaussianInitialization(0, 0.1));
    policyNetwork.Add(new Linear<>(3, 128));
    policyNetwork.Add(new ReLULayer<>());
    policyNetwork.Add(new Linear<>(128, 1));
    policyNetwork.Add(new TanHLayer<>());

    FFN<EmptyLoss<>, GaussianInitialization>
        qNetwork(EmptyLoss<>(), GaussianInitialization(0, 0.1));
    qNetwork.Add(new Linear<>(3+1, 128));
    qNetwork.Add(new ReLULayer<>());
    qNetwork.Add(new Linear<>(128, 1));

    // Set up Soft actor-critic agent.
    SAC<Pendulum, decltype(qNetwork), decltype(policyNetwork), AdamUpdate>
        agent(config, qNetwork, policyNetwork, replayMethod);

    converged = testAgent<decltype(agent)>(agent, -900, 500, 10);
    if (converged)
      break;
  }
  BOOST_REQUIRE(converged);
}

BOOST_AUTO_TEST_SUITE_END();


#include "hoCuNDArray_utils.h"
#include "radial_utilities.h"
#include "cuNDArray_fileio.h"
#include "cuNDArray_math.h"
#include "imageOperator.h"
#include "identityOperator.h"
#include "hoPartialDerivativeOperator.h"
#include "hoCuConebeamProjectionOperator.h"
#include "cuConvolutionOperator.h"
#include "hoCuNDArray_blas.h"
#include "hoCuNDArray_elemwise.h"
#include "hoCuNDArray_blas.h"
#include "cgSolver.h"
#include "CBCT_acquisition.h"
#include "complext.h"
#include "encodingOperatorContainer.h"
#include "vector_td_io.h"
#include "hoPartialDerivativeOperator.h"
#include "hoCuGPBBSolver.h"
#include "hoCuTvOperator.h"
#include "hoCuTvPicsOperator.h"
#include "hoCuNCGSolver.h"
#include "hoCuCgDescentSolver.h"
#include "hoNDArray_utils.h"
#include "hoCuPartialDerivativeOperator.h"
#include "cuTvOperator.h"
#include "cuTv1dOperator.h"
#include "cuTvPicsOperator.h"
#include "CBSubsetOperator.h"
#include "osSPSSolver.h"
#include "osMOMSolver.h"
#include "osMOMSolverD.h"
#include "osMOMSolverD2.h"
#include "hoCuNCGSolver.h"
#include "osMOMSolverD3.h"
#include "osMOMSolverL1.h"
#include "osMOMSolverF.h"
#include "osAHZCSolver.h"
#include "ADMMSolver.h"
#include <iostream>
#include <algorithm>
#include <sstream>
#include <math_constants.h>
#include <boost/program_options.hpp>
#include <boost/make_shared.hpp>
#include <GPUTimer.h>
#include "cuSolverUtils.h"
#include "osPDsolver.h"
#include "osLALMSolver.h"
#include "osLALMSolver2.h"
#include "cuATrousOperator.h"
#include "hdf5_utils.h"
#include "cuEdgeATrousOperator.h"
#include "cuDCTOperator.h"
#include "cuDCTDerivativeOperator.h"
#include "dicomWriter.h"
#include "conebeam_projection.h"
#include "weightingOperator.h"
#include "hoNDArray_math.h"
#include "cuNCGSolver.h"
#include "CT_acquisition.h"
#include "hoCuOSOMSolver.h"

using namespace std;
using namespace Gadgetron;

namespace po = boost::program_options;

boost::shared_ptr<cuNDArray<float> > calculate_prior(boost::shared_ptr<CBCT_binning>  binning,boost::shared_ptr<CBCT_acquisition> ps, hoCuNDArray<float>& projections, std::vector<size_t> is_dims, floatd3 imageDimensions){
	std::cout << "Calculating FDK prior" << std::endl;
	boost::shared_ptr<CBCT_binning> binning_pics=binning->get_3d_binning();
	std::vector<size_t> is_dims3d = is_dims;
	is_dims3d.pop_back();
	boost::shared_ptr< hoCuConebeamProjectionOperator >
	Ep( new hoCuConebeamProjectionOperator() );
	Ep->setup(ps,binning_pics,imageDimensions);
	Ep->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());
	Ep->set_domain_dimensions(&is_dims3d);
	Ep->set_use_filtered_backprojection(true);
	boost::shared_ptr<hoCuNDArray<float> > prior3d(new hoCuNDArray<float>(&is_dims3d));
	Ep->mult_MH(&projections,prior3d.get());

	hoCuNDArray<float> tmp_proj(*ps->get_projections());
	Ep->mult_M(prior3d.get(),&tmp_proj);
	float s = dot(ps->get_projections().get(),&tmp_proj)/dot(&tmp_proj,&tmp_proj);
	*prior3d *= s;
	boost::shared_ptr<cuNDArray<float> > prior(new cuNDArray<float>(*expand( prior3d.get(), is_dims.back() )));
	std::cout << "Prior complete" << std::endl;
	return prior;
}

boost::shared_ptr<cuNDArray<float>> calculate_weightImage(boost::shared_ptr<CBCT_binning>  binning,boost::shared_ptr<CBCT_acquisition> ps, hoCuNDArray<float>& ho_projections, std::vector<size_t> is_dims, floatd3 imageDimensions){

	cuNDArray<float> projections(ho_projections);
	boost::shared_ptr<CBCT_binning> binning_pics=binning->get_3d_binning();
	std::vector<size_t> is_dims3d = is_dims;
	is_dims3d.pop_back();
	auto Ep = boost::make_shared<cuConebeamProjectionOperator>();
	auto ps2 = boost::make_shared<CBCT_acquisition>(boost::shared_ptr<hoCuNDArray<float>>(),ps->get_geometry());
	Ep->setup(ps2,binning_pics,imageDimensions);
	Ep->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());
	Ep->set_domain_dimensions(&is_dims3d);
	//Ep->set_use_filtered_backprojection(true);
	Ep->offset_correct(&projections);
	//Ep->mult_MH(&projections,prior3d.get());
	//cgSolver<hoCuNDArray<float>> solv;
	cuNCGSolver<float> solv;
	solv.set_non_negativity_constraint(true);
	solv.set_encoding_operator(Ep);
	solv.set_max_iterations(10);
	auto prior3d = solv.solve(&projections);
	//auto prior3d = boost::make_shared<cuNDArray<float>>(is_dims3d);
	write_nd_array(prior3d.get(),"fdk.real");
	cuNDArray<float> tmp_proj(projections);
	clear(&tmp_proj);
	Ep->mult_M(prior3d.get(),&tmp_proj);
	//float s = dot(ps->get_projections().get(),&tmp_proj)/dot(&tmp_proj,&tmp_proj);
	//std::cout << "Scaling " << s << std::endl;
/*
	 //Ep->offset_correct(&tmp_proj);
	//tmp_proj *= s;
	write_nd_array(&tmp_proj,"projtmp.real");
	tmp_proj -= projections;
	tmp_proj *= float(-1);
	abs_inplace(&tmp_proj);

	write_nd_array(ps->get_projections().get(),"proj.real");
	write_nd_array(&tmp_proj,"projdiff.real");
	std::cout << "Proj size ";
	auto pdims = *tmp_proj.get_dimensions();
	for (auto p : pdims ) std::cout << p << " ";
	std::cout << std::endl;
	//Ep->set_use_filtered_backprojection(false);
	//Ep->mult_MH(&tmp_proj,prior3d.get());
	//solv.set_non_negativity_constraint(false);
	prior3d = solv.solve(&tmp_proj);
	//abs_inplace(prior.get());
	std::cout << "Prior complete" << std::endl;
	*/
	return prior3d;
}

int main(int argc, char** argv)
{

	string acquisition_filename;
	string outputFile;
	uintd3 imageSize;
	floatd3 voxelSize;
	int device;
	unsigned int downsamples;
	unsigned int iterations;
	unsigned int subsets;
	float rho,tau;
	float tv_weight,pics_weight, wavelet_weight,huber,sigma,dct_weight;
	float tv_4d;
    bool use_non_negativity;
	int reg_iter;

	po::options_description desc("Allowed options");

	desc.add_options()
    				("help", "produce help message")
    				("acquisition,a", po::value<string>(&acquisition_filename)->default_value("acquisition.hdf5"), "Acquisition data")
    				("samples,n",po::value<unsigned int>(),"Number of samples per ray")
    				("output,f", po::value<string>(&outputFile)->default_value("reconstruction.hdf5"), "Output filename")
    				("size,s",po::value<uintd3>(&imageSize)->default_value(uintd3(512,512,1)),"Image size in pixels")
    				("binning,b",po::value<string>(),"Binning file for 4d reconstruction")
    				("SAG","Use exact SAG correction if present")
    				("voxelSize,v",po::value<floatd3>(&voxelSize)->default_value(floatd3(0.488f,0.488f,1.0f)),"Voxel size in mm")
    				("dimensions,d",po::value<floatd3>(),"Image dimensions in mm. Overwrites voxelSize.")
    				("iterations,i",po::value<unsigned int>(&iterations)->default_value(10),"Number of iterations")
    				("device",po::value<int>(&device)->default_value(0),"Number of the device to use (0 indexed)")
    				("downsample,D",po::value<unsigned int>(&downsamples)->default_value(0),"Downsample projections this factor")
    				("subsets,u",po::value<unsigned int>(&subsets)->default_value(10),"Number of subsets to use")
    				("TV",po::value<float>(&tv_weight)->default_value(0),"Total variation weight in spatial dimensions")
							("TV4D",po::value<float>(&tv_4d)->default_value(0),"Total variation weight in temporal dimensions")
    				("PICS",po::value<float>(&pics_weight)->default_value(0),"PICS weight")
    				("Wavelet,W",po::value<float>(&wavelet_weight)->default_value(0),"Weight of the wavelet operator")
    				("Huber",po::value<float>(&huber)->default_value(0),"Huber weight")
    				("use_prior","Use an FDK prior")
    				("use_non_negativity",po::value<bool>(&use_non_negativity)->default_value(true),"Prevent image from having negative attenuation")
    				("sigma",po::value<float>(&sigma)->default_value(0.1),"Sigma for billateral filter")
    				("DCT",po::value<float>(&dct_weight)->default_value(0),"DCT regularization")
    				("3D","Only use binning for selecting valid projections")
							("tau",po::value<float>(&tau)->default_value(1e-5),"Tau value for solver")
							("reg_iter",po::value<int>(&reg_iter)->default_value(2))
    				;

	po::variables_map vm;
	po::store(po::parse_command_line(argc, argv, desc), vm);
	po::notify(vm);

	if (vm.count("help")) {
		cout << desc << "\n";
		return 1;
	}

	std::stringstream command_line_string;
	std::cout << "Command line options:" << std::endl;
	for (po::variables_map::iterator it = vm.begin(); it != vm.end(); ++it){
		boost::any a = it->second.value();
		command_line_string << it->first << ": ";
		if (a.type() == typeid(std::string)) command_line_string << it->second.as<std::string>();
		else if (a.type() == typeid(int)) command_line_string << it->second.as<int>();
		else if (a.type() == typeid(unsigned int)) command_line_string << it->second.as<unsigned int>();
		else if (a.type() == typeid(float)) command_line_string << it->second.as<float>();
		else if (a.type() == typeid(vector_td<float,3>)) command_line_string << it->second.as<vector_td<float,3> >();
		else if (a.type() == typeid(vector_td<int,3>)) command_line_string << it->second.as<vector_td<int,3> >();
		else if (a.type() == typeid(vector_td<unsigned int,3>)) command_line_string << it->second.as<vector_td<unsigned int,3> >();
        else if (a.type() == typeid(bool)) command_line_string << it->second.as<bool>();
		else command_line_string << "Unknown type" << std::endl;
		command_line_string << std::endl;
	}
	std::cout << command_line_string.str();

	cudaSetDevice(device);
	cudaDeviceReset();

	//Really weird stuff. Needed to initialize the device?? Should find real bug.
	cudaDeviceManager::Instance()->lockHandle();
	cudaDeviceManager::Instance()->unlockHandle();

	boost::shared_ptr<CBCT_acquisition> ps(new CBCT_acquisition());
	ps->load(acquisition_filename);
	ps->get_geometry()->print(std::cout);
	ps->downsample(downsamples);

	float SDD = ps->get_geometry()->get_SDD();
	float SAD = ps->get_geometry()->get_SAD();

	boost::shared_ptr<CBCT_binning> binning(new CBCT_binning());
	if (vm.count("binning")){
		std::cout << "Loading binning data" << std::endl;
		binning->load(vm["binning"].as<string>());
		if (vm.count("3D"))
			binning = binning->get_3d_binning();
	} else binning->set_as_default_3d_bin(ps->get_projections()->get_size(2));
	binning->print(std::cout);

	floatd3 imageDimensions;
	if (vm.count("dimensions")){
		imageDimensions = vm["dimensions"].as<floatd3>();
		voxelSize = imageDimensions/imageSize;
	}
	else imageDimensions = voxelSize*imageSize;

	float lengthOfRay_in_mm = norm(imageDimensions);
	unsigned int numSamplesPerPixel = 3;
	float minSpacing = min(voxelSize)/numSamplesPerPixel;

	unsigned int numSamplesPerRay;
	if (vm.count("samples")) numSamplesPerRay = vm["samples"].as<unsigned int>();
	else numSamplesPerRay = ceil( lengthOfRay_in_mm / minSpacing );

	float step_size_in_mm = lengthOfRay_in_mm / numSamplesPerRay;
	size_t numProjs = ps->get_projections()->get_size(2);
	size_t needed_bytes = 2 * prod(imageSize) * sizeof(float);
	std::vector<size_t> is_dims = to_std_vector((uint64d3)imageSize);

	std::cout << "IS dimensions " << is_dims[0] << " " << is_dims[1] << " " << is_dims[2] << std::endl;
	std::cout << "Image size " << imageDimensions << std::endl;

	is_dims.push_back(binning->get_number_of_bins());

	//osLALMSolver<cuNDArray<float>> solver;
	hoCuOSMOMSolver<float> solver;
	//osMOMSolverL1<cuNDArray<float>> solver;
	//osAHZCSolver<cuNDArray<float>> solver;
	//osMOMSolverF<cuNDArray<float>> solver;
	//ADMMSolver<cuNDArray<float>> solver;
	solver.set_dump(false);

	//hoCuCgDescentSolver<float> solver;

	//osSPSSolver<hoNDArray<float>> solver;
	//hoCuNCGSolver<float> solver;
	//solver.set_domain_dimensions(&is_dims);
	solver.set_max_iterations(iterations);
	solver.set_output_mode(osSPSSolver<cuNDArray<float>>::OUTPUT_VERBOSE);
	solver.set_tau(tau);
	solver.set_non_negativity_constraint(use_non_negativity);
	solver.set_huber(huber);
	solver.set_reg_steps(reg_iter);
	//solver.set_rho(rho);

  if (tv_weight > 0){

  	auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,4>>(0);
  	Dx->set_weight(tv_weight);
  	Dx->set_domain_dimensions(&is_dims);
  	Dx->set_codomain_dimensions(&is_dims);

  	auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,4>>(1);
  	Dy->set_weight(tv_weight);
  	Dy->set_domain_dimensions(&is_dims);
  	Dy->set_codomain_dimensions(&is_dims);

  	auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,4>>(2);
  	Dz->set_weight(tv_weight);
  	Dz->set_domain_dimensions(&is_dims);
  	Dz->set_codomain_dimensions(&is_dims);

  	solver.add_regularization_group({Dx,Dy,Dz});
    if (tv_4d > 0) {
        auto Dt = boost::make_shared<cuPartialDerivativeOperator<float, 4>>(3);
        Dt->set_weight(tv_4d);
        Dt->set_domain_dimensions(&is_dims);
        Dt->set_codomain_dimensions(&is_dims);
        solver.add_regularization_operator(Dt);
    }
/*
	auto projections = *ps->get_projections();
  	auto prior_weight = calculate_weightImage(binning,ps,projections,is_dims,imageDimensions);
  	//sqrt_inplace(prior_weight.get());
  	std::cout << "Prior min " << min(prior_weight.get()) << std::endl;
  	//*prior_weight -= min(prior_weight.get());
  	*prior_weight /= asum(prior_weight.get())/prior_weight->get_number_of_elements();
  	*prior_weight -= max(prior_weight.get());
  	*prior_weight *= float(-1);
  	//clamp_min(prior_weight.get(),float(1e-2));
  	//reciprocal_inplace(prior_weight.get());

  	write_nd_array(prior_weight.get(),"prior.real");
  	//cudaDeviceReset();
  	auto Wt = boost::make_shared<weightingOperator<cuNDArray<float>>>(prior_weight,Dt);
	Wt->set_weight(tv_weight);
  	Wt->set_domain_dimensions(&is_dims);
  	Wt->set_codomain_dimensions(&is_dims);
*/
  	//lver.add_regularization_operator(Dt);

  }
  /*
  if (pics_weight > 0){

  	auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,4>>(0);
  	Dx->set_weight(pics_weight);
  	Dx->set_domain_dimensions(&is_dims);
  	Dx->set_codomain_dimensions(&is_dims);

  	auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,4>>(1);
  	Dy->set_weight(pics_weight);
  	Dy->set_domain_dimensions(&is_dims);
  	Dy->set_codomain_dimensions(&is_dims);

  	auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,4>>(2);
  	Dz->set_weight(pics_weight);
  	Dz->set_domain_dimensions(&is_dims);
  	Dz->set_codomain_dimensions(&is_dims);

  	solver.add_regularization_group({Dx,Dy,Dz},prior);

  }*/

/*
	if (tv_weight > 0){

		auto Dx = boost::make_shared<cuDCTDerivativeOperator<float>>(0);
		Dx->set_weight(tv_weight);
		Dx->set_domain_dimensions(&is_dims);
		Dx->set_codomain_dimensions(&is_dims);

		auto Dy = boost::make_shared<cuDCTDerivativeOperator<float>>(1);
		Dy->set_weight(tv_weight);
		Dy->set_domain_dimensions(&is_dims);
		Dy->set_codomain_dimensions(&is_dims);

		auto Dz = boost::make_shared<cuDCTDerivativeOperator<float>>(2);
		Dz->set_weight(tv_weight);
		Dz->set_domain_dimensions(&is_dims);
		Dz->set_codomain_dimensions(&is_dims);

		solver.add_regularization_group({Dx,Dy,Dz});

	}

*/

	if (dct_weight > 0){
		auto dctOp = boost::make_shared<cuDCTOperator<float>>();
		dctOp->set_domain_dimensions(&is_dims);
		dctOp->set_weight(dct_weight);
		solver.add_regularization_operator(dctOp);
	}

	auto E = boost::make_shared<CBSubsetOperator<cuNDArray> >(subsets);

	//E->setup(ps,binning,imageDimensions);
	E->setup(ps,binning,imageDimensions);
	E->set_domain_dimensions(&is_dims);
	E->set_codomain_dimensions(ps->get_projections()->get_dimensions().get());

	solver.set_encoding_operator(E);

	//auto projections = boost::make_shared<cuNDArray<float>>(*ps->get_projections());
	hoCuNDArray<float> projections = *ps->get_projections();
	ps->set_projections(boost::shared_ptr<hoCuNDArray<float>>());
	std::cout << "Projection norm:" << nrm2(&projections) << std::endl;
	//E->set_use_offset_correction(false);

	//boost::shared_ptr<cuNDArray<bool>> mask;
	//mask = E->calculate_mask(projections,0.03f);
	//ps->set_projections(boost::shared_ptr<hoCuNDArray<float>>()); //Clear projections from host memory.
	{
		auto cu_proj = cuNDArray<float>(projections);
		E->offset_correct(&cu_proj);
		cu_proj.to_host(&projections);
	}
	//E->set_mask(mask);
	std::cout << "Projection norm:" << nrm2(&projections) << std::endl;

	//solver.set_damping(1e-6);

	/*
    boost::shared_ptr<hoCuNDArray<float> > prior;

  if (vm.count("use_prior")) {
  	prior = calculate_prior(binning,ps,projections,is_dims,imageDimensions);
  	solver.set_x0(prior);
  }
	 */
	boost::shared_ptr<hoNDArray<float>> result;
	{
		GPUTimer tim("Solver");
		result = solver.solve(&projections);
	}
//	global_timer.reset();
	std::cout << "Penguin" << nrm2(result.get()) << std::endl;

	std::cout << "Result sum " << asum(result.get()) << std::endl;

	//apply_mask(result.get(),mask.get());

	std::cout << "Result sum " << asum(result.get()) << std::endl;
	//saveNDArray2HDF5(result.get(),outputFile,imageDimensions,vector_td<float,3>(0),command_line_string.str(),iterations);

	saveNDArray2HDF5(result.get(),outputFile,imageDimensions,floatd3(0,0,0),command_line_string.str(),iterations);
//	write_nd_array(result.get(),"reconstruction.real");
	//write_dicom(result.get(),command_line_string.str(),imageDimensions);

}


#ifndef __VSIM_GEOMETRY_HPP__
#define __VSIM_GEOMETRY_HPP__

#include <string>
#include <vector>
#include <memory>
#include <map>
#include <Eigen/Core>

#include <vsim/env/base_element.hpp>

namespace vsim {

struct Geometry ;
typedef std::shared_ptr<Geometry> GeometryPtr ;

// Geometry structure associates a mesh instance with material.

struct Geometry: public BaseElement {

    Geometry() = default ;
    virtual ~Geometry() {}
};

struct BoxGeometry: public Geometry {
    Eigen::Vector3f half_extents_ ;
};

typedef std::shared_ptr<BoxGeometry> BoxGeometryPtr ;

struct PlaneGeometry: public Geometry {
    Eigen::Vector4f coeffs_ ;
};

typedef std::shared_ptr<PlaneGeometry> PlaneGeometryPtr ;

struct SphereGeometry: public Geometry {
    float radius_ ;
};

struct CylinderGeometry: public Geometry {
    float radius_, height_ ;
};

struct ConeGeometry: public Geometry {
    float radius_, height_ ;
};

static const int MAX_MESH_TEXTURES = 4 ;

class Mesh ;
typedef std::shared_ptr<Mesh> MeshPtr ;

// mesh data of triangular mesh

class Mesh: public Geometry {
public:

    enum PrimitiveType { Triangles, Lines, Points } ;

    // vertex attribute arrays, they can be of different size if attributes (e.g. normals) are shared for multiple vertices (e.g. per face normals)
    std::vector<Eigen::Vector3f> vertices_, normals_, colors_ ;
    std::vector<Eigen::Vector2f> tex_coords_[MAX_MESH_TEXTURES] ;

    // triplets of vertex attribute indices corresponding to the triangles of the mesh (all same size)
    std::vector<uint32_t> vertex_indices_, normal_indices_, color_indices_, tex_coord_indices_[MAX_MESH_TEXTURES] ;

    PrimitiveType ptype_ ;

    void load(const std::string &fpath) {}

    // a 3d cube

    static MeshPtr createWireCube(float sz) ;
    static MeshPtr createSolidCube(float sz) ;

    static MeshPtr createWireSphere(float radius, size_t slices, size_t stacks) ;
    static MeshPtr createSolidSphere(float radius, size_t slices, size_t stacks) ;

    // the base of the cone is on (0, 0, 0) aligned with the z-axis and pointing towards positive z

    static MeshPtr createWireCone(float radius, float height, size_t slices, size_t stacks) ;
    static MeshPtr createSolidCone(float radius, float height, size_t slices, size_t stacks) ;

    static MeshPtr createWireCylinder(float radius, float height, size_t slices, size_t stacks) ;
    static MeshPtr createSolidCylinder(float radius, float height, size_t slices, size_t stacks) ;
};

}
#endif