//==============================================================================
//         Copyright 2003 - 2013   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2013   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#include <nt2/include/constants/i.hpp>

#include <boost/dispatch/functor/meta/call.hpp>
#include <nt2/sdk/functor/meta/call.hpp>
#include <nt2/sdk/unit/tests/relation.hpp>
#include <nt2/sdk/unit/tests/type_expr.hpp>
#include <complex>
#include <nt2/sdk/complex/complex.hpp>
#include <nt2/sdk/unit/module.hpp>
#include <boost/simd/sdk/config.hpp>

NT2_TEST_CASE_TPL ( i_real,  BOOST_SIMD_REAL_TYPES)
{
  typedef typename std::complex<T> cT;

  // specific values tests
   NT2_TEST_EQUAL(nt2::I<cT>(),  cT(0,1) );
}


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY 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 of the License,
// or (at your option) any later version.
//
// FlexibleSUSY 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 FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

#ifndef FIXED_POINT_ITERATOR_H
#define FIXED_POINT_ITERATOR_H

#include <iostream>
#include <limits>
#include <string>
#include <utility>
#include <Eigen/Core>
#include <gsl/gsl_sys.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_multiroots.h>

#include "logger.hpp"
#include "wrappers.hpp"
#include "error.hpp"
#include "ewsb_solver.hpp"
#include "gsl_utils.hpp"
#include "gsl_vector.hpp"

namespace flexiblesusy {

namespace fixed_point_iterator {

class Convergence_tester_absolute {
public:
   explicit Convergence_tester_absolute(double precision_ = 1.0e-2)
      : precision(precision_)
   {}

   std::string name() const { return "Convergence_tester_absolute"; }

   /**
    * Test whether the absolute value of the residual, defined by
    * \f$|a-b| = \sqrt{\sum_i (a_i - b_i)^2}\f$,
    * is less than the set precision.
    *
    * @param a GSL vector
    * @param b GSL vector
    * @return GSL error code (GSL_SUCCESS or GSL_CONTINUE)
    */
   int operator()(const GSL_vector& a, const GSL_vector& b) const {
      if (a.size() != b.size()) throw SetupError("Error: vectors have different size.");

      const auto dimension = a.size();
      double residual = 0.;

      if (precision < 0.)
         GSL_ERROR("absolute tolerance is negative", GSL_EBADTOL);

      for (std::size_t i = 0; i < dimension; ++i)
         residual += Sqr(a[i] - b[i]);

      residual = Sqrt(residual);

      return (residual < precision ? GSL_SUCCESS : GSL_CONTINUE);
   }

private:
   double precision;                 ///< precision goal
};

class Convergence_tester_relative {
public:
   explicit Convergence_tester_relative(double precision_ = 1.0e-2)
      : precision(precision_)
   {}

   std::string name() const { return "Convergence_tester_relative"; }

   /**
    * Test whether the relative difference is less than the set
    * precision. The relative difference test used here is carried out
    * by applying \a MaxRelDiff to each element of the vector.
    *
    * @param a GSL vector
    * @param b GSL vector
    * @return GSL error code (GSL_SUCCESS or GSL_CONTINUE)
    */
   int operator()(const GSL_vector& a, const GSL_vector& b) const {
      if (a.size() != b.size()) throw SetupError("Error: vectors have different size.");

      const auto dimension = a.size();
      double rel_diff = 0.;

      if (precision < 0.)
         GSL_ERROR("relative tolerance is negative", GSL_EBADTOL);

      for (std::size_t i = 0; i < dimension; ++i) {
         rel_diff = MaxRelDiff(a[i], b[i]);

         if (rel_diff > precision)
            return GSL_CONTINUE;
      }

      return GSL_SUCCESS;
   }

private:
   double precision;                 ///< precision goal
};

template <std::size_t dimension>
class Convergence_tester_tadpole {
public:
   using Vector_t = Eigen::Matrix<double,dimension,1>;
   using Function_t = std::function<Vector_t(const Vector_t&)>;

   Convergence_tester_tadpole(double precision_,
                              const Function_t& tadpole_function_)
      : precision(precision_)
      , tadpole_function(tadpole_function_)
   {}

   std::string name() const { return "Convergence_tester_tadpole"; }

   /**
    * Test whether the relative difference is less than the set
    * precision. The relative difference test used here is carried out
    * by applying \a MaxRelDiff to each element of the vector. If the
    * relative difference is below the precision, it is tested whether
    * the tadpoles are below the precision. If the tadpoles are larger
    * than the precision, GSL_CONTINUE is returned.
    *
    * @param a GSL vector
    * @param b GSL vector
    * @return GSL error code (GSL_SUCCESS or GSL_CONTINUE)
    */
   int operator()(const GSL_vector& a, const GSL_vector& b) const {
      if (a.size() != b.size()) throw SetupError("Error: vectors have different size.");

      if (precision < 0.)
         GSL_ERROR("relative tolerance is negative", GSL_EBADTOL);

      const double max_rel_diff =
         MaxRelDiff(to_eigen_vector(a), to_eigen_vector(b));

      if (max_rel_diff > precision)
         return GSL_CONTINUE;

      static const double eps = 10*std::pow(10., -std::numeric_limits<double>::digits10);

      if (max_rel_diff < eps)
         return GSL_SUCCESS;

      return check_tadpoles(a);
   }

private:
   double precision;                 ///< precision goal
   const Function_t tadpole_function; ///< function to calculate tadpole

   int check_tadpoles(const GSL_vector& x) const {
      const GSL_vector t(to_GSL_vector(tadpole_function(to_eigen_vector(x))));
      return gsl_multiroot_test_residual(t.raw(), precision);
   }
};

} // namespace fixed_point_iterator

/**
 * @class Fixed_point_iterator
 * @brief Does fixed point iteration
 * @author Dylan Harries, Alexander Voigt
 * @tparam dimension dimension of function
 * @tparam Convergence_tester function for relative comparison
 *    of subsequent iteration steps
 *
 * The user has to provide the function (of which a fixed point should
 * be found) of the type \a Function_t. This function gets as
 * arguments a Eigen vector of length \a dimension and returns a
 * vector with the next point.
 *
 * @note The standard relative convergence criterion
 * \f$\text{MaxRelDiff}(x_{n+1}, x_{n}) < \text{precision}\f$ is not
 * very good: The iteration might converge slowly.  This means, that
 * subsequent steps are very close to each other, but \f$x_n\f$ might
 * not be close to the true fixed point.
 *
 * @todo implement check for no progress towards solution
 */
template <std::size_t dimension, class Convergence_tester = fixed_point_iterator::Convergence_tester_relative>
class Fixed_point_iterator : public EWSB_solver {
public:
   using Vector_t = Eigen::Matrix<double,dimension,1>;
   using Function_t = std::function<Vector_t(const Vector_t&)>;

   Fixed_point_iterator() = default;
   template <typename F>
   Fixed_point_iterator(F&&, std::size_t, const Convergence_tester&);
   virtual ~Fixed_point_iterator() = default;

   template <typename F>
   void set_function(F&& f) { function = std::forward<F>(f); }
   void set_max_iterations(std::size_t n) { max_iterations = n; }
   int find_fixed_point(const Eigen::VectorXd&);

   // EWSB_solver interface methods
   virtual std::string name() const override { return "Fixed_point_iterator<" + convergence_tester.name() + ">"; }
   virtual int solve(const Eigen::VectorXd&) override;
   virtual Eigen::VectorXd get_solution() const override;

private:
   std::size_t max_iterations{100};         ///< maximum number of iterations
   GSL_vector xn{dimension};                ///< current iteration point
   GSL_vector fixed_point{dimension};       ///< vector of fixed point estimate
   Function_t function{nullptr};            ///< function defining fixed point
   Convergence_tester convergence_tester{}; ///< convergence tester

   int fixed_point_iterator_iterate();
   void print_state(std::size_t) const;
   static int gsl_function(const gsl_vector*, void*, gsl_vector*);
};

/**
 * Constructor
 *
 * @param function_ pointer to the function to find fixed point for
 * @param max_iterations_ maximum number of iterations
 * @param convergence_tester_ convergence tester
 */
template <std::size_t dimension, class Convergence_tester>
template <typename F>
Fixed_point_iterator<dimension,Convergence_tester>::Fixed_point_iterator(
   F&& function_,
   std::size_t max_iterations_,
   const Convergence_tester& convergence_tester_
)
   : max_iterations(max_iterations_)
   , function(std::forward<F>(function_))
   , convergence_tester(convergence_tester_)
{
}

/**
 * Start the iteration
 *
 * @param start starting point
 *
 * @return GSL error code (GSL_SUCCESS if fixed point found)
 */
template <std::size_t dimension, class Convergence_tester>
int Fixed_point_iterator<dimension,Convergence_tester>::find_fixed_point(
   const Eigen::VectorXd& start
)
{
   if (!function)
      throw SetupError("Fixed_point_iterator: function not callable");

   int status;
   std::size_t iter = 0;

#ifndef ENABLE_DEBUG
   gsl_set_error_handler_off();
#endif

   fixed_point = xn = to_GSL_vector(start);

#ifdef ENABLE_VERBOSE
   print_state(iter);
#endif

   do {
      iter++;
      status = fixed_point_iterator_iterate();

#ifdef ENABLE_VERBOSE
      print_state(iter);
#endif

      if (status)   // check if iterator has problems
         break;

      status = convergence_tester(fixed_point, xn);

   } while (status == GSL_CONTINUE && iter < max_iterations);

   VERBOSE_MSG("\t\t\tFixed_point_iterator status = "
               << gsl_strerror(status));

   return status;
}

/**
 * Perform a single step of the fixed point iteration
 *
 * @return GSL error code
 */
template <std::size_t dimension, class Convergence_tester>
int Fixed_point_iterator<dimension,Convergence_tester>::fixed_point_iterator_iterate()
{
   xn = fixed_point;

   void* parameters = &function;

   int status = gsl_function(xn.raw(), parameters, fixed_point.raw());

   if (status != GSL_SUCCESS)
      return GSL_EBADFUNC;

   // For safety, include a check for nans or infs here (which
   // should be sufficient for now)
   if (!is_finite(fixed_point))
      GSL_ERROR("update point is not finite", GSL_EBADFUNC);

   return GSL_SUCCESS;
}

/**
 * Print state of the fixed point iterator
 *
 * @param iteration iteration number
 */
template <std::size_t dimension, class Convergence_tester>
void Fixed_point_iterator<dimension,Convergence_tester>::print_state(std::size_t iteration) const
{
   VERBOSE_MSG("\t\t\tIteration n = " << iteration
               << ": x_{n} = " << xn
               << ", x_{n+1} = " << fixed_point);
}

template <std::size_t dimension, class Convergence_tester>
int Fixed_point_iterator<dimension,Convergence_tester>::gsl_function(const gsl_vector* x, void* params, gsl_vector* f)
{
   if (!is_finite(x)) {
      gsl_vector_set_all(f, std::numeric_limits<double>::max());
      return GSL_EDOM;
   }

   Function_t* fun = static_cast<Function_t*>(params);
   int status = GSL_SUCCESS;
   const Vector_t arg(to_eigen_vector(x));
   auto result = arg;

   try {
      result = (*fun)(arg);
      status = GSL_SUCCESS;
   } catch (const flexiblesusy::Error&) {
      status = GSL_EDOM;
   }

   copy(result, f);

   return status;
}

template <std::size_t dimension, class Convergence_tester>
int Fixed_point_iterator<dimension,Convergence_tester>::solve(const Eigen::VectorXd& start)
{
   return (find_fixed_point(start) == GSL_SUCCESS ?
           EWSB_solver::SUCCESS : EWSB_solver::FAIL);
}

template <std::size_t dimension, class Convergence_tester>
Eigen::VectorXd Fixed_point_iterator<dimension,Convergence_tester>::get_solution() const
{
   return to_eigen_vector(fixed_point);
}

} // namespace flexiblesusy

