#include "autoopt/dual.hpp"

#include <gtest/gtest.h>

#include "autoopt/derivative.hpp"

using namespace autoopt;

TEST(DualTest, BasicOperations) {
  dual<double> a(2.0, 1.0);  // a = 2.0, da/dx = 1.0
  dual<double> b(3.0, 0.0);  // b = 3.0, db/dx = 0.0

  dual<double> c = a + b;
  EXPECT_DOUBLE_EQ(c._x, 5.0);
  EXPECT_DOUBLE_EQ(c._dx, 1.0);

  dual<double> d = a * b;
  EXPECT_DOUBLE_EQ(d._x, 6.0);
  EXPECT_DOUBLE_EQ(d._dx, 3.0);

  dual<double> e = a / b;
  EXPECT_DOUBLE_EQ(e._x, 2.0 / 3.0);
  EXPECT_DOUBLE_EQ(e._dx, 1.0 / 3.0);
}

TEST(DualTest, StandardFunctions) {
  dual<double> a(0.5, 1.0);  // a = 0.5, da/dx = 1.0

  dual<double> b = std::sin(a);
  EXPECT_DOUBLE_EQ(b._x, std::sin(0.5));
  EXPECT_DOUBLE_EQ(b._dx, std::cos(0.5));

  dual<double> c = std::exp(a);
  EXPECT_DOUBLE_EQ(c._x, std::exp(0.5));
  EXPECT_DOUBLE_EQ(c._dx, std::exp(0.5));

  dual<double> d = std::log(a);
  EXPECT_DOUBLE_EQ(d._x, std::log(0.5));
  EXPECT_DOUBLE_EQ(d._dx, 1.0 / 0.5);
}

TEST(DualTest, DerivativeFunction) {
  auto func = []<typename T>(const T& x) { return std::sin(x) * std::exp(x); };

  for (double val : {0.0, 0.5, 1.0, 2.0}) {
    double deriv = derivative(func, val);
    double expected = (std::cos(val) + std::sin(val)) * std::exp(val);
    EXPECT_DOUBLE_EQ(deriv, expected);
  }
}

TEST(DualTest, GradientFunction) {
  auto func = []<typename T>(const Eigen::VectorX<T>& x) {
    return x(0) * x(0) + std::sin(x(1));
  };

  Eigen::VectorX<double> point(2);
  point << 1.0, 0.0;
  Eigen::VectorX<double> grad = gradient(func, point);

  EXPECT_DOUBLE_EQ(grad(0), 2.0 * point(0));      // d/dx1
  EXPECT_DOUBLE_EQ(grad(1), std::cos(point(1)));  // d/dx2
}

TEST(DualTest, JacobianFunction) {
  auto func = []<typename T>(const Eigen::VectorX<T>& x) {
    Eigen::VectorX<T> y(2);
    y << x(0) * x(0), std::sin(x(1));
    return y;
  };

  Eigen::VectorX<double> point(2);
  point << 1.0, 0.0;
  auto jacob = jacobian<double>(func, point);

  EXPECT_DOUBLE_EQ(jacob(0, 0), 2.0 * point(0));      // d(f1)/d(x1)
  EXPECT_DOUBLE_EQ(jacob(0, 1), 0.0);                 // d(f1)/d(x2)
  EXPECT_DOUBLE_EQ(jacob(1, 0), 0.0);                 // d(f2)/d(x1)
  EXPECT_DOUBLE_EQ(jacob(1, 1), std::cos(point(1)));  // d(f2)/d(x2)
}

TEST(DualTest, HessianFunction) {
  auto func = []<typename T>(const Eigen::VectorX<T>& x) {
    return x(0) * x(0) + x(1) * x(1);
  };

  Eigen::VectorX<double> point(2);
  point << 1.0, 2.0;
  auto hess = hessian<double>(func, point);

  EXPECT_DOUBLE_EQ(hess(0, 0), 2.0);  // d²f/dx1²
  EXPECT_DOUBLE_EQ(hess(0, 1), 0.0);  // d²f/dx1dx2
  EXPECT_DOUBLE_EQ(hess(1, 0), 0.0);  // d²f/dx2dx1
  EXPECT_DOUBLE_EQ(hess(1, 1), 2.0);  // d²f/dx2²
}

struct opti_func {
  std::vector<Eigen::Vector2<double>> test_data;

  // loss function
  template <typename T>
  T operator()(const Eigen::VectorX<T>& params) const {
    T sum = T(0);
    for (const auto& data_point : test_data) {
      T x = T{data_point(0)};
      T y_true = T{data_point(1)};
      T a = params(0);
      T b = params(1);
      T c = params(2);
      T y_pred = a * x * x + b * x + c;
      T error = y_pred - y_true;
      sum = sum + error * error;
    }
    return sum / T(test_data.size());
  }
};

TEST(DualTest, OptimizationFunction) {
  opti_func f;
  f.test_data = {
      {0.0, 4.0},
      {1.0, 1.0},
      {2.0, 0.0},
      {3.0, 1.0},
      {4.0, 4.0},
  };

  Eigen::VectorX<double> params(3);
  params << 1.0, -4.0, 4.0;
  
  auto grad = gradient<double>(f, params);

  EXPECT_DOUBLE_EQ(grad(0), 0.0);  // dL/da
  EXPECT_DOUBLE_EQ(grad(1), 0.0);  // dL/db
  EXPECT_DOUBLE_EQ(grad(2), 0.0);  // dL/dc

  auto hess = hessian<double>(f, params);
  for (std::size_t i = 0; i < 3; ++i) {
    for (std::size_t j = 0; j < 3; ++j) {
      EXPECT_GE(hess(i, j), 0.0);  // Hessian should be positive semi-definite
    }
  }
}