paintutilsPlus

-- Computercraft paintutilsPlus, shortened to pup
-- Adds image manipulation functions to paintutils
---------------------------
-- Variable Declarations --
---------------------------

-- Create table populated with existing paintutils
-- Allows overwriting paintutils, use of any unchanged functions, and potential integration with other libraries
local pup = {}
for key,val in pairs(paintutils) do
    pup[key] = val
end

pup.config = {
    restoreCursor = true, -- After drawing, attempt to reset cursor
    restoreColour = true, -- After drawing, attempt to reset colours
    overlayResizes = true, -- When overlaying an image, resize img if topImage does not fit
    padColour = 0,         -- When images need to be padded, use this colour (colours[name] code, 0 is 'transparent')
    bgColour = 0,           -- When bgImage not defined, use this colour (colours[name] code, 0 is 'transparent')
    mutate = false,         -- When using setColour or setPixel, change original image
    drawToScreen = true,    -- When creating a box, draw it immediately to screen
    fileOverwriting = false,-- When saving an image, allow files to be overwritten
}

-- Create a table to store all colour codes for easy lookup
-- Avoids expensive math.log(code,2) checks
local colourLookup = {}
local hexLookup = {}
for i = 0,15 do
    colourLookup[2^i] = string.format("%x",i)
    hexLookup[string.format("%x",i)] = 2^i
end

-- Create nfp metatable holder
local nfp = {}
nfp.__index = nfp
-- Create bimg metatable holder
local bimg = {}
bimg.__index = bimg

----------------------
-- Helper Functions --
----------------------

-- For safely editing an image table without mutations
-- This returns a reference-free copy of a given image
local function clone(image)
    local clone = {}
    for y = 1,#image do
        clone[y] = {}
        for x = 1,#image[y] do
            clone[y][x] = image[y][x]
        end
    end
    return clone
end

