#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 std::array<T, 2>& x) {
    return x[0] * x[0] + std::sin(x[1]);
  };

  std::array<double, 2> point = {1.0, 0.0};
  std::array<double, 2> 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 std::array<T, 2>& x) {
    return std::array<T, 2>{x[0] * x[0], std::sin(x[1])};
  };

  std::array<double, 2> point = {1.0, 0.0};
  auto jacob = jacobian<double, 2, 2>(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 std::array<T, 2>& x) {
    return x[0] * x[0] + x[1] * x[1];
  };

  std::array<double, 2> point = {1.0, 2.0};
  auto hess = hessian<double, 2>(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<std::array<double, 2>> test_data;

  // loss function
  template <typename T>
  T operator()(const std::array<T, 3>& 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},
  };

  std::array<double, 3> params = {1.0, -4.0, 4.0};
  
  auto grad = gradient<double, 3>(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, 3>(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
    }
  }
}