#pragma once
#include "vde/sdf/sdf_primitives.h"
#include "vde/sdf/sdf_operations.h"
#include "vde/core/point.h"
#include <cmath>
#include <functional>

namespace vde::sdf {

using core::Point3D;
using core::Vector3D;

// ═══════════════════════════════════════════════
//  Finite-Difference Gradient
// ═══════════════════════════════════════════════

/// Gradient of SDF at point p (normal direction)
/// Computed via central finite differences: (f(p+h) - f(p-h)) / (2h)
/// @param f SDF function (primitives or evaluate from tree)
/// @param p evaluation point
/// @param h step size (default 1e-6)
[[nodiscard]] inline Vector3D gradient(
    const std::function<double(const Point3D&)>& f,
    const Point3D& p, double h = 1e-6)
{
    Vector3D g;
    double inv_2h = 0.5 / h;
    g.x() = (f(Point3D(p.x() + h, p.y(), p.z())) - f(Point3D(p.x() - h, p.y(), p.z()))) * inv_2h;
    g.y() = (f(Point3D(p.x(), p.y() + h, p.z())) - f(Point3D(p.x(), p.y() - h, p.z()))) * inv_2h;
    g.z() = (f(Point3D(p.x(), p.y(), p.z() + h)) - f(Point3D(p.x(), p.y(), p.z() - h))) * inv_2h;
    return g;
}

// ═══════════════════════════════════════════════
//  Gradient Result Struct
// ═══════════════════════════════════════════════

/// Forward-mode AD: compute SDF value AND gradient in one pass
/// Returns {sdf_value, gradient_vector}
struct GradResult {
    double value;
    Vector3D grad;  // df/dx, df/dy, df/dz
};

[[nodiscard]] inline GradResult evaluate_with_gradient(
    const std::function<double(const Point3D&)>& f,
    const Point3D& p, double h = 1e-6)
{
    return GradResult{f(p), gradient(f, p, h)};
}

// ═══════════════════════════════════════════════
//  Parameter Gradients (Finite Difference)
// ═══════════════════════════════════════════════

/// Compute df/dradius for a sphere at point p
[[nodiscard]] inline double sphere_radius_gradient(const Point3D& p, double radius, double h = 1e-6)
{
    return (sphere(p, radius + h) - sphere(p, radius - h)) / (2.0 * h);
}

/// Compute df/d(extents[i]) for a box (axis 0=x, 1=y, 2=z)
[[nodiscard]] inline double box_extent_gradient(const Point3D& p, const Point3D& extents,
                                                  int axis, double h = 1e-6)
{
    Point3D ext_plus = extents;
    Point3D ext_minus = extents;
    ext_plus[axis] += h;
    ext_minus[axis] -= h;
    return (box(p, ext_plus) - box(p, ext_minus)) / (2.0 * h);
}

/// Compute df/d(major_radius) for a torus
[[nodiscard]] inline double torus_major_gradient(const Point3D& p, double major_r, double minor_r,
                                                   double h = 1e-6)
{
    return (torus(p, major_r + h, minor_r) - torus(p, major_r - h, minor_r)) / (2.0 * h);
}

/// Compute df/d(minor_radius) for a torus
[[nodiscard]] inline double torus_minor_gradient(const Point3D& p, double major_r, double minor_r,
                                                   double h = 1e-6)
{
    return (torus(p, major_r, minor_r + h) - torus(p, major_r, minor_r - h)) / (2.0 * h);
}

/// Compute df/dheight for cylinder
[[nodiscard]] inline double cylinder_height_gradient(const Point3D& p, double radius, double height,
                                                       double h = 1e-6)
{
    return (cylinder(p, radius, height + h) - cylinder(p, radius, height - h)) / (2.0 * h);
}

/// Compute df/dradius for cylinder
[[nodiscard]] inline double cylinder_radius_gradient(const Point3D& p, double radius, double height,
                                                       double h = 1e-6)
{
    return (cylinder(p, radius + h, height) - cylinder(p, radius - h, height)) / (2.0 * h);
}

// ═══════════════════════════════════════════════
//  Analytical Gradients (Exact)
// ═══════════════════════════════════════════════

/// Exact gradient of sphere: radial unit vector (p - center) / |p - center|
/// For sphere centered at origin: p / |p|
[[nodiscard]] inline Vector3D sphere_gradient_analytic(const Point3D& p, double /*radius*/)
{
    double n = p.norm();
    if (n < 1e-30) return Vector3D::Zero();
    return p / n;
}

/// Exact gradient of box: sign(p) component-wise for exterior points
/// For interior points, the gradient is the unit normal of the nearest face
[[nodiscard]] inline Vector3D box_gradient_analytic(const Point3D& p, const Point3D& half_extents)
{
    Vector3D g = Vector3D::Zero();
    if (std::abs(p.x()) > half_extents.x()) g.x() = (p.x() > 0.0 ? 1.0 : -1.0);
    if (std::abs(p.y()) > half_extents.y()) g.y() = (p.y() > 0.0 ? 1.0 : -1.0);
    if (std::abs(p.z()) > half_extents.z()) g.z() = (p.z() > 0.0 ? 1.0 : -1.0);

    // For interior: gradient points toward nearest face
    if (g.norm() < 1e-30 && half_extents.norm() > 1e-30) {
        // Find the dimension where the point is closest to a face
        Point3D q = Point3D(
            half_extents.x() - std::abs(p.x()),
            half_extents.y() - std::abs(p.y()),
            half_extents.z() - std::abs(p.z())
        );
        // Smallest q component = closest face
        if (q.x() <= q.y() && q.x() <= q.z()) g.x() = (p.x() >= 0.0 ? 1.0 : -1.0);
        else if (q.y() <= q.x() && q.y() <= q.z()) g.y() = (p.y() >= 0.0 ? 1.0 : -1.0);
        else g.z() = (p.z() >= 0.0 ? 1.0 : -1.0);
    }
    return g;
}

/// Exact gradient of torus
/// SDF: sqrt((sqrt(px^2 + pz^2) - major)^2 + py^2) - minor
[[nodiscard]] inline Vector3D torus_gradient_analytic(const Point3D& p, double major_r, double minor_r)
{
    double rho = std::sqrt(p.x() * p.x() + p.z() * p.z());
    if (rho < 1e-30) {
        // On the Y axis: gradient points radially outward in XZ plane
        return Vector3D(0.0, (p.y() >= 0.0 ? 1.0 : -1.0), 0.0);
    }
    double inner = std::sqrt((rho - major_r) * (rho - major_r) + p.y() * p.y());
    if (inner < 1e-30) return Vector3D::Zero();
    double d_rho = (rho - major_r) / inner;
    return Vector3D(
        d_rho * p.x() / rho,
        p.y() / inner,
        d_rho * p.z() / rho
    );
}

/// Exact gradient of cylinder (infinite along Y): normalize(p.x, 0, p.z)
[[nodiscard]] inline Vector3D cylinder_gradient_analytic(const Point3D& p, double /*radius*/)
{
    double rho = std::sqrt(p.x() * p.x() + p.z() * p.z());
    if (rho < 1e-30) return Vector3D::UnitY(); // degenerate case: along the axis
    return Vector3D(p.x() / rho, 0.0, p.z() / rho);
}

/// Exact gradient of plane: normal
[[nodiscard]] inline Vector3D plane_gradient_analytic(const Vector3D& normal)
{
    return normal;
}

/// Exact gradient of capsule: gradient of segment distance
[[nodiscard]] inline Vector3D capsule_gradient_analytic(const Point3D& p,
                                                          const Point3D& a, const Point3D& b,
                                                          double radius)
{
    Point3D pa = p - a;
    Point3D ba = b - a;
    double ba_sq = ba.squaredNorm();
    if (ba_sq < 1e-30) {
        // Degenerate to sphere
        double n = pa.norm();
        if (n < 1e-30) return Vector3D::Zero();
        return pa / n;
    }
    double h = std::clamp(pa.dot(ba) / ba_sq, 0.0, 1.0);
    Point3D closest = a + ba * h;
    Point3D delta = p - closest;
    double d = delta.norm();
    if (d < 1e-30) return Vector3D::Zero();
    return delta / d;
}

// ═══════════════════════════════════════════════
//  Chain Rule for CSG Operations
// ═══════════════════════════════════════════════

/// Gradient after union: gradient of the child that is closer
[[nodiscard]] inline Vector3D chain_union(const Vector3D& grad_a, const Vector3D& grad_b,
                                            double d1, double d2)
{
    return (d1 < d2) ? grad_a : grad_b;
}

/// Gradient after intersection: gradient of the child that is farther
[[nodiscard]] inline Vector3D chain_intersection(const Vector3D& grad_a, const Vector3D& grad_b,
                                                   double d1, double d2)
{
    return (d1 > d2) ? grad_a : grad_b;
}

/// Gradient after difference: max(-d1, d2)
[[nodiscard]] inline Vector3D chain_difference(const Vector3D& grad_a, const Vector3D& grad_b,
                                                 double d1, double d2)
{
    return (-d1 > d2) ? Vector3D(-grad_a.x(), -grad_a.y(), -grad_a.z()) : grad_b;
}

/// Gradient after smooth union with full chain rule
/// f = d1*(1-h) + d2*h - k*h*(1-h) where h = clamp(0.5+0.5*(d2-d1)/k, 0, 1)
[[nodiscard]] inline Vector3D chain_smooth_union(const Vector3D& grad_a, const Vector3D& grad_b,
                                                   double d1, double d2, double k)
{
    if (k <= 0.0) return chain_union(grad_a, grad_b, d1, d2);
    double h = std::clamp(0.5 + 0.5 * (d2 - d1) / k, 0.0, 1.0);
    if (h <= 0.0) return grad_a;
    if (h >= 1.0) return grad_b;

    // Full derivatives w.r.t. d1, d2:
    // ∂f/∂d1 = 2 - 3h,  ∂f/∂d2 = 3h - 1
    double w1 = 2.0 - 3.0 * h;
    double w2 = 3.0 * h - 1.0;
    return grad_a * w1 + grad_b * w2;
}

// ═══════════════════════════════════════════════
//  Chain Rule for Domain Transforms
// ═══════════════════════════════════════════════

/// Chain through translation: gradient unchanged (isometric)
[[nodiscard]] inline GradResult chain_translate(const GradResult& child_result, const Point3D& /*offset*/)
{
    return child_result;
}

/// Chain through rotation around Y axis (inverse rotation applied to gradient)
/// op_rotate applies R(-θ) to the point, so chain applies R(θ) to the gradient
[[nodiscard]] inline GradResult chain_rotate(const GradResult& child_result, double angle_rad)
{
    GradResult result = child_result;
    double c = std::cos(angle_rad);
    double s = std::sin(angle_rad);
    // R(θ) = [[c, 0, -s], [0, 1, 0], [s, 0, c]]
    result.grad = Vector3D(
        c * child_result.grad.x() - s * child_result.grad.z(),
        child_result.grad.y(),
        s * child_result.grad.x() + c * child_result.grad.z()
    );
    return result;
}

/// Chain through non-uniform scale: divide gradient components by scale factors
[[nodiscard]] inline GradResult chain_scale(const GradResult& child_result, const Point3D& factors)
{
    GradResult result = child_result;
    result.grad = Vector3D(
        child_result.grad.x() / factors.x(),
        child_result.grad.y() / factors.y(),
        child_result.grad.z() / factors.z()
    );
    return result;
}

/// Chain through twist around Y axis
/// Twist: T(p) rotates (x,z) by +amount*p.y
/// Chain: J_T^T * child_grad where J_T includes y-dependence of the rotation angle
[[nodiscard]] inline GradResult chain_twist(const GradResult& child_result,
                                              double amount, const Point3D& p)
{
    double angle = amount * p.y();
    double c = std::cos(angle);
    double s = std::sin(angle);

    // Forward-twisted position
    double tp_x = p.x() * c - p.z() * s;
    double tp_z = p.x() * s + p.z() * c;

    GradResult result = child_result;

    // Spatial part: inverse rotation R(-angle) applied to child gradient
    double gx = c * child_result.grad.x() + s * child_result.grad.z();
    double gz = -s * child_result.grad.x() + c * child_result.grad.z();

    // Y-component: extra term from angle's dependence on p.y
    // ∂/∂y correction: amount * (tp_x * gz_child - tp_z * gx_child)
    double gy = child_result.grad.y() + amount * (tp_x * child_result.grad.z()
                                                 - tp_z * child_result.grad.x());

    result.grad = Vector3D(gx, gy, gz);
    return result;
}

/// Chain through repeat: gradient unchanged (periodic domain)
[[nodiscard]] inline GradResult chain_repeat(const GradResult& child_result,
                                               const Point3D& /*cell*/, const Point3D& /*p*/)
{
    return child_result;
}

// ═══════════════════════════════════════════════
//  Parameter Gradient Vector (All Shape Params)
// ═══════════════════════════════════════════════

/// Parameter gradient struct: holds all possible shape-parameter sensitivities
struct ParamGrad {
    double d_radius = 0.0;
    double d_height = 0.0;
    Point3D d_extents = Point3D::Zero();
    double d_major = 0.0;
    double d_minor = 0.0;
    double d_angle = 0.0;
    Vector3D d_normal = Vector3D::Zero();
    double d_offset = 0.0;
    Point3D d_translate = Point3D::Zero();
    double d_rotate = 0.0;
    Point3D d_scale = Point3D::Zero();
};

/// Compute parameter gradients for a sphere primitive
[[nodiscard]] inline ParamGrad sphere_param_grad(const Point3D& p, double radius, double h = 1e-6)
{
    ParamGrad pg;
    pg.d_radius = sphere_radius_gradient(p, radius, h);
    return pg;
}

/// Compute parameter gradients for a box primitive
[[nodiscard]] inline ParamGrad box_param_grad(const Point3D& p, const Point3D& half_extents,
                                                double h = 1e-6)
{
    ParamGrad pg;
    pg.d_extents.x() = box_extent_gradient(p, half_extents, 0, h);
    pg.d_extents.y() = box_extent_gradient(p, half_extents, 1, h);
    pg.d_extents.z() = box_extent_gradient(p, half_extents, 2, h);
    return pg;
}

/// Compute parameter gradients for a cylinder primitive
[[nodiscard]] inline ParamGrad cylinder_param_grad(const Point3D& p, double radius, double height,
                                                     double h = 1e-6)
{
    ParamGrad pg;
    pg.d_radius = cylinder_radius_gradient(p, radius, height, h);
    pg.d_height = cylinder_height_gradient(p, radius, height, h);
    return pg;
}

} // namespace vde::sdf