#!/usr/bin/luajit
-- Decompose a dilation kernel for mathematical
-- morphology.
-- Rather, compute a dilation kernel from a
-- decomposition.
-- Dilatiom ((+)) is sup; erosion ((-)) is
-- inf.  The kernel is added in dilation,
-- subtracted in erosion.
-- I realize that this is maybe the wrong approach
-- for emitting code or analyzing costs; for that
-- we'd want to maintain the flow graph.
local meta = {__index={}}

-- Adds a point to a kernel, normally during creation
function meta.__index.add(kern, x, y, val)
    if kern[x] == nil then kern[x] = {} end
    local col = kern[x]
    local old = col[y]
    if old == nil or old < val then
        col[y] = val
    end
end

-- Get a point from a kernel, normally nil
function meta.__index.at(kern, x, y)
    local col = kern[x]
    if col then return col[y] end
end

-- Iterator over (x,y,level) triples
function meta.__index.triples(this)
    local x, y
    x = next(this)

    return function()
        if x == nil then return nil end
        y = next(this[x], y)

        while y == nil do -- next column
            x = next(this, x)
            if x == nil then return nil end
            y = next(this[x])
        end

        return x, y, this[x][y]
    end
end

-- Appends strings to a list
local function append(tbl, ...)
    for _, v in ipairs({...}) do
        table.insert(tbl, tostring(v))
    end
end

-- Return a textual (x, y): v listing
function meta.__index.dump(kern)
    local rv = {'{'}
    for x, y, v in kern:triples() do
        append(rv, "(", x, ", ", y, "): ", v, "; ")
    end
    append(rv, '}')
    return table.concat(rv, '')
end

-- ASCII art plot of support with no levels
function meta.__tostring(this)
    local xmin, xmax, ymin, ymax
    for x, y in this:triples() do
        if xmin == nil or xmin > x then xmin = x end
        if ymin == nil or ymin > y then ymin = y end
        if xmax == nil or xmax < x then xmax = x end
        if ymax == nil or ymax < y then ymax = y end
    end

    if xmin == nil then return "(empty)" end

    local rv = {}

    for y = ymax, ymin, -1 do
        append(rv, '    ')
        for x = xmin, xmax do
            local c = '.'
            if x == 0 and y == 0 then
                if this:at(x, y) then 
                    c = '#'
                else
                    c = '*'
                end
            else
                if this:at(x, y) then c = '@' end
            end
            append(rv, c)
            --append(rv, c)
        end
        append(rv, '\n')
    end

    return table.concat(rv)
end

-- Translate a kernel by {x, y}
function meta.__call(this, idx)
    local x, y = unpack(idx)
    local result = kern()
    for xi, yi, v in this:triples() do
        result:add(x+xi, y+yi, v)
    end
    return result
end

-- Add a constant to the level of a kernel
function meta.__add(this, level)
    local result = kern()
    for xi, yi, v in this:triples() do
        result:add(xi, yi, v+level)
    end
    return result
end

-- Union (pointwise max) of two kernels
function meta.__concat(this, other)
    local result = kern()

    for x, y, v in this:triples() do
        result:add(x, y, v)
    end

    for x, y, v in other:triples() do
        result:add(x, y, v)
    end

    return result
end

-- Helper for the common case of dilating a kernel
-- by a two-point kernel.
function meta.__index.X(k, p)
    return k..k(p)
end

-- Constructor for kernel objects
function kern()
    return setmetatable({}, meta)
end


-- An alternative implementation that remembers
-- the dataflow graph.
--
-- The idea is that we want to be able to compose
-- the graph with the same :X, .., and call
-- operations I've been using, but not only
-- compute the resulting sum kernel, but also
-- things like how many min operations are needed
-- per pixel to apply it to an image, and maybe
-- even generate Lua or C code to do it.
--
-- Metatable for composition operations:
local gmeta = {__index={}}
-- Four variants: zero impulse, translate,
-- constant add, and union.
local k0 = kern()
k0:add(0, 0, 0)



-- Some stupid test code

print(kern(), kern()[3], kern(){3,4})

