/* ----------------------------------------------------------------------------
 * GTDynamics Copyright 2020, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * See LICENSE for the license information
 * -------------------------------------------------------------------------- */

/**
 * @file  main.cpp
 * @brief Trajectory optimization for a legged robot with contacts.
 * @author Alejandro Escontrela
 */

#include <gtdynamics/dynamics/DynamicsGraph.h>
#include <gtdynamics/dynamics/OptimizerSetting.h>
#include <gtdynamics/factors/MinTorqueFactor.h>
#include <gtdynamics/universal_robot/Robot.h>
#include <gtdynamics/universal_robot/sdf.h>
#include <gtdynamics/utils/initialize_solution_utils.h>
#include <gtsam/base/Value.h>
#include <gtsam/base/Vector.h>
#include <gtsam/linear/NoiseModel.h>
#include <gtsam/nonlinear/LevenbergMarquardtOptimizer.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/slam/PriorFactor.h>

#include <boost/algorithm/string/join.hpp>
#include <boost/optional.hpp>
#include <fstream>
#include <iostream>
#include <string>
#include <utility>

using namespace gtdynamics;

int main(int argc, char** argv) {
  // Load the Vision 60 quadruped by Ghost robotics:
  // https://youtu.be/wrBNJKZKg10
  auto vision60 =
      CreateRobotFromFile(kUrdfPath + std::string("/vision60.urdf"));

  // Env parameters.
  gtsam::Vector3 gravity = (gtsam::Vector(3) << 0, 0, -9.8).finished();
  double mu = 2.0;

  // Contact points at feet.
  ContactPoints contact_points;
  contact_points.emplace(
      "lower0", ContactPoint{gtsam::Point3(0.14, 0, 0), 0});
  contact_points.emplace(
      "lower1", ContactPoint{gtsam::Point3(0.14, 0, 0), 1});
  contact_points.emplace(
      "lower2", ContactPoint{gtsam::Point3(0.14, 0, 0), 2});
  contact_points.emplace(
      "lower3", ContactPoint{gtsam::Point3(0.14, 0, 0), 3});

  // Specify optimal control problem parameters.
  double T = 3.0;                                     // Time horizon (s.)
  double dt = 1. / 240;                               // Time step (s.)
  int t_steps = static_cast<int>(std::ceil(T / dt));  // Timesteps.

  double sigma_dynamics = 1e-5;    // Variance of dynamics constraints.
  double sigma_objectives = 1e-3;  // Variance of additional objectives.

  // Specify boundary conditions for base and joints.
  gtsam::Pose3 base_pose_init = vision60.link("body")->wTcom();
  gtsam::Vector6 base_twist_init = gtsam::Z_6x1,
                 base_twist_final = gtsam::Z_6x1,
                 base_accel_init = gtsam::Z_6x1,
                 base_accel_final = gtsam::Z_6x1;
  gtsam::Vector joint_angles_init = gtsam::Vector::Zero(12),
                joint_vels_init = gtsam::Vector::Zero(12),
                joint_accels_init = gtsam::Vector::Zero(12),
                joint_vels_final = gtsam::Vector::Zero(12),
                joint_accels_final = gtsam::Vector::Zero(12);

  // Specify target poses we want to reach and at what times
  // we want to reach them.
  std::vector<gtsam::Pose3> des_poses;
  std::vector<double> des_poses_t;
  auto des_pose_nm = gtsam::noiseModel::Isotropic::Sigma(6, sigma_objectives);

  des_poses.push_back(gtsam::Pose3(gtsam::Rot3::RzRyRx(M_PI / 8, 0.0, 0.0),
                                   gtsam::Point3(0, 0, 0.1)));
  des_poses_t.push_back(0.75);

  des_poses.push_back(gtsam::Pose3(gtsam::Rot3::RzRyRx(-M_PI / 8, 0.0, 0.0),
                                   gtsam::Point3(0, 0, 0.1)));
  des_poses_t.push_back(1.5);

  des_poses.push_back(
      gtsam::Pose3(gtsam::Rot3::RzRyRx(0.0, M_PI / 10, M_PI / 8),
                   gtsam::Point3(0.0, 0, 0.1)));
  des_poses_t.push_back(2.25);

  des_poses.push_back(
      gtsam::Pose3(gtsam::Rot3::RzRyRx(0.0, -M_PI / 10, -M_PI / 8),
                   gtsam::Point3(0.0, 0, 0.1)));
  des_poses_t.push_back(3.0);

  // Build the trajectory factor graph and add boundary condition and goal
  // pose factors.
  auto opt = OptimizerSetting();
  opt.bp_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.bv_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.ba_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.p_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.v_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.a_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.f_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.fa_cost_model = gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics);
  opt.t_cost_model = gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.cp_cost_model = gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.cfriction_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.cv_cost_model = gtsam::noiseModel::Isotropic::Sigma(3, sigma_dynamics);
  opt.ca_cost_model = gtsam::noiseModel::Isotropic::Sigma(3, sigma_dynamics);
  opt.planar_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(3, sigma_dynamics);
  opt.prior_q_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.prior_qv_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.prior_qa_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.prior_t_cost_model =
      gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.q_col_cost_model = gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.v_col_cost_model = gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  opt.time_cost_model = gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics);
  auto graph_builder = DynamicsGraph(opt, gravity);
  gtsam::NonlinearFactorGraph graph = graph_builder.trajectoryFG(
      vision60, t_steps, dt, CollocationScheme::Trapezoidal, contact_points,
      mu);

  auto base_link = vision60.link("body");
  gtsam::NonlinearFactorGraph objective_factors;

  // Add certain poses to be reached.
  for (size_t i = 0; i < des_poses.size(); i++)
    objective_factors.addPrior(
        internal::PoseKey(base_link->id(),
                          static_cast<int>(std::ceil(des_poses_t[i] / dt))),
        des_poses[i], des_pose_nm);

  // Add base boundary conditions to FG.
  objective_factors.addPrior(
      internal::PoseKey(base_link->id(), 0), base_pose_init,
      gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics));
  objective_factors.addPrior<gtsam::Vector6>(
      internal::TwistKey(base_link->id(), 0), base_twist_init,
      gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics));
  objective_factors.addPrior<gtsam::Vector6>(
      internal::TwistAccelKey(base_link->id(), 0), base_accel_init,
      gtsam::noiseModel::Isotropic::Sigma(6, sigma_dynamics));
  objective_factors.addPrior<gtsam::Vector6>(
      internal::TwistKey(base_link->id(), t_steps), base_twist_final,
      gtsam::noiseModel::Isotropic::Sigma(6, sigma_objectives));
  objective_factors.addPrior<gtsam::Vector6>(
      internal::TwistAccelKey(base_link->id(), t_steps), base_accel_final,
      gtsam::noiseModel::Isotropic::Sigma(6, sigma_objectives));

  // Add joint boundary conditions to FG.
  for (auto&& joint : vision60.joints()) {
    objective_factors.addPrior(
        internal::JointAngleKey(joint->id(), 0), 0.0,
        gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics));
    objective_factors.addPrior(
        internal::JointVelKey(joint->id(), 0), 0.0,
        gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics));
    objective_factors.addPrior(
        internal::JointAccelKey(joint->id(), 0), 0.0,
        gtsam::noiseModel::Isotropic::Sigma(1, sigma_dynamics));
    objective_factors.addPrior(
        internal::JointVelKey(joint->id(), t_steps), 0.0,
        gtsam::noiseModel::Isotropic::Sigma(1, sigma_objectives));
    objective_factors.addPrior(
        internal::JointAccelKey(joint->id(), t_steps), 0.0,
        gtsam::noiseModel::Isotropic::Sigma(1, sigma_objectives));
  }

  // Add min torque objectives.
  for (int t = 0; t <= t_steps; t++) {
    for (auto&& joint : vision60.joints())
      objective_factors.add(MinTorqueFactor(
          internal::TorqueKey(joint->id(), t),
          gtsam::noiseModel::Gaussian::Covariance(gtsam::I_1x1)));
  }
  graph.add(objective_factors);

  // Initialize solution.
  gtsam::Values init_vals;
  std::string initialization_technique = "inverse_kinematics";
  if (initialization_technique == "interp")
    // TODO(aescontrela): Figure out why the linearly interpolated initial
    // trajectory fails to optimize. My initial guess is that the optimizer has
    // a difficult time optimizing the trajectory when the initial solution lies
    // in the infeasible region. This would make sense if I were using an IPM to
    // solve this problem...
    init_vals = InitializeSolutionInterpolationMultiPhase(
        vision60, "body", base_pose_init, des_poses, des_poses_t, dt, 0.0,
        contact_points);
  else if (initialization_technique == "zeros")
    init_vals = ZeroValuesTrajectory(vision60, t_steps, 0, 0.0, contact_points);
  else if (initialization_technique == "inverse_kinematics")
    init_vals = InitializeSolutionInverseKinematics(
        vision60, "body", base_pose_init, des_poses, des_poses_t, dt, 0.0,
        contact_points);

  gtsam::LevenbergMarquardtParams params;
  params.setVerbosityLM("SUMMARY");
  gtsam::LevenbergMarquardtOptimizer optimizer(graph, init_vals, params);
  gtsam::Values results = optimizer.optimize();

  gtsam::Pose3 optimized_pose_init = Pose(results, base_link->id(), 0);
  gtsam::Pose3 optimized_pose_final =
      Pose(results, base_link->id(), t_steps - 1);

  std::cout << "Optimized Pose init trans: "
            << optimized_pose_init.translation()
            << "\n\tinit rot:" << optimized_pose_init.rotation().rpy()
            << std::endl;
  std::cout << "Optimized Pose final trans: "
            << optimized_pose_final.translation()
            << "\n\tfinal rot:" << optimized_pose_final.rotation().rpy()
            << std::endl;

  auto joint_vals_init = graph_builder.jointAnglesMap(vision60, results, 0);
  auto joint_vals_final =
      graph_builder.jointAnglesMap(vision60, results, t_steps);

  std::cout << "Joint vals init" << std::endl;
  for (auto&& jval : joint_vals_init)
    std::cout << "\t" << jval.first << ": " << jval.second << "," << std::endl;
  std::cout << "Joint vals final" << std::endl;
  for (auto&& jval : joint_vals_final)
    std::cout << "\t" << jval.first << ": " << jval.second << "," << std::endl;

  // Log the joint angles, velocities, accels, torques, and current goal pose.
  std::vector<std::string> jnames;
  for (auto&& joint : vision60.joints()) jnames.push_back(joint->name());
  std::string jnames_str = boost::algorithm::join(jnames, ",");
  std::ofstream traj_file;
  traj_file.open("traj.csv");
  // angles, vels, accels, torques.
  traj_file << jnames_str << "," << jnames_str << "," << jnames_str << ","
            << jnames_str << ",gol_x"
            << ",gol_y"
            << ",gol_z"
            << ",gol_qx"
            << ",gol_qy"
            << ",gol_qz"
            << ",gol_qw"
            << "\n";
  for (int t = 0; t <= t_steps; t++) {
    std::vector<std::string> vals;
    for (auto&& joint : vision60.joints())
      vals.push_back(std::to_string(JointAngle(results, joint->id(), t)));
    for (auto&& joint : vision60.joints())
      vals.push_back(std::to_string(JointVel(results, joint->id(), t)));
    for (auto&& joint : vision60.joints())
      vals.push_back(std::to_string(JointAccel(results, joint->id(), t)));
    for (auto&& joint : vision60.joints())
      vals.push_back(std::to_string(Torque(results, joint->id(), t)));

    for (size_t i = 0; i < des_poses.size(); i++) {
      gtsam::Pose3 dp = des_poses[i];
      if (t <= static_cast<int>(std::round(des_poses_t[i] / dt))) {
        vals.push_back(std::to_string(dp.x()));
        vals.push_back(std::to_string(dp.y()));
        vals.push_back(std::to_string(dp.z()));
        vals.push_back(std::to_string(dp.rotation().toQuaternion().x()));
        vals.push_back(std::to_string(dp.rotation().toQuaternion().y()));
        vals.push_back(std::to_string(dp.rotation().toQuaternion().z()));
        vals.push_back(std::to_string(dp.rotation().toQuaternion().w()));
        break;
      }
    }

    std::string vals_str = boost::algorithm::join(vals, ",");
    traj_file << vals_str << "\n";
  }
  traj_file.close();

  return 0;
}


