#pragma once

#include <array>

#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, std::size_t N, class Func>
std::array<T, N> gradient(Func&& f, const std::array<T, N>& x) {
  std::array<T, N> grad{};
  std::array<dual<T>, N> dual_x{};
  for (std::size_t i = 0; i < N; ++i) {
    dual_x[i] = dual<T>(x[i], T(0));
  }
  for (std::size_t i = 0; i < N; ++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, std::size_t N, std::size_t M>
using matrix_t = std::array<std::array<T, M>, N>;

template <typename T, std::size_t N, std::size_t M, class Func>
matrix_t<T, M, N> jacobian(Func&& f, const std::array<T, N>& x) {
  matrix_t<T, M, N> jacob{};
  std::array<dual<T>, N> dual_x{};
  for (std::size_t i = 0; i < N; ++i) {
    dual_x[i] = dual<T>(x[i], T(0));
  }
  for (std::size_t i = 0; i < N; ++i) {
    dual_x[i]._dx = T(1);
    std::array<dual<T>, M> dual_y = f(dual_x);
    for (std::size_t j = 0; j < M; ++j) {
      jacob[j][i] = dual_y[j]._dx;
    }
    dual_x[i]._dx = T(0);
  }
  return jacob;
}

template <typename T, std::size_t N, class Func>
matrix_t<T, N, N> hessian(Func&& f, const std::array<T, N>& x) {
  auto helper_func = [&f]<typename U>(const std::array<U, N>& y) {
    return gradient<U, N>(f, y);
  };
  return jacobian<T, N, N>(helper_func, x);
}

} // namespace autoopt