#include "lantern.hpp"
#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <new>
#include <numbers>

namespace {
using namespace lantern;
bool finite(Vec3 v, double bound = 1000) {
    return std::isfinite(v.x) && std::isfinite(v.y) && std::isfinite(v.z) &&
           std::abs(v.x) <= bound && std::abs(v.y) <= bound && std::abs(v.z) <= bound;
}
double dot(Vec3 a, Vec3 b) { return a.x*b.x + a.y*b.y + a.z*b.z; }
bool valid_ray(const Ray& ray) {
    return finite(ray.origin) && finite(ray.direction, 1) && std::abs(dot(ray.direction, ray.direction) - 1) < 1e-10;
}
double component(Vec3 v, unsigned axis) { return axis == 0 ? v.x : axis == 1 ? v.y : v.z; }
bool slab(const Ray& ray, const Aabb& bounds, double nearest) {
    double low = 0, high = nearest;
    for (unsigned axis = 0; axis < 3; ++axis) {
        const double origin = component(ray.origin, axis), direction = component(ray.direction, axis);
        const double minimum = component(bounds.low, axis), maximum = component(bounds.high, axis);
        if (direction == 0) { // Includes signed zero. Never divide a parallel direction.
            if (origin < minimum || origin > maximum) return false;
            continue;
        }
        double a = (minimum - origin) / direction, b = (maximum - origin) / direction;
        if (a > b) std::swap(a, b);
        low = std::max(low, a); high = std::min(high, b);
        if (high < low) return false; // Tangency is conservatively accepted.
    }
    return true;
}
std::optional<Hit> sphere_hit(const Ray& ray, const Sphere& sphere, double nearest) {
    const Vec3 oc{ray.origin.x-sphere.center.x, ray.origin.y-sphere.center.y, ray.origin.z-sphere.center.z};
    const double a = dot(ray.direction, ray.direction);
    const double center_t = -dot(oc, ray.direction)/a;
    // Subtracting half_b^2 and a*c loses a tiny perpendicular miss when the
    // origin is far along the ray. Compute the closest-point vector instead;
    // fma keeps the projection subtraction to one rounding per component.
    const Vec3 perpendicular{std::fma(center_t,ray.direction.x,oc.x),
                             std::fma(center_t,ray.direction.y,oc.y),
                             std::fma(center_t,ray.direction.z,oc.z)};
    const double radial = sphere.radius*sphere.radius-dot(perpendicular,perpendicular);
    if (radial < 0) return {};
    const double root = std::sqrt(radial/a);
    double distance = center_t-root;
    if (distance <= 1e-5) distance = center_t+root;
    if (distance <= 1e-5 || distance > nearest) return {};
    return Hit{sphere.id, distance};
}
void retain(std::optional<Hit>& best, std::optional<Hit> candidate) {
    if (candidate && (!best || candidate->distance < best->distance ||
        (candidate->distance == best->distance && candidate->id < best->id))) best = candidate;
}
} // namespace