#endif


// Copyright Oleg Maximenko 2014.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//
// See http://github.com/svgpp/svgpp for library home page.

#pragma once

#include <boost/mpl/if.hpp>
#include <boost/type_traits/is_same.hpp>
#include <svgpp/config.hpp>
#include <svgpp/definitions.hpp>
#include <svgpp/detail/adapt_context.hpp>
#include <svgpp/parser/value_parser_fwd.hpp>
#include <svgpp/parser/external_function/parse_length.hpp>
#if !defined(SVGPP_USE_EXTERNAL_LENGTH_PARSER)
# include <svgpp/parser/external_function/parse_length_impl.hpp>
# include <svgpp/parser/grammar/length.hpp>
#endif
#include <svgpp/parser/detail/common.hpp>
#include <svgpp/parser/detail/finite_function_iterator.hpp>
#include <svgpp/parser/detail/parse_list_iterator.hpp>
#include <svgpp/parser/detail/value_parser_parameters.hpp>
#include <svgpp/traits/length_dimension_by_attribute.hpp>

namespace svgpp
{

template<SVGPP_TEMPLATE_ARGS>
struct value_parser<tag::type::length, SVGPP_TEMPLATE_ARGS_PASS>
{
  template<class AttributeTag, class Context, class AttributeValue, class PropertySource>
  static bool parse(AttributeTag tag, Context & context, AttributeValue const & attribute_value,
                                    PropertySource source)
  {
    typedef typename traits::length_dimension_by_attribute<AttributeTag>::type direction_t;
    typedef typename boost::range_const_iterator<AttributeValue>::type iterator_t;
    typedef detail::value_parser_parameters<Context, SVGPP_TEMPLATE_ARGS_PASS> args_t;
    typedef typename boost::parameter::parameters<
      boost::parameter::optional<tag::length_policy>
    >::template bind<SVGPP_TEMPLATE_ARGS_PASS>::type args2_t;
    typedef typename detail::unwrap_context<Context, tag::length_policy> length_policy_context;
    typedef typename length_policy_context::template bind<args2_t>::type length_policy_t;

    typename length_policy_t::length_factory_type & length_factory
      = length_policy_t::length_factory(length_policy_context::get(context));
    iterator_t it = boost::begin(attribute_value), end = boost::end(attribute_value);
    typename length_policy_t::length_factory_type::length_type value;
    if (detail::parse_length<direction_t, PropertySource>(length_factory, it, end, value)
      && it == end)
    {
      args_t::value_events_policy::set(args_t::value_events_context::get(context), tag, source, value);
      return true;
    }
    else
    {
      return args_t::error_policy::parse_failed(args_t::error_policy_context::get(context), tag, attribute_value);
    }
  }
};

template<SVGPP_TEMPLATE_ARGS>
struct value_parser<tag::type::list_of<tag::type::length>, SVGPP_TEMPLATE_ARGS_PASS>
{
  template<class AttributeTag, class Context, class AttributeValue, class PropertySource>
  static bool parse(AttributeTag tag, Context & context, AttributeValue const & attribute_value,
                                    PropertySource property_source)
  {
    typedef typename traits::length_dimension_by_attribute<AttributeTag>::type direction_t;
    typedef typename boost::range_const_iterator<AttributeValue>::type iterator_t;
    return parseT<
      detail::length_grammar_tag,
      AttributeTag, Context, AttributeValue, PropertySource, direction_t>
      (tag, context, attribute_value, property_source);
  }

protected:
  template<class LengthGrammarTag, class AttributeTag, class Context, class AttributeValue, class PropertySource, class Direction>
  static bool parseT(AttributeTag tag, Context & context, AttributeValue const & attribute_value,
                                    PropertySource property_source)
  {
    typedef typename boost::range_const_iterator<AttributeValue>::type iterator_t;
    typedef detail::value_parser_parameters<Context, SVGPP_TEMPLATE_ARGS_PASS> args_t;
    typedef typename boost::parameter::parameters<
      boost::parameter::optional<tag::length_policy>
    >::template bind<SVGPP_TEMPLATE_ARGS_PASS>::type args2_t;
    typedef typename detail::unwrap_context<Context, tag::length_policy> length_policy_context;
    typedef typename length_policy_context::template bind<args2_t>::type length_policy_t;
    typedef typename length_policy_t::length_factory_type length_factory_t;
    typedef qi::rule<iterator_t, typename length_factory_t::length_type()> length_rule_t;
    typedef detail::comma_wsp_rule_no_skip<iterator_t> separator_t;
    typedef detail::parse_list_iterator<
      typename length_factory_t::length_type,
      iterator_t,
      length_rule_t,
      separator_t
    > parse_list_iterator_t;
    typedef detail::finite_function_iterator<parse_list_iterator_t> output_iterator_t;

    length_factory_t & length_factory = length_policy_t::length_factory(length_policy_context::get(context));
    length_rule_t length_rule;
#ifdef SVGPP_USE_EXTERNAL_LENGTH_PARSER
    length_rule %= detail::get_length_rule<
      Direction,
      iterator_t,
      PropertySource
    >(LengthGrammarTag(), length_factory)(boost::phoenix::ref(length_factory));
#else
    SVGPP_STATIC_IF_SAFE const typename boost::mpl::if_<
      boost::is_same<LengthGrammarTag, detail::length_grammar_tag>,
      length_grammar<
        PropertySource,
        iterator_t,
        length_factory_t,
        Direction,
        typename length_factory_t::number_type
      >,
      percentage_or_length_css_grammar<
        iterator_t,
        length_factory_t,
        typename length_factory_t::number_type
      >
    >::type length_grammar;
    length_rule %= length_grammar(boost::phoenix::ref(length_factory));
#endif

    SVGPP_STATIC_IF_SAFE const separator_t separator_grammar;
    parse_list_iterator_t parse_list(
      boost::begin(attribute_value), boost::end(attribute_value),
      length_rule,
      separator_grammar);

    args_t::value_events_policy::set(args_t::value_events_context::get(context), tag, property_source,
      boost::make_iterator_range(output_iterator_t(parse_list), output_iterator_t()));
    if (parse_list.error())
    {
      return args_t::error_policy::parse_failed(args_t::error_policy_context::get(context), tag, attribute_value);
    }
    else
      return true;
  }
};

}


// Copyright © 2016-2021 Thomas Nagler and Thibault Vatter
//
// This file is part of the vinecopulib library and licensed under the terms of
// the MIT license. For a copy, see the LICENSE file in the root directory of
// vinecopulib or https://vinecopulib.github.io/vinecopulib/.

#pragma once

#include <memory>

#include <Eigen/Dense>
#include <vinecopulib/bicop/family.hpp>

namespace vinecopulib {
//! @brief An abstract class for bivariate copula families.
//!
//! This class is used in the implementation underlying the Bicop class.
//! Users should not use AbstractBicop or derived classes directly, but
//! always work with the Bicop interface.
class AbstractBicop
{
  friend class Bicop;

public:
  virtual ~AbstractBicop() = 0;

protected:
  // Factories
  static std::shared_ptr<AbstractBicop> create(
    BicopFamily family = BicopFamily::indep,
    const Eigen::MatrixXd& parameters = Eigen::MatrixXd());

  // Getters and setters
  BicopFamily get_family() const;

  std::string get_family_name() const;

  double get_loglik() const;

  void set_loglik(const double loglik = NAN);

  void set_var_types(const std::vector<std::string>& var_types);

  virtual Eigen::MatrixXd get_parameters() const = 0;

  virtual Eigen::MatrixXd get_parameters_lower_bounds() const = 0;

  virtual Eigen::MatrixXd get_parameters_upper_bounds() const = 0;

  virtual void set_parameters(const Eigen::MatrixXd& parameters) = 0;

  // Virtual methods
  virtual void fit(const Eigen::MatrixXd& data,
                   std::string method,
                   double mult,
                   const Eigen::VectorXd& weights) = 0;

  virtual double get_npars() const = 0;

  virtual void set_npars(const double& npars) = 0;

  virtual double parameters_to_tau(const Eigen::MatrixXd& parameters) = 0;

  virtual void flip() = 0;

  // following are virtual so they can be overriden by KernelBicop
  virtual Eigen::VectorXd pdf(const Eigen::MatrixXd& u);