#ifndef ___INVERSECOLORPSACEADJUSTEMENT_H___
#define ___INVERSECOLORPSACEADJUSTEMENT_H___
#include <iostream>
#include <math.h>
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <stdio.h>
#include <cstdlib>
// #include <opencv/cv.h>
// #include <opencv/highgui.h>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

/**
 * @brief Class to bring back the RGB color space from the oRGB
 *
 */
class Inverse {

 public:
  Inverse(void) = default;

/**
 * @brief contains a matrix of rotation, that we will use to rotate a vector in the next function.
 *
 * @param angle
 * @return Eigen::Matrix3d
 */
    Eigen::Matrix3d rotatePoint(double angle);

// /**
//  * @brief based on the angle taken from the vector of the loaded image, it rotates that vector using the previous rotatePoint function
//  *
//  * @param img
//  * @return cv::Mat
//  */
    cv::Mat fullRotation(cv::Mat img1);

/**
 * @brief function that returns a vector after it has been multiplied with a specific matrix given because of the formula. Also it normalizes the values
 *
 * @param img
 * @return cv::Mat
 */
    cv::Mat setlinearImage(cv::Mat img1);
};

#endif


// Boost.Geometry - gis-projections (based on PROJ4)

// Copyright (c) 2008-2015 Barend Gehrels, Amsterdam, the Netherlands.

// This file was modified by Oracle on 2017, 2018.
// Modifications copyright (c) 2017-2018, Oracle and/or its affiliates.
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle.

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

// This file is converted from PROJ4, http://trac.osgeo.org/proj
// PROJ4 is originally written by Gerald Evenden (then of the USGS)
// PROJ4 is maintained by Frank Warmerdam
// PROJ4 is converted to Boost.Geometry by Barend Gehrels

// Last updated version of proj: 5.0.0

// Original copyright notice:

// This implements the Quadrilateralized Spherical Cube (QSC) projection.
// Copyright (c) 2011, 2012  Martin Lambers <marlam@marlam.de>

// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

// The QSC projection was introduced in:
// [OL76]
// E.M. O'Neill and R.E. Laubscher, "Extended Studies of a Quadrilateralized
// Spherical Cube Earth Data Base", Naval Environmental Prediction Research
// Facility Tech. Report NEPRF 3-76 (CSC), May 1976.

// The preceding shift from an ellipsoid to a sphere, which allows to apply
// this projection to ellipsoids as used in the Ellipsoidal Cube Map model,
// is described in
// [LK12]
// M. Lambers and A. Kolb, "Ellipsoidal Cube Maps for Accurate Rendering of
// Planetary-Scale Terrain Data", Proc. Pacfic Graphics (Short Papers), Sep.
// 2012

// You have to choose one of the following projection centers,
// corresponding to the centers of the six cube faces:
// phi0 = 0.0, lam0 = 0.0       ("front" face)
// phi0 = 0.0, lam0 = 90.0      ("right" face)
// phi0 = 0.0, lam0 = 180.0     ("back" face)
// phi0 = 0.0, lam0 = -90.0     ("left" face)
// phi0 = 90.0                  ("top" face)
// phi0 = -90.0                 ("bottom" face)
// Other projection centers will not work!

// In the projection code below, each cube face is handled differently.
// See the computation of the face parameter in the ENTRY0(qsc) function
// and the handling of different face values (FACE_*) in the forward and
// inverse projections.

// Furthermore, the projection is originally only defined for theta angles
// between (-1/4 * PI) and (+1/4 * PI) on the current cube face. This area
// of definition is named AREA_0 in the projection code below. The other
// three areas of a cube face are handled by rotation of AREA_0.

#ifndef BOOST_GEOMETRY_PROJECTIONS_QSC_HPP
#define BOOST_GEOMETRY_PROJECTIONS_QSC_HPP

#include <boost/core/ignore_unused.hpp>
#include <boost/geometry/util/math.hpp>

#include <boost/geometry/srs/projections/impl/base_static.hpp>
#include <boost/geometry/srs/projections/impl/base_dynamic.hpp>
#include <boost/geometry/srs/projections/impl/projects.hpp>
#include <boost/geometry/srs/projections/impl/factory_entry.hpp>

namespace boost { namespace geometry
{

namespace srs { namespace par4
{
    struct qsc {}; // Quadrilateralized Spherical Cube

}} //namespace srs::par4

namespace projections
{
    #ifndef DOXYGEN_NO_DETAIL
    namespace detail { namespace qsc
    {

            /* The six cube faces. */
            enum face_type {
                face_front  = 0,
                face_right  = 1,
                face_back   = 2,
                face_left   = 3,
                face_top    = 4,
                face_bottom = 5
            };

            template <typename T>
            struct par_qsc
            {
                face_type face;
                T a_squared;
                T b;
                T one_minus_f;
                T one_minus_f_squared;
            };

            static const double epsilon10 = 1.e-10;

            /* The four areas on a cube face. AREA_0 is the area of definition,
             * the other three areas are counted counterclockwise. */
            enum area_type {
                area_0 = 0,
                area_1 = 1,
                area_2 = 2,
                area_3 = 3
            };

            /* Helper function for forward projection: compute the theta angle
             * and determine the area number. */
            template <typename T>
            inline T qsc_fwd_equat_face_theta(T const& phi, T const& y, T const& x, area_type *area)
            {
                static const T fourth_pi = detail::fourth_pi<T>();
                static const T half_pi = detail::half_pi<T>();
                static const T pi = detail::pi<T>();

                T theta;
                if (phi < epsilon10) {
                    *area = area_0;
                    theta = 0.0;
                } else {
                    theta = atan2(y, x);
                    if (fabs(theta) <= fourth_pi) {
                        *area = area_0;
                    } else if (theta > fourth_pi && theta <= half_pi + fourth_pi) {
                        *area = area_1;
                        theta -= half_pi;
                    } else if (theta > half_pi + fourth_pi || theta <= -(half_pi + fourth_pi)) {
                        *area = area_2;
                        theta = (theta >= 0.0 ? theta - pi : theta + pi);
                    } else {
                        *area = area_3;
                        theta += half_pi;
                    }
                }
                return theta;
            }

            /* Helper function: shift the longitude. */
            template <typename T>
            inline T qsc_shift_lon_origin(T const& lon, T const& offset)
            {
                static const T pi = detail::pi<T>();
                static const T two_pi = detail::two_pi<T>();

                T slon = lon + offset;
                if (slon < -pi) {
                    slon += two_pi;
                } else if (slon > +pi) {
                    slon -= two_pi;
                }
                return slon;
            }

            /* Forward projection, ellipsoid */

