#pragma once

#include <Eigen/Eigen>

#include "autoopt/dual.hpp"

namespace autoopt {

template <typename T, class Func>
T derivative(Func&& f, const T& x) {
  dual<T> a(x, T(1));
  dual<T> b = f(a);
  return b._dx;
}

template <typename T, class Func>
Eigen::VectorX<T> gradient(Func&& f, const Eigen::VectorX<T>& x) {
  Eigen::VectorX<T> grad{x.size()};
  Eigen::VectorX<dual<T>> dual_x{x.size()};
  for (int i = 0; i < x.size(); ++i) {
    dual_x(i) = dual<T>(x(i), T(0));
  }
  for (int i = 0; i < x.size(); ++i) {
    dual_x(i)._dx = T(1);
    dual<T> dual_y = f(dual_x);
    grad(i) = dual_y._dx;
    dual_x(i)._dx = T(0);
  }
  return grad;
}

template <typename T, class Func>
Eigen::MatrixX<T> jacobian(Func&& f, const Eigen::VectorX<T>& x) {
  Eigen::MatrixX<T> jacob(f(x).size(), x.size());
  Eigen::VectorX<dual<T>> dual_x(x.size());
  for (int i = 0; i < x.size(); ++i) {
    dual_x(i) = dual<T>(x(i), T(0));
  }
  for (int i = 0; i < x.size(); ++i) {
    dual_x(i)._dx = T(1);
    Eigen::VectorX<dual<T>> dual_y = f(dual_x);
    for (int j = 0; j < dual_y.size(); ++j) {
      jacob(j, i) = dual_y(j)._dx;
    }
    dual_x(i)._dx = T(0);
  }
  return jacob;
}

template <typename T, class Func>
Eigen::MatrixX<T> hessian(Func&& f, const Eigen::VectorX<T>& x) {
  auto helper_func = [&f]<typename U>(const Eigen::VectorX<U>& y) {
    return gradient<U>(f, y);
  };
  return jacobian<T>(helper_func, x);
}

} // namespace autoopt