#include <gtest/gtest.h>
#include "vde/sdf/sdf_torch.h"
#include "vde/sdf/sdf_gradient.h"
#include "vde/sdf/sdf_tree.h"
#include "vde/sdf/sdf_primitives.h"
#include "vde/core/point.h"
#include <random>
#include <cmath>
#include <vector>

using namespace vde::sdf;
using vde::core::Point3D;

constexpr double EPS = 1e-5;

// ═══════════════════════════════════════════════════
//  Helper: generate random points
// ═══════════════════════════════════════════════════

std::vector<double> random_points(int n) {
    std::mt19937 rng(42);
    std::uniform_real_distribution<double> dist(-3.0, 3.0);
    std::vector<double> pts(3 * n);
    for (int i = 0; i < 3 * n; ++i) {
        pts[i] = dist(rng);
    }
    return pts;
}

// ═══════════════════════════════════════════════════
//  evaluate_batch
// ═══════════════════════════════════════════════════

TEST(SdfTorchBridge, EvaluateBatch_MatchesSequential) {
    auto root = SdfNode::sphere(2.0);
    const int n = 100;
    auto pts_arr = random_points(n);

    std::vector<double> batch_dist(n);
    evaluate_batch(root, pts_arr.data(), n, batch_dist.data());

    for (int i = 0; i < n; ++i) {
        Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]);
        double expected = evaluate(root, p);
        EXPECT_NEAR(batch_dist[i], expected, EPS);
    }
}

TEST(SdfTorchBridge, EvaluateBatch_CSGTree) {
    auto a = SdfNode::sphere(1.5);
    auto b = SdfNode::box(Point3D(0.8, 0.8, 0.8));
    auto root = SdfNode::op_union(a, b);
    const int n = 50;
    auto pts_arr = random_points(n);

    std::vector<double> batch_dist(n);
    evaluate_batch(root, pts_arr.data(), n, batch_dist.data());

    for (int i = 0; i < n; ++i) {
        Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]);
        double expected = evaluate(root, p);
        EXPECT_NEAR(batch_dist[i], expected, EPS);
    }
}

// ═══════════════════════════════════════════════════
//  gradient_batch
// ═══════════════════════════════════════════════════

TEST(SdfTorchBridge, GradientBatch_MatchesSequential) {
    auto root = SdfNode::sphere(2.0);
    const int n = 100;
    auto pts_arr = random_points(n);

    std::vector<double> batch_grad(3 * n);
    gradient_batch(root, pts_arr.data(), n, batch_grad.data());

    auto fn = [&](const Point3D& p) { return evaluate(root, p); };
    double h = 1e-6;

    for (int i = 0; i < n; ++i) {
        Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]);

        // Sequential gradient via central FD
        double dfdx = (fn(Point3D(p.x()+h, p.y(), p.z())) -
                       fn(Point3D(p.x()-h, p.y(), p.z()))) / (2.0*h);
        double dfdy = (fn(Point3D(p.x(), p.y()+h, p.z())) -
                       fn(Point3D(p.x(), p.y()-h, p.z()))) / (2.0*h);
        double dfdz = (fn(Point3D(p.x(), p.y(), p.z()+h)) -
                       fn(Point3D(p.x(), p.y(), p.z()-h))) / (2.0*h);

        EXPECT_NEAR(batch_grad[i*3],     dfdx, 1e-4);
        EXPECT_NEAR(batch_grad[i*3+1],   dfdy, 1e-4);
        EXPECT_NEAR(batch_grad[i*3+2],   dfdz, 1e-4);
    }
}

TEST(SdfTorchBridge, GradientBatch_Box) {
    auto root = SdfNode::box(Point3D(2.0, 1.0, 3.0));
    const int n = 50;
    auto pts_arr = random_points(n);

    std::vector<double> batch_grad(3 * n);
    gradient_batch(root, pts_arr.data(), n, batch_grad.data());

    // Check that gradient at origin (inside box) is zero or close
    Point3D origin(0, 0, 0);
    Vector3D g = gradient([&](const Point3D& p) { return evaluate(root, p); }, origin);
    // Gradient near origin should point outward (or be near-internal max)
    (void)g;  // just verifying no crash

    for (int i = 0; i < n; ++i) {
        // All gradients should be finite
        EXPECT_TRUE(std::isfinite(batch_grad[i*3]));
        EXPECT_TRUE(std::isfinite(batch_grad[i*3+1]));
        EXPECT_TRUE(std::isfinite(batch_grad[i*3+2]));
    }
}

// ═══════════════════════════════════════════════════
//  evaluate_with_gradient_batch
// ═══════════════════════════════════════════════════

TEST(SdfTorchBridge, EvaluateWithGradientBatch_BothCorrect) {
    auto root = SdfNode::sphere(1.5);
    const int n = 50;
    auto pts_arr = random_points(n);

    std::vector<double> batch_dist(n);
    std::vector<double> batch_grad(3 * n);
    evaluate_with_gradient_batch(root, pts_arr.data(), n,
                                  batch_dist.data(), batch_grad.data());

    // Compare distances against sequential
    std::vector<double> seq_dist(n);
    evaluate_batch(root, pts_arr.data(), n, seq_dist.data());
    for (int i = 0; i < n; ++i) {
        EXPECT_NEAR(batch_dist[i], seq_dist[i], EPS);
    }

    // Compare gradients against sequential
    std::vector<double> seq_grad(3 * n);
    gradient_batch(root, pts_arr.data(), n, seq_grad.data());
    for (int i = 0; i < 3 * n; ++i) {
        EXPECT_NEAR(batch_grad[i], seq_grad[i], 1e-10);
    }
}

TEST(SdfTorchBridge, EvaluateWithGradientBatch_CSGTree) {
    // Smooth union of sphere and box
    auto sphere = SdfNode::sphere(1.0);
    auto box_node = SdfNode::box(Point3D(0.8, 0.8, 0.8));
    auto root = SdfNode::smooth_union(sphere, box_node, 0.2);
    const int n = 30;
    auto pts_arr = random_points(n);

    std::vector<double> batch_dist(n);
    std::vector<double> batch_grad(3 * n);
    evaluate_with_gradient_batch(root, pts_arr.data(), n,
                                  batch_dist.data(), batch_grad.data());

    for (int i = 0; i < n; ++i) {
        EXPECT_TRUE(std::isfinite(batch_dist[i]));
        EXPECT_TRUE(std::isfinite(batch_grad[i*3]));
        EXPECT_TRUE(std::isfinite(batch_grad[i*3+1]));
        EXPECT_TRUE(std::isfinite(batch_grad[i*3+2]));
    }
}