  virtual Eigen::VectorXd cdf(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::VectorXd hfunc1(const Eigen::MatrixXd& u);

  virtual Eigen::VectorXd hfunc2(const Eigen::MatrixXd& u);

  Eigen::VectorXd hinv1(const Eigen::MatrixXd& u);

  Eigen::VectorXd hinv2(const Eigen::MatrixXd& u);

  virtual Eigen::VectorXd pdf_raw(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::VectorXd hfunc1_raw(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::VectorXd hfunc2_raw(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::VectorXd hinv1_raw(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::VectorXd hinv2_raw(const Eigen::MatrixXd& u) = 0;

  virtual Eigen::MatrixXd tau_to_parameters(const double& tau) = 0;
  Eigen::MatrixXd no_tau_to_parameters(const double&);

  // Misc methods
  Eigen::VectorXd hinv1_num(const Eigen::MatrixXd& u);

  Eigen::VectorXd hinv2_num(const Eigen::MatrixXd& u);

  Eigen::VectorXd pdf_c_d(const Eigen::MatrixXd& u);

  Eigen::VectorXd pdf_d_d(const Eigen::MatrixXd& u);

  double loglik(const Eigen::MatrixXd& u,
                const Eigen::VectorXd weights = Eigen::VectorXd());

  // Data members
  BicopFamily family_;
  double loglik_{ NAN };
  std::vector<std::string> var_types_{ "c", "c" };
};

//! A shared pointer to an object of class AbstracBicop.
typedef std::shared_ptr<AbstractBicop> BicopPtr;
}

#include <vinecopulib/bicop/implementation/abstract.ipp>


// __BEGIN_LICENSE__
// Copyright (C) 2006-2011 United States Government as represented by
// the Administrator of the National Aeronautics and Space Administration.
// All Rights Reserved.
// __END_LICENSE__

/// \file results.cc
///

/************************************************************************
 *     File: results.cc
 *     Date: December 2008
 *       By: Ara Nefian and Michael Broxton
 *      For: NASA Ames Research Center, Intelligent Mechanisms Group
 * Function: Main program for the measuring the performance and accuracy of the stereo pipeline
 ************************************************************************/
#include <boost/algorithm/string.hpp>
#include <boost/program_options.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/convenience.hpp>
#include <boost/filesystem/fstream.hpp>
using namespace boost;
namespace po = boost::program_options;
namespace fs = boost::filesystem;

#include <vw/Core.h>
#include <vw/Image.h>
#include <vw/Math.h>
#include <vw/FileIO.h>
#include <vw/Camera.h>
#include <vw/Stereo.h>

using namespace vw;
using namespace vw::math;
using namespace vw::camera;
using namespace vw::stereo;

#include <asp/Core/StereoSettings.h>
#include <asp/Sessions.h>

// Support for ISIS image files
#if defined(ASP_HAVE_PKG_ISIS) && ASP_HAVE_PKG_ISIS == 1
#include <asp/IsisIO/DiskImageResourceIsis.h>
#endif

// Boost headers
#include <boost/thread/xtime.hpp>
// Posix time is not fully supported in the version of Boost for RHEL
// Workstation 4
#ifdef __APPLE__
#include <boost/date_time/posix_time/posix_time.hpp>
#else
#include <ctime>
#endif

using namespace std;

//***********************************************************************
// MAIN
//***********************************************************************
int main(int argc, char* argv[]) {

  string true_prefix = argv[1];
  string pred_prefix = argv[2];

  //string true_prefix = "/Users/anefian/projects/stereopipeline/data/apollo/apollo-crop-small/results/apollo-crop-small";
  //string pred_prefix = "/Users/anefian/projects/stereopipeline/data/apollo/apollo-crop-small/results/apollo-crop-small";

  string true_disp_map_h_filename;
  string true_disp_map_v_filename;
  string pred_disp_map_h_filename;
  string pred_disp_map_v_filename;
  string good_pixels_map_filename;

  true_disp_map_h_filename = true_prefix+"-R-H.tif";
  true_disp_map_v_filename = true_prefix+"-R-V.tif";
  pred_disp_map_h_filename = pred_prefix+"-R-H.tif";
  pred_disp_map_v_filename = pred_prefix+"-R-V.tif";
  good_pixels_map_filename = pred_prefix+"-GoodPixelMap.tif";

  DiskImageView<PixelGray<float> > true_disp_map_h (true_disp_map_h_filename);
  DiskImageView<PixelGray<float> > true_disp_map_v (true_disp_map_v_filename);
  DiskImageView<PixelGray<float> > pred_disp_map_h (pred_disp_map_h_filename);
  DiskImageView<PixelGray<float> > pred_disp_map_v (pred_disp_map_v_filename);
  DiskImageView<PixelRGB<uint8> >  good_pixels_map (good_pixels_map_filename);

  //compute the average coverage measure
  int img_size = true_disp_map_h.rows() * true_disp_map_h.cols();
  int num_good_pixels = 0;
  int i, j;

  for (i = 0; i < good_pixels_map.rows(); i++){
    for (j = 0; j < good_pixels_map.cols(); j++){

      if (good_pixels_map(i,j).r()==200){
          num_good_pixels++;
      }
    }
  }

  float avg_cover = num_good_pixels/(float) img_size;
  printf("Average Coverage = %f\n", avg_cover);

  //compute the average error measure
  float avg_error_h = 0;

  for (i = 0; i < pred_disp_map_h.rows(); i++){
    for (j = 0; j < pred_disp_map_h.cols(); j++){
      float tmp = pred_disp_map_h(i,j).v()-true_disp_map_h(i,j).v();
       avg_error_h = avg_error_h + tmp*tmp;
    }
  }

  avg_error_h = avg_error_h/num_good_pixels;
  printf("avg_error_h = %f\n", avg_error_h);

  float avg_error_v = 0;

  for (i = 0; i < pred_disp_map_h.rows(); i++){
    for (j = 0; j < pred_disp_map_h.cols(); j++){
      float tmp = pred_disp_map_v(i,j).v()-true_disp_map_v(i,j).v();
       avg_error_v = avg_error_v + tmp*tmp;
    }
  }

  avg_error_v = avg_error_v/num_good_pixels;
  printf("avg_error_v = %f\n", avg_error_v);
  float avg_error = 0.5*(avg_error_v + avg_error_h);
  printf("avg_error = %f\n", avg_error);

  return(EXIT_SUCCESS);
}


#include <std_msgs/Float32MultiArray.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/PointCloud.h>
#include <sensor_msgs/point_cloud_conversion.h>
#include <geometry_msgs/Transform.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <tf/LinearMath/Matrix3x3.h>
#include <tf_conversions/tf_eigen.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl_ros/point_cloud.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/common/common_headers.h>
#include <pcl/features/normal_3d.h>
#include <pcl/features/normal_3d_omp.h>
#include <pcl/features/fpfh.h>
#include <pcl/features/fpfh_omp.h>
#include <pcl/keypoints/sift_keypoint.h>
#include <pcl/io/pcd_io.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/registration/correspondence_estimation.h>
#include <pcl/registration/correspondence_rejection_sample_consensus.h>
#include <pcl/registration/transformation_estimation_svd.h>
#include <pcl/filters/passthrough.h>
#include <pcl/registration/icp.h>
#include <pcl/console/parse.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/voxel_grid.h>
#include <iostream>
#include <thread>
#include <chrono>
#include <string>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/LU>
#include <Eigen/Eigen>
#include <Eigen/Eigenvalues>
#include <Eigen/SVD>
#include <Eigen/Dense>
#include <Eigen/StdVector>

pcl::PointCloud<pcl::PointXYZ>::Ptr ref_cloud_ptr(new pcl::PointCloud<pcl::PointXYZ>);
pcl::PointCloud<pcl::PointXYZ>::Ptr ref_main_cloud_ptr(new pcl::PointCloud<pcl::PointXYZ>);
pcl::PointCloud<pcl::PointXYZ>::Ptr scene_cloud_ptr(new pcl::PointCloud<pcl::PointXYZ>);
pcl::PointCloud<pcl::PointXYZ>::Ptr final_output(new pcl::PointCloud<pcl::PointXYZ>);
// Reference to Rough Transformation matrix coming from Pose Estimation using DL
Eigen::Matrix4f Rough_T_matrix_from_pose_est = Eigen::Matrix4f::Identity();

// Parameters for SIFT Computation
const float min_scale =  0.01f;
const int n_octaves = 8;
const int n_scales_per_octave = 4;
const float min_contrast =  0.001f;
double leaf_size = 0.005f;
unsigned int text_id = 0;
void keyboardEventOccurred (const pcl::visualization::KeyboardEvent &event,
                            void* viewer_void)
{
  pcl::visualization::PCLVisualizer *viewer = static_cast<pcl::visualization::PCLVisualizer *> (viewer_void);
  if (event.getKeySym () == "r" && event.keyDown ())
  {
    std::cout << "r was pressed => removing all text" << std::endl;

    char str[512];
    for (unsigned int i = 0; i < text_id; ++i)
    {
      sprintf (str, "text#%03d", i);
      viewer->removeShape (str);
    }
    text_id = 0;
  }
}

void mouseEventOccurred (const pcl::visualization::MouseEvent &event,
                         void* viewer_void)
{
  pcl::visualization::PCLVisualizer *viewer = static_cast<pcl::visualization::PCLVisualizer *> (viewer_void);
  if (event.getButton () == pcl::visualization::MouseEvent::LeftButton &&
      event.getType () == pcl::visualization::MouseEvent::MouseButtonRelease)
  {
    std::cout << "Left mouse button released at position (" << event.getX () << ", " << event.getY () << ")" << std::endl;

    char str[512];
    sprintf (str, "text#%03d", text_id ++);
    viewer->addText ("clicked here", event.getX (), event.getY (), str);
  }
}

void cloud_cb(const sensor_msgs::PointCloud2ConstPtr& input)
{
  pcl::fromROSMsg(*input, *scene_cloud_ptr);
}
pcl::visualization::PCLVisualizer::Ptr simpleVis (pcl::PointCloud<pcl::PointXYZ>::ConstPtr cloud)
{
  // --------------------------------------------
  // -----Open 3D viewer and add point cloud-----
  // --------------------------------------------
  pcl::visualization::PCLVisualizer::Ptr viewer (new pcl::visualization::PCLVisualizer ("3D Viewer"));
  viewer->setBackgroundColor (0, 0, 0);
  viewer->addPointCloud<pcl::PointXYZ> (cloud, "sample cloud");
  viewer->setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 1, "sample cloud");
  viewer->addCoordinateSystem (1.0);
  viewer->initCameraParameters ();
  return (viewer);
}

pcl::PointCloud<pcl::PointNormal>::Ptr estimateNormals(pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud_ptr,double p_normal_radius)
{
  double normal_radius = p_normal_radius;
  pcl::NormalEstimationOMP<pcl::PointXYZ, pcl::PointNormal> ne;
  pcl::PointCloud<pcl::PointNormal>::Ptr cloud_normals(new pcl::PointCloud<pcl::PointNormal>);
  pcl::search::KdTree<pcl::PointXYZ>::Ptr tree_n(new pcl::search::KdTree<pcl::PointXYZ>());
  std::cout<<"Started to compute normals in the function"<<std::endl;
  ne.setInputCloud(input_cloud_ptr);
  ne.setSearchMethod(tree_n);
  ne.setRadiusSearch(normal_radius);
  ne.compute(*cloud_normals);
  std::cout<<"Normals were computed"<<std::endl;
  for(int i = 0; i<cloud_normals->points.size(); ++i)
  {
    cloud_normals->points[i].x = input_cloud_ptr->points[i].x;
    cloud_normals->points[i].y = input_cloud_ptr->points[i].y;
    cloud_normals->points[i].z = input_cloud_ptr->points[i].z;
  }
  std::cout<<"Cloud normals are "<<cloud_normals->points.size()<<std::endl;
  return cloud_normals;
}

pcl::PointCloud<pcl::PointXYZ>::Ptr estimateKeypointsSIFT(pcl::PointCloud<pcl::PointNormal>::Ptr input_cloud_normal_ptr)
{
  pcl::SIFTKeypoint<pcl::PointNormal, pcl::PointWithScale> sift;
  pcl::PointCloud<pcl::PointWithScale> result;
  pcl::search::KdTree<pcl::PointNormal>::Ptr tree(new pcl::search::KdTree<pcl::PointNormal>());

  sift.setSearchMethod(tree);
  sift.setScales(min_scale, n_octaves, n_scales_per_octave);
  sift.setMinimumContrast(min_contrast);
  sift.setInputCloud(input_cloud_normal_ptr);
  sift.compute(result);

  // Copying the pointwithscale to pointxyz so as visualize the cloud
  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_temp_ptr(new pcl::PointCloud<pcl::PointXYZ>);
  copyPointCloud(result, *cloud_temp_ptr);
  std::cout << "SIFT keypoints in the cloud are " << cloud_temp_ptr->points.size () << std::endl;
  return cloud_temp_ptr;
}

pcl::PointCloud<pcl::FPFHSignature33>::Ptr fpfhExtractor(pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud_filtered_ptr, pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud_ptr, pcl::PointCloud<pcl::PointNormal>::Ptr input_normals_ptr, double p_fpfh_radius)
{
  pcl::FPFHEstimationOMP<pcl::PointXYZ, pcl::PointNormal, pcl::FPFHSignature33> fpfh;
  for (int i = 0; i < input_normals_ptr->points.size(); i++)
  {
    if (!pcl::isFinite<pcl::PointNormal>(input_normals_ptr->points[i]))
    {
      PCL_WARN("normals[%d] is not finite\n", i);
    }
  }
  double fpfh_radius = p_fpfh_radius;
  fpfh.setInputCloud(input_cloud_ptr);
  fpfh.setSearchSurface(input_cloud_filtered_ptr);
  fpfh.setInputNormals(input_normals_ptr);
  pcl::search::KdTree<pcl::PointXYZ>::Ptr tree (new pcl::search::KdTree<pcl::PointXYZ>());
  fpfh.setSearchMethod(tree);
  // Output datasets
  pcl::PointCloud<pcl::FPFHSignature33>::Ptr fpfhs(new pcl::PointCloud<pcl::FPFHSignature33>());
  // Radius here should be bigger than that radius used in estimation of normals
  fpfh.setRadiusSearch(fpfh_radius);
  fpfh.compute(*fpfhs);
  std::cout<<"FPFHs points size = "<<fpfhs->points.size()<<std::endl;
  std::cout<<"Input keypoint clouds size = "<<input_cloud_ptr->points.size()<<std::endl;
  return fpfhs;
}

void callbackMatrix(const geometry_msgs::Transform::ConstPtr& msg)
{
    Eigen::Matrix<double,3,3> tmat;
    float x = msg->translation.x;
    float y = msg->translation.y;
    float z = msg->translation.z;
    float beta_x = msg->rotation.x;
    float beta_y = msg->rotation.y;
    float beta_z = msg->rotation.z;
    float beta_w = msg->rotation.w;
    tf::Quaternion quat = tf::Quaternion(beta_x,beta_y,beta_z,beta_w);

    tf::Matrix3x3 rot(quat);
    //rot.getRotation(quat);
    tf::matrixTFToEigen(rot,tmat);
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
          std::cout<<tmat(i,j)<<endl;
          Rough_T_matrix_from_pose_est(i,j) = tmat(i,j);
        }
        std::cout<<"-----"<<std::endl;
    }
    Rough_T_matrix_from_pose_est(0,3) = x;
    Rough_T_matrix_from_pose_est(1,3) = y;
    Rough_T_matrix_from_pose_est(2,3) = z;
    Rough_T_matrix_from_pose_est(3,0) = 0;
    Rough_T_matrix_from_pose_est(3,1) = 0;
    Rough_T_matrix_from_pose_est(3,2) = 0;
    Rough_T_matrix_from_pose_est(3,3) = 1;
}

int main (int argc, char** argv)
{
  int cycle = 0;
  const std::string PCD_PATH = "/home/dlar/catkin_ws/body_visible_rotated.pcd";
  const std::string PCD_MAIN_PATH = "/home/dlar/catkin_ws/body_visible_rotated.pcd";

  // Initialize ROS
  ros::init (argc, argv, "point_cloud_register");
  ros::NodeHandle nh;
  ros::Rate r(20);
  // Create a ROS subscriber for the input point cloud from the sensor
  ros::Subscriber sub = nh.subscribe<sensor_msgs::PointCloud2>("/camera/depth_registered/points", 1, cloud_cb);
  ros::Publisher pub = nh.advertise<sensor_msgs::PointCloud2> ("xtion/depth_registered/registered_points", 20);
  ros::Subscriber T_matrix_sub = nh.subscribe<geometry_msgs::Transform>("/transformation_wrt_depth_optical_frame", 1, callbackMatrix);
  // Ros Message
  sensor_msgs::PointCloud2 output;
  // Import reference point clouds data
  if (pcl::io::loadPCDFile<pcl::PointXYZ> (PCD_PATH, *ref_cloud_ptr) == -1) //* load the file
  {
    PCL_ERROR ("Couldn't read file reference point cloud. \n");
    return (-1);
  }
  if (pcl::io::loadPCDFile<pcl::PointXYZ> (PCD_MAIN_PATH, *ref_main_cloud_ptr) == -1) //* load the file
  {
    PCL_ERROR ("Couldn't read file main reference point cloud \n");
    return (-1);
  }

  // Filtering Reference Cloud ---------------------------------------
  pcl::PointCloud<pcl::PointXYZ>::Ptr ref_cloud_filtered_ptr(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::VoxelGrid<pcl::PointXYZ> sor;
  sor.setInputCloud (ref_cloud_ptr);
  sor.setLeafSize (leaf_size, leaf_size, leaf_size);
  sor.filter (*ref_cloud_filtered_ptr);

  pcl::PointCloud<pcl::PointNormal>::Ptr ref_cloud_normal_ptr(new pcl::PointCloud<pcl::PointNormal>);
  pcl::PointCloud<pcl::PointNormal>::Ptr scene_cloud_normal_ptr(new pcl::PointCloud<pcl::PointNormal>);
  pcl::PointCloud<pcl::PointXYZ>::Ptr ref_SIFT_keypoints_ptr(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::PointCloud<pcl::PointXYZ>::Ptr scene_SIFT_keypoints_ptr(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::PointCloud<pcl::FPFHSignature33>::Ptr ref_fpfh_ptr(new pcl::PointCloud<pcl::FPFHSignature33>());
  pcl::PointCloud<pcl::FPFHSignature33>::Ptr scene_fpfh_ptr(new pcl::PointCloud<pcl::FPFHSignature33>());

  std::cout<<"Started to compute normals and keypoints of the reference object."<<std::endl;
  // Setting normal and fpfh search radius
  double normal_radius_ = 0.02f;
  double fpfh_radius_ = 0.02f;
  ref_cloud_normal_ptr = estimateNormals(ref_cloud_filtered_ptr,normal_radius_);
  ref_SIFT_keypoints_ptr = estimateKeypointsSIFT(ref_cloud_normal_ptr);
  ref_fpfh_ptr = fpfhExtractor(ref_cloud_filtered_ptr,ref_SIFT_keypoints_ptr,ref_cloud_normal_ptr,fpfh_radius_);
  float depth_limit;
  float height_limit;
  float width_limit;

  // Rough Transformation matrix
  Eigen::Matrix4f T_matrix = Eigen::Matrix4f::Identity();
  pcl::PointCloud<pcl::PointXYZ>::Ptr scene_cloud_transformed(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::PointCloud<pcl::PointXYZ>::Ptr scene_cloud_transformed_dl(new pcl::PointCloud<pcl::PointXYZ>);
  pcl::PointCloud<pcl::PointXYZ>::Ptr aligned_ref_cloud (new pcl::PointCloud<pcl::PointXYZ>);

  while(ros::ok())
  {
    cycle++;

    if(cycle>1)
    {

        // First filter scene cloud data ---------------------------------
        // ...removing distant points
        pcl::PassThrough<pcl::PointXYZ> pass;
        pass.setInputCloud (scene_cloud_ptr);
        pass.setFilterFieldName ("z");
        depth_limit = Rough_T_matrix_from_pose_est(2,3);
        pass.setFilterLimits (0, depth_limit+0.15);
        pass.filter (*scene_cloud_ptr);
        pass.setFilterFieldName("x");
        width_limit = Rough_T_matrix_from_pose_est(0,3);
        pass.setFilterLimits(width_limit-0.40,width_limit+0.40);
        pass.filter(*scene_cloud_ptr);
        pass.setFilterFieldName("y");
        height_limit = Rough_T_matrix_from_pose_est(1,3);
        pass.setFilterLimits(height_limit-0.20, height_limit+0.75);
        pass.filter(*scene_cloud_ptr);
        pass.setFilterFieldName("z");
        pass.setNegative (true);
        pass.setFilterLimits(0,depth_limit-0.05);
        pass.filter(*scene_cloud_ptr);
        // Downsampling scene point cloud
        pcl::VoxelGrid<pcl::PointXYZ> sor;
        sor.setInputCloud (scene_cloud_ptr);
        sor.setLeafSize (leaf_size, leaf_size, leaf_size);
        sor.filter (*scene_cloud_ptr);

        // Remove Statistical Outliers -------------------------------------
        pcl::StatisticalOutlierRemoval<pcl::PointXYZ> sor_statistic_removal;
        sor_statistic_removal.setInputCloud (scene_cloud_ptr);
        sor_statistic_removal.setMeanK (50);
        sor_statistic_removal.setStddevMulThresh(1.0);
        sor_statistic_removal.filter (*scene_cloud_ptr);

        pcl::io::savePCDFileASCII("data_scene_filtered_SIFT.pcd", *scene_cloud_ptr);
        auto start = std::chrono::high_resolution_clock::now();
        // Estimate normal, find keypoints and features ------------------
        std::cout<<"Started to compute scene cloud normals."<<std::endl;
        scene_cloud_normal_ptr = estimateNormals(scene_cloud_ptr,normal_radius_);
        scene_SIFT_keypoints_ptr = estimateKeypointsSIFT(scene_cloud_normal_ptr);
        scene_fpfh_ptr = fpfhExtractor(scene_cloud_ptr,scene_SIFT_keypoints_ptr,scene_cloud_normal_ptr,fpfh_radius_);
        // Find Correspondencses -----------------------------------------
        pcl::registration::CorrespondenceEstimation<pcl::FPFHSignature33, pcl::FPFHSignature33> feature_matching_est;
        pcl::CorrespondencesPtr correspondences(new pcl::Correspondences());
        feature_matching_est.setInputSource(scene_fpfh_ptr);
        feature_matching_est.setInputTarget(ref_fpfh_ptr);
        feature_matching_est.determineCorrespondences(*correspondences);
        std::cout<<"Size of correspondences = "<< correspondences->size()<<std::endl;
        // RANSAC false matching removal  --------------------------------
        pcl::registration::CorrespondenceRejectorSampleConsensus<pcl::PointXYZ> rejector_sac;
        pcl::CorrespondencesPtr correspondences_filtered(new pcl::Correspondences());
        rejector_sac.setInputSource(scene_SIFT_keypoints_ptr);
        rejector_sac.setInputTarget(ref_SIFT_keypoints_ptr);
        rejector_sac.setInlierThreshold(1.0);
        rejector_sac.setMaximumIterations(1000);
        rejector_sac.setRefineModel(true);
        rejector_sac.setInputCorrespondences(correspondences);
        rejector_sac.getCorrespondences(*correspondences_filtered);
        std::cout<<"Size of correspondences after outlier correspondence removal = "<< correspondences_filtered->size()<<std::endl;
        // Motion Estimation using SVD -----------------------------------
        pcl::registration::TransformationEstimationSVD<pcl::PointXYZ, pcl::PointXYZ> rough_T_est;
        rough_T_est.estimateRigidTransformation(*scene_SIFT_keypoints_ptr, *ref_SIFT_keypoints_ptr, *correspondences_filtered, T_matrix);
        // NOTE: T_matrix coming from SVD. Rough_T_matrix_from_pose_est coming from Pose Estimation
      if(cycle<3)
      {
        std::cout << "Estimated Transform:" << std::endl << T_matrix << std::endl;
        pcl::transformPointCloud(*scene_cloud_ptr, *scene_cloud_transformed, T_matrix);
        pcl::io::savePCDFileASCII("data_init_trans_sift.pcd", *scene_cloud_transformed);
      }
      auto end = std::chrono::high_resolution_clock::now();
      auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
	    std::cout << "PCL GPU duration: " << duration.count() << std::endl;
        /*pcl::transformPointCloud(*scene_cloud_ptr, *scene_cloud_transformed_dl, Rough_T_matrix_from_pose_est.inverse());
        pcl::io::savePCDFileASCII("shekil_data_dl.pcd", *scene_cloud_transformed_dl);
        std::cout << "Rought Transformation coming from Pose Estimation" << std::endl << Rough_T_matrix_from_pose_est << std::endl;
        */
      // Iterative Closest Point Algorithm -----------------------------
      pcl::IterativeClosestPoint<pcl::PointXYZ,pcl::PointXYZ> icp;
      icp.setInputSource(scene_cloud_transformed);
      icp.setInputTarget(ref_cloud_filtered_ptr);
      icp.setTransformationEpsilon(1e-8);
      icp.setMaximumIterations(100000);
      icp.align(*aligned_ref_cloud);
      std::cout << "has converged:" << icp.hasConverged() << " score: " <<
      icp.getFitnessScore() << std::endl;
      std::cout << icp.getFinalTransformation()<< std::endl;

      // Transform reference model into scene points using general transformation inverse
      Eigen::Matrix4f icp_transformation = icp.getFinalTransformation();
      if(cycle<3)
      {
          Eigen::Matrix4f gen_transformation = icp_transformation*T_matrix;
          pcl::VoxelGrid<pcl::PointXYZ> sor_main;
          sor_main.setInputCloud (ref_main_cloud_ptr);
          sor_main.setLeafSize (0.005, 0.005, 0.005);
          sor_main.filter (*ref_main_cloud_ptr);
          pcl::transformPointCloud (*ref_main_cloud_ptr, *final_output, gen_transformation.inverse());
      }
      else
      {
        pcl::transformPointCloud (*final_output, *final_output, icp_transformation.inverse());
      }

    }

    std::cout<<"Computation is done"<<std::endl;
    pcl::toROSMsg(*final_output,output);
    output.header.frame_id = "camera_depth_optical_frame";
    output.header.stamp = ros::Time::now();
    pub.publish (output);

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

  }
  return 0;
}


#include "Test.h"
#include <Eigen/Dense>
#include "Robot.h"
#include "Body/Body.h"
#include "Utils.h"
#include "SimEnvGenerator.h"
#include "Joint/Joint.h"
#include "Body/BodyCuboid.h"
#include "Body/BodySphere.h"
#include "Joint/JointFixed.h"
#include "Joint/JointSphericalExp.h"
#include "Force/ForceGeneralPrimitiveContact.h"

using namespace redmax;

Eigen::VectorXd Test::random_generated_vector(int len, Eigen::VectorXd lower_bound, Eigen::VectorXd upper_bound) {
    Eigen::VectorXd res(len);
    for (int i = 0;i < len;++i) {
        double a = (double)rand() / RAND_MAX;
        res[i] = a * lower_bound[i] + (1.0 - a) * upper_bound[i];
    }
    return res;
}

Eigen::VectorXd Test::random_generated_vector(int len, double lower_bound, double upper_bound) {
    Eigen::VectorXd res(len);
    for (int i = 0;i < len;++i) {
        double a = (double)rand() / RAND_MAX;
        res[i] = a * lower_bound + (1.0 - a) * upper_bound;
    }
    return res;
}

Eigen::MatrixXd Test::random_generated_matrix(int n, int m, double lower_bound, double upper_bound) {
    Eigen::MatrixXd res(n, m);
    for (int i = 0;i < n;++i)
        for (int j = 0;j < m;++j) {
            double a = (double)rand() / RAND_MAX;
            res(i, j) = a * lower_bound + (1.0 - a) * upper_bound;
        }
    return res;
}

void Test::test_xml_reader() {
    Simulation* sim = new Simulation("../../assets/finger_torque.xml");
    // Simulation* sim = new Simulation("../../assets/rss_two_finger_assemble.xml");

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;

    sim->init();

    sim->print_ctrl_info();

    sim->print_design_params_info();

    auto qdot_init = sim->get_qdot_init();
    qdot_init[sim->_ndof_r - 2] = 5.;
    sim->set_qdot_init(qdot_init);

    sim->reset(false);

    int t0 = clock();
    sim->forward(300, false, false);

    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    sim->replay();
}

void Test::test_speed() {
    Simulation* sim = new Simulation("../../assets/rss_finger_flip.xml");

    sim->init();

    sim->reset();

    int t0 = clock();
    for (int i = 0;i < 5000;i++)
        sim->forward(1, false, false);

    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;
}

// for passive simulation without actuators
void Test::test_forward() {
    // Simulation* sim = SimEnvGenerator::createSinglePendulumTest();
    // Simulation* sim = SimEnvGenerator::createMultiPendulumTest(4);
    // Simulation* sim = SimEnvGenerator::createPrismaticTest();
    // Simulation* sim = SimEnvGenerator::createFree2DTest();
    // Simulation* sim = SimEnvGenerator::createGroundContactTest("BDF1");
    // Simulation* sim = SimEnvGenerator::createBoxContactTest("BDF1");
    // Simulation* sim = SimEnvGenerator::createSinglePendulumObjTest();
    // Simulation* sim = SimEnvGenerator::createSphereGroundContactTest("BDF2");
    // Simulation* sim = SimEnvGenerator::createFree3DEulerTest();
    // Simulation* sim = SimEnvGenerator::createFree3DExpTest();
    Simulation* sim = SimEnvGenerator::createTorqueFingerFlickDemo();

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;

    sim->init();

    int num_steps = 1000;

    int t0 = clock();
    sim->reset();
    sim->forward(num_steps, false);
    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay();
}

void Test::test_derivative() {
    // Simulation* sim = SimEnvGenerator::createSinglePendulumTest();
    // Simulation* sim = SimEnvGenerator::createMultiPendulumTest(4);
    // Simulation* sim = SimEnvGenerator::createPrismaticTest();
    // Simulation* sim = SimEnvGenerator::createFree2DTest();
    // Simulation* sim = SimEnvGenerator::createGroundContactTest();
    // Simulation* sim = SimEnvGenerator::createBoxContactTest();
    Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/finger_mesh.xml");

    sim->init();

    // srand(time(0));
    // sim->set_q(Test::random_generated_vector(sim->_ndof_r, -1., 1.));
    // sim->set_qdot(Test::random_generated_vector(sim->_ndof_r, -1., 1.));

    // sim->test_derivative();

    // test JointSphericalExp Derivatives
    {
        JointSphericalExp* joint = new JointSphericalExp(sim, 0, nullptr, Matrix3::Identity(), Vector3::Zero());
        joint->test_derivatives();
        delete joint;
    }

    // test BodySphere Derivatives
    {
        JointFixed* joint = new JointFixed(sim, 0, nullptr, Matrix3::Identity(), Vector3::Zero());
        BodySphere* body_sphere = new BodySphere(nullptr, joint, 2., Matrix3::Identity(), Vector3::Zero(), 1.);
        body_sphere->test_collision_derivatives();
        delete joint;
        delete body_sphere;
    }

    // test BodyCuboid Derivatives
    {
        JointFixed* joint = new JointFixed(sim, 0, nullptr, Matrix3::Identity(), Vector3::Zero());
        BodyCuboid* body_cuboid = new BodyCuboid(nullptr, joint, Vector3(1., 1., 1.), Matrix3::Identity(), Vector3::Zero(), 1.);
        body_cuboid->test_collision_derivatives();
        delete joint;
        delete body_cuboid;
    }

    // test general-primitive contact Derivatives
    // {
    //     Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/finger_mesh.xml");

    //     dynamic_cast<ForceGeneralPrimitiveContact*>(sim->_robot->_forces[2])->test_derivatives();

    //     delete sim;
    // }
}

void Test::test_design_derivative() {
    // Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/double_pendulum_design_1.xml");
    // Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/double_pendulum_design.xml");
    // Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/double_pendulum_contact_design.xml");
    Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/abstract_finger.xml");
    // Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/new_finger.xml");

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;

    sim->init();

    sim->print_ctrl_info();

    sim->print_design_params_info();

    VectorX design_params = sim->get_design_params();

    sim->set_design_params(design_params);

    sim->reset(true, true);

    int t0 = clock();
    dtype f = 0.;
    int num_steps = 100;
    for (int ii = 0;ii < num_steps;ii ++) {
        sim->forward(1, false, true);
        VectorX var = sim->get_variables();
        f += var.sum();
    }
    sim->_backward_info.set_flags(false, false, true, false);
    sim->_backward_info._df_dq0 = VectorX::Zero(sim->_ndof_r);
    sim->_backward_info._df_dqdot0 = VectorX::Zero(sim->_ndof_r);
    sim->_backward_info._df_dq = VectorX::Zero(sim->_ndof_r * num_steps);
    sim->_backward_info._df_dvar = VectorX::Ones(sim->_ndof_var * num_steps);
    sim->_backward_info._df_dp = VectorX::Zero(sim->_ndof_p);
    sim->backward();
    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    design_params = sim->get_design_params();

    std::cerr << "design params = " << design_params.transpose() << std::endl;
    std::cerr << "f = " << f << std::endl;
    std::cerr << "df_dp = " << sim->_backward_results._df_dp.transpose() << std::endl;
    VectorX df_dp = sim->_backward_results._df_dp;

    // dtype eps = 1e-2;
    // for (int trial = 0;trial < 10;trial++) {
    //     VectorX df_dp_fd = VectorX::Zero(sim->_ndof_p);
    //     for (int i = 0;i < sim->_ndof_p;i++) {
    //         // if (i % 12 < 9) {
    //         //     df_dp(i) = 0;
    //         //     continue;
    //         // }
    //         VectorX design_params_pos = design_params;
    //         design_params_pos[i] += eps;
    //         sim->set_design_params(design_params_pos);
    //         sim->reset(false);
    //         dtype f_pos = 0.;
    //         for (int ii = 0;ii < num_steps;ii++) {
    //             sim->forward(1, false, false);
    //             VectorX var = sim->get_variables();
    //             f_pos += var.sum();
    //         }
    //         df_dp_fd(i) = (f_pos - f) / eps;
    //     }

    //     // std::cerr << "df_dp_fd = " << df_dp_fd.transpose() << std::endl;

    //     print_error_full("df_dp", df_dp, df_dp_fd);
    //     eps /= 10.;
    // }

    sim->replay();
}

void Test::test_control_derivative() {
    Simulation* sim = new Simulation("/home/eanswer/Workspace/Projects/RoboticHand/code/RoboticHandDesign/RedMax/assets/abstract_finger.xml");

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;

    sim->init();

    sim->print_ctrl_info();

    sim->reset(true, false);

    int t0 = clock();
    dtype f = 0.;
    int num_steps = 1;
    VectorX action = VectorX::Zero(num_steps * sim->_ndof_u);
    action << -3.65349773,  2.90075299,  2.12093891,  2.06099197;
    VectorX u = action;
    for (int i = 0;i < u.size();i++)
        u(i) = tanh(u(i));

    for (int ii = 0;ii < num_steps;ii ++) {
        sim->set_u(u.segment(ii * sim->_ndof_u, (ii + 1) * sim->_ndof_u));
        sim->forward(1, false, false);
        VectorX var = sim->get_variables();
        f += var.sum();
    }
    sim->_backward_info.set_flags(false, false, false, true);
    sim->_backward_info._df_dq = VectorX::Zero(sim->_ndof_r * num_steps);
    sim->_backward_info._df_du = VectorX::Zero(sim->_ndof_u * num_steps);
    sim->_backward_info._df_dvar = VectorX::Ones(sim->_ndof_var * num_steps);
    // sim->_backward_info._df_du = VectorX::Ones(sim->_ndof_u * num_steps);
    // sim->_backward_info._df_dvar = VectorX::Zero(sim->_ndof_var * num_steps);
    sim->backward();
    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    std::cerr << "f = " << f << std::endl;

    VectorX df_du = sim->_backward_results._df_du;
    VectorX df_daction = df_du;
    for (int i = 0;i < df_daction.size();i++) {
        df_daction[i] *= (1. - tanh(action(i)) * tanh(action(i)));
    }

    std::cerr << "df_du: " << df_du.transpose() << std::endl;
    {
        dtype eps = 1e-2;
        for (int trial = 0;trial < 10;trial++) {
            VectorX df_du_fd = VectorX::Zero(sim->_ndof_u * num_steps);
            for (int i = 0;i < num_steps * sim->_ndof_u;i++) {
                VectorX u_pos = u;
                u_pos(i) += eps;

                sim->set_u(u_pos);
                sim->reset(false, false);
                dtype f_pos = 0.;
                for (int ii = 0;ii < num_steps;ii++) {
                    sim->set_u(u_pos.segment(ii * sim->_ndof_u, (ii + 1) * sim->_ndof_u));
                    sim->forward(1, false, false);
                    VectorX var = sim->get_variables();
                    f_pos += var.sum();
                }
                // f_pos = u_pos.sum();
                df_du_fd(i) = (f_pos - f) / eps;
            }
            print_error_full("eps = " + to_string(eps) + ", df_du", df_du, df_du_fd);
            eps /= 10.;
        }
    }

    std::cerr << "df_daction: " << df_daction.transpose() << std::endl;

    {
        dtype eps = 1e-2;
        for (int trial = 0;trial < 10;trial++) {
            VectorX df_daction_fd = VectorX::Zero(sim->_ndof_u * num_steps);
            for (int i = 0;i < num_steps * sim->_ndof_u;i++) {
                VectorX action_pos = action;
                action_pos(i) += eps;
                VectorX u_pos = action_pos;
                for (int j = 0;j < u_pos.size();j++)
                    u_pos(j) = tanh(u_pos(j));

                sim->set_u(u_pos);
                sim->reset(false, false);
                dtype f_pos = 0.;
                for (int ii = 0;ii < num_steps;ii++) {
                    sim->set_u(u_pos.segment(ii * sim->_ndof_u, (ii + 1) * sim->_ndof_u));
                    sim->forward(1, false, false);
                    VectorX var = sim->get_variables();
                    f_pos += var.sum();
                }
                // f_pos = u_pos.sum();
                df_daction_fd(i) = (f_pos - f) / eps;
            }
            print_error_full("eps = " + to_string(eps) + ", df_daction", df_daction, df_daction_fd);
            eps /= 10.;
        }
    }

    sim->replay();
}

void Test::test_torque_finger() {
    Simulation* sim = SimEnvGenerator::createTorqueFingerDemo("BDF1");

    sim->init();

    int num_steps = 50;
    int ndof_r = sim->_ndof_r;
    int ndof_u = sim->_ndof_u;

    VectorX ctrl_min, ctrl_max;
    sim->get_ctrl_range(ctrl_min, ctrl_max);

    Vector3 q_goal(0., constants::pi / 2., constants::pi / 4.);
    Vector3 P_q = Vector3(10., 1.5, 1.5) * 1e5;

    int t0 = clock();
    sim->reset(false);

    for (int i = 0;i < num_steps;i++) {
        VectorX q = sim->get_q();
        VectorX error = q_goal - q;
        VectorX u = error.cwiseProduct(P_q);

        sim->set_u(u);

        sim->forward(1, false);
    }

    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay();
}

void Test::test_torque_finger_flick() {
    Simulation* sim = SimEnvGenerator::createTorqueFingerFlickDemo("BDF1");

    sim->init();

    int num_steps = 1000;
    int ndof_r = sim->_ndof_r;
    int ndof_u = sim->_ndof_u;

    VectorX ctrl_min, ctrl_max;
    sim->get_ctrl_range(ctrl_min, ctrl_max);

    Vector3 q_goal(0., 0., 0.);
    Vector3 P_q = Vector3(10., 1.5, 1.5) * 1e5;

    int t0 = clock();
    Vector6 q0 = (Vector6() << 0., constants::pi / 2., constants::pi / 4., 0., 0., 0.).finished();
    sim->set_q_init(q0);

    sim->reset(false);

    for (int i = 0;i < num_steps;i++) {
        VectorX q = sim->get_q();
        VectorX error = q_goal - q.head(3);
        VectorX u = error.cwiseProduct(P_q);

        sim->set_u(u);

        sim->forward(1, false);
    }

    int t1 = clock();

    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay();
}

void Test::test_pendulum_optimization() {
    Simulation* sim = SimEnvGenerator::createMultiPendulumTest(2, "BDF2");

    sim->init();

    sim->_viewer_options->_loop = false;
    sim->_viewer_options->_infinite = true;
    bool test_gradient = false;

    int num_steps = 100;

    Vector2 q_goal(constants::pi / 2., constants::pi / 4.);

    Vector2 q0(0., 0.);
    Vector2 qdot0(0., 0.);
    for (int opt_iter = 0;opt_iter < 2000;opt_iter ++) {
        sim->set_state_init(q0, qdot0);

        sim->reset(true);
        sim->forward(num_steps, false);

        VectorX q = sim->get_q();
        VectorX qdot = sim->get_qdot();

        if (opt_iter == 0) { // render the initial guess
            sim->replay();
        }

        // objective f = |q - q_goal|^2
        dtype f = (q - q_goal).dot(q - q_goal);
        std::cerr << "iter = " << opt_iter << ", f = " << f << std::endl;

        // set backward info
        sim->_backward_info.set_flags(true, true, false, false);

        int ndof_r = sim->_ndof_r;

        // set terminal derivatives
        sim->_backward_info._df_dq0 = VectorX::Zero(ndof_r);
        sim->_backward_info._df_dqdot0 = VectorX::Zero(ndof_r);
        sim->_backward_info._df_dq = VectorX::Zero(ndof_r * num_steps);
        sim->_backward_info._df_dq.tail(ndof_r) = 2. * (q - q_goal);

        // backward
        sim->backward();

        // test gradient
        if (test_gradient) {
            dtype eps = 0.01;
            for (int ii = 0;ii < 10;ii ++) {
                VectorX df_dq0_fd(ndof_r);
                for (int i = 0;i < ndof_r;i++) {
                    VectorX q0_pos = q0;
                    q0_pos(i) += eps;
                    sim->set_state_init(q0_pos, qdot0);
                    sim->reset(false);
                    sim->forward(num_steps, false);
                    VectorX q_pos = sim->get_q();
                    dtype f_pos = (q_pos - q_goal).dot(q_pos - q_goal);

                    VectorX q0_neg = q0;
                    q0_neg(i) -= eps;
                    sim->set_state_init(q0_neg, qdot0);
                    sim->reset(false);
                    sim->forward(num_steps, false);
                    VectorX q_neg = sim->get_q();
                    dtype f_neg = (q_neg - q_goal).dot(q_neg - q_goal);

                    df_dq0_fd(i) = (f_pos - f_neg) / eps / 2.;
                }
                VectorX df_dqdot0_fd(ndof_r);
                for (int i = 0;i < ndof_r;i++) {
                    VectorX qdot0_pos = qdot0;
                    qdot0_pos(i) += eps;
                    sim->set_state_init(q0, qdot0_pos);
                    sim->reset(false);
                    sim->forward(num_steps, false);
                    VectorX q_pos = sim->get_q();
                    dtype f_pos = (q_pos - q_goal).dot(q_pos - q_goal);

                    VectorX qdot0_neg = qdot0;
                    qdot0_neg(i) -= eps;
                    sim->set_state_init(q0, qdot0_neg);
                    sim->reset(false);
                    sim->forward(num_steps, false);
                    VectorX q_neg = sim->get_q();
                    dtype f_neg = (q_neg - q_goal).dot(q_neg - q_goal);

                    df_dqdot0_fd(i) = (f_pos - f_neg) / eps / 2.;
                }

                std::cerr << "eps = " << eps << std::endl;
                print_error("df_dq0", sim->_backward_results._df_dq0, df_dq0_fd);
                print_error("df_dqdot0", sim->_backward_results._df_dqdot0, df_dqdot0_fd);

                eps /= 2.;
            }
        }

        dtype step = 0.02;
        q0 -= step * sim->_backward_results._df_dq0;
        qdot0 -= step * sim->_backward_results._df_dqdot0;
    }

    sim->set_state_init(q0, qdot0);

    sim->reset(false);
    sim->forward(num_steps, false);

    std::cerr << "q0 = " << q0.transpose() << ", qdot0 = " << qdot0.transpose() << std::endl;
    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay(); // render the optimized values
}

void Test::test_torque_finger_optimization() {
    Simulation* sim = SimEnvGenerator::createTorqueFingerDemo("BDF2");

    sim->init();

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;
    bool test_gradient = false;

    int num_steps = 50;
    int ndof_r = sim->_ndof_r;
    int ndof_u = sim->_ndof_u;

    VectorX ctrl_min, ctrl_max;
    sim->get_ctrl_range(ctrl_min, ctrl_max);

    Vector3 q_goal(0., constants::pi / 2., constants::pi / 4.);

    VectorX u = VectorX::Zero(ndof_u * num_steps);

    int t0 = clock();
    for (int opt_iter = 0;opt_iter < 1000;opt_iter ++) {
        sim->reset(true);
        std::vector<VectorX> qs;
        qs.clear();
        for (int i = 0;i < num_steps;i++) {
            sim->set_u(u.segment(ndof_u * i, ndof_u));
            sim->forward(1, false);
            qs.push_back(sim->get_q());
        }

        if (opt_iter == 0) { // render the initial guess
            sim->replay();
        }

        // objective f = sum(|q(i) - q_goal|^2)
        dtype f = 0.;
        for (int i = 0;i < num_steps;i++) {
            f += (qs[i] - q_goal).dot(qs[i] - q_goal);
        }
        std::cerr << "iter = " << opt_iter << ", f = " << f << std::endl;

        // set backward info
        sim->_backward_info.set_flags(false, false, false, true);

        // set terminal derivatives
        sim->_backward_info._df_du = VectorX::Zero(ndof_u * num_steps);
        sim->_backward_info._df_dq = VectorX::Zero(ndof_r * num_steps);
        for (int i = 0;i < num_steps;i++) {
            sim->_backward_info._df_dq.segment(ndof_r * i, ndof_r) = 2. * (qs[i] - q_goal);
        }

        // backward
        sim->backward();

        // test gradient
        if (test_gradient) {
            // dtype eps = 1e-7;
            dtype eps = 0.01;
            for (int ii = 0;ii < 10;ii++) {
                VectorX df_du_fd(ndof_u * num_steps);
                for (int k = 0;k < ndof_u * num_steps;k++) {
                    VectorX u_pos = u;
                    u_pos(k) += eps;
                    sim->reset(false);
                    dtype f_pos = 0.;
                    for (int i = 0;i < num_steps;i++) {
                        sim->set_u(u_pos.segment(ndof_u * i, ndof_u));
                        sim->forward(1, false);
                        VectorX q = sim->get_q();
                        f_pos += (q - q_goal).dot(q - q_goal);
                    }
                    df_du_fd(k) = (f_pos - f) / eps;
                }

                std::cerr << "eps = " << eps << std::endl;
                print_error("df_du", sim->_backward_results._df_du, df_du_fd);

                eps /= 2.;
            }
        }

        dtype step = 0.05;
        u -= step * sim->_backward_results._df_du;
        // for (int i = 0;i < num_steps;i++)
        //     for (int j = 0;j < ndof_u;j++)
        //         u(i * ndof_u + j) = min(ctrl_max(j), max(ctrl_min(j), u(i * ndof_u + j)));
    }

    int t1 = clock();
    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    sim->reset(false);
    for (int i = 0;i < num_steps;i++) {
        sim->set_u(u.segment(ndof_u * i, ndof_u));
        sim->forward(1, false);
    }

    sim->replay();
}

void Test::test_torque_finger_flick_optimization() {
    Simulation* sim = SimEnvGenerator::createTorqueFingerFlickDemo("BDF2");

    sim->init();

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;
    bool test_gradient = false;

    int num_steps = 1000;
    int ndof_r = sim->_ndof_r;
    int ndof_u = sim->_ndof_u;

    VectorX ctrl_min, ctrl_max;
    sim->get_ctrl_range(ctrl_min, ctrl_max);

    Vector6 q0 = (Vector6() << 0., constants::pi / 2., constants::pi / 4., 0., 0., 0.).finished();
    sim->set_q_init(q0);

    dtype x_goal = 5.5;

    VectorX u = VectorX::Zero(ndof_u * num_steps);

    // initial guess
    Vector3 q_goal(0., 0., 0.);
    Vector3 P_q = Vector3(10., 2.0, 3.0);

    sim->reset(false);

    for (int i = 0;i < num_steps;i++) {
        VectorX q = sim->get_q();
        VectorX error = q_goal - q.head(3);
        VectorX ui = error.cwiseProduct(P_q);
        // ui.setZero();
        // for (int j = 0;j < ui.size();j++)
        //     ui[j] = min(1., max(ui[j], -1.));

        u.segment(ndof_u * i, ndof_u) = ui;

        sim->set_u(ui);

        sim->forward(1, false);
    }

    std::cerr << "u = " << u.head(10).transpose() << std::endl;
    std::cerr << u.maxCoeff() << ", " << u.minCoeff() << std::endl;
    std::cerr << "u.norm = " << u.norm() << std::endl;

    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;
    int t0 = clock();
    for (int opt_iter = 0;opt_iter < 1;opt_iter ++) {
        sim->reset(true);
        for (int i = 0;i < num_steps;i++) {
            sim->set_u(u.segment(ndof_u * i, ndof_u));
            sim->forward(1, false);
        }

        VectorX q = sim->get_q();

        if (opt_iter == 0) { // render the initial guess
            sim->replay();
        }

        // objective f = sum(|q(3) - x_goal|^2)
        dtype f = (q(3) - x_goal) * (q(3) - x_goal);

        std::cerr << "iter = " << opt_iter << ", f = " << f << std::endl;

        // set backward info
        sim->_backward_info.set_flags(false, false, false, true);

        // set terminal derivatives
        sim->_backward_info._df_du = VectorX::Zero(ndof_u * num_steps);
        sim->_backward_info._df_dq = VectorX::Zero(ndof_r * num_steps);
        sim->_backward_info._df_dq(ndof_r * (num_steps - 1) + 3) = 2. * (q(3) - x_goal);

        // backward
        sim->backward();

        // // test gradient
        // if (test_gradient) {
        //     dtype eps = 1e-7;
        //     VectorX df_du_fd(ndof_u * num_steps);
        //     for (int k = 0;k < ndof_u * num_steps;k++) {
        //         VectorX u_pos = u;
        //         u_pos(k) += eps;
        //         sim->reset(false);
        //         dtype f_pos = 0.;
        //         for (int i = 0;i < num_steps;i++) {
        //             sim->set_u(u_pos.segment(ndof_u * i, ndof_u));
        //             sim->forward(1, false);
        //             VectorX q_pos = sim->get_q();
        //             f_pos += (q_pos(3) - x_goal) * (q_pos(3) - x_goal);
        //         }
        //         df_du_fd(k) = (f_pos - f) / eps;
        //     }

        //     // print_error("df_du", sim->_backward_results._df_du, df_du_fd);
        //     std::cerr << "df_du_fd = " << df_du_fd.head(10).transpose() << std::endl;
        // }

        std::cerr << "df_du = " << sim->_backward_results._df_du.head(10).transpose() << std::endl;
        std::cerr << "df_du = " << sim->_backward_results._df_du.norm() << std::endl;
        dtype step = 0.2;
        u -= step * sim->_backward_results._df_du;
        // for (int i = 0;i < num_steps;i++)
        //     for (int j = 0;j < ndof_u;j++)
        //         u(i * ndof_u + j) = min(1., max(-1., u(i * ndof_u + j)));
    }

    int t1 = clock();
    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    sim->reset(false);
    for (int i = 0;i < num_steps;i++) {
        sim->set_u(u.segment(ndof_u * i, ndof_u));
        sim->forward(1, false);
    }

    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay();
}

void Test::test_ground_contact_optimization() {
    Simulation* sim = SimEnvGenerator::createGroundContactTest("BDF2");

    // target x,y position
    Vector2 pos_goal(1.0, 0.5);

    // Add virtual goal
    JointFixed* joint = new JointFixed(sim, 10, nullptr, Matrix3::Identity(), Vector3(pos_goal(0), pos_goal(1), 0.));
    BodyCuboid* target = new BodyCuboid(sim, joint, Vector3(0.2, 0.2, 0.2), Matrix3::Identity(), Vector3::Zero(), 1.);
    target->set_color(Vector3(0., 1., 0.));

    sim->_robot->_root_joints.push_back(joint);
    sim->_robot->init();

    // start optimization
    sim->init();

    sim->_viewer_options->_loop = true;
    sim->_viewer_options->_infinite = true;
    bool test_gradient = false;

    int num_steps = 1000;
    int ndof_r = sim->_ndof_r;

    // optimize the initial state for cube 1 (i.e. qdot[0]~qdot[2])
    Vector3 q0 = (Vector3() << -2, 2, 0).finished();
    Vector3 qdot0 = (Vector3() << 5, 50, 50).finished();

    int t0 = clock();
    for (int opt_iter = 0;opt_iter < 200;opt_iter ++) {
        sim->set_state_init(q0, qdot0);

        sim->reset(true);
        sim->forward(num_steps, false);

        VectorX q = sim->get_q();

        if (opt_iter == 0) {// render the initial guess
            sim->replay();
        }

        // objective f = |q[0:2] - pos_goal|^2
        dtype f = (q.head(2) - pos_goal).dot(q.head(2) - pos_goal);
        std::cerr << "iter = " << opt_iter << ", f = " << f << std::endl;

        // set backward info
        sim->_backward_info.set_flags(false, true, false, false);

        int ndof_r = sim->_ndof_r;

        // set terminal derivatives
        sim->_backward_info._df_dqdot0 = VectorX::Zero(ndof_r);
        sim->_backward_info._df_dq = VectorX::Zero(ndof_r * num_steps);
        sim->_backward_info._df_dq.segment(ndof_r * (num_steps - 1), 2) = 2. * (q.head(2) - pos_goal);

        // backward
        sim->backward();

        // test gradient
        if (test_gradient) {
            dtype eps = 1e-7;
            VectorX df_dqdot0_fd(ndof_r);
            for (int i = 0;i < ndof_r;i++) {
                VectorX qdot0_pos = qdot0;
                qdot0_pos(i) += eps;
                sim->set_state_init(q0, qdot0_pos);
                sim->reset(false);
                sim->forward(num_steps, false);
                VectorX q_pos = sim->get_q();
                dtype f_pos = (q_pos.head(2) - pos_goal).dot(q_pos.head(2) - pos_goal);
                df_dqdot0_fd(i) = (f_pos - f) / eps;
            }

            print_error("df_dqdot0", sim->_backward_results._df_dqdot0, df_dqdot0_fd);
        }

        dtype step = 0.002;
        qdot0.head(3) -= step * sim->_backward_results._df_dqdot0.head(3);

        qdot0(0) = min((dtype)50., max((dtype)-50., qdot0(0)));
        qdot0(1) = min((dtype)50., max((dtype)-50., qdot0(1)));
        qdot0(2) = min((dtype)50., max((dtype)-50., qdot0(2)));

        std::cerr << "df_dqdot0 = " << sim->_backward_results._df_dqdot0.transpose() << std::endl;
        std::cerr << "qdot0 = " << qdot0.head(3).transpose() << std::endl;
    }

    int t1 = clock();
    std::cerr << "time = " << (t1 - t0) / 1000.0 << "ms" << std::endl;

    sim->set_state_init(q0, qdot0);

    sim->reset(false);
    sim->forward(num_steps, false);

    std::cerr << "qdot0 = " << qdot0.transpose() << std::endl;
    std::cerr << "q = " << sim->get_q().transpose() << std::endl;

    sim->replay();
}

void Test::test_box_contact_optimization() {
    Simulation* sim = SimEnvGenerator::createBoxContactTest("BDF2");

    // target x position for the cube 2 (i.e. target q[3])
    dtype x_goal = 3.5;

    // Add virtual goal
    JointFixed* joint = new JointFixed(sim, 10, nullptr, Matrix3::Identity(), Vector3(x_goal + 3., 0.5, 0.));
    BodyCuboid* target = new BodyCuboid(sim, joint, Vector3(0.2, 0.2, 0.2), Matrix3::Identity(), Vector3::Zero(), 1.);
    target->set_color(Vector3(0., 1., 0.));

    sim->_robot->_root_joints.push_back(joint);
    sim->_robot->init();

    sim->init();

    sim->_viewer_options->_loop = false;
    sim->_viewer_options->_infinite = true;
    bool test_gradient = false;

    int num_steps = 400;
    int ndof_r = sim->_ndof_r;

    // optimize the initial state for cube 1 (i.e. qdot[0]~qdot[2])
    Vector6 q0 = (Vector6() << 0., 0., 0., 0., 0., 0.).finished();
    Vector6 qdot0 = (Vector6() << 35., 0., 0., 0., 0., 0.).finished();
    for (int opt_iter = 0;opt_iter < 200;opt_iter ++) {
        sim->set_state_init(q0, qdot0);

        sim->reset(true);
        sim->forward(num_steps, false);

        VectorX q = sim->get_q();

        // if (opt_iter == 0) {// render the initial guess
        //     std:cerr << "q = " << q.transpose() << std::endl;
        //     sim->replay();
        // }

        // objective f = |q[3] - x_goal|^2
        dtype f = (q[3] - x_goal) * (q[3] - x_goal);
        std::cerr << "iter = " << opt_iter << ", f = " << f << std::endl;

        // set backward info
        sim->_backward_info.set_flags(false, true, false, false);

        int ndof_r = sim->_ndof_r;

        // set terminal derivatives
        sim->_backward_info._df_dq0 = VectorX::Zero(ndof_r);
        sim->_backward_info._df_dqdot0 = VectorX::Zero(ndof_r);
        sim->_backward_info._df_dq = VectorX::Zero(ndof_r * num_steps);
        sim->_backward_info._df_dq[ndof_r * (num_steps - 1) + 3] = 2. * (q[3] - x_goal);

        // backward
        sim->backward();

        // test gradient
        if (test_gradient) {
            dtype eps = 1e-7;
            VectorX df_dq0_fd(ndof_r);
            for (int i = 0;i < ndof_r;i++) {
                VectorX q0_pos = q0;
                q0_pos(i) += eps;
                sim->set_state_init(q0_pos, qdot0);
                sim->reset(false);
                sim->forward(num_steps, false);
                VectorX q_pos = sim->get_q();
                dtype f_pos = (q_pos[3] - x_goal) * (q_pos[3] - x_goal);
                df_dq0_fd(i) = (f_pos - f) / eps;
            }
            VectorX df_dqdot0_fd(ndof_r);
            for (int i = 0;i < ndof_r;i++) {
                VectorX qdot0_pos = qdot0;
                qdot0_pos(i) += eps;
                sim->set_state_init(q0, qdot0_pos);
                sim->reset(false);
                sim->forward(num_steps, false);
                VectorX q_pos = sim->get_q();
                dtype f_pos = (q_pos[3] - x_goal) * (q_pos[3] - x_goal);
                df_dqdot0_fd(i) = (f_pos - f) / eps;
            }

            // print_error("df_dq0", sim->_backward_results._df_dq0, df_dq0_fd);
            print_error("df_dqdot0", sim->_backward_results._df_dqdot0, df_dqdot0_fd);
        }

        dtype step = 0.5;
        qdot0(0) -= step * sim->_backward_results._df_dqdot0(0);
        qdot0(1) -= step * sim->_backward_results._df_dqdot0(1);

        // qdot0(0) = min(60., max(30., qdot0(0)));

        std::cerr << "df_dqdot0 = " << sim->_backward_results._df_dqdot0.head(3).transpose() << std::endl;
        std::cerr << "qdot0 = " << qdot0.head(3).transpose() << std::endl;
    }

    sim->set_state_init(q0, qdot0);

    sim->reset(false);
    sim->forward(num_steps, false);

    std::cerr << "q0 = " << q0.transpose() << ", qdot0 = " << qdot0.transpose() << std::endl;
    std::cerr << "q = " << sim->get_q().transpose() << ", qdot = " << sim->get_qdot().transpose() << std::endl;

    sim->replay();
}


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

/**
 * \file mpc_solver_XUopt.hpp
 * \brief
 * \details
 *
 * \author Autonomous Space Robotics Lab (ASRL)
 */
#pragma once

#include <algorithm>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include <Eigen/Cholesky>
#include <Eigen/Dense>
#include <unsupported/Eigen/SparseExtra>

#include <vtr_path_tracker/robust_mpc/optimization/mpc_solver_base.hpp>

const int SOLN_OK = 1;
const int UNCERT_EXCEEDS_CONSTRAINTS = 2;

namespace vtr {
namespace path_tracker {

class MpcSolverXUopt : public MpcSolverBase {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  /** \brief Constructor */
  MpcSolverXUopt();

  /** \brief Destructor */
  ~MpcSolverXUopt();

  /// \todo: add documentation

  /** \brief This reset occurs once per timestep */
  void reset_solver_specific() override;
  MpcNominalModel::model_state_t *select_x_pred(int index) override;
  MpcNominalModel::model_trajectory_t *select_x_pred() override;
  void set_desired_speed_ctrl(int &index, const double &speed,
                              const double &ctrl) override;
  void set_desired_speed(int &index, const double &speed) override;
  void get_desired_ctrl(int &index) override;
  void post_process_x_pred(int &index) override;

  /** \brief This computation occurs for each iteration of the solver */
  void compute_solver_update(const local_path_t &local_path,
                             int iteration) override;

  void get_u(Eigen::MatrixXf &u_mat) override;
  void get_v(Eigen::MatrixXf &v_mat) override;
  void get_w(Eigen::MatrixXf &w_mat) override;
  void set_cmd_km1(const float &u_km1_in, const float &v_km1_in) override;
  void set_exp_consts(float c_in, float c_out);
  float c_in, c_out;

 private:
  // Timing variables
  std::vector<float> norm_ds_opt;
  float step_size_;
  int iterations_small_step;

  // The variable mu is used by constrained optimization and varies as a
  // function of iteration
  float mu_value;
  int mu_index;
  void update_mu_value(float &mu_value);

  // Optimization matrix sizes & indices
  int size_x_vec, size_u_vec, size_v_vec, size_y_vec, size_z_vec, size_w_vec;
  int x_offset_index, u_offset_index, v_offset_index, y_offset_index;
  int z_lb_offset_index, z_ub_offset_index, w_lb_offset_index,
      w_ub_offset_index;

  // Intermediate optimization variables
  Eigen::MatrixXf x_ub_lims, x_lb_lims;
  Eigen::MatrixXf x_ub_abc, x_lb_abc;
  mtx_triplet_list_t weight_Q_triplet, weight_Ru_triplet, weight_Rv_triplet;
  Eigen::SparseMatrix<float, 0> weight_Q;  // weight_Ru, weight_Rv;
  Eigen::MatrixXf weight_Ru_diag, weight_Rv_diag, dI_t_Rdu_dI, dI_t_Rdv_dI,
      weight_u0, weight_v0;
  int solution_result;
  Eigen::MatrixXf x_minus_fx_, gx_minus_w_;

  Eigen::MatrixXf exp_bw, f_1, df_1, d2f_1, d2_cost_w;

  // Final optimization variables: opt_update = - J_grad_L_mtx^(-1)*grad_L;
  Eigen::MatrixXf s_bar, ds_bar;
  Eigen::MatrixXf grad_L;
  Eigen::SparseMatrix<float> J_grad_L_mtx;
  mtx_triplet_list_t J_grad_L_triplet_list;
  std::vector<bool> J_grad_L_mtx_entries_flg;
  std::vector<int> J_grad_L_mtx_indices;
  Eigen::MatrixXf grad_L_im1, ds_opt_im1;

  /** \brief Utility function
   * Finds the sign of a number
   * \todo duplicated in utilities.h
   */
  float getSign(float number);

  // Optimization algorithm helper functions
  void compute_mtx_offsets_and_sizes();
  void reset_solver_variables();
  void compute_weight_mtxs();
  void extract_x(const local_path_t &local_path);
  void compute_constraints_V2(const local_path_t &local_path);
  void extract_grad_L();
  void extract_J_grad_L(int iteration);

  void insert_triplet_list(mtx_triplet_list_t &triplet_list_out,
                           mtx_triplet_list_t &triplet_list_in, float &mult,
                           int &i, int &j, float sign = 1.0f);
  void insert_mtx(mtx_triplet_list_t &triplet_list_out,
                  const Eigen::MatrixXf &gradients_in, int &i, int &j,
                  float sign = 1.0f, bool transpose = false);
  void insert_triplet(mtx_triplet_list_t &triplet_list_out, const int &i,
                      const int &j, float v_ij, float sign = 1.0f);

  void insert_triplet_list(Eigen::MatrixXf &mtx,
                           mtx_triplet_list_t &triplet_list_in, float &mult,
                           int &i, int &j, float sign = 1.0f);
  void insert_mtx(Eigen::MatrixXf &mtx, const Eigen::MatrixXf &gradients_in,
                  int &i, int &j, float sign = 1.0f, bool transpose = false);
  void insert_triplet(Eigen::MatrixXf &mtx, const int &i, const int &j,
                      float v_ij, float sign = 1.0f);

  void populate_mtx(Eigen::SparseMatrix<float, 0> &mtx_sm,
                    mtx_triplet_list_t &triplet_list);
  void reset_index_list(mtx_triplet_list_t &triplet_list_in);
  void compute_limiting_step_size(float &step_size, float &s_km1, float &ds,
                                  float lb, float ub);
};

}  // namespace path_tracker
}  // namespace vtr