            // template class, using CRTP to implement forward/inverse
            template <typename T, typename Parameters>
            struct base_qsc_ellipsoid
                : public base_t_fi<base_qsc_ellipsoid<T, Parameters>, T, Parameters>
            {
                par_qsc<T> m_proj_parm;

                inline base_qsc_ellipsoid(const Parameters& par)
                    : base_t_fi<base_qsc_ellipsoid<T, Parameters>, T, Parameters>(*this, par)
                {}

                // FORWARD(e_forward)
                // Project coordinates from geographic (lon, lat) to cartesian (x, y)
                inline void fwd(T& lp_lon, T& lp_lat, T& xy_x, T& xy_y) const
                {
                    static const T fourth_pi = detail::fourth_pi<T>();
                    static const T half_pi = detail::half_pi<T>();
                    static const T pi = detail::pi<T>();

                        T lat, lon;
                        T theta, phi;
                        T t, mu; /* nu; */
                        area_type area;

                        /* Convert the geodetic latitude to a geocentric latitude.
                         * This corresponds to the shift from the ellipsoid to the sphere
                         * described in [LK12]. */
                        if (this->m_par.es != 0.0) {
                            lat = atan(this->m_proj_parm.one_minus_f_squared * tan(lp_lat));
                        } else {
                            lat = lp_lat;
                        }

                        /* Convert the input lat, lon into theta, phi as used by QSC.
                         * This depends on the cube face and the area on it.
                         * For the top and bottom face, we can compute theta and phi
                         * directly from phi, lam. For the other faces, we must use
                         * unit sphere cartesian coordinates as an intermediate step. */
                        lon = lp_lon;
                        if (this->m_proj_parm.face == face_top) {
                            phi = half_pi - lat;
                            if (lon >= fourth_pi && lon <= half_pi + fourth_pi) {
                                area = area_0;
                                theta = lon - half_pi;
                            } else if (lon > half_pi + fourth_pi || lon <= -(half_pi + fourth_pi)) {
                                area = area_1;
                                theta = (lon > 0.0 ? lon - pi : lon + pi);
                            } else if (lon > -(half_pi + fourth_pi) && lon <= -fourth_pi) {
                                area = area_2;
                                theta = lon + half_pi;
                            } else {
                                area = area_3;
                                theta = lon;
                            }
                        } else if (this->m_proj_parm.face == face_bottom) {
                            phi = half_pi + lat;
                            if (lon >= fourth_pi && lon <= half_pi + fourth_pi) {
                                area = area_0;
                                theta = -lon + half_pi;
                            } else if (lon < fourth_pi && lon >= -fourth_pi) {
                                area = area_1;
                                theta = -lon;
                            } else if (lon < -fourth_pi && lon >= -(half_pi + fourth_pi)) {
                                area = area_2;
                                theta = -lon - half_pi;
                            } else {
                                area = area_3;
                                theta = (lon > 0.0 ? -lon + pi : -lon - pi);
                            }
                        } else {
                            T q, r, s;
                            T sinlat, coslat;
                            T sinlon, coslon;

                            if (this->m_proj_parm.face == face_right) {
                                lon = qsc_shift_lon_origin(lon, +half_pi);
                            } else if (this->m_proj_parm.face == face_back) {
                                lon = qsc_shift_lon_origin(lon, +pi);
                            } else if (this->m_proj_parm.face == face_left) {
                                lon = qsc_shift_lon_origin(lon, -half_pi);
                            }
                            sinlat = sin(lat);
                            coslat = cos(lat);
                            sinlon = sin(lon);
                            coslon = cos(lon);
                            q = coslat * coslon;
                            r = coslat * sinlon;
                            s = sinlat;

                            if (this->m_proj_parm.face == face_front) {
                                phi = acos(q);
                                theta = qsc_fwd_equat_face_theta(phi, s, r, &area);
                            } else if (this->m_proj_parm.face == face_right) {
                                phi = acos(r);
                                theta = qsc_fwd_equat_face_theta(phi, s, -q, &area);
                            } else if (this->m_proj_parm.face == face_back) {
                                phi = acos(-q);
                                theta = qsc_fwd_equat_face_theta(phi, s, -r, &area);
                            } else if (this->m_proj_parm.face == face_left) {
                                phi = acos(-r);
                                theta = qsc_fwd_equat_face_theta(phi, s, q, &area);
                            } else {
                                /* Impossible */
                                phi = theta = 0.0;
                                area = area_0;
                            }
                        }

                        /* Compute mu and nu for the area of definition.
                         * For mu, see Eq. (3-21) in [OL76], but note the typos:
                         * compare with Eq. (3-14). For nu, see Eq. (3-38). */
                        mu = atan((12.0 / pi) * (theta + acos(sin(theta) * cos(fourth_pi)) - half_pi));
                        // TODO: (cos(mu) * cos(mu)) could be replaced with sqr(cos(mu))
                        t = sqrt((1.0 - cos(phi)) / (cos(mu) * cos(mu)) / (1.0 - cos(atan(1.0 / cos(theta)))));
                        /* nu = atan(t);        We don't really need nu, just t, see below. */

                        /* Apply the result to the real area. */
                        if (area == area_1) {
                            mu += half_pi;
                        } else if (area == area_2) {
                            mu += pi;
                        } else if (area == area_3) {
                            mu += half_pi + pi;
                        }

                        /* Now compute x, y from mu and nu */
                        /* t = tan(nu); */
                        xy_x = t * cos(mu);
                        xy_y = t * sin(mu);
                }
                /* Inverse projection, ellipsoid */

                // INVERSE(e_inverse)
                // Project coordinates from cartesian (x, y) to geographic (lon, lat)
                inline void inv(T& xy_x, T& xy_y, T& lp_lon, T& lp_lat) const
                {
                    static const T half_pi = detail::half_pi<T>();
                    static const T pi = detail::pi<T>();

                        T mu, nu, cosmu, tannu;
                        T tantheta, theta, cosphi, phi;
                        T t;
                        int area;

                        /* Convert the input x, y to the mu and nu angles as used by QSC.
                         * This depends on the area of the cube face. */
                        nu = atan(sqrt(xy_x * xy_x + xy_y * xy_y));
                        mu = atan2(xy_y, xy_x);
                        if (xy_x >= 0.0 && xy_x >= fabs(xy_y)) {
                            area = area_0;
                        } else if (xy_y >= 0.0 && xy_y >= fabs(xy_x)) {
                            area = area_1;
                            mu -= half_pi;
                        } else if (xy_x < 0.0 && -xy_x >= fabs(xy_y)) {
                            area = area_2;
                            mu = (mu < 0.0 ? mu + pi : mu - pi);
                        } else {
                            area = area_3;
                            mu += half_pi;
                        }

                        /* Compute phi and theta for the area of definition.
                         * The inverse projection is not described in the original paper, but some
                         * good hints can be found here (as of 2011-12-14):
                         * http://fits.gsfc.nasa.gov/fitsbits/saf.93/saf.9302
                         * (search for "Message-Id: <9302181759.AA25477 at fits.cv.nrao.edu>") */
                        t = (pi / 12.0) * tan(mu);
                        tantheta = sin(t) / (cos(t) - (1.0 / sqrt(2.0)));
                        theta = atan(tantheta);
                        cosmu = cos(mu);
                        tannu = tan(nu);
                        cosphi = 1.0 - cosmu * cosmu * tannu * tannu * (1.0 - cos(atan(1.0 / cos(theta))));
                        if (cosphi < -1.0) {
                            cosphi = -1.0;
                        } else if (cosphi > +1.0) {
                            cosphi = +1.0;
                        }

                        /* Apply the result to the real area on the cube face.
                         * For the top and bottom face, we can compute phi and lam directly.
                         * For the other faces, we must use unit sphere cartesian coordinates
                         * as an intermediate step. */
                        if (this->m_proj_parm.face == face_top) {
                            phi = acos(cosphi);
                            lp_lat = half_pi - phi;
                            if (area == area_0) {
                                lp_lon = theta + half_pi;
                            } else if (area == area_1) {
                                lp_lon = (theta < 0.0 ? theta + pi : theta - pi);
                            } else if (area == area_2) {
                                lp_lon = theta - half_pi;
                            } else /* area == AREA_3 */ {
                                lp_lon = theta;
                            }
                        } else if (this->m_proj_parm.face == face_bottom) {
                            phi = acos(cosphi);
                            lp_lat = phi - half_pi;
                            if (area == area_0) {
                                lp_lon = -theta + half_pi;
                            } else if (area == area_1) {
                                lp_lon = -theta;
                            } else if (area == area_2) {
                                lp_lon = -theta - half_pi;
                            } else /* area == area_3 */ {
                                lp_lon = (theta < 0.0 ? -theta - pi : -theta + pi);
                            }
                        } else {
                            /* Compute phi and lam via cartesian unit sphere coordinates. */
                            T q, r, s;
                            q = cosphi;
                            t = q * q;
                            if (t >= 1.0) {
                                s = 0.0;
                            } else {
                                s = sqrt(1.0 - t) * sin(theta);
                            }
                            t += s * s;
                            if (t >= 1.0) {
                                r = 0.0;
                            } else {
                                r = sqrt(1.0 - t);
                            }
                            /* Rotate q,r,s into the correct area. */
                            if (area == area_1) {
                                t = r;
                                r = -s;
                                s = t;
                            } else if (area == area_2) {
                                r = -r;
                                s = -s;
                            } else if (area == area_3) {
                                t = r;
                                r = s;
                                s = -t;
                            }
                            /* Rotate q,r,s into the correct cube face. */
                            if (this->m_proj_parm.face == face_right) {
                                t = q;
                                q = -r;
                                r = t;
                            } else if (this->m_proj_parm.face == face_back) {
                                q = -q;
                                r = -r;
                            } else if (this->m_proj_parm.face == face_left) {
                                t = q;
                                q = r;
                                r = -t;
                            }
                            /* Now compute phi and lam from the unit sphere coordinates. */
                            lp_lat = acos(-s) - half_pi;
                            lp_lon = atan2(r, q);
                            if (this->m_proj_parm.face == face_right) {
                                lp_lon = qsc_shift_lon_origin(lp_lon, -half_pi);
                            } else if (this->m_proj_parm.face == face_back) {
                                lp_lon = qsc_shift_lon_origin(lp_lon, -pi);
                            } else if (this->m_proj_parm.face == face_left) {
                                lp_lon = qsc_shift_lon_origin(lp_lon, +half_pi);
                            }
                        }

                        /* Apply the shift from the sphere to the ellipsoid as described
                         * in [LK12]. */
                        if (this->m_par.es != 0.0) {
                            int invert_sign;
                            T tanphi, xa;
                            invert_sign = (lp_lat < 0.0 ? 1 : 0);
                            tanphi = tan(lp_lat);
                            xa = this->m_proj_parm.b / sqrt(tanphi * tanphi + this->m_proj_parm.one_minus_f_squared);
                            lp_lat = atan(sqrt(this->m_par.a * this->m_par.a - xa * xa) / (this->m_proj_parm.one_minus_f * xa));
                            if (invert_sign) {
                                lp_lat = -lp_lat;
                            }
                        }
                }

                static inline std::string get_name()
                {
                    return "qsc_ellipsoid";
                }

            };

            // Quadrilateralized Spherical Cube
            template <typename Parameters, typename T>
            inline void setup_qsc(Parameters& par, par_qsc<T>& proj_parm)
            {
                static const T fourth_pi = detail::fourth_pi<T>();
                static const T half_pi = detail::half_pi<T>();

                /* Determine the cube face from the center of projection. */
                if (par.phi0 >= half_pi - fourth_pi / 2.0) {
                    proj_parm.face = face_top;
                } else if (par.phi0 <= -(half_pi - fourth_pi / 2.0)) {
                    proj_parm.face = face_bottom;
                } else if (fabs(par.lam0) <= fourth_pi) {
                    proj_parm.face = face_front;
                } else if (fabs(par.lam0) <= half_pi + fourth_pi) {
                    proj_parm.face = (par.lam0 > 0.0 ? face_right : face_left);
                } else {
                    proj_parm.face = face_back;
                }
                /* Fill in useful values for the ellipsoid <-> sphere shift
                 * described in [LK12]. */
                if (par.es != 0.0) {
                    proj_parm.a_squared = par.a * par.a;
                    proj_parm.b = par.a * sqrt(1.0 - par.es);
                    proj_parm.one_minus_f = 1.0 - (par.a - proj_parm.b) / par.a;
                    proj_parm.one_minus_f_squared = proj_parm.one_minus_f * proj_parm.one_minus_f;
                }
            }

    }} // namespace detail::qsc
    #endif // doxygen

    /*!
        \brief Quadrilateralized Spherical Cube projection
        \ingroup projections
        \tparam Geographic latlong point type
        \tparam Cartesian xy point type
        \tparam Parameters parameter type
        \par Projection characteristics
         - Azimuthal
         - Spheroid
        \par Example
        \image html ex_qsc.gif
    */
    template <typename T, typename Parameters>
    struct qsc_ellipsoid : public detail::qsc::base_qsc_ellipsoid<T, Parameters>
    {
        inline qsc_ellipsoid(const Parameters& par) : detail::qsc::base_qsc_ellipsoid<T, Parameters>(par)
        {
            detail::qsc::setup_qsc(this->m_par, this->m_proj_parm);
        }
    };

    #ifndef DOXYGEN_NO_DETAIL
    namespace detail
    {

        // Static projection
        BOOST_GEOMETRY_PROJECTIONS_DETAIL_STATIC_PROJECTION(srs::par4::qsc, qsc_ellipsoid, qsc_ellipsoid)

        // Factory entry(s)
        template <typename T, typename Parameters>
        class qsc_entry : public detail::factory_entry<T, Parameters>
        {
            public :
                virtual base_v<T, Parameters>* create_new(const Parameters& par) const
                {
                    return new base_v_fi<qsc_ellipsoid<T, Parameters>, T, Parameters>(par);
                }
        };

        template <typename T, typename Parameters>
        inline void qsc_init(detail::base_factory<T, Parameters>& factory)
        {
            factory.add_to_factory("qsc", new qsc_entry<T, Parameters>);
        }

    } // namespace detail
    #endif // doxygen

} // namespace projections

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_PROJECTIONS_QSC_HPP


/* Copyright 2016 Guillem Pascual */

#pragma once

#include "defs/intrusive.hpp"
#include <boost/intrusive_ptr.hpp>

#include "defs/common.hpp"

INCL_NOWARN
#include <glm/glm.hpp>
INCL_WARN

class Bezier;
class MapAwareEntity;
class Packet;

class MovementGenerator
{
public:
    virtual ~MovementGenerator();

    virtual boost::intrusive_ptr<Packet> packet() = 0;
    virtual glm::vec3 update(MapAwareEntity* owner, float elapsed) = 0;
    virtual bool hasNext() = 0;
};

class RandomMovement : public MovementGenerator
{
public:
    RandomMovement();
    virtual ~RandomMovement();

    boost::intrusive_ptr<Packet> packet() override;
    glm::vec3 update(MapAwareEntity* owner, float elapsed) override;
    bool hasNext() override;

private:
    bool _hasPoint;

    Bezier* _bezier;

    float _t;
    glm::vec2 _previous;
};

class Bezier
{
public:
    Bezier(glm::vec2 a, glm::vec2 b, glm::vec2 c, glm::vec2 d) :
        _a(a), _b(b), _c(c), _d(d)
    {}

    virtual ~Bezier();

    virtual glm::vec2 next(float t) = 0;
    virtual float increase(float t, float speed) = 0;

