# Based on the paper:
#   T. Várady, P. Benkő, G. Kós: Implicit surfaces revisited - I-patches.
#     Geometric Modelling, pp. 323-335, Springer, 2001.

module IPatch

using LinearAlgebra

export plane, sphere, cylinder, cone, union, ipatch, write_volume

"""
    plane(p, n)

Returns an implicit representation for the plane
interpolating `p` and having normal vector `n`.
"""
plane(p, n) = q -> dot(q - p, n)

"""
    sphere(p, r)

Returns an implicit representation for the sphere
with center `p` and a radius of `r`.
"""
sphere(p, r) = q -> norm(q - p) - r

"""
    cylinder(p, n, r)

Returns an implicit representation for the cylinder
with an axis going through `p` in direction `n`, and with radius `r`.
"""
cylinder(p, n, r) = q -> r - norm((q - p) - n * dot(q - p, n))

"""
    cone(v, p, r)

Returns an implicit representation for the cone
with an axis going through `v` and `p`.
The apex is at `v`, and the radius at `p` is `r`.
"""
function cone(v, p, r)
    d = p - v
    n = normalize(d)
    sinphi2 = r^2 / (r^2 + dot(d, d))
    sinphi = sqrt(sinphi2)
    cosphi = sqrt(1 - sinphi2)
    function (q)
        qh = dot(q - v, n)
        qr = norm(q - v - n * qh)
        qr * cosphi - abs(qh) * sinphi
    end
end

"""
    union(s1, s2)

Returns the implicit representation of the union
of the two (implicit) surfaces `s1` and `s2`.
"""
union(s1, s2) = q -> s1(q) * s2(q)


# I-patch & file output

"""
    ipatch(constraints, point, [exponent, weights])

`constraints` is a vector of `(primary, boundary)` tuples,
where both `primary` and `boundary` are implicit surfaces.
The intersection of each tuple defines one side of the I-patch,
which will interpolate these with G^1 continuity (or higher,
when `exponent` is larger than 2).

`point` (a 3D vector) is to be interpolated by the patch,
which implicitly defines the weight of the final term.
The weight of all other terms are 1 by default,
but these can be changed with the `weights` argument.
"""
function ipatch(constraints, point, exponent = 2, weights = ones(length(constraints)))
    boundaries = Set([b for (_,b) in constraints])
    function ipatch_eval(q)
        patch = 0.0
        n = length(constraints)
        for i in 1:n
            rest = setdiff(boundaries, [constraints[i][2]])
            bs = mapreduce(b -> b(q), *, rest)
            patch += constraints[i][1](q) * bs^exponent * weights[i]
        end
        fullness = mapreduce(b -> b(q), *, boundaries)^exponent
        (patch, fullness)
    end
    pc, fc = ipatch_eval(point)
    wc = pc / fc
    function (q)
        patch, fullness = ipatch_eval(q)
        patch - fullness * wc
    end
end

"""
    write_volume(voxels, filename)

Writes the 3D matrix `voxels` into `filename.raw`
with the metaimage header `filename.mhd`.
"""
function write_volume(voxels, filename)
    n = size(voxels)
    open("$(filename).mhd", "w") do f
        println(f, "NDims = 3")
        println(f, "DimSize = $(n[1]) $(n[2]) $(n[3])")
        println(f, "ElementSize = 1.0 1.0 1.0")
        println(f, "ElementSpacing = 1.0 1.0 1.0")
        println(f, "ElementType = MET_DOUBLE")
        println(f, "ElementByteOrderMSB = False") # LittleEndian
        println(f, "ElementDataFile = $(filename).raw")
    end
    open("$(filename).raw", "w") do f
        for val in voxels
            write(f, val)
        end
    end
end


# Examples

