#include <gtest/gtest.h>

#include <autoopt/ellipse.hpp>
#include <autoopt/optimization_problem.hpp>
#include <autoopt/util.hpp>
#include <autoopt/btls.hpp>
#include <iomanip>
#include <iostream>

using namespace autoopt;

TEST(Ellipse, Slope) {
  ellipse<double> e{100, 1000, deg2rad(1.0)};  // entrance angle 1 degree
  quadric<double> q = e.to_quadric();

  EXPECT_NEAR(q.slope_at(-10), -0.0010305116165301856, 1e-9);
  EXPECT_NEAR(q.slope_at(0), 0.0, 1e-9);
  EXPECT_NEAR(q.slope_at(10), 0.00090001261192696272, 1e-9);
}

TEST(Ellipse, ParamGradient) {
  std::vector<std::pair<double, double>> data_points = {
      {-10.0, -0.001}, {0.0, 0.0}, {10, 0.0009}};

  Eigen::VectorX<double> params(4);
  params << 100, 1000, deg2rad(1.0), 0.0;

  auto loss_func = [&data_points]<typename T>(const Eigen::VectorX<T>& p) {
    ellipse<T> e{T{p(0)}, T{p(1)}, T{p(2)}};
    quadric<T> q = e.to_quadric().rotated_by(T{p(3)});
    T loss = T{0};
    for (const auto& [x, y_true] : data_points) {
      T y_pred = q.slope_at(T{x});
      T error = y_pred - T{y_true};
      loss = loss + error * error;
    }
    return loss / T(data_points.size());
  };

  auto_diff_optimization_problem<double, decltype(loss_func)> problem(loss_func, params);

  auto grad = problem.gradient(params);

  std::cout << grad << std::endl;

  EXPECT_NEAR(grad(0), -2.0789313126683308e-10, 1e-15);  // d/d(left_arm)
  EXPECT_NEAR(grad(1), -1.7464984353858657e-12, 1e-15);  // d/d(right_arm)
  EXPECT_NEAR(grad(2), 1.2013025455499119e-06, 1e-15);   // d/d(entrance_angle)
  EXPECT_NEAR(grad(3), -2.0332702665822054e-05, 1e-15);  // d/d(rotation_angle)

  auto hess = problem.hessian(params);

  std::cout << hess << std::endl;

  Eigen::VectorX<double> params_delta(4);
  params_delta << 1.0, 1.0, deg2rad(0.1), deg2rad(0.1);

  params(0) += 0.9;  // left_arm

  // log barrier
  log_barrier_optimization_problem<double> log_barrier_problem(
      problem,
      params_delta,
      1e-3);
  
  auto log_barrier_grad = log_barrier_problem.gradient(params);

  std::cout << "Log barrier gradient:" << std::endl;
  std::cout << log_barrier_grad << std::endl;

  btls(problem);
}