    inline glm::vec2 start() { return _a; }
    inline glm::vec2 startOffset() { return _b; }
    inline glm::vec2 endOffset() { return _c; }
    inline glm::vec2 end() { return _d; }

protected:
    glm::vec2 calc(float t)
    {
        float u = 1.0f - t;
        float tt = t * t;
        float uu = u * u;
        float uuu = uu * u;
        float ttt = tt * t;

        glm::vec2 p = uuu * _a;     // first term
        p += 3 * uu * t * _b;       // second term
        p += 3 * u * tt * _c;       // third term
        p += ttt * _d;              // fourth term

        return p;
    }

protected:
    glm::vec2 _a;
    glm::vec2 _b;
    glm::vec2 _c;
    glm::vec2 _d;
};

template <int len>
class SegmentedBezier : public Bezier
{
public:
    SegmentedBezier(glm::vec2 a, glm::vec2 b, glm::vec2 c, glm::vec2 d) :
        Bezier(a, b, c, d),
        _length(0)
    {
        _segments[0] = 0;

        glm::vec2 origin = calc(0);
        for (int i = 1; i <= len; i += 1)
        {
            glm::vec2 p = calc(i * (1.0f / len));
            _length += glm::distance(p, origin);
            _segments[i] = _length;
            origin = p;
        }
    }

    float increase(float t, float speed)
    {
        return speed / _length;
    }

    glm::vec2 next(float t)
    {
        return calc(map(t));
    }

protected:
    float map(float t)
    {
        float targetLength = t * _segments[len];
        int low = 0;
        int high = len;
        int index = 0;

        while (low < high)
        {
            index = low + (high - low) / 2;
            if (_segments[index] < targetLength)
            {
                low = index + 1;
            }
            else
            {
                high = index;
            }
        }

        if (_segments[index] > targetLength)
        {
            index--;
        }

        float lengthBefore = _segments[index];
        if (lengthBefore == targetLength)
        {
            return index / static_cast<float>(len);
        }

        return (index + (targetLength - lengthBefore) /
            static_cast<float>(_segments[index + 1] - lengthBefore)) / static_cast<float>(len);
    }

protected:
    float _segments[len + 1];
    float _length;
};

class DerivativeBezier : public Bezier
{
public:
    DerivativeBezier(glm::vec2 a, glm::vec2 b, glm::vec2 c, glm::vec2 d) :
        Bezier(a, b, c, d)
    {
        _dv1 = -3.0f * a + 9.0f * b - 9.0f * c + 3.0f * d;
        _dv2 = 6.0f * a - 12.0f * b + 6.0f * c;
        _dv3 = -3.0f * a + 3.0f * b;
    }

    float increase(float t, float speed)
    {
        return speed / glm::length(t * t * _dv1 + t * _dv2 + _dv3);
    }

