# Line "drawing" iterator
struct LineIterator{N}
    start :: CartesianIndex{N}
    ϕ     :: Float64
    θ     :: Float64
end

RandomLineIterator(start) =
    LineIterator(start, 2π*rand(Float64), π*rand(Float64) - π/2)

Base.IteratorSize(::LineIterator) = Base.IsInfinite()
Base.iterate(iter :: LineIterator) = iter.start, 0.1

function Base.iterate(iter :: LineIterator{2}, r :: Float64)
    r    = r + sqrt(2)
    x, y = iter.start[1], iter.start[2]
    ϕ    = iter.ϕ

    xn = x + r*cos(ϕ) |> floor |> Int
    yn = y + r*sin(ϕ) |> floor |> Int
    return CartesianIndex(xn, yn), r
end

function Base.iterate(iter :: LineIterator{3}, r :: Float64)
    r       = r + sqrt(3)
    x, y, z = iter.start[1], iter.start[2], iter.start[3]
    ϕ       = iter.ϕ
    θ       = iter.θ

    xn = x + r*cos(θ)*cos(ϕ) |> floor |> Int
    yn = y + r*cos(θ)*sin(ϕ) |> floor |> Int
    zn = z + r*sin(θ)        |> floor |> Int
    return CartesianIndex(xn, yn, zn), r
end

# Theoretical two-point functions for the void phase for overlapping
# spheres with centers generated by Poisson process

function s2_theory(r, R, λ, :: Val{2})
    tmp = r/(2R)
    tmp2 = (r > 2R) ? 2 : 2/π*(π + tmp*sqrt(1 - tmp^2) - acos(tmp))
    η = λ * π * R^2
    return exp(-η*abs(tmp2))
end

function s2_theory(r, R, λ, :: Val{3})
    tmp = r/R
    tmp2 = (r > 2R) ? 2 : 1 + 3/4*tmp - 1/16*tmp^3
    η = λ * 4/3 * π * R^3
    return exp(-η*tmp2)
end

# Search for the best parameters for random spheres
function spheres_parameters(target :: AbstractArray{T, N},
                            R0     :: Float64,
                            λ0     :: Float64) where {T, N}
    @. model(x, p) = s2_theory(x, p[1], p[2], Val(N))
    s2 = Directional.s2(target, 0) |> mean
    r = range(0.0; length = length(s2))
    fit = curve_fit(model, r, s2, [R0, λ0])
    return coef(fit)
end

function draw_sphere!(array  :: AbstractArray{T, N},
                      center :: CartesianIndex{N},
                      R      :: Float64) where {T, N}
    Ri = R |> round |> Int

    sphere_indices = CartesianIndices(Tuple(-Ri:Ri for i in 1:N))
    foreach(sphere_indices) do point
        dist = mapreduce(x -> x^2, +, Tuple(point))
        coord = point + center
        if dist < R^2 && checkbounds(Bool, array, coord)
            @inbounds array[coord] = 1
        end
    end
end

function generate_spheres(shape :: NTuple{N, Int},
                          T     :: Type,
                          R     :: Float64,
                          λ     :: Float64) where N
    nspheres = pois_rand(λ * prod(shape))
    spheres  = zeros(T, shape)
    indices  = CartesianIndices(spheres)

    for n in 1:nspheres
        draw_sphere!(spheres, rand(indices), R)
    end

    return spheres
end