print(k0:at(0,0), '0,1', k0:at(0,1),
    '1,0', k0:at(1,0), 'ok')
kx = (k0{3,4}..k0..k0{-1,-1}+2)
print(kx:dump())
print(kx)

-- Approximate a 35-pixel-long diagonal line
print(k0:X{1,1}:X{2,1}:X{4,2}:X{8,5}:X{16,10}:X{3,2})

-- Approximate an L2 SDF
sdf0 = k0..(k0{-1,0}..k0{1,0}..
            k0{0,-1}..k0{0,1})
            + -1
print(sdf0:dump())
print(sdf0)

sdf1 = sdf0..(sdf0{-1,-1}..sdf0{-1,1}
            ..sdf0{1,-1}..sdf0{1,1})
            + -1.414
print(sdf1, sdf1:dump())

sdf2 = sdf1..(sdf1{-2,0}..sdf1{2,0}
            ..sdf1{0,-2}..sdf1{0,2})
            + -2
print(sdf2, sdf2:dump())

sdf3 = sdf2..(sdf2{-1,-2}..sdf2{-1,2}
            ..sdf2{-2,-1}..sdf2{2,-1}
            ..sdf2{1,-2}..sdf2{1,2}
            ..sdf2{-2,1}..sdf2{2,1})
            + -2.236

print(sdf3, sdf3:dump())
-- That may be kind of a dumb approach in the sense
-- that if we allow sequences of moves of the same
-- size we get estimates for a larger area, but
-- they would ve worse estimates.

-- Serra's hexagon example, but on a square grid.
print(k0:X{1,0}:X{2,0}:X{-1,1}:X{-2,2}:X{1,1}:X{2,2})
-- Or:
print(k0:X{1,0}:X{2,0}:X{0,1}:X{0,2}:X{1,1}:X{2,2})


-- 23-pixel-square checkerboard, example from Urbach
-- & Wilkinson 02008.
hc = k0:X{2,0}:X{4,0}:X{8,0}:X{6,0}
vc = hc:X{0,2}:X{0,4}:X{0,8}:X{0,6}
print(vc..vc{1,1}..vc{2,2}..vc{2,0}..vc{0,2})

-- Better alternative, only 12 maxes per pixel:
purina = k0..k0{-1,-1}:X{2,0}:X{0,2}
print(purina)
hp = purina:X{2,0}:X{4,0}:X{8,0}:X{6,0}
vp = hp    :X{0,2}:X{0,4}:X{0,8}:X{0,6}
print(vp)

-- Some pen nib shapes of interest for stroking
-- letterforms
print(k0:X{-1,0}:X{1,1}:X{2,2})
print(k0:X{-1,0}:X{1,1}:X{1,-1})
print(k0:X{-1,0}:X{1,1}:X{1,-1}:X{0,-1})
print(k0{-1,0}:X{1,1}:X{1,-1}..k0)
print(k0:X{1,1}:X{3,0})

-- 5-high diamond
k2 = k0:X{1,1}:X{-1,1}
print(k2{0,-1}..k2{0,-2}:X{1,1}:X{-1,1})

-- 5-high circle in 7 operations:
k3 = k0:X{1,0}:X{-1,0}
r5 = k3{-1,0}:X{2,0}:X{0,1}:X{0,-1}
print(k3{0,2}:X{0,-4}..r5)

-- or in only 6 operations as an octagon:
print(k0:X{1,0}:X{1,0}:X{1,1}:X{0,1}:X{0,1}:X{-1,1})

-- Urbach–Wilkinson needs 8:
k2 = k0:X{1,0}
k3 = k2:X{1,0}
k5 = k2:X{2,0}:X{1,0}
print(k3..k5{-1,1}..k5{-1,2}..k5{-1,3}..k3{0,4})

-- a similar checkerboard:
print(k0:X{1,1}:X{-1,1}:X{2,0}:X{0,2})

-- a horizontal fuzzy line
print(k0:X{3,0}:X{5,0}:X{7,0}:X{11,0}:X{1,1}:X{1,-1})