    glm::vec2 next(float t)
    {
        return calc(t);
    }

protected:
    glm::vec2 _dv1;
    glm::vec2 _dv2;
    glm::vec2 _dv3;
};


/* -*- c++ -*- */
/*
 * Copyright 2015 Free Software Foundation, Inc.
 *
 * This file is part of GNU Radio
 *
 * GNU Radio is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 *
 * GNU Radio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Radio; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include <gnuradio/fec/tpc_encoder.h>
#include <gnuradio/fec/tpc_common.h>
#include <gnuradio/fec/generic_encoder.h>

#include <math.h>
#include <boost/assign/list_of.hpp>
#include <volk/volk.h>
#include <sstream>
#include <stdio.h>
#include <vector>

#include <algorithm>            // for std::reverse
#include <string.h>             // for memcpy

namespace gr {
namespace fec {
generic_encoder::sptr
tpc_encoder::make(std::vector<int> row_polys, std::vector<int> col_polys, int krow, int kcol, int bval, int qval)
{
    return generic_encoder::sptr(new tpc_encoder(row_polys, col_polys, krow, kcol, bval, qval));
}

tpc_encoder::tpc_encoder (std::vector<int> row_polys, std::vector<int> col_polys, int krow, int kcol, int bval, int qval)
    : d_rowpolys(row_polys), d_colpolys(col_polys), d_krow(krow), d_kcol(kcol),
      d_bval(bval), d_qval(qval)
{
    // first we operate on data chunks of get_input_size()
    // TODO: should we verify this and throw an error if it doesn't match?  YES
    // hwo do we do that?

    // calculate the input and output sizes
    inputSize = (d_krow*d_kcol - (d_bval+d_qval));
    rowEncoder_K = ceil(log(d_rowpolys[0])/log(2));    // rowEncoder_K is the constraint length of the row encoder polynomial
    rowEncoder_n = d_rowpolys.size();
    rowEncoder_m = rowEncoder_K - 1;
    colEncoder_K = ceil(log(d_colpolys[0])/log(2));    // colEncoder_K is the constraint length of the col encoder polynomial
    colEncoder_n = d_colpolys.size();
    colEncoder_m = colEncoder_K - 1;

    outputSize = ((d_krow+rowEncoder_m)*rowEncoder_n)*((d_kcol+colEncoder_m)*colEncoder_n) - d_bval;

    fp = NULL;
    //DEBUG_PRINT("inputSize=%d outputSize=%d\n", inputSize, outputSize);
    //fp = fopen("c_encoder_output.txt", "w");

    // resize internal matrices
    rowNumStates = 1 << (rowEncoder_K-1);       // 2^(row_mm)
    colNumStates = 1 << (colEncoder_K-1);       // 2^(col_mm)
    rowOutputs.resize(2, std::vector<int>(rowNumStates,0));
    rowNextStates.resize(2, std::vector<int>(rowNumStates,0));
    colOutputs.resize(2, std::vector<int>(colNumStates,0));
    colNextStates.resize(2, std::vector<int>(colNumStates,0));;

    rowTail.resize(rowNumStates, 0);
    colTail.resize(colNumStates, 0);

    // precalculate the state transition matrix for the row polynomial
    tpc_common::precomputeStateTransitionMatrix_RSCPoly(rowNumStates, d_rowpolys, rowEncoder_K, rowEncoder_n,
                                                        rowOutputs, rowNextStates);

    // calculate the tail for the row
    tpc_common::rsc_tail(rowTail, d_rowpolys, rowNumStates, rowEncoder_m);

    // precalculate the state transition matrix for the column polynomial
    tpc_common::precomputeStateTransitionMatrix_RSCPoly(colNumStates, d_colpolys, colEncoder_K, colEncoder_n,
                                                        colOutputs, colNextStates);
    // calculate the tail for the col
    tpc_common::rsc_tail(colTail, d_colpolys, colNumStates, colEncoder_m);

    // pre-allocate memory we use for encoding
    inputSizeWithPad = d_bval+d_qval+inputSize;
    inputWithPad.resize(inputSizeWithPad, 0);

    numRowsToEncode = inputSizeWithPad/d_krow;      // this should be OK w/ integer division -- TODO: check this?
    rowToEncode.resize(d_krow,0);
    rowEncoded_block.resize(d_krow+(rowEncoder_m*rowEncoder_n), 0);
    rowEncodedBits.resize(d_kcol, std::vector<uint8_t>(rowEncoder_m*rowEncoder_n,0) );

    numColsToEncode = d_krow+(rowEncoder_m*rowEncoder_n);
    colToEncode.resize(d_kcol,0);
    colEncoded_block.resize(d_kcol+(colEncoder_m*colEncoder_n), 0);
    colEncodedBits.resize(d_krow+(rowEncoder_m*rowEncoder_n), std::vector<uint8_t>(colEncoder_m*colEncoder_n,0) );
}

int tpc_encoder::get_output_size() {
    return outputSize;
}

int tpc_encoder::get_input_size() {
    return inputSize;
}

void tpc_encoder::block_conv_encode( std::vector<uint8_t> &output,
                               std::vector<uint8_t> input,
                               std::vector< std::vector<int> > transOutputVec,
                               std::vector< std::vector<int> > transNextStateVec,
                               std::vector<int> tail,
                               size_t KK,
                               size_t nn)
{
    size_t outsym, ii, jj;
    int state = 0;
    size_t LL = input.size();

    std::vector<int> binVec(nn,0);

    // encode data bits one bit at a time
    for (ii=0; ii<LL; ii++) {

        // determine the output symbol
        outsym = transOutputVec[(int)input[ii]][state];
        // determine the next state
        state = transNextStateVec[(int)input[ii]][state];

        // Convert symbol to a binary vector
        tpc_common::itob( binVec, outsym, nn );

        // Assign to output : TODO: investigate using memcpy for this?
        for (jj=0;jj<nn;jj++) {
            output[nn*ii+jj] = binVec[jj];
        }
    }

    // encode tail
    for (ii=LL;ii<LL+KK-1;ii++) {

        // determine the output symbol
        outsym = transOutputVec[tail[state]][state];
        // determine the next state
        state = transNextStateVec[tail[state]][state];

        // Convert symbol to a binary vector
        tpc_common::itob( binVec, outsym, nn );

        // Assign to output : TODO: investigate using memcpy for this?
        for (jj=0;jj<nn;jj++) {
            output[nn*ii+jj] = binVec[jj];
        }
    }
}

void tpc_encoder::generic_work(void *inBuffer, void *outBuffer) {
    const uint8_t *in = (const uint8_t *) inBuffer;
    uint8_t *out = (uint8_t *) outBuffer;

    size_t ii, jj;     // indexing var

    //DEBUG_PRINT_UCHAR_ARRAY(in, inputSize);

    // TODO: probably a better way to do this than memcpy?
    memcpy(&inputWithPad[d_bval+d_qval], in, sizeof(unsigned char)*inputSize);

    //DEBUG_PRINT("Input with Pad -->\n");
    //DEBUG_PRINT_UCHAR_ARRAY(&inputWithPad[0], inputSizeWithPad);
    //DEBUG_PRINT_F(fp, "Input with Pad -->\n");
    //DEBUG_PRINT_UCHAR_ARRAY_F(fp, &inputWithPad[0], inputSizeWithPad);

    // encode the row data
    for(ii=0; ii<numRowsToEncode; ii++) {
        // populate rowToEncode
        memcpy(&rowToEncode[0], &inputWithPad[ii*d_krow], sizeof(unsigned char)*d_krow);

        //DEBUG_PRINT("Encoding row=[%d] -->\n",ii);
        //DEBUG_PRINT_UCHAR_ARRAY(&rowToEncode[0], d_krow);
        //DEBUG_PRINT_F(fp, "Encoding row=[%d] -->\n",ii);
        //DEBUG_PRINT_UCHAR_ARRAY_F(fp, &rowToEncode[0], d_krow);

        // encode it
        block_conv_encode( rowEncoded_block,
                           rowToEncode,
                           rowOutputs,
                           rowNextStates,
                           rowTail,
                           rowEncoder_K,
                           rowEncoder_n);

        //DEBUG_PRINT("Row Encoded Block=[%d] -->\n",ii);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR(&rowEncoded_block[0], tmp);
        //DEBUG_PRINT_F(fp, "Row Encoded Block=[%d] -->\n",ii);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR_F(fp, &rowEncoded_block[0], tmp);

        // store only the encoded bits, b/c we read out the data in a special way
        memcpy(&rowEncodedBits[ii][0], &rowEncoded_block[d_krow], sizeof(uint8_t)*(rowEncoder_m*rowEncoder_n));

//         DEBUG_PRINT("Row Encoded Bits");
//         tmp = rowEncoder_m*rowEncoder_n;
//         DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR(&rowEncodedBits[ii][0], tmp);
    }

    // encode the column data
    size_t numDataColsToEncode = d_krow;
    size_t numCheckColsToEncode = numColsToEncode-numDataColsToEncode;
    for(ii=0; ii<numDataColsToEncode; ii++) {
        // populate colToEncode
        for(jj=0; jj<d_kcol; jj++) {
            colToEncode[jj] = inputWithPad[jj*d_krow+ii];
        }

        //DEBUG_PRINT("Encoding col=[%d] -->\n",ii);
        //DEBUG_PRINT_UCHAR_ARRAY(&colToEncode[0], d_kcol);
        //DEBUG_PRINT_F(fp, "Encoding col=[%d] -->\n",ii);
        //DEBUG_PRINT_UCHAR_ARRAY_F(fp, &colToEncode[0], d_kcol);

        // encode it
        block_conv_encode( colEncoded_block,
                           colToEncode,
                           colOutputs,
                           colNextStates,
                           colTail,
                           colEncoder_K,
                           colEncoder_n);

        //DEBUG_PRINT("Col Encoded Block=[%d] -->\n",ii);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR(&colEncoded_block[0], tmp);
        //DEBUG_PRINT_F(fp, "Col Encoded Block=[%d] -->\n",ii);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR_F(fp, &colEncoded_block[0], tmp);

        // store only the encoded bits, b/c we read the data out in a special way
        memcpy(&colEncodedBits[ii][0], &colEncoded_block[d_kcol], sizeof(uint8_t)*(colEncoder_m*colEncoder_n));

//         DEBUG_PRINT("Col Encoded Bits");
//         tmp = colEncoder_m*colEncoder_n;
//         DEBUG_PRINT_FLOAT_ARRAY(&colEncodedBits[ii][0], tmp);
    }

    // encode checks on checks (encode the row-encoded bits)
    for(ii=0; ii<numCheckColsToEncode; ii++) {
        // populate colToEncode
        for(jj=0; jj<d_kcol; jj++) {
            colToEncode[jj] = rowEncodedBits[jj][ii];       // indexing is weird b/c of the way we declared the vector :(
        }

        //DEBUG_PRINT("Encoding col=[%d] -->\n",ii+numDataColsToEncode);
        //DEBUG_PRINT_UCHAR_ARRAY(&colToEncode[0], d_kcol);
        //DEBUG_PRINT_F(fp, "Encoding col=[%d] -->\n",ii+numDataColsToEncode);
        //DEBUG_PRINT_UCHAR_ARRAY_F(fp, &colToEncode[0], d_kcol);

        // encode it
        block_conv_encode( colEncoded_block,
                           colToEncode,
                           colOutputs,
                           colNextStates,
                           colTail,
                           colEncoder_K,
                           colEncoder_n);

        //DEBUG_PRINT("Col Encoded Block=[%d] -->\n",ii+numDataColsToEncode);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR(&colEncoded_block[0], tmp);

        //DEBUG_PRINT_F(fp, "Col Encoded Block=[%d] -->\n",ii+numDataColsToEncode);
        //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR_F(fp, &colEncoded_block[0], tmp);

        // store only the encoded bits, b/c we read the data out in a special way
        memcpy(&colEncodedBits[ii+numDataColsToEncode][0], &colEncoded_block[d_kcol], sizeof(uint8_t)*(colEncoder_m*colEncoder_n));

//         DEBUG_PRINT("Col Encoded Bits");
//         tmp = colEncoder_m*colEncoder_n;
//         DEBUG_PRINT_FLOAT_ARRAY(&colEncodedBits[ii][0], tmp);
    }

    unsigned char* inputDataPtr;
    uint8_t *outputDataPtr = out;

    int curRowInRowEncodedBits = 0;
    // read out the data along the rows into the "out" array

    // skip B zeros & do the first row
    inputDataPtr =  &inputWithPad[d_bval];
    if(d_bval > d_krow){ throw std::runtime_error("bval must be < krow"); }
    size_t firstRowRemainingBits = d_krow-d_bval;
    for(ii=0; ii<firstRowRemainingBits; ii++) {
        *outputDataPtr++ = (uint8_t)(*inputDataPtr++);
    }

    // copy the encoded bits
    memcpy(outputDataPtr, &rowEncodedBits[curRowInRowEncodedBits++][0],
            sizeof(uint8_t)*(rowEncoder_m*rowEncoder_n));

    outputDataPtr += (rowEncoder_m*rowEncoder_n);

    // copy out the rest of the data
    for(ii=1; ii<d_kcol; ii++) {        // ii starts at 1, b/c we already did idx=0
        // copy systematic bits
        for(jj=0; jj<d_krow; jj++) {
            *outputDataPtr++ = (uint8_t)(*inputDataPtr++);
        }

        // copy the encoded bits
        memcpy(outputDataPtr, &rowEncodedBits[curRowInRowEncodedBits++][0],
               sizeof(uint8_t)*(rowEncoder_m*rowEncoder_n));

        outputDataPtr += (rowEncoder_m*rowEncoder_n);
    }

    // copy the encoded cols
    for(ii=0; ii<(colEncoder_m*colEncoder_n); ii++) {
        // copy checks
        for(jj=0; jj<d_krow; jj++) {
            *outputDataPtr++ = colEncodedBits[jj][ii];
        }
        int kk = jj;
        // copy checks on checks
        for(jj=0; jj<(rowEncoder_m*rowEncoder_n); jj++) {
            *outputDataPtr++ = colEncodedBits[kk++][ii];
        }
    }

    //DEBUG_PRINT("Output\n");
    //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR(out, outputSize);
    //DEBUG_PRINT_F(fp, "Output\n");
    //DEBUG_PRINT_FLOAT_ARRAY_AS_UCHAR_F(fp, out, outputSize);
}

tpc_encoder::~tpc_encoder()
{
    if(fp)
      fclose(fp);
}

}
}


#include "ExplorationPlanner/planner/ParticleSet.h"
#include <boost/random/uniform_01.hpp>
#include <cmath>
#include <sstream>

void ParticleSet::weightParticles(const std::vector<TrajectoryValueHandler>& h)
{
    assert( h.size() == getNumberOfParticles() );
    // Get min+max rewards
    double rmin = std::numeric_limits<double>::max();
    double rmax = -std::numeric_limits<double>::max();
    for ( auto it = h.begin(), iend = h.end(); it != iend; ++it)
    {
        rmin = std::min(rmin, it->getMinSumOfDiscRew() );
        rmax = std::max(rmax, it->getMaxSumOfDiscRew() );
    }

    double scaleFactor = (rmax - rmin);
    if (rmax == rmin)
        scaleFactor = 1;

    // Now we can scale the particles
    for ( unsigned int i = 0; i < getNumberOfParticles(); ++i )
    {
        double w = h[i].getWeightFactor(rmin, scaleFactor);
        particles_[i].setWeight( w * particles_[i].getWeight()  );
    }
    normalizeWeights();
}

void ParticleSet::resample(RNG& rng)
{
    if (getNumberOfParticles() <= 1)
        return;

    const double interval = getSumOfWeights() / getNumberOfParticles();
    boost::random::uniform_01<double> d;
    double target = interval * rng(d); //Initial target weight

    // Compute resampled vector
    std::vector<Particle> resampledParticles(getNumberOfParticles(), particles_.front() );
    double cumw = 0;
    unsigned int n = 0;
    for (auto it=particles_.begin(), iend=particles_.end(); it != iend; ++it)
    {
        cumw += it->getWeight();
        while (cumw > target)
        {
            resampledParticles[n++] = *it; // Postfix increment is crucial.
            target += interval;
        }
    }

    // Assign the new particles and normalize weights
    particles_ = resampledParticles;
    normalizeWeights();
}

Particle ParticleSet::getMaxWeightParticle() const
{
    Particle p;
    double w = -std::numeric_limits<double>::max();
    for (auto it = particles_.begin(), iend = particles_.end(); it != iend; ++it)
    {
        double cw = it->getWeight();
        if (w < cw)
        {
            w = cw;
            p = *it;
        }
    }
    return p;
}

double ParticleSet::getEffectiveNumberOfParticles() const
{
    double w_sq = 0.0;
    for (auto it = particles_.begin(), iend = particles_.end(); it != iend; ++it)
        w_sq += std::pow(it->getWeight(), 2);
    return (1.0/w_sq);
}

void ParticleSet::normalizeWeights()
{
    double w = getSumOfWeights();
    for (auto it = particles_.begin(), iend = particles_.end(); it != iend; ++it)
        it->setWeight( it->getWeight() / w  );
}

double ParticleSet::getSumOfWeights() const
{
    double w = 0.0;
    for (auto it = particles_.begin(), iend = particles_.end(); it != iend; ++it)
        w += it->getWeight();
    return w;
}

std::string ParticleSet::Print() const
{
    std::stringstream ss;
    const unsigned int sz = particles_.size();
    ss << "Particle set with " << sz << " particles\n";
    for (unsigned int i = 0; i < sz; ++i)
        ss << "Particle " << i << ":\n" << particles_[i].Print() << "\n";
    return ss.str();
}


//==================================================================================================
/**
  Copyright 2016 NumScale SAS
  Copyright 2016 J.T.Lapreste

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

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/detail/traits.hpp>
#include <boost/simd/meta/as_logical.hpp>
#include <boost/simd/function/genmask.hpp>
#include <boost/simd/function/hmsb.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;

  BOOST_DISPATCH_OVERLOAD_IF( all_
                            , (typename A0, typename X)
                            , (detail::is_native<X>)
                            , bd::cpu_
                            , bs::pack_<bd::fundamental_<A0>, X>
                            )
  {
    BOOST_FORCEINLINE bool operator()(const A0& a0) const BOOST_NOEXCEPT
    {
      using count = std::integral_constant<std::size_t, (1ull << A0::static_size)-1>;
      return hmsb(genmask(a0)) == count::value;
    }
  };
} } }

#endif


#include <iostream>
#include <ceres/ceres.h>
#include <ceres/rotation.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <algorithm>
#include <glog/logging.h>
#include <iterator>
#include <unordered_set>
#include <cmath>
#include <math.h>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>

#define EPS T(0.00001)

using namespace pybind11::literals;
namespace py = pybind11;
using namespace Eigen;

// use 13 dof features to do PCA
struct ReprojectionResidualFixedFNewPCA {
    ReprojectionResidualFixedFNewPCA(double *locs_2d_in, Matrix <double, 13, 5> &pca_comps_in,
                                     Matrix <double, 13, 1> &mean_shape_in, double cad_to_meter_in, double *img_size_in,
                                     double focal_in, int *valid_mask_per_obj_, double reprojection_scale_factor_in, double plane_scale_factor_in)
            : pca_comps(pca_comps_in), mean_shape(mean_shape_in), locs_2d(locs_2d_in),
            cad_to_meter(cad_to_meter_in), img_size(img_size_in), focal(focal_in),
            valid_mask_per_obj(valid_mask_per_obj_), reprojection_scale_factor(reprojection_scale_factor_in),
            plane_scale_factor(plane_scale_factor_in){}

    Matrix <double, 13, 5> &pca_comps;
    Matrix <double, 13, 1> &mean_shape;
    double *locs_2d;
    double cad_to_meter;
    double *img_size; // h, w
    double focal;
    int *valid_mask_per_obj;  // length of 12: 0 or 1
    double reprojection_scale_factor;
    double plane_scale_factor;

    template<typename T>
    bool operator()(const T *const rvec, const T *const tvec,
                    const T *const pca_coeff, const T *const plane, T *residual) const {
        // return the reprojection error
        Matrix <T, 5, 1> pca_coeff_;
        pca_coeff_<<pca_coeff[0], pca_coeff[1], pca_coeff[2], pca_coeff[3], pca_coeff[4];
        Matrix <T, 13, 1> shape_3d = mean_shape.cast <T>() + pca_comps.cast <T>() * pca_coeff_;

        Matrix <T, 12, 4> shape_full;
        T base_front_length = shape_3d[0] * T(cad_to_meter);
        T base_back_length = shape_3d[1] * T(cad_to_meter);
        T base_half_width = shape_3d[2] * T(cad_to_meter);
        T base_height = shape_3d[3] * T(cad_to_meter);
        T top_front_length = shape_3d[4] * T(cad_to_meter);
        T top_back_length = shape_3d[5] * T(cad_to_meter);
        T top_half_width = shape_3d[6] * T(cad_to_meter);
        T top_height = shape_3d[7] * T(cad_to_meter);
        T mid_front_length = shape_3d[8] * T(cad_to_meter);
        T mid_back_length = shape_3d[9] * T(cad_to_meter);
        T mid_front_height = shape_3d[10] * T(cad_to_meter);
        T mid_back_height = shape_3d[11] * T(cad_to_meter);
        T mid_half_width = shape_3d[12] * T(cad_to_meter);

        //# base four keypoints
        shape_full.row(0) = Matrix <T, 4, 1>(base_front_length, base_height, -base_half_width, T(1));
        shape_full.row(1) = Matrix <T, 4, 1>(base_front_length, base_height, base_half_width, T(1));
        shape_full.row(2) = Matrix <T, 4, 1>(base_back_length, base_height, -base_half_width, T(1));
        shape_full.row(3) = Matrix <T, 4, 1>(base_back_length, base_height, base_half_width, T(1));

        //# middle four keypoints
        shape_full.row(4) = Matrix <T, 4, 1>(mid_front_length, mid_front_height, -mid_half_width, T(1));
        shape_full.row(5) = Matrix <T, 4, 1>(mid_front_length, mid_front_height, mid_half_width, T(1));
        shape_full.row(6) = Matrix <T, 4, 1>(mid_back_length, mid_back_height, -mid_half_width, T(1));
        shape_full.row(7) = Matrix <T, 4, 1>(mid_back_length, mid_back_height, mid_half_width, T(1));

        //# top four keypoints
        shape_full.row(8) = Matrix <T, 4, 1>(top_front_length, top_height, -top_half_width, T(1));
        shape_full.row(9) = Matrix <T, 4, 1>(top_front_length, top_height, top_half_width, T(1));
        shape_full.row(10) = Matrix <T, 4, 1>(top_back_length, top_height, -top_half_width, T(1));
        shape_full.row(11) = Matrix <T, 4, 1>(top_back_length, top_height, top_half_width, T(1));
        Matrix <T, 4, 12> P = shape_full.transpose();

        Matrix <T, 3, 3> R;
        Matrix <T, 4, 4> Rt;
        T theta = sqrt(rvec[0] * rvec[0] + rvec[1] * rvec[1] + rvec[2] * rvec[2]);
        Matrix <T, 3, 3> omega;
        omega<<T(0), -rvec[2], rvec[1], rvec[2], T(0), -rvec[0], -rvec[1], rvec[0], T(0);
        R = (sin(theta) / theta) * omega + ((T(1) - cos(theta)) / (theta * theta)) * (omega * omega);
        R(0, 0) += T(1), R(1, 1) += T(1), R(2, 2) += T(1);

        /*
        Matrix<T,1,3> rvec1;
        rvec1(0,0) = T(1.0)*rvec[0];
        rvec1(0,1) = T(-1.0)*rvec[1];
        rvec1(0,2) = T(-1.0)*rvec[2];
        Rt.setZero();
        ceres::AngleAxisToRotationMatrix(rvec1.data(), R.data());
        */
        Rt(0, 0) = R(0, 0), Rt(0, 1) = R(0, 1), Rt(0, 2) = R(0, 2);
        Rt(1, 0) = R(1, 0), Rt(1, 1) = R(1, 1), Rt(1, 2) = R(1, 2);
        Rt(2, 0) = R(2, 0), Rt(2, 1) = R(2, 1), Rt(2, 2) = R(2, 2);
        Rt(0, 3) = tvec[0], Rt(1, 3) = tvec[1], Rt(2, 3) = tvec[2];
        Rt(3, 3) = T(1);

