#pragma once

#include <cmath>

namespace autoopt {

template <typename T>
struct dual {
  T _x;
  T _dx;

  constexpr dual(T x = T{0}, T dx = T{0}) : _x(x), _dx(dx) {}

  // allow arbitrary deeply-nested dual construction
  template <typename U>
  constexpr dual(U x) : _x(T(x)), _dx(T(0)) {}
};

template <typename T>
constexpr dual<T> operator+(const dual<T>& a, const dual<T>& b) {
  return dual<T>(a._x + b._x, a._dx + b._dx);
}

template <typename T>
constexpr dual<T> operator-(const dual<T>& a) {
  return dual<T>(-a._x, -a._dx);
}

template <typename T>
constexpr dual<T> operator-(const dual<T>& a, const dual<T>& b) {
  return dual<T>(a._x - b._x, a._dx - b._dx);
}

template <typename T>
constexpr dual<T> operator*(const dual<T>& a, const dual<T>& b) {
  return dual<T>(a._x * b._x, a._x * b._dx + a._dx * b._x);
}

template <typename T>
constexpr dual<T> operator/(const dual<T>& a, const dual<T>& b) {
  return dual<T>(a._x / b._x, (a._dx * b._x - a._x * b._dx) / (b._x * b._x));
}

template <typename T>
constexpr bool operator==(const dual<T>& a, const dual<T>& b) {
  return a._x == b._x;
}

template <typename T>
constexpr bool operator!=(const dual<T>& a, const dual<T>& b) {
  return a._x != b._x;
}

template <typename T>
constexpr bool operator<(const dual<T>& a, const dual<T>& b) {
  return a._x < b._x;
}

template <typename T>
constexpr bool operator<=(const dual<T>& a, const dual<T>& b) {
  return a._x <= b._x;
}

template <typename T>
constexpr bool operator>(const dual<T>& a, const dual<T>& b) {
  return a._x > b._x;
}

template <typename T>
constexpr bool operator>=(const dual<T>& a, const dual<T>& b) {
  return a._x >= b._x;
}

}  // namespace autoopt

namespace std {

using autoopt::dual;

// forward declarations of standard functions

template <typename T>
constexpr dual<T> abs(const dual<T>& a);

template <typename T>
constexpr dual<T> exp(const dual<T>& a);

template <typename T>
constexpr dual<T> log(const dual<T>& a);

template <typename T>
constexpr dual<T> pow(const dual<T>& a, const dual<T>& b);

template <typename T>
constexpr dual<T> sqrt(const dual<T>& a);

// forward declarations of trigonometric functions

template <typename T>
constexpr dual<T> sin(const dual<T>& a);

template <typename T>
constexpr dual<T> cos(const dual<T>& a);

template <typename T>
constexpr dual<T> tan(const dual<T>& a);

template <typename T>
constexpr dual<T> asin(const dual<T>& a);

template <typename T>
constexpr dual<T> acos(const dual<T>& a);

template <typename T>
constexpr dual<T> atan(const dual<T>& a);

template <typename T>
constexpr dual<T> atan2(const dual<T>& y, const dual<T>& x);

// forward declarations of hyperbolic functions

template <typename T>
constexpr dual<T> sinh(const dual<T>& a);

template <typename T>
constexpr dual<T> cosh(const dual<T>& a);

template <typename T>
constexpr dual<T> tanh(const dual<T>& a);

template <typename T>
constexpr dual<T> asinh(const dual<T>& a);

template <typename T>
constexpr dual<T> acosh(const dual<T>& a);

template <typename T>
constexpr dual<T> atanh(const dual<T>& a);

// standard functions

template <typename T>
constexpr dual<T> abs(const dual<T>& a) {
  return dual<T>(std::abs(a._x), (a._x >= T(0) ? T(1) : T(-1)) * a._dx);
}

template <typename T>
constexpr dual<T> exp(const dual<T>& a) {
  T exp_x = std::exp(a._x);
  return dual<T>(exp_x, exp_x * a._dx);
}

template <typename T>
constexpr dual<T> log(const dual<T>& a) {
  return dual<T>(std::log(a._x), (T(1) / a._x) * a._dx);
}

template <typename T>
constexpr dual<T> pow(const dual<T>& a, const dual<T>& b) {
  return std::exp(b * std::log(a));
}

template <typename T>
constexpr dual<T> sqrt(const dual<T>& a) {
  T sqrt_x = std::sqrt(a._x);
  return dual<T>(sqrt_x, (T(1) / (T(2) * sqrt_x)) * a._dx);
}

// trigonometric functions

template <typename T>
constexpr dual<T> sin(const dual<T>& a) {
  return dual<T>(std::sin(a._x), std::cos(a._x) * a._dx);
}

template <typename T>
constexpr dual<T> cos(const dual<T>& a) {
  return dual<T>(std::cos(a._x), -std::sin(a._x) * a._dx);
}

template <typename T>
constexpr dual<T> tan(const dual<T>& a) {
  T cos_x = std::cos(a._x);
  return dual<T>(std::tan(a._x), (T(1) / (cos_x * cos_x)) * a._dx);
}

template <typename T>
constexpr dual<T> asin(const dual<T>& a) {
  return dual<T>(std::asin(a._x), (T(1) / std::sqrt(T(1) - a._x * a._x)) * a._dx);
}

template <typename T>
constexpr dual<T> acos(const dual<T>& a) {
  return dual<T>(std::acos(a._x), (-T(1) / std::sqrt(T(1) - a._x * a._x)) * a._dx);
}

template <typename T>
constexpr dual<T> atan(const dual<T>& a) {
  return dual<T>(std::atan(a._x), (T(1) / (T(1) + a._x * a._x)) * a._dx);
}

template <typename T>
constexpr dual<T> atan2(const dual<T>& y, const dual<T>& x) {
  return dual<T>(std::atan2(y._x, x._x),
                 (x._x * y._dx - y._x * x._dx) / (x._x * x._x + y._x * y._x));
}

// hyperbolic functions

template <typename T>
constexpr dual<T> sinh(const dual<T>& a) {
  return dual<T>(std::sinh(a._x), std::cosh(a._x) * a._dx);
}

template <typename T>
constexpr dual<T> cosh(const dual<T>& a) {
  return dual<T>(std::cosh(a._x), std::sinh(a._x) * a._dx);
}

template <typename T>
constexpr dual<T> tanh(const dual<T>& a) {
  T cosh_x = std::cosh(a._x);
  return dual<T>(std::tanh(a._x), (T(1) / (cosh_x * cosh_x)) * a._dx);
}

template <typename T>
constexpr dual<T> asinh(const dual<T>& a) {
  return dual<T>(std::asinh(a._x), (T(1) / std::sqrt(a._x * a._x + T(1))) * a._dx);
}

template <typename T>
constexpr dual<T> acosh(const dual<T>& a) {
  return dual<T>(std::acosh(a._x), (T(1) / std::sqrt(a._x * a._x - T(1))) * a._dx);
}

template <typename T>
constexpr dual<T> atanh(const dual<T>& a) {
  return dual<T>(std::atanh(a._x), (T(1) / (T(1) - a._x * a._x)) * a._dx);
}

}  // namespace std