#pragma once

#include <array>
#include <cstddef>

#include "autoopt/derivative.hpp"

namespace autoopt {

template <typename T, size_t N>
struct optimization_problem {
  virtual std::array<T, N> initial_guess() = 0;

  virtual T objective(const std::array<T, N>& params) = 0;
  virtual std::array<T, N> gradient(const std::array<T, N>& params) = 0;
  virtual matrix_t<T, N, N> hessian(const std::array<T, N>& params) = 0;
};

template <typename T, size_t N, class Func>
struct auto_diff_optimization_problem : public optimization_problem<T, N> {
  Func _objective_func;
  std::array<T, N> _initial_guess;

  auto_diff_optimization_problem(
      Func objective_func, std::array<T, N> initial_guess = std::array<T, N>{})
      : _objective_func(objective_func), _initial_guess(initial_guess) {}

  std::array<T, N> initial_guess() override { return _initial_guess; }

  T objective(const std::array<T, N>& params) override {
    return _objective_func(params);
  }

  std::array<T, N> gradient(const std::array<T, N>& params) override {
    return autoopt::gradient<T, N>(_objective_func, params);
  }

  matrix_t<T, N, N> hessian(const std::array<T, N>& params) override {
    return autoopt::hessian<T, N>(_objective_func, params);
  }
};

template <typename T, size_t N>
struct log_barrier_optimization_problem
    : public optimization_problem<T, N> {
  optimization_problem<T, N>& _base_problem;
  std::array<T, N> _delta;
  T _barrier_strength;

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

  std::array<T, N> initial_guess() override {
    return _base_problem.initial_guess();
  }

  T objective(const std::array<T, N>& params) override {
    T base_obj = _base_problem.objective(params);
    T barrier = barrier_term(params);
    return base_obj + barrier;
  }

  std::array<T, N> gradient(const std::array<T, N>& params) override {
    auto base_grad = _base_problem.gradient(params);
    std::array<T, N> barrier_grad = autoopt::gradient<T, N>(
        [this]<typename U>(const std::array<U, N>& p) {
          return barrier_term<U>(p);
        },
        params);
    std::array<T, N> total_grad;
    for (size_t i = 0; i < N; ++i) {
      total_grad[i] = base_grad[i] + barrier_grad[i];
    }
    return total_grad;
  }

  matrix_t<T, N, N> hessian(const std::array<T, N>& params) override {
    auto base_hess = _base_problem.hessian(params);
    matrix_t<T, N, N> barrier_hess = autoopt::hessian<T, N>(
        [this]<typename U>(const std::array<U, N>& p) {
          return barrier_term<U>(p);
        },
        params);
    matrix_t<T, N, N> total_hess;
    for (size_t i = 0; i < N; ++i) {
      for (size_t j = 0; j < N; ++j) {
        total_hess[i][j] = base_hess[i][j] + barrier_hess[i][j];
      }
    }
    return total_hess;
  }

private:
  template <typename U>
  U barrier_term(const std::array<U, N>& params) {
    U barrier = U{0};
    for (size_t i = 0; i < N; ++i) {
      U lb = _base_problem.initial_guess()[i] - _delta[i];
      U ub = _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