#pragma once

#include <cstddef>

#include "autoopt/derivative.hpp"

namespace autoopt {

template <typename T>
struct optimization_problem {
  virtual Eigen::VectorX<T>& initial_guess() = 0;
  virtual Eigen::VectorX<T>& x() = 0;

  virtual T objective(const Eigen::VectorX<T>& params) = 0;
  virtual Eigen::VectorX<T> gradient(const Eigen::VectorX<T>& params) = 0;
  virtual Eigen::MatrixX<T> hessian(const Eigen::VectorX<T>& params) = 0;
};

template <typename T, class Func>
struct auto_diff_optimization_problem : public optimization_problem<T> {
  Func _objective_func;
  Eigen::VectorX<T> _initial_guess;
  Eigen::VectorX<T> _x;

  auto_diff_optimization_problem(
      Func objective_func, Eigen::VectorX<T> initial_guess = Eigen::VectorX<T>{})
      : _objective_func(objective_func), _initial_guess(initial_guess), _x(initial_guess) {}

  Eigen::VectorX<T>& initial_guess() override { return _initial_guess; }

  Eigen::VectorX<T>& x() override { return _x; }

  T objective(const Eigen::VectorX<T>& params) override {
    return _objective_func(params);
  }

  Eigen::VectorX<T> gradient(const Eigen::VectorX<T>& params) override {
    return autoopt::gradient<T>(_objective_func, params);
  }

  Eigen::MatrixX<T> hessian(const Eigen::VectorX<T>& params) override {
    return autoopt::hessian<T>(_objective_func, params);
  }
};

template <typename T>
struct log_barrier_optimization_problem
    : public optimization_problem<T> {
  optimization_problem<T>& _base_problem;
  Eigen::VectorX<T> _delta;
  T _barrier_strength;

  log_barrier_optimization_problem(
      optimization_problem<T>& base_problem,
      Eigen::VectorX<T> delta,
      T barrier_strength = T{1e-3})
      : _base_problem(base_problem),
        _delta(delta),
        _barrier_strength(barrier_strength) {}

  Eigen::VectorX<T>& initial_guess() override {
    return _base_problem.initial_guess();
  }

  Eigen::VectorX<T>& x() override {
    return _base_problem.x();
  }

  T objective(const Eigen::VectorX<T>& params) override {
    T base_obj = _base_problem.objective(params);
    T barrier = barrier_term(params);
    return base_obj + barrier;
  }

  Eigen::VectorX<T> gradient(const Eigen::VectorX<T>& params) override {
    auto base_grad = _base_problem.gradient(params);
    Eigen::VectorX<T> barrier_grad = autoopt::gradient<T>(
        [this]<typename U>(const Eigen::VectorX<U>& p) {
          return barrier_term<U>(p);
        },
        params);
    Eigen::VectorX<T> total_grad(params.size());
    total_grad = base_grad + barrier_grad;
    return total_grad;
  }

  Eigen::MatrixX<T> hessian(const Eigen::VectorX<T>& params) override {
    auto base_hess = _base_problem.hessian(params);
    Eigen::MatrixX<T> barrier_hess = autoopt::hessian<T>(
        [this]<typename U>(const Eigen::VectorX<U>& p) {
          return barrier_term<U>(p);
        },
        params);
    Eigen::MatrixX<T> total_hess(params.size(), params.size());
    total_hess = base_hess + barrier_hess;
    return total_hess;
  }

private:
  template <typename U>
  U barrier_term(const Eigen::VectorX<U>& params) {
    U barrier = U{0};
    for (int i = 0; i < params.size(); ++i) {
      U lb = U{_base_problem.initial_guess()(i) - _delta(i)};
      U ub = U{_base_problem.initial_guess()(i) + _delta(i)};
      barrier = barrier + std::log(params(i) - lb) + std::log(ub - params(i));
    }
    return -U{_barrier_strength} * barrier;
  }
};

}  // namespace autoopt