        Matrix <T, 3, 4> I;
        I.setIdentity();

        Matrix <T, 3, 3> K;
        K<<T(focal), T(0), T(img_size[1] / 2.0), T(0), T(focal), T(img_size[0] / 2.0), T(0), T(0), T(1);

        Matrix <T, 3, 12> P_cam = I * Rt * P;
        Matrix <T, 3, 12> p_homo = K * P_cam;

        float valid_sum = 0.0f;
        for (int i = 0; i<12; i++) {
            valid_sum += float(valid_mask_per_obj[i]);
        }

        T scale_inlier = T(1.0);
        T scale_outlier = T(1.0);
        T scale_factor = T(reprojection_scale_factor) / (T(valid_sum) * scale_inlier + T(12 - valid_sum) * scale_outlier);
        T p[24];
        for (int i = 0; i<12; i++) {
            p[2 * i] = p_homo(0, i) / p_homo(2, i);
            p[2 * i + 1] = p_homo(1, i) / p_homo(2, i);
        }

        T *locs_2d_est = p; // 12 x 2
        for (int i = 0; i<24; i++) {
            if (valid_mask_per_obj[i / 2] == 1)
                residual[i] = scale_factor * scale_inlier * (T(locs_2d[i]) - locs_2d_est[i]);
            else
                residual[i] = scale_factor * scale_outlier * (T(locs_2d[i]) - locs_2d_est[i]);
        }

        // four wheels should lie on the same plane
        T plane_scale = T(plane_scale_factor);
        for (int i = 0; i<4; i++) {
            T bottom_plane_dis = abs(plane[0] * P_cam(0, 0 + i) + plane[1] * P_cam(1, 0 + i)
                                     + plane[2] * P_cam(2, 0 + i) + plane[3]);
            bottom_plane_dis = bottom_plane_dis / sqrt(plane[0] * plane[0] + plane[1] * plane[1]
                                                       + plane[2] * plane[2]);
            residual[24 + i] = plane_scale * bottom_plane_dis;
        }

        return true;
    }
};