namespace lantern {
std::uint64_t Rng::next() noexcept {
    // SplitMix64 (Sebastiano Vigna, public-domain implementation; see NOTICE).
    // Unsigned wraparound is intentional and defined by C++.
    std::uint64_t value = (state_ += 0x9e3779b97f4a7c15ULL);
    value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL;
    value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL;
    return value ^ (value >> 31);
}
double Rng::uniform() noexcept { return static_cast<double>(next() >> 11) * 0x1.0p-53; }

std::expected<DirectionSample, Error> cosine_hemisphere(double u, double v) {
    if (!std::isfinite(u) || !std::isfinite(v) || u < 0 || u >= 1 || v < 0 || v >= 1)
        return std::unexpected(Error{"Sampler inputs must be finite in [0, 1)."});
    const double radius = std::sqrt(u), angle = 2 * std::numbers::pi * v;
    const double z = std::sqrt(1 - u);
    return DirectionSample{{radius*std::cos(angle), radius*std::sin(angle), z}, z/std::numbers::pi};
}

std::expected<void, Error> Moments::add(double value) {
    if (!std::isfinite(value) || std::abs(value) > 1e12 || count_ >= max_samples)
        return std::unexpected(Error{"Sample value or count exceeds the statistics budget."});
    ++count_;
    const double delta = value - mean_;
    mean_ += delta / static_cast<double>(count_);
    m2_ += delta * (value - mean_);
    return {};
}
std::expected<double, Error> cosine_integral(std::size_t samples, std::uint64_t seed) {
    if (samples == 0 || samples > max_samples) return std::unexpected(Error{"Sample count must be in 1–1048576."});
    Rng rng{seed}; Moments moments;
    for (std::size_t i = 0; i < samples; ++i) {
        // Uniform hemisphere z has density 1/(2*pi) per steradian.
        if (auto added = moments.add(2*std::numbers::pi*rng.uniform()); !added) return std::unexpected(added.error());
    }
    return moments.mean(); // Analytic answer: integral cos(theta) d omega = pi.
}
std::expected<double, Error> power_weight(double a, double b) {
    if (!std::isfinite(a) || !std::isfinite(b) || a < 0 || b < 0 || a > 1e6 || b > 1e6 || a+b == 0)
        return std::unexpected(Error{"PDFs must be bounded nonnegative with nonzero total."});
    const double scale = std::max(a,b);
    a /= scale; b /= scale; // Avoid squaring tiny PDFs into an all-zero denominator.
    return a*a / (a*a + b*b); // One sample from each strategy; same PDF measure.
}

std::expected<bool, Error> box_hit(const Ray& ray, const Aabb& box) {
    if (!valid_ray(ray) || !finite(box.low, 1100) || !finite(box.high, 1100) ||
        box.low.x > box.high.x || box.low.y > box.high.y || box.low.z > box.high.z)
        return std::unexpected(Error{"Box query requires bounded ordered bounds and a finite unit ray."});
    return slab(ray, box, std::numeric_limits<double>::infinity());
}
std::expected<SceneIndex, Error> SceneIndex::build(std::span<const Sphere> input) {
    if (input.size() > max_spheres) return std::unexpected(Error{"Index exceeds the 128-sphere budget."});
    for (std::size_t i = 0; i < input.size(); ++i) {
        if (!finite(input[i].center) || !std::isfinite(input[i].radius) || input[i].radius < 0.01 || input[i].radius > 100)
            return std::unexpected(Error{"Sphere geometry is outside the finite bounded domain."});
        for (std::size_t j = 0; j < i; ++j)
            if (input[i].id == input[j].id) return std::unexpected(Error{"Sphere identifiers must be unique."});
    }
    try {
        SceneIndex result;
        result.spheres_.assign(input.begin(), input.end());
        if (!input.empty()) {
            result.nodes_.reserve(input.size()*2 - 1); // Bounded before multiplication.
            result.build_node(0, input.size());
        }
        return result;
    } catch (const std::bad_alloc&) { return std::unexpected(Error{"Index allocation failed."}); }
}
std::size_t SceneIndex::build_node(std::size_t first, std::size_t count) {
    Aabb bounds{{1100,1100,1100}, {-1100,-1100,-1100}};
    for (std::size_t i = first; i < first+count; ++i) {
        const auto& sphere = spheres_[i]; const auto c = sphere.center; const auto r = sphere.radius;
        bounds.low = {std::min(bounds.low.x,c.x-r), std::min(bounds.low.y,c.y-r), std::min(bounds.low.z,c.z-r)};
        bounds.high = {std::max(bounds.high.x,c.x+r), std::max(bounds.high.y,c.y+r), std::max(bounds.high.z,c.z+r)};
    }
    const auto index = nodes_.size(); nodes_.push_back({bounds, first, count, 0, 0});
    if (count <= 2) return index;
    const Vec3 extent{bounds.high.x-bounds.low.x, bounds.high.y-bounds.low.y, bounds.high.z-bounds.low.z};
    const unsigned axis = extent.y > extent.x ? (extent.z > extent.y ? 2U : 1U) : (extent.z > extent.x ? 2U : 0U);
    std::sort(spheres_.begin()+static_cast<std::ptrdiff_t>(first), spheres_.begin()+static_cast<std::ptrdiff_t>(first+count),
        [axis](const Sphere& a, const Sphere& b) { const auto x=component(a.center,axis), y=component(b.center,axis); return x<y || (x==y && a.id<b.id); });
    // Equal-count splitting makes recursion logarithmic, including tied centroids.
    const auto left = build_node(first,count/2), right = build_node(first+count/2,count-count/2);
    nodes_[index].left=left; nodes_[index].right=right; nodes_[index].count=0;
    return index;
}
std::expected<std::optional<Hit>, Error> SceneIndex::linear_closest(const Ray& ray) const {
    if (!valid_ray(ray)) return std::unexpected(Error{"Query requires a finite bounded unit ray."});
    std::optional<Hit> best;
    for (const auto& sphere : spheres_) retain(best,sphere_hit(ray,sphere,best ? best->distance : std::numeric_limits<double>::infinity()));
    return best;
}
std::expected<std::optional<Hit>, Error> SceneIndex::closest(const Ray& ray) const {
    if (!valid_ray(ray)) return std::unexpected(Error{"Query requires a finite bounded unit ray."});
    std::optional<Hit> best;
    if (nodes_.empty()) return best;
    std::array<std::size_t,max_spheres*2> stack{}; std::size_t used=1;
    while (used > 0) {
        const auto& node = nodes_[stack[--used]];
        if (!slab(ray,node.bounds,best ? best->distance : std::numeric_limits<double>::infinity())) continue;
        if (node.count != 0) {
            for (std::size_t i=node.first; i<node.first+node.count; ++i)
                retain(best,sphere_hit(ray,spheres_[i],best ? best->distance : std::numeric_limits<double>::infinity()));
        } else {
            if (used+2 > stack.size()) return std::unexpected(Error{"Traversal stack budget exceeded."});
            stack[used++]=node.right; stack[used++]=node.left;
        }
    }
    return best;
}
} // namespace lantern