#pragma once

#include <Eigen/Eigen>
#include <autoopt/btls.hpp>
#include <autoopt/optimization_problem.hpp>
#include <cmath>
#include <cstddef>
#include <optional>
#include <utility>
#include <vector>

#include "autoopt/util.hpp"

namespace autoopt {

struct test_data {
  Eigen::VectorXd _xs;
  Eigen::VectorXd _ys;
  double _center;

  size_t size() const { return _xs.size(); }

  std::pair<double, double> operator[](size_t i) const {
    return {_xs(i) - _center, _ys(i)};
  };
};

struct curve_base {
  Eigen::VectorXd initial;
  Eigen::VectorXd delta;
  std::vector<const char*> param_names;
  bool flip = false;

  std::pair<Eigen::VectorXd, double> fit(this auto&& self,
                                         const test_data& measurements) {
    auto opt_func = [&]<typename T>(const Eigen::VectorX<T>& params) {
      auto grapher = self.template grapher_from_params<T>(params);
      T error{0.0};
      for (size_t i = 0; i < measurements.size(); i++) {
        auto [x, y] = measurements[i];
        T distance = T{self.flip ? -1.0 : 1.0} *
                         std::atan(grapher.slope_at(T{self.flip ? -x : x})) -
                     T{arcsec2rad(y)};
        error = error + distance * distance;
      }
      return error / T(measurements.size());
    };

    auto_diff_optimization_problem problem(opt_func, self.initial);

    log_barrier_optimization_problem lb_problem(problem, self.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();
    return {fitted_params, std::sqrt(problem.objective(fitted_params))};
  }

  Eigen::VectorXd plot(this auto&& self, const test_data& measurements,
                       std::optional<Eigen::VectorXd> params = {}) {
    auto grapher = self.template grapher_from_params<double>(
        params.value_or(self.initial));
    Eigen::VectorXd ys(measurements.size());
    for (size_t i = 0; i < measurements.size(); i++) {
      double x = measurements[i].first;
      ys(i) = rad2arcsec((self.flip ? -1.0 : 1.0) *
                         atan(grapher.slope_at(self.flip ? -x : x)));
    }
    return ys;
  }
};

}  // namespace autoopt