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

Eigen::VectorX<double> autoopt::fit_ellipse(
    const std::vector<std::pair<double, double>>& data,
    const Eigen::VectorX<double>& initial_params,
    const Eigen::VectorX<double>& delta) {
  auto opt_func = [&]<typename T>(const Eigen::VectorX<T>& params) {
    ellipse<T> ellip(params(0), params(1), params(2));
    quadric<T> quad = ellip.to_quadric().rotated_by(params(3));
    T error = T(0);
    for (const auto& [x, y] : data) {
      T y_fit = quad.slope_at(T(x));
      error = error + (y_fit - T(y)) * (y_fit - T(y));
    }
    return error / T(data.size());
  };

  auto_diff_optimization_problem problem(opt_func, initial_params);

  log_barrier_optimization_problem lb_problem(problem, delta, 1e-4);

  while (lb_problem._barrier_strength > 1e-20) {
    btls(lb_problem);
    lb_problem._barrier_strength *= 1e-2;
  }

  Eigen::VectorX<double> fitted_params = problem.x();

  std::cout << "Fitted parameters: " << std::setprecision(10)
            << fitted_params.transpose() << std::endl;

  double obj_value = problem.objective(fitted_params);
  std::cout << "Objective value: " << obj_value << std::endl;

  // rms in radians
  std::cout << "RMS error: " << std::sqrt(obj_value) << " radians" << std::endl;
  // rms in arcsec
  double rms_arcsec = std::sqrt(obj_value) * (3600.0 * 180.0 / M_PI);
  std::cout << "RMS error: " << rms_arcsec << " arcsec" << std::endl;

  return fitted_params;
}

Eigen::VectorX<double> autoopt::fit_parabola(
    const std::vector<std::pair<double, double>>& data,
    const Eigen::VectorX<double>& initial_params,
    const Eigen::VectorX<double>& delta) {
  auto opt_func = [&]<typename T>(const Eigen::VectorX<T>& params) {
    parabola<T> parab(params(0), params(1));
    quadric<T> quad = parab.to_quadric().rotated_by(params(2));
    T error = T(0);
    for (const auto& [x, y] : data) {
      T y_fit = quad.slope_at(T(x));
      error = error + (y_fit - T(y)) * (y_fit - T(y));
    }
    return error / T(data.size());
  };

  auto_diff_optimization_problem problem(opt_func, initial_params);

  log_barrier_optimization_problem lb_problem(problem, delta, 1e-4);

  while (lb_problem._barrier_strength > 1e-20) {
    btls(lb_problem);
    lb_problem._barrier_strength *= 1e-2;
  }

  Eigen::VectorX<double> fitted_params = problem.x();

  std::cout << "Fitted parameters: " << std::setprecision(10)
            << fitted_params.transpose() << std::endl;

  double obj_value = problem.objective(fitted_params);
  std::cout << "Objective value: " << obj_value << std::endl;

  // rms in radians
  std::cout << "RMS error: " << std::sqrt(obj_value) << " radians" << std::endl;
  // rms in arcsec
  double rms_arcsec = std::sqrt(obj_value) * (3600.0 * 180.0 / M_PI);
  std::cout << "RMS error: " << rms_arcsec << " arcsec" << std::endl;

  return fitted_params;
}