-- For identifying the format of a given image
-- Returns either "nfp", "bimg", or nil
local function getFormat(image)
    -- check format matches basic structure
    if type(image) ~= "table" then return nil end
    if #image == 0 or type(image[1]) ~= "table" or not image[1][1] then return nil end

    -- check if format matches nfp structure
    local isNFP = true
    for _, row in ipairs(image) do
        for _, pixel in ipairs(row) do
            if type(pixel) ~= "number" then
                isNFP = false
            end
        end
        if not isNFP then break end
    end
    if isNFP then return "nfp" end

    -- check if format matches bimg structure
    for _, row in ipairs(image) do
        if type(row) ~= "table"
        or type(row[1]) ~= "string"
        or type(row[2]) ~= "string"
        or type(row[3]) ~= "string"
        or math.max(#row[1], #row[2], #row[3]) ~= math.min(#row[1], #row[2], #row[3]) then
            return nil
        end
    end
    return "bimg"
end

---------------------
-- Setup Functions --
---------------------

-- For loading an image from a file
-- Detects format, returns image table+meta or nil
function pup.loadImage(path)
    -- check file exists
    if not fs.exists(path) then error("File ["..path.."] does not exist",2) end

    -- try to load as NFP
    local success, img = pcall(paintutils.loadImage,path)
    if success then
        return setmetatable(img,nfp)
    end

    -- try to load as BIMG
    local file = fs.open(path,"r")
    local data = file.readAll()
    file.close()
    pass, img = pcall(textutils.unserialise, data)
    if pass and getFormat(img) == "bimg" then
        return setmetatable(img,bimg)
    end

    -- if neither format, return nil
    return nil
end

-- For loading an image from memory
-- Detects format, returns image table+meta or nil
function pup.parseImage(image)
    -- try to parse as NFP
    local success, img = pcall(paintutils.parseImage,image)
    if success then
        return setmetatable(img,nfp)
    end

    -- check if format matches BIMG
    if getFormat(image) == "bimg" then
        return setmetatable(image,bimg)
    end

    -- if neither format, return nil
    return nil
end

-- For creating a new image (single colour box)
-- Returns nfp table, optionally renders to screen
function pup.drawFilledBox(x,y,x2,y2,colour,render)
    if render == nil then render = pup.config.drawToScreen end
    local bgC = term.getBackgroundColour()
    local xPos,yPos = term.getCursorPos()
    local clr = tonumber(colour) or colours[colour] or pup.config.bgColour
    local x1 = math.min(x,x2)
    x2 = math.max(x,x2)
    local y1 = math.min(y,y2)
    y2 = math.max(y,y2)
    local width = math.abs(x2-x1)+1
    local height = math.abs(y2-y1)+1
    local img = {}

    for y = 1,height do
        img[y] = {}
        for x = 1,width do
            img[y][x] = clr
        end
    end

    if render then pup.drawImage(img,x1,y1) end
    if pup.config.restoreCursor then term.setCursorPos(xPos,yPos) end
    if pup.config.restoreColour then term.setBackgroundColour(bgC) end

    return setmetatable(img, nfp)
end

-- For creating a new image (single colour border)
-- Returns nfp table, optionally renders to screen
function pup.drawBox(x,y,x2,y2,colour,render)
    if render == nil then render = pup.config.drawToScreen end
    local x1 = math.min(x,x2)
    x2 = math.max(x,x2)
    local y1 = math.min(y,y2)
    y2 = math.max(y,y2)
    local width = math.abs(x2-x1)
    local height = math.abs(y2-y1)
    local img = pup.drawFilledBox(x1,y1,x2,y2,colour,false)

    for y = 2,height do
        for x = 2, width do
            img[y][x] = pup.config.padColour
        end
    end
    if render then pup.drawImage(img,x1,y1) end
    return img
end

-- Inspect an image to get the current canvas size
-- If lines have different sizes, width will always be the longest
function nfp:getSize()
    local imgHeight = #self
    local imgWidth = 0
    for _,line in ipairs(self) do
        if #line > imgWidth then
            imgWidth = #line
        end
    end
    return imgWidth, imgHeight
end

----------------------------
-- Manipulation Functions --
----------------------------

-- Use a table to control padding or cropping per-side
-- Table format: {left = numA, right = numB, top = numC, bottom = numD}
-- Positive values increase that side, negative values decrease it
-- Can be used for translation (e.g. left -1 right +1 shifts the image 1 pixel left)
function nfp:adjustCanvas(adjustTable, padColour)
    adjustTable = adjustTable or {}

    local left   = tonumber(adjustTable.left)   or 0
    local right  = tonumber(adjustTable.right)  or 0
    local top    = tonumber(adjustTable.top)    or 0
    local bottom = tonumber(adjustTable.bottom) or 0

    local oldW, oldH = self:getSize()

    local newW = oldW + left + right
    local newH = oldH + top + bottom

    if newW <= 0 or newH <= 0 then
        error("Attempt to adjust canvas to non-positive dimensions", 2)
    end

    padColour = tonumber(padColour) or colours[padColour] or pup.config.padColour

    local newSelf = {}

    for y = 1, newH do
        newSelf[y] = {}
        for x = 1, newW do
            newSelf[y][x] = padColour
        end
    end

    for row = 1, oldH do
        for pixel = 1, oldW do
            local x = pixel + left
            local y = row + top

            if x >= 1 and x <= newW
            and y >= 1 and y <= newH then
                local val = self[row] and self[row][pixel]
                if val then
                    newSelf[y][x] = val
                end
            end
        end
    end

    return setmetatable(newSelf, nfp)
end

-- Adjust the canvas size to a set width/height
-- Will add/subtract pixels from right side or bottom
function nfp:resizeCanvas(width, height, bgColour)
    local oldW, oldH = self:getSize()
    width = tonumber(width) or oldW
    height = tonumber(height) or oldH
    return self:adjustCanvas({
            right  = width  - oldW,
            bottom = height - oldH
        }, bgColour)
end

-- Draw an image on top of another
-- Position is relative to underlying image
-- If new image does not fit, resizes canvas
-- If above but resize = false, crops top image
function nfp:overlay(topImage,xPos,yPos,resize)
    xPos = xPos or 1
    yPos = yPos or 1
    if resize == nil then resize = pup.config.overlayResizes end
    local selfW, selfH = self:getSize()
    local topW,topH = topImage:getSize()
    local newSelf = nil
    if topW+xPos-1 > selfW or topH+yPos-1 > selfH then
        if resize == true then
            newSelf = self:resizeCanvas(
                math.max(selfW,topW+xPos-1),
                math.max(selfH,topH+yPos-1),
                pup.config.padColour
            )
            selfW,selfH = newSelf:getSize()
        end
    end
    newSelf = newSelf or clone(self)
    for y = yPos, yPos + topH - 1 do
        for x = xPos, xPos + topW - 1 do
            local ty = y - yPos + 1
            local tx = x - xPos + 1

            if newSelf[y]
            and newSelf[y][x]
            and topImage[ty]
            and topImage[ty][tx]
            and topImage[ty][tx] ~= 0 then
                newSelf[y][x] = topImage[ty][tx]
            end
        end
    end
    return setmetatable(newSelf, nfp)
end

-- Replace all instances of a given colour with another
function nfp:setColour(oldColour,newColour,mutate)
    oldColour = tonumber(oldColour) or colours[oldColour]
    newColour = tonumber(newColour) or colours[newColour]
    local newSelf = mutate and self or clone(self)
    for lk,line in ipairs(newSelf) do
        for ck,c in ipairs(line) do
            if c == oldColour then
                newSelf[lk][ck] = newColour
            end
        end
    end
    return setmetatable(newSelf, nfp)
end

-- Replace a pixel at given location with a new colour
-- if mutate = true then edit the image instead of returning a new one
function nfp:setPixel(x,y,newColour,mutate)
    newColour = tonumber(newColour) or colours[newColour]
    local newSelf = mutate and self or clone(self)
    if not (newSelf[y] and newSelf[y][x]) then
        error("Attempt to set non-existent pixel",2)
    end
    newSelf[y][x] = newColour
    return setmetatable(newSelf,nfp)
end

-- Mirror an image along x and/or y axis ("x"/"y"/"xy")
function nfp:mirror(axis)
    local mirrorX = axis == "x" or axis == "xy"
    local mirrorY = axis == "y" or axis == "xy"
    local newSelf = clone(self)
    if mirrorX == true then
        local changed = {}
        for k1,l in ipairs(newSelf) do
            changed[k1] = {}
            for k2,p in ipairs(l) do
                changed[k1][#l+1-k2] = p
            end
        end
        newSelf = changed
    end
    if mirrorY == true then
        local changed = {}
        for k,l in ipairs(newSelf) do
            changed[#newSelf+1-k] = l
        end
        newSelf = changed
    end
    return setmetatable(newSelf,nfp)
end

-- Flip an image 90 degrees clockwise
function nfp:rotate()
    local newSelf = {}
    for y = 1,#self do
        for x = 1,#(self[y] or {}) do
            newSelf[x] = newSelf[x] or {}
            newSelf[x][#self-y+1] = self[y][x]
        end
    end
    return setmetatable(newSelf, nfp)
end

-- Set pixel ratio 1:factor, enlarging the image
function nfp:enlarge(factor)
    factor = math.floor(tonumber(factor) or 2)
    if factor < 1 then error("Factor must be >= 1",2) end
    local newSelf = {}
    for lk,l in ipairs(self) do
        table.insert(newSelf, {})
        for ck,c in ipairs(l) do
            for i = 1,factor do
                table.insert(newSelf[#newSelf], c)
            end
        end
        for n = 1,factor-1 do
            table.insert(newSelf,clone({newSelf[#newSelf]})[1])
        end
    end
    return setmetatable(newSelf, nfp)
end

--------------------------
-- Formatting Functions --
--------------------------

-- Convert NFP image to blit format
-- Pixels are converted 1:1
function nfp:blitImage()
    local bg = colourLookup[pup.config.bgColour] or colourLookup[term.getBackgroundColour()]
    local blit = {}
    for lk,l in ipairs(self) do
        if not blit[lk] then blit[lk] = {""} end
        for _,pixel in ipairs(l) do
            local clr = colourLookup[pixel] or bg
            blit[lk][1] = blit[lk][1]..clr
        end
        table.insert(blit[lk],1,string.rep(bg,#l))
        table.insert(blit[lk],1,string.rep(" ",#l))
    end
    return setmetatable(blit,bimg)
end

-- Convert NFP image to upscaled blit format
-- Every 2x3 cluster of pixels becomes 1 character
function nfp:upscale(bgColour)
    local blit = {}

    local w = 0
    for _,l in ipairs(self) do
        if #l > w then w = #l end
    end
    w = w + w%2
    local h = #self
    if h%3 > 0 then
        h = h + (3-(h%3))
    end

    local bgC = tonumber(bgColour) or colours[bgColour] or pup.config.bgColour

    local val, v1, v2, pxl, charVal = 0,nil,nil,nil,nil


    for cl = 1,h,3 do
        blit[#blit + 1] = {"","",""}
        for c = 1,w,2 do
            val = 0
            v1 = nil
            v2 = nil
            charVal = nil
            for pl = 1,3 do
                for p = 1,2 do
                    if (not self[(cl-1)+pl]) or (not self[(cl-1)+pl][(c-1)+p]) then
                        pxl = 0
                    else
                        pxl = self[(cl-1)+pl][(c-1)+p]
                    end
                    if not v1 then v1 = pxl
                    elseif v1 ~= pxl then v2 = pxl end
                    if v1 == pxl then
                        val = val + 2^((((pl-1)*2)+p)-1)
                    end
                end
            end
            if not v2 then v2 = v1 end
            if val > 31 then
                charVal = 128+63-val
                local tmp = v1
                v1 = v2
                v2 = tmp
            else
                charVal = 128+val
            end
            v1 = colourLookup[v1] or colourLookup[bgC] or colourLookup[term.getBackgroundColour()]
            v2 = colourLookup[v2] or colourLookup[bgC] or colourLookup[term.getBackgroundColour()]
            blit[#blit][1] = blit[#blit][1]..string.char(charVal)
            blit[#blit][2] = blit[#blit][2]..v1
            blit[#blit][3] = blit[#blit][3]..v2
        end
    end
    return setmetatable(blit,bimg)
end

-- Convert a 1:1 blit image to NFP format
-- This allows manipulation functions
function bimg:nfpImage()
    local newSelf = {}
    for y,row in ipairs(self) do
        newSelf[y] = {}
        for x = 1,#row[3] do
            newSelf[y][x] = hexLookup[row[3]:sub(x,x)]
        end
    end
    return setmetatable(newSelf, nfp)
end

----------------------
-- Output Functions --
----------------------

-- Draw function for BIMG format
function bimg:drawImage(xPos,yPos)
    local cx,cy = term.getCursorPos()
    xPos = xPos or cx
    yPos = yPos or cy
    for k,l in ipairs(self) do
        term.setCursorPos(xPos, yPos+k-1)
        term.blit(table.unpack(l))
    end
    if pup.config.restoreCursor then
        term.setCursorPos(cx,cy)
    end
    return self
end

-- Reset cursor and colours after drawImage
-- Otherwise, as with original paintutils
function pup.drawImage(image,x,y)
    local oldX, oldY = term.getCursorPos()
    local bg, fg = term.getBackgroundColour(),term.getTextColour()
    x = x or oldX
    y = y or oldY
    paintutils.drawImage(image,x,y)
    if pup.config.restoreColour then
        term.setBackgroundColour(bg)
        term.setTextColour(fg)
    end
    if pup.config.restoreCursor then
        term.setCursorPos(oldX,oldY)
    end
    return image
end
-- Add chainable drawImage function to metatable
function nfp:drawImage(x,y)
    return pup.drawImage(self,x,y)
end

-- Save function for blit format
function bimg:saveImage(path)
    if type(path) ~= "string" or (not pup.config.fileOverwriting and fs.exists(path)) then
        error("Image save location ["..tostring(path).."] invalid",2)
    end
    local file = fs.open(path,"w")
    file.write(textutils.serialise(self))
    file.close()
end

-- Save function for NFP format
function nfp:saveImage(path)
    local buffer = {}
    if type(path) ~= "string" or (not pup.config.fileOverwriting and fs.exists(path)) then
        error("File name ["..tostring(path).."] invalid.",2)
    end
    for _,line in ipairs(self) do
        for _,char in ipairs(line) do
            if char > 0 then
                char = colourLookup[char]
            else
                char = " "
            end
            buffer[#buffer+1] = char
        end
        buffer[#buffer+1] = "\n"
    end
    local file = fs.open(path,"w")
    file.write(table.concat(buffer))
    file.close()
end

return pup