py::object optimize_pca_per_vehicle_new_pca(double focal, py::buffer rt_, py::buffer pca_, int num_obj, int num_veh,
                                            py::buffer pca_comps_, py::buffer mean_shape_, py::buffer locs_2d_,
                                            double cad_to_meter, py::buffer img_size_, py::buffer plane_,
                                            int num_iter, py::buffer mapping_, py::buffer kp_valid_mask_,
                                            double reprojection_scale_factor, double plane_scale_factor) {
    auto buf0 = rt_.request();    // num_obj x 6
    double *rt = static_cast<double *>(buf0.ptr);
    auto buf1 = pca_.request();  // num_veh x 5
    double *pca_start = static_cast<double *>(buf1.ptr);
    auto buf2 = pca_comps_.request();
    double *pca_comps = static_cast<double *>(buf2.ptr);
    auto buf3 = mean_shape_.request();
    double *mean_shape = static_cast<double *>(buf3.ptr);
    auto buf4 = locs_2d_.request();
    double *locs_2d_start = static_cast<double *>(buf4.ptr); // n x 12 x 2
    auto buf5 = img_size_.request();
    double *img_size = static_cast<double *>(buf5.ptr);
    auto buf6 = plane_.request();
    double *plane0 = static_cast<double *>(buf6.ptr);
    auto buf7 = mapping_.request();
    int *mapping = static_cast<int *>(buf7.ptr);
    auto buf8 = kp_valid_mask_.request();
    int *valid_mask = static_cast<int *>(buf8.ptr);

    double *locs_2d, *rvec, *tvec, *pca_coeff;

    Matrix<double, 5, 13>pca_comps_in = Map<Matrix<double, 5, 13>>(pca_comps);
    Matrix <double, 13, 5> pca_comps_in_t = pca_comps_in.transpose();
    Matrix<double, 13, 1>mean_shape_in = Map<Matrix<double, 13, 1>>(mean_shape);
    ceres::Problem problem;
    for (int i = 0; i<num_obj; i++) {
        rvec = rt + 6 * i;
        tvec = rvec + 3;
        int pca_index = *(mapping + i);
        pca_coeff = pca_start + pca_index * 5;
        locs_2d = locs_2d_start + i * 24;
        int *valid_mask_per_obj = valid_mask + i * 12;
        ceres::CostFunction *cost_function;
        cost_function = new ceres::AutoDiffCostFunction <ReprojectionResidualFixedFNewPCA, 28, 3, 3, 5, 4>(
                new ReprojectionResidualFixedFNewPCA(locs_2d,
                                                     pca_comps_in_t,
                                                     mean_shape_in,
                                                     cad_to_meter,
                                                     img_size,
                                                     focal,
                                                     valid_mask_per_obj,
                                                     reprojection_scale_factor, plane_scale_factor)
        );
        problem.AddResidualBlock(cost_function, NULL, rvec, tvec, pca_coeff, plane0);
        for (int j = 0; j<5; j++) {
            problem.SetParameterLowerBound(pca_coeff, j, -0.15);
            problem.SetParameterUpperBound(pca_coeff, j, 0.15);
        }
    }

    ceres::Solver::Options options;
    options.max_num_iterations = num_iter;
    options.minimizer_progress_to_stdout = false;
    options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;  //ceres::SPARSE_SCHUR;  //ceres::DENSE_SCHUR;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    std::cout<<summary.FullReport()<<"\n";
    // return optimized values
    py::array_t <double> x_rt = py::array_t <double>(
            py::buffer_info(
                    rt,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (num_obj * 6)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );

    py::array_t <double> x_pca = py::array_t <double>(
            py::buffer_info(
                    pca_start,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (num_veh * 5)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );

    py::array_t <double> plane = py::array_t <double>(
            py::buffer_info(
                    plane0,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (4)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );
    py::list outputs;
    outputs.append(x_rt);
    outputs.append(x_pca);
    outputs.append(plane);
    return outputs;
}

PYBIND11_PLUGIN(ceres_reconstruct) {
        py::module m("ceres_reconstruct", "Python bindings to the Ceres-Solver minimizer.");
        // google::InitGoogleLogging("ceres_reconstruct");
        m.def("optimize_pca_per_vehicle_new_pca", &optimize_pca_per_vehicle_new_pca, "Optimizes with pca and rt with ransac");
        return m.ptr();
}


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 *    Notes
 *      The test cases in this unit test should be more extensive.
 *
 */

#include <boost/bind.hpp>
#include <boost/make_shared.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/unit_test.hpp>

#include "Tudat/Mathematics/RootFinders/newtonRaphson.h"
#include "Tudat/Mathematics/RootFinders/UnitTests/testFunction1.h"
#include "Tudat/Mathematics/RootFinders/UnitTests/testFunction2.h"
#include "Tudat/Mathematics/RootFinders/UnitTests/testFunction3.h"
#include "Tudat/Mathematics/RootFinders/UnitTests/testFunctionWithLargeRootDifference.h"
#include "Tudat/Mathematics/RootFinders/UnitTests/testFunctionWithZeroRoot.h"

namespace tudat
{
namespace unit_tests
{

BOOST_AUTO_TEST_SUITE( testsuite_rootfinders )

using namespace tudat;
using namespace root_finders;
using namespace root_finders::termination_conditions;

//! Check if Newton-Raphson converges on test function #1 (TestFunction1).
BOOST_AUTO_TEST_CASE( test_newtonRaphson_testFunction1 )
{
    // Create object containing the test functions.
    std::shared_ptr< TestFunction1 > testFunction = std::make_shared< TestFunction1 >( 1 );

    // The termination condition.
    NewtonRaphson::TerminationFunction terminationConditionFunction =
            std::bind( &RootAbsoluteToleranceTerminationCondition< double >::checkTerminationCondition,
                         std::make_shared< RootAbsoluteToleranceTerminationCondition< double > >(
                             testFunction->getTrueRootAccuracy( ) ), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5 );

    // Test Newton-Raphson object.
    NewtonRaphson newtonRaphson( terminationConditionFunction );

    // Let Newton-Raphson search for the root.
    const double root = newtonRaphson.execute( testFunction, testFunction->getInitialGuess( ) );

    // Check if the result is within the requested accuracy.
    BOOST_CHECK_CLOSE_FRACTION( root, testFunction->getTrueRootLocation( ), 1.0e-15 );
    BOOST_CHECK_LT( testFunction->evaluate( root ), testFunction->getTrueRootAccuracy( ) );
}

//! Check if Newton-Raphson converges on test function #2 (TestFunction2).
BOOST_AUTO_TEST_CASE( test_newtonRaphson_testFunction2 )
{
    // Create object containing the test functions.
    std::shared_ptr< TestFunction2 > testFunction = std::make_shared< TestFunction2 >( 1 );

    // The termination condition.
    NewtonRaphson::TerminationFunction terminationConditionFunction =
            std::bind( &RootAbsoluteToleranceTerminationCondition< double >::checkTerminationCondition,
                         std::make_shared< RootAbsoluteToleranceTerminationCondition< double > >(
                             testFunction->getTrueRootAccuracy( ) ), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5 );

    // Test Newton-Raphson object.
    NewtonRaphson newtonRaphson( terminationConditionFunction );

    // Let Newton-Raphson search for the root.
    const double root = newtonRaphson.execute( testFunction, testFunction->getInitialGuess( ) );

    // Check if the result is within the requested accuracy.
    BOOST_CHECK_CLOSE_FRACTION( root, testFunction->getTrueRootLocation( ), 1.0e-15 );
    BOOST_CHECK_LT( testFunction->evaluate( root ), testFunction->getTrueRootAccuracy( ) );
}

//! Check if Newton-Raphson converges on test function #3 (TestFunction3).
BOOST_AUTO_TEST_CASE( test_newtonRaphson_testFunction3 )
{
    // Create object containing the test functions.
    std::shared_ptr< TestFunction3 > testFunction = std::make_shared< TestFunction3 >( 1 );

    // The termination condition.
    NewtonRaphson::TerminationFunction terminationConditionFunction =
            std::bind( &RootAbsoluteToleranceTerminationCondition< double >::checkTerminationCondition,
                         std::make_shared< RootAbsoluteToleranceTerminationCondition< double > >(
                             testFunction->getTrueRootAccuracy( ) ), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5 );

    // Test Newton-Raphson object.
    NewtonRaphson newtonRaphson( terminationConditionFunction );

    // Let Newton-Raphson search for the root.
    const double root = newtonRaphson.execute( testFunction, testFunction->getInitialGuess( ) );

    // Check if the result is within the requested accuracy.
    BOOST_CHECK_CLOSE_FRACTION( root, testFunction->getTrueRootLocation( ), 1.0e-15 );
    BOOST_CHECK_LT( testFunction->evaluate( root ), testFunction->getTrueRootAccuracy( ) );
}

//! Check if Newton-Raphson converges on function with large root difference
//! (testFunctionWithLargeRootDifference).
// Not the best test case. Inheritance from old code.
BOOST_AUTO_TEST_CASE( test_newtonRaphson_testFunctionWithLargeRootDifference )
{
    // Declare tolerance.
    const double tolerance = 1.0e-10;

    // Declare expected roots.
    const double expectedRootLowCase = 1.00000000793634;
    const double expectedRootHighCase = 7937.3386333591;

    // Create objects containing the test functions. Values were obtained during a limit case
    // gravity assist calculation (while evaluating Cassini-1 trajectory).
    std::shared_ptr< TestFunctionWithLargeRootDifference > testFunctionLowCase =
            std::make_shared< TestFunctionWithLargeRootDifference >
            ( 1, -3.24859999867635e18, -3248600.0, 1.5707963267949 );
    std::shared_ptr< TestFunctionWithLargeRootDifference > testFunctionHighCase =
            std::make_shared< TestFunctionWithLargeRootDifference >
            ( 1, -3248600.0, -3.24859999867635e18, 1.5707963267949 );

    // The termination condition.
    NewtonRaphson::TerminationFunction terminationConditionFunction
            = std::bind( &RootRelativeToleranceTerminationCondition< >::checkTerminationCondition,
                           std::make_shared< RootRelativeToleranceTerminationCondition< > >(
                               1.0e-10 ), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5 );

    // Make the Newton-Raphson object.
    NewtonRaphson newtonRaphson( terminationConditionFunction );

    // Let Newton-Raphson search for the root for both cases.
    double rootLowCase = newtonRaphson.execute( testFunctionLowCase, 1.0 + 1.0e-10 );
    double rootHighCase = newtonRaphson.execute( testFunctionHighCase, 1.0 + 1.0e-2 );

    // Check if the result is within the requested accuracy.
    BOOST_CHECK_CLOSE_FRACTION( rootLowCase, expectedRootLowCase, tolerance );
    BOOST_CHECK_CLOSE_FRACTION( rootHighCase, expectedRootHighCase, tolerance );
}

//! Check if Newton-Raphson converges on function with zero root (testFunctionWithZeroRoot).
// Not the best test case. Inheritance from old code. Not really relevant anymore. The basic
// idea is that Newton-Raphson should work for both a function that becomes zero, as well as for a
// function that does not become zero. A better case should be written.
BOOST_AUTO_TEST_CASE( test_newtonRaphson_testFunctionWithZeroRoot )
{
    // Create object containing the test functions.
    std::shared_ptr< TestFunctionWithZeroRoot > testFunction =
            std::make_shared< TestFunctionWithZeroRoot >( 1 );

    // The termination condition.
    NewtonRaphson::TerminationFunction terminationConditionFunction
            = std::bind(
                &RootAbsoluteOrRelativeToleranceTerminationCondition< >::
                checkTerminationCondition,
                std::make_shared< RootAbsoluteOrRelativeToleranceTerminationCondition< > >(
                    1.0e-308, 1.0e-15 ), std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5 );

    // Test Newton-Raphson object.
    NewtonRaphson newtonRaphson( terminationConditionFunction );

    // Let Newton-Raphson search for the root.
    const double root = newtonRaphson.execute( testFunction, testFunction->getInitialGuess( ) );

    // Check if the result is within the requested accuracy.
    BOOST_CHECK_SMALL( root, 1.0e-150 );
    BOOST_CHECK_LT( testFunction->evaluate( root ), testFunction->getTrueRootAccuracy( ) );
}

BOOST_AUTO_TEST_SUITE_END( ) // testsuite_rootfinders

} // namespace unit_tests
} // namespace tudat


#include <Python.h>
#include <pybind11/stl.h>
#include <pybind11/eigen.h>
#include <arrow/python/pyarrow.h>
#include <arrow/python/platform.h>
#include <arrow/api.h>
#include <factors/continuous/LinearGaussianCPD.hpp>
#include <factors/continuous/CKDE.hpp>
#include <factors/discrete/DiscreteFactor.hpp>
#include <models/BayesianNetwork.hpp>
#include <dataset/dataset.hpp>
#include <util/bit_util.hpp>
#include <util/math_constants.hpp>
#include <util/arrow_macros.hpp>
#include <Eigen/Dense>
#include <learning/parameters/mle_base.hpp>
#include <boost/math/distributions/normal.hpp>

namespace py = pybind11;

using arrow::Type;
using Eigen::Matrix, Eigen::Dynamic, Eigen::Array;

using boost::math::normal_distribution;
using dataset::DataFrame;
using factors::discrete::DiscreteFactorType;
using learning::parameters::MLE;
using models::BayesianNetworkBase, models::ConditionalBayesianNetworkBase;
using util::pi, util::one_div_root_two;

typedef std::shared_ptr<arrow::Array> Array_ptr;

namespace factors::continuous {

std::shared_ptr<Factor> LinearGaussianCPDType::new_factor(const BayesianNetworkBase& m,
                                                          const std::string& variable,
                                                          const std::vector<std::string>& evidence,
                                                          py::args args,
                                                          py::kwargs kwargs) const {
    for (const auto& e : evidence) {
        if (m.node_type(e) == DiscreteFactorType::get()) {
            return generic_new_factor<CLinearGaussianCPD>(variable, evidence, args, kwargs);
        }
    }

    return generic_new_factor<LinearGaussianCPD>(variable, evidence, args, kwargs);
}

std::shared_ptr<Factor> LinearGaussianCPDType::new_factor(const ConditionalBayesianNetworkBase& m,
                                                          const std::string& variable,
                                                          const std::vector<std::string>& evidence,
                                                          py::args args,
                                                          py::kwargs kwargs) const {
    for (const auto& e : evidence) {
        if (m.node_type(e) == DiscreteFactorType::get()) {
            return generic_new_factor<CLinearGaussianCPD>(variable, evidence, args, kwargs);
        }
    }

    return generic_new_factor<LinearGaussianCPD>(variable, evidence, args, kwargs);
}

LinearGaussianCPD::LinearGaussianCPD(std::string variable, std::vector<std::string> evidence)
    : Factor(variable, evidence), m_fitted(false), m_beta(), m_variance(-1){};

LinearGaussianCPD::LinearGaussianCPD(std::string variable,
                                     std::vector<std::string> evidence,
                                     VectorXd beta,
                                     double variance)
    : Factor(variable, evidence), m_fitted(true), m_variance(variance) {
    if (static_cast<size_t>(beta.rows()) != (evidence.size() + 1)) {
        throw std::invalid_argument(
            "Wrong number of beta parameters. Beta vector size: " + std::to_string(beta.size()) +
            ". Expected beta vector size: " + std::to_string(evidence.size() + 1) + ".");
    }

    if (variance <= 0) {
        throw std::invalid_argument("Variance must be a positive value.");
    }

    m_beta = beta;
};

void LinearGaussianCPD::fit(const DataFrame& df) {
    MLE<LinearGaussianCPD> mle;

    auto params = mle.estimate(df, this->variable(), this->evidence());

    m_beta = params.beta;
    m_variance = params.variance;
    m_fitted = true;
}

template <typename ArrowType>
Matrix<typename ArrowType::c_type, Dynamic, 1> logl_impl(const DataFrame& df,
                                                         const VectorXd& beta,
                                                         double variance,
                                                         const std::string& var,
                                                         const std::vector<std::string>& evidence) {
    using CType = typename ArrowType::c_type;
    using ArrayVecType = Array<CType, Dynamic, 1>;

    double inv_std = 1 / std::sqrt(variance);

    ArrayVecType logl;
    auto var_array = df.to_eigen<false, ArrowType, false>(var);
    if (evidence.empty()) {
        logl = -0.5 * (inv_std * (var_array->array() - beta(0))).square();
    } else {
        ArrayVecType means = ArrayVecType::Constant(df->num_rows(), beta(0));
        int idx = 1;
        for (auto it = evidence.begin(); it != evidence.end(); ++it, ++idx) {
            auto ev_array = df.to_eigen<false, ArrowType, false>(*it);
            means += static_cast<CType>(beta[idx]) * ev_array->array();
        }

        logl = -0.5 * (inv_std * (var_array->array() - means)).square();
    }

    logl += -0.5 * std::log(variance) - 0.5 * std::log(2 * pi<CType>);
    return logl.matrix();
}

template <typename ArrowType>
Matrix<typename ArrowType::c_type, Dynamic, 1> logl_impl_null(const DataFrame& df,
                                                              const VectorXd& beta,
                                                              double variance,
                                                              const std::string& var,
                                                              const std::vector<std::string>& evidence) {
    using CType = typename ArrowType::c_type;

    auto logl = logl_impl<ArrowType>(df, beta, variance, var, evidence);
    auto combined_bitmap = df.combined_bitmap(var, evidence);
    auto bitmap_data = combined_bitmap->data();
    auto logl_ptr = logl.data();

    for (int i = 0; i < df->num_rows(); ++i) {
        if (!util::bit_util::GetBit(bitmap_data, i)) logl_ptr[i] = util::nan<CType>;
    }

    return logl;
}

template <typename ArrowType>
double slogl_impl(const DataFrame& df,
                  const VectorXd& beta,
                  double variance,
                  const std::string& var,
                  const std::vector<std::string>& evidence) {
    return logl_impl<ArrowType>(df, beta, variance, var, evidence).sum();
}

template <typename ArrowType>
double slogl_impl_null(const DataFrame& df,
                       const VectorXd& beta,
                       double variance,
                       const std::string& var,
                       const std::vector<std::string>& evidence) {
    auto logl = logl_impl<ArrowType>(df, beta, variance, var, evidence);

    auto combined_bitmap = df.combined_bitmap(var, evidence);
    auto bitmap_data = combined_bitmap->data();
    auto logl_ptr = logl.data();

    double accum = 0;
    for (int i = 0; i < df->num_rows(); ++i) {
        if (util::bit_util::GetBit(bitmap_data, i)) accum += logl_ptr[i];
    }

    return accum;
}

template <typename ArrowType>
Matrix<typename ArrowType::c_type, Dynamic, 1> cdf_impl(const DataFrame& df,
                                                        const VectorXd& beta,
                                                        double variance,
                                                        const std::string& var,
                                                        const std::vector<std::string>& evidence) {
    using CType = typename ArrowType::c_type;
    using ArrayVecType = Array<CType, Dynamic, 1>;

    CType inv_std = static_cast<CType>(1. / std::sqrt(variance));
    auto var_ptr = df.to_eigen<false, ArrowType, false>(var);
    auto var_array = *var_ptr;

    ArrayVecType t(df->num_rows());
    if (evidence.empty()) {
        for (int64_t i = 0, end = df->num_rows(); i < end; ++i) {
            t(i) = std::erfc((beta(0) - var_array(i)) * inv_std * one_div_root_two<CType>);
        }
    } else {
        ArrayVecType means = ArrayVecType::Constant(df->num_rows(), beta[0]);

        int idx = 1;
        for (auto it = evidence.begin(); it != evidence.end(); ++it, ++idx) {
            auto ev_array = df.to_eigen<false, ArrowType, false>(*it);
            means += static_cast<CType>(beta[idx]) * ev_array->array();
        }

        for (int64_t i = 0, end = df->num_rows(); i < end; ++i) {
            t(i) = std::erfc((means(i) - var_array(i)) * inv_std * one_div_root_two<CType>);
        }
    }

    return (0.5 * t).matrix();
}

template <typename ArrowType>
Matrix<typename ArrowType::c_type, Dynamic, 1> cdf_impl_null(const DataFrame& df,
                                                             const VectorXd& beta,
                                                             double variance,
                                                             const std::string& var,
                                                             const std::vector<std::string>& evidence) {
    using CType = typename ArrowType::c_type;
    using ArrayVecType = Array<CType, Dynamic, 1>;

    CType inv_std = static_cast<CType>(1. / std::sqrt(variance));
    auto combined_bitmap = df.combined_bitmap(var, evidence);
    auto bitmap_data = combined_bitmap->data();

    auto var_ptr = df.to_eigen<false, ArrowType, false>(var);
    auto var_array = *var_ptr;

    ArrayVecType t(df->num_rows());
    if (evidence.empty()) {
        for (int64_t i = 0, end = df->num_rows(); i < end; ++i) {
            if (util::bit_util::GetBit(bitmap_data, i))
                t(i) = std::erfc((beta(0) - var_array(i)) * inv_std * one_div_root_two<CType>);
            else
                t(i) = util::nan<CType>;
        }
    } else {
        auto non_null = util::bit_util::non_null_count(combined_bitmap, df->num_rows());
        ArrayVecType means = ArrayVecType::Constant(non_null, beta[0]);

        int idx = 1;
        for (auto it = evidence.begin(); it != evidence.end(); ++it, ++idx) {
            auto ev_array = df.to_eigen<false, ArrowType>(combined_bitmap, *it);
            means += static_cast<CType>(beta[idx]) * ev_array->array();
        }

        for (int64_t i = 0, k = 0, end = df->num_rows(); i < end; ++i) {
            if (util::bit_util::GetBit(bitmap_data, i))
                t(i) = std::erfc((means(k++) - var_array(i)) * inv_std * one_div_root_two<CType>);
            else
                t(i) = util::nan<CType>;
        }
    }

    return (0.5 * t).matrix();
}

VectorXd LinearGaussianCPD::logl(const DataFrame& df) const {
    check_fitted();
    switch (df.same_type(this->variable(), this->evidence())->id()) {
        case Type::DOUBLE: {
            if (df.null_count(this->variable(), this->evidence()) == 0)
                return logl_impl<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
            else
                return logl_impl_null<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
        }
        case Type::FLOAT: {
            if (df.null_count(this->variable(), this->evidence()) == 0) {
                auto t = logl_impl<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
                return t.template cast<double>();
            } else {
                auto t = logl_impl_null<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
                return t.template cast<double>();
            }
        }
        default:
            throw std::invalid_argument("Wrong data type to compute logl. [double] or [float] data is expected.");
    }
}

double LinearGaussianCPD::slogl(const DataFrame& df) const {
    check_fitted();
    switch (df.same_type(this->variable(), this->evidence())->id()) {
        case Type::DOUBLE: {
            if (df.null_count(this->variable(), this->evidence()) == 0)
                return slogl_impl<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
            else
                return slogl_impl_null<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
        }
        case Type::FLOAT: {
            if (df.null_count(this->variable(), this->evidence()) == 0)
                return slogl_impl<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
            else
                return slogl_impl_null<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
        }
        default:
            throw std::invalid_argument("Wrong data type to compute slogl. [double] or [float] data is expected.");
    }
}

VectorXd LinearGaussianCPD::cdf(const DataFrame& df) const {
    check_fitted();
    switch (df.same_type(this->variable(), this->evidence())->id()) {
        case Type::DOUBLE: {
            if (df.null_count(this->variable(), this->evidence()) == 0)
                return cdf_impl<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
            else
                return cdf_impl_null<arrow::DoubleType>(df, m_beta, m_variance, this->variable(), this->evidence());
        }
        case Type::FLOAT: {
            if (df.null_count(this->variable(), this->evidence()) == 0) {
                auto t = cdf_impl<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
                return t.template cast<double>();
            } else {
                auto t = cdf_impl_null<arrow::FloatType>(df, m_beta, m_variance, this->variable(), this->evidence());
                return t.template cast<double>();
            }
        }
        default:
            throw std::invalid_argument("Wrong data type to compute cdf. [double] or [float] data is expected.");
    }
}

Array_ptr LinearGaussianCPD::sample(int n, const DataFrame& evidence_values, unsigned int seed) const {
    if (n < 0) {
        throw std::invalid_argument("n should be a non-negative number");
    }

    check_fitted();
    arrow::NumericBuilder<arrow::DoubleType> builder;
    RAISE_STATUS_ERROR(builder.Resize(n));

    std::mt19937 rng{seed};
    std::normal_distribution<> normal(m_beta(0), std::sqrt(m_variance));

    for (auto i = 0; i < n; ++i) {
        builder.UnsafeAppend(normal(rng));
    }

    std::shared_ptr<arrow::DoubleArray> out;
    RAISE_STATUS_ERROR(builder.Finish(&out));

    if (!this->evidence().empty()) {
        if (!evidence_values.has_columns(this->evidence()))
            throw std::domain_error("Evidence values not present for sampling.");

        auto out_values = reinterpret_cast<double*>(out->values()->mutable_data());
        const auto& e = this->evidence();
        for (size_t j = 0; j < e.size(); ++j) {
            auto evidence = evidence_values->GetColumnByName(e[j]);

            switch (evidence->type_id()) {
                case Type::DOUBLE: {
                    auto dwn_evidence = std::static_pointer_cast<arrow::DoubleArray>(evidence);
                    auto raw_evidence = dwn_evidence->raw_values();

                    for (auto i = 0; i < n; ++i) {
                        out_values[i] += m_beta(j + 1) * raw_evidence[i];
                    }
                    break;
                }
                case Type::FLOAT: {
                    auto dwn_evidence = std::static_pointer_cast<arrow::FloatArray>(evidence);
                    auto raw_evidence = dwn_evidence->raw_values();

                    for (auto i = 0; i < n; ++i) {
                        out_values[i] += m_beta(j + 1) * raw_evidence[i];
                    }
                    break;
                }
                default: {
                    throw std::invalid_argument("Wrong data type \"" + evidence->type()->ToString() +
                                                "\" for LinearGaussianCPD parent data.");
                }
            }
        }
    }

    return out;
}

std::string LinearGaussianCPD::ToString() const {
    std::stringstream stream;
    stream << std::fixed << std::setprecision(3);
    if (!this->evidence().empty()) {
        const auto& e = this->evidence();
        stream << "[LinearGaussianCPD] P(" << this->variable() << " | " << e[0];
        for (size_t i = 1; i < e.size(); ++i) {
            stream << ", " << e[i];
        }

        if (m_fitted) {
            stream << ") = N(" << m_beta(0);
            for (size_t i = 0; i < e.size(); ++i) {
                stream << " + " << m_beta(i + 1) << "*" << e[i];
            }
            stream << ", " << m_variance << ")";
        } else {
            stream << ") not fitted";
        }
    } else {
        if (m_fitted)
            stream << "[LinearGaussianCPD] P(" << this->variable() << ") = N(" << m_beta(0) << ", " << m_variance
                   << ")";
        else
            stream << "[LinearGaussianCPD] P(" << this->variable() << ") not fitted";
    }
    return stream.str();
}

py::tuple LinearGaussianCPD::__getstate__() const {
    return py::make_tuple(this->variable(), this->evidence(), m_fitted, m_beta, m_variance);
}

LinearGaussianCPD LinearGaussianCPD::__setstate__(py::tuple& t) {
    if (t.size() != 5) throw std::runtime_error("Not valid LinearGaussianCPD.");

    LinearGaussianCPD cpd(t[0].cast<std::string>(), t[1].cast<std::vector<std::string>>());

    cpd.m_fitted = t[2].cast<bool>();
    cpd.m_beta = t[3].cast<VectorXd>();
    cpd.m_variance = t[4].cast<double>();

    return cpd;
}

}  // namespace factors::continuous


//==================================================================================================
/**
  Copyright 2016 Numscale SAS

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

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/function/raw.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd =  boost::dispatch;
  namespace bs =  boost::simd;

  BOOST_DISPATCH_OVERLOAD ( rec_
                          , (typename A0)
                          , bs::sse2_
                          , bs::raw_tag
                          , bs::pack_<bd::double_<A0>, bs::sse_>
                          )
  {
    BOOST_FORCEINLINE A0 operator()(raw_tag const&, const A0 & a0) const BOOST_NOEXCEPT
    {
      return  _mm_cvtps_pd(_mm_rcp_ps(_mm_cvtpd_ps(a0))); //The maximum error for this approximation is 1.5e-12
    }
  };

} } }

#endif