--- Utility library for creating palette swaps of LÖVE-based images.

local palette_swap = {}

--------------------------------------------------------------------------------
-- Main work functions

local function get_palette_of_image_data (data)
    assert(data)
    --
    local palette  = {}
    local has_seen = {}
    --
    for x = 0, data:getWidth() - 1 do
        for y = 0, data:getHeight() - 1 do
            local r, g, b, a  =  data:getPixel(x, y)
            if not has_seen[r]       then  has_seen[r]    = {}  end
            if not has_seen[r][g]    then  has_seen[r][g] = {}  end
            if not has_seen[r][g][b] then
                local color = {r,g,b,a}
                has_seen[r][g][b]    =  color
                palette[#palette+1]  =  color
            end
        end
    end
    --
    return palette
end

local function change_palette_to_dim_tree (change_palette)
    assert(type(change_palette) == 'table')
    --
    local tree  =  {}
    --
    for color_from, color_to in pairs(change_palette) do
        assert(type(color_from) == 'table')
        assert(type(color_to)   == 'table')
        --
        local r, g, b  =  unpack(color_from)
        if not tree[r]       then  tree[r]    = {}  end
        if not tree[r][g]    then  tree[r][g] = {}  end
        --
        assert(not tree[r][g][b])
        tree[r][g][b]    =  color_to
    end
    --
    return tree
end

local function apply_change_palette (data, change_palette)
    assert(data)
    assert(change_palette)
    --
    local change_tree  =  change_palette_to_dim_tree(change_palette)
    local function map (_x, _y, r, g, b, a)
        local c  =  change_tree[r][g][b]
        return c[1], c[2], c[3], a
    end
    --
    local new_data = data:clone()
          new_data:mapPixel(map)
    --
    return new_data
end

local function change_palette_from_map (map_func, palette)
    local change_palette = {}
    for i = 1, #palette do
        local from  =  palette[i]
        local to    =  {[4] = 1, map_func(unpack(from))}
        change_palette[from] = to
    end
    return change_palette
end

--------------------------------------------------------------------------------
-- API

palette_swap.get_palette_from_image_data  = get_palette_of_image_data

local function concat_functions (first_fn, ...)
    --
    local funcs = {first_fn, ...}

    -- Identity function.
    if #funcs == 0 then
        return function (...) return ... end

    -- Only one function
    elseif #funcs <= 1 then
        return first_fn
    end

    -- Generic
    return function (...)
        local val = {...}
        for i = 1, #funcs do
            val = {funcs[i](unpack(val))}
        end
        return unpack(val)
    end
end

function palette_swap.apply_color_map_to_image_data (img_data, ...)
    assert(img_data)
    --
    local map_fn  =  concat_functions(...)
    --
    return apply_change_palette(img_data, change_palette_from_map(map_fn, get_palette_of_image_data(img_data)))
end

--------------------------------------------------------------------------------
-- Example color maps

palette_swap.color_map = {}

function palette_swap.color_map.invert (r,g,b)
    return 1 - r, 1 - g, 1 - b
end

local CHANNEL_ORDERS = {
    [0] = {'r', 'g', 'b'},
    [1] = {'g', 'b', 'r'},
    [2] = {'b', 'r', 'g'},
    [3] = {'r', 'b', 'g'},
    [4] = {'b', 'g', 'r'},
    [5] = {'g', 'r', 'b'},
    [6] = {'b', 'g', 'r'},
    [7] = {'b', 'g', 'r'}, -- TODO
    [8] = {'b', 'g', 'r'}, -- TODO
    [9] = {'r', 'g', 'b'}, -- TODO
}

function palette_swap.color_map.switch_channel (mode)
    local channel_order = CHANNEL_ORDERS[mode]
    local func_s = ("return function (r, g, b)  return %s, %s, %s  end"):format(unpack(channel_order))
    return loadstring (func_s)()
end

local function load_colors ()

    -- Only need to import colors when needed
    local colors = require 'colors'
    assert(type(colors) == 'table')
    assert(type(colors.rgb1_to_rgb255) == 'function')
    assert(type(colors.rgb_to_hsl)     == 'function')
    assert(type(colors.hsl_to_rgb1)    == 'function')

    function palette_swap.color_map.increase_saturation (amount)
        -- Create function
        return function (r, g, b)
            local hsl = colors.rgb_to_hsl(colors.rgb1_to_rgb255(r, g, b))
            --
            hsl[2] = math.min(1, math.max(0, hsl[2] + amount))
            --
            return unpack(colors.hsl_to_rgb1(hsl))
        end
    end

    function palette_swap.color_map.rotate_hue (rad)
        return function (r, g, b)
            local hsl = colors.rgb_to_hsl(colors.rgb1_to_rgb255(r, g, b))
            --
            hsl[1] = (hsl[1] + rad) % 1
            --
            return unpack(colors.hsl_to_rgb1(hsl))
        end
    end
end

for _, key in ipairs {'rotate_hue', 'increase_saturation'} do
    palette_swap.color_map[key]  =  function (...)
        load_colors()
        return palette_swap.color_map[key](...)
    end
end

--------------------------------------------------------------------------------
-- Return

return palette_swap