# Cube pocket
function test1(res)
    # Note:
    # This is an example from the original paper, which does _not_ work.
    # In Section 2, paragraph 7, the paper says that when different primary surfaces
    # have the same bounding functions, each has to be used only once, which is
    # implemented in the `ipatch` function, but this does not solve the problem,
    # actually it makes it worse: now it doesn't even interpolate the positions.
    # (For the real solution, see `test1b`.)
    center = [0.5,0.5,0.5]
    a1 = plane([1,0,0], [1,0,0])
    a2 = plane([0,1,0], [0,1,0])
    a3 = plane([0,0,1], [0,0,1])
    b1 = plane(center, [-1,0,0])
    b2 = plane(center, [0,-1,0])
    b3 = plane(center, [0,0,-1])
    auxiliary = [0.7, 0.7, 0.7]
    patch = ipatch([(a1, b2), (a1, b3), (a2, b3), (a2, b1), (a3, b1), (a3, b2)], auxiliary)
    voxels = Array{Float64}(undef, res, res, res)
    r = range(0.5, stop=1, length=res)
    for i in 1:res, j in 1:res, k in 1:res
        voxels[i,j,k] = patch([r[i], r[j], r[k]])
    end
    write_volume(voxels, "/tmp/ipatch-test")
end

# Cube pocket - working version
function test1b(res, fullness = 0.75)
    # Here we use the same data, but now primary surfaces with the same bounding functions
    # are united (using `union`, which is just multiplication), and everything works,
    # since the source of the problem is eliminated.
    # For two-sided surfaces, this method reverts to functional splines; for other
    # surfaces it is a hybrid between functional splines and (normal) I-patches.
    center = [0.5,0.5,0.5]
    a1 = plane([1,0,0], [1,0,0])
    a2 = plane([0,1,0], [0,1,0])
    a3 = plane([0,0,1], [0,0,1])
    b1 = plane(center, [-1,0,0])
    b2 = plane(center, [0,-1,0])
    b3 = plane(center, [0,0,-1])
    a12 = union(a1, a2)
    a13 = union(a1, a3)
    a23 = union(a2, a3)
    auxiliary = repeat([fullness], 3)
    patch = ipatch([(a12, b3), (a13, b2), (a23, b1)], auxiliary)
    voxels = Array{Float64}(undef, res, res, res)
    r = range(0.5, stop=1, length=res)
    for i in 1:res, j in 1:res, k in 1:res
        voxels[i,j,k] = patch([r[i], r[j], r[k]])
    end
    write_volume(voxels, "/tmp/ipatch-test")
end

# Sphere octant
function test2(res, fullness = 0.6)
    o = [0,0,0]
    s = sphere(o, 1)
    b1 = plane(o, [1,0,0])
    b2 = plane(o, [0,1,0])
    b3 = plane(o, [0,0,1])
    auxiliary = [fullness, fullness, fullness]
    patch = ipatch([(s, b1), (s, b2), (s, b3)], auxiliary)
    voxels = Array{Float64}(undef, res, res, res)
    r = range(0, stop=1, length=res)
    for i in 1:res, j in 1:res, k in 1:res
        voxels[i,j,k] = patch([r[i], r[j], r[k]])
    end
    write_volume(voxels, "/tmp/ipatch-test")
end

# Two-sided
function test3(res, fullness = 3)
    o = [0,0,0]
    s1 = sphere(o, 3)
    s2 = sphere([0,4,0], 5)
    p1 = plane(o, [0,0,1])
    p2 = plane(o, normalize([0,1,1]))
    auxiliary = [0.0, 2.0, fullness]
    patch = ipatch([(s1, p1), (s2, p2)], auxiliary)
    voxels = Array{Float64}(undef, res, res, res)
    r = range(-3, stop=3, length=res)
    for i in 1:res, j in 1:res, k in 1:res
        voxels[i,j,k] = patch([r[i], r[j], r[k]])
    end
    write_volume(voxels, "/tmp/ipatch-test")
end

# Cylinder & cone
function test4(res)
    o = [0,0,0]
    c1 = cylinder(o, [1,0,0], 0.5)
    c2 = cone([0,0.2,0], [0,0,1], 0.6)
    p1 = plane([-0.7,0,0], [1,0,0])
    p2 = plane([0.7,0,0], [1,0,0])
    p3 = plane([0,0,0], [1,0,0])
    auxiliary = [-0.4, 0.0, 0.5]
    patch = ipatch([(c1, p1), (c1, p2), (c2, p3)], auxiliary)
    voxels = Array{Float64}(undef, res, res, res)
    r = range(-1, stop=1, length=res)
    for i in 1:res, j in 1:res, k in 1:res
        voxels[i,j,k] = c2([r[i], r[j], r[k]])
    end
    write_volume(voxels, "/tmp/ipatch-test")
end

end