using SpecialFunctions


"""
    mc_volume(i, xs, nmc, searcher)

Estimate the area and volume of the `i`-th Voronoi cell from the Voronoi Diagram generated by `xs`
by a Monte Carlo estimate from random rays
"""
function mc_volume(i, xs, nmc=1000, searcher = Raycast(xs))
    y, δy, V, A = mc_integrate(x->1, i, xs, nmc, 0, searcher)
    A, V
end

"""
    mc_volumes(xs::Points, nmc=1000)

Estimate the areas and volumes of the Voronoi Cells generated by `xs` using `nmc` Monte Carlo samples.

# Returns
- `SparseMatrix`: the areas of the common boundaries of two cells
- `Vector`: the volumes of each cell
"""
function mc_volumes(xs::Points, nmc=1000)
    V = zeros(length(xs))
    A = spzeros(length(xs), length(xs))

    searcher = Raycast(xs)

    for i in 1:length(xs)
        a, v = mc_volume(i, xs, nmc, searcher)
        V[i] = v
        A[:, i] = a
    end
    #A = (A + A') / 2
    A, V
end

"""
    mc_volumes(sig::Vertices, xs::Points, nmc=1000)

In the case when the simplicial complex is already known this information can be used
to speed up the Monte-Carlo sampling by restricting the search space
"""
function mc_volumes(σ::Vertices, xs::Points, nmc=1000)
    V = zeros(length(xs))
    A = spzeros(length(xs), length(xs))

    searcher = Raycast(xs)
    neigh = neighbors(σ)

    for i in 1:length(xs)
        ix = [i; neigh[i]]
        a, v = mc_volume(1, xs[ix], nmc, RaycastBruteforce())
        V[i] = v
        A[ix, i] = a
    end
    #A = (A + A') / 2
    A, V
end

"""
    neighbors(sig::Vertices) --> Dict{Int, Vector{Int}}
Compute the neighbors of all generators.
"""
neighbors(σ::Vertices) = neighbors(keys(σ))
function neighbors(σ)
    d = Dict{Int, Vector{Int}}()
    for s in σ
        for i in s
            v = get!(d, i, Vector{Int}())
            append!(v, s)
        end
    end
    for i in keys(d)
        unique!(d[i])
        filter!(x->x != i, d[i])
    end
    return d
end

"""
    mc_integrate(f::Function, i::Int, xs::Points, nmc=1000, nmc2=1000, searcher=Raycast(xs))
Integrate function `f` over cell `i` and its boundary using `nmc` rays per cell and `nmc2` points per ray for the volume integral.

# Returns:
- `Vf::Real`: the volume integral of `f`
- `Af::SparseVector`: Af[j] is the surface integral of `f` over the intersection between cells i and j.
- `V::Real`: the volume of cell `i`
- `A::SparseVector`: A[j] is the surface area of the intersection between cells i and j.
"""
function mc_integrate(f::Function, i::Int, xs::Vector, nmc=1000, nmc2=1000, searcher = Raycast(xs))
    x = xs[i]
    d = length(x)

    # volume/area computations are basically for free, leaving them commented for now
    V = 0.
    A = spzeros(length(xs))

    y = 0.
    δy = spzeros(length(xs))

    for _ in 1:nmc
        u = normalize(randn(d))
        (j, t) = raycast([i], x, u, xs, searcher)

        V += t^d
        for _ in 1:nmc2
            r = t * rand()
            x′ = x + u * r
            y += f(x′) * r^(d-1) * t
        end

        if t < Inf
            j = j[1] == i ? j[2] : j[1]
            normal = normalize(xs[j] - xs[i])
            dA = t ^ (d-1) / abs(dot(normal, u))
            A[j] += dA
            δy[j] += dA * f(x + t*u)
        end
    end

    c_vol = pi^(d/2) / gamma(d/2 + 1)
    c_area = d * c_vol

    V *= c_vol / nmc
    A *= c_area / nmc

    y *= c_area / nmc / nmc2
    δy *= c_area / nmc

    y, δy, V, A
end

function mc_integrate(f, σ::Vertices, xs::Points, nmc, nmc2)
    return mc_integrate_i((x,i)->f(x), σ, xs, nmc, nmc2)
end

# MC integration using f(x, i) function which is also passed the index of the cell
function mc_integrate_i(f, σ::Vertices, xs::Points, nmc, nmc2)
    n = length(xs)
    ys = zeros(n)
    ∂ys = zeros(n)

    neigh = neighbors(σ)

    for i in 1:n
        ix = [i; neigh[i]]
        y, ∂y, V, A = mc_integrate(x->f(x, i), 1, xs[ix], nmc, nmc2, RaycastBruteforce())
        ys[i] = y
        ∂ys[i] = sum(∂y)
    end
    return ys, ∂ys
end
