Fragment Shader API

-- Rev 2

local sin   = math.sin
local cos   = math.cos
local tan   = math.tan
local abs   = math.abs
local min   = math.min
local max   = math.max
local floor = math.floor
local sqrt  = math.sqrt
local clock = os.clock

local term_setCursorPos = term.setCursorPos
local term_blit         = term.blit
local table_concat      = table.concat

local STRATEGY_FULL_CELL = 0
local STRATEGY_HEX_CELL = 1
local STRATEGY_TWO_CELL = 2

local function vec2(x, y)
    return { x = x, y = y }
end

local function vec3(x, y, z)
    return { x = x, y = y, z = z }
end

local function vec4(x, y, z, w)
    return { x = x, y = y, z = z, w = w }
end

local function length2(v)
    return (v.x^2 + v.y^2)^0.5
end

local function length3(v)
    return (v.x^2 + v.y^2 + v.z^2)^0.5
end

local function length4(v)
    return (v.x^2 + v.y^2 + v.z^2 + v.w^2)^0.5
end

local function normalize2(v)
    return div2s(v, length2(v))
end

local function normalize3(v)
    return div3s(v, length3(v))
end

local function normalize4(v)
    return div4s(v, length4(v))
end

local function distance2(v1, v2)
    return ((v1.x - v2.x)^2 + (v1.y - v2.y)^2)^0.5
end

local function distance3(v1, v2)
    return ((v1.x - v2.x)^2 + (v1.y - v2.y)^2 + (v1.z - v2.z)^2)^0.5
end

local function distance4(v1, v2)
    return ((v1.x - v2.x)^2 + (v1.y - v2.y)^2 + (v1.z - v2.z)^2 + (v1.w - v2.w)^2)^0.5
end

local function dot2(v1, v2)
    return v1.x*v2.x + v1.y*v2.y
end

local function dot3(v1, v2)
    return v1.x*v2.x + v1.y*v2.y + v1.z*v2.z
end

local function dot4(v1, v2)
    return v1.x*v2.x + v1.y*v2.y + v1.z*v2.z + v1.w*v2.w
end

local function add2(v1, v2)
    return vec2(v1.x+v2.x, v1.y+v2.y)
end

local function add3(v1, v2)
    return vec3(v1.x+v2.x, v1.y+v2.y, v1.z+v2.z)
end

local function add4(v1, v2)
    return vec4(v1.x+v2.x, v1.y+v2.y, v1.z+v2.z, v1.w+v2.w)
end

local function sub2(v1, v2)
    return vec2(v1.x-v2.x, v1.y-v2.y)
end

local function sub3(v1, v2)
    return vec3(v1.x-v2.x, v1.y-v2.y, v1.z-v2.z)
end

local function sub4(v1, v2)
    return vec4(v1.x-v2.x, v1.y-v2.y, v1.z-v2.z, v1.w-v2.w)
end

local function mul2(v1, v2)
    return vec2(v1.x*v2.x, v1.y*v2.y)
end

local function mul3(v1, v2)
    return vec3(v1.x*v2.x, v1.y*v2.y, v1.z*v2.z)
end

local function mul4(v1, v2)
    return vec4(v1.x*v2.x, v1.y*v2.y, v1.z*v2.z, v1.w*v2.w)
end

local function div2(v1, v2)
    return vec2(v1.x/v2.x, v1.y/v2.y)
end

local function div3(v1, v2)
    return vec3(v1.x/v2.x, v1.y/v2.y, v1.z/v2.z)
end

local function div4(v1, v2)
    return vec4(v1.x/v2.x, v1.y/v2.y, v1.z/v2.z, v1.w/v2.w)
end

local function add2s(v1, n)
    return vec2(v1.x+n, v1.y+n)
end

local function add3s(v1, n)
    return vec3(v1.x+n, v1.y+n, v1.z+n)
end

local function add4s(v1, n)
    return vec4(v1.x+n, v1.y+n, v1.z+n, v1.w+n)
end

local function sub2s(v1, n)
    return vec2(v1.x-n, v1.y-n)
end

local function sub3s(v1, n)
    return vec3(v1.x-n, v1.y-n, v1.z-n)
end

local function sub4s(v1, n)
    return vec4(v1.x-n, v1.y-n, v1.z-n, v1.w-n)
end

local function mul2s(v1, n)
    return vec2(v1.x*n, v1.y*n)
end

local function mul3s(v1, n)
    return vec3(v1.x*n, v1.y*n, v1.z*n)
end

local function mul4s(v1, n)
    return vec4(v1.x*n, v1.y*n, v1.z*n, v1.w*n)
end

local function div2s(v1, n)
    return vec2(v1.x/n, v1.y/n)
end

local function div3s(v1, n)
    return vec3(v1.x/n, v1.y/n, v1.z/n)
end

local function div4s(v1, n)
    return vec4(v1.x/n, v1.y/n, v1.z/n, v1.w/n)
end

local function mat2mulvec2(m, v)
    return vec2(m[1] * v.x + m[2] * v.y,
                m[3] * v.x + m[4] * v.y)
end

local function mat3mulvec3(m, v)
    return vec3(m[1] * v.x + m[2] * v.y + m[3] * v.z,
                m[4] * v.x + m[5] * v.y + m[6] * v.z,
                m[7] * v.x + m[8] * v.y + m[9] * v.z)
end

local function getRotMat2(a)
    return { cos(a), -sin(a), sin(a), cos(a) }
end

local function getRotMat3x(a)
    return {    1,          0,          0,
                0,          cos(a),     -sin(a),
                0,          sin(a),     cos(a)  }
end

local function getRotMat3y(a)
    return {    cos(a),     0,          sin(a),
                0,          1,          0,
                -sin(a),    0,          cos(a)  }
end

local function getRotMat3z(a)
    return {    cos(a),     -sin(a),    0,
                sin(a),     cos(a),     0,
                0,          0,          1       }
end

local function fract(n)
    return n - floor(n)
end

local function mix(x, y, a)
    return x * (1 - a) + y * a
end

local function clamp(a, s, l)
    return a < s and s or a > l and l or a
end

local function smoothstep(edge0, edge1, x)
    local t = clamp((x - edge0) / (edge1 - edge0), 0, 1)
    return t * t * (3 - 2 * t)
end

local function lab2xyz(lab)
    local y = (lab.x + 16.0) / 116.0
    local y3 = y^3.0
    local x = (lab.y / 500.0) + y
    local x3 = x^3.0
    local z = y - (lab.z / 200.0)
    local z3 = z^3.0

    if y3 > 0.008856 then
        y = y3
    else
        y = (y - (16.0 / 116.0)) / 7.787
    end
    if x3 > 0.008856 then
        x = x3
    else
        x = (x - (16.0 / 116.0)) / 7.787
    end

    if z3 > 0.008856 then
        z = z3
    else
        z = (z - (16.0 / 116.0)) / 7.787
    end

    return vec3(x * 95.0429, y * 100.0, z * 108.8900)
end

local Mi =  {{ 3.2406, -1.5372, -0.4986},
             {-0.9689,  1.8758,  0.0415},
             { 0.0557, -0.2040,  1.0570}}

local function xyz2rgb(xyz)
    local x = xyz.x / 100.0
    local y = xyz.y / 100.0
    local z = xyz.z / 100.0

    local r = (x * Mi[1][1]) + (y * Mi[1][2]) + (z * Mi[1][3])
    local g = (x * Mi[2][1]) + (y * Mi[2][2]) + (z * Mi[2][3])
    local b = (x * Mi[3][1]) + (y * Mi[3][2]) + (z * Mi[3][3])

    -- assume sRGB
    if r > 0.0031308 then
        r = ((1.055 * r^(1.0 / 2.4)) - 0.055)
    else
        r = (r * 12.92)
    end
    if g > 0.0031308 then
        g = ((1.055 * g^(1.0 / 2.4)) - 0.055)
    else
        g = (g * 12.92)
    end
    if b > 0.0031308 then
        b = ((1.055 * b^(1.0 / 2.4)) - 0.055)
    else
        b = (b * 12.92)
    end

    r = (r < 0) and 0 or r
    g = (g < 0) and 0 or g
    b = (b < 0) and 0 or b

    return vec3(r, g, b)
end

local function lab2rgb(lab)
    return xyz2rgb(lab2xyz(lab))
end

local function rgb2lab(color)
    -- Original code: http://www.brucelindbloom.com
    -- Ported to Lua by Yevano.
    local r, g, b = color.x, color.y, color.z
    local X, Y, Z, fx, fy, fz, xr, yr, z
    local Ls, as, bs
    local eps = 216 / 24389
    local k = 24389 / 27

    local Xr = 0.964221  -- reference white D50
    local Yr = 1
    local Zr = 0.825211

    -- assuming sRGB (D65)
    if r <= 0.04045 then
        r = r / 12
    else
        r = ((r + 0.055) / 1.055)^2.4
    end

    if g <= 0.04045 then
        g = g / 12
    else
        g = ((g + 0.055) / 1.055)^2.4
    end

    if b <= 0.04045 then
        b = b / 12
    else
        b = ((b + 0.055) / 1.055)^2.4
    end

    X = 0.436052025 * r + 0.385081593 * g + 0.143087414 * b
    Y = 0.222491598 * r + 0.71688606 * g + 0.060621486 * b
    Z = 0.013929122 * r + 0.097097002 * g + 0.71418547 * b

    -- XYZ to Lab
    xr = X / Xr
    yr = Y / Yr
    zr = Z / Zr

    if xr > eps then
        fx = xr^(1 / 3)
    else
        fx = ((k * xr + 16.) / 116.)
    end

    if yr > eps then
        fy = yr^(1 / 3)
    else
        fy = ((k * yr + 16.) / 116.)
    end

    if zr > eps then
        fz = zr^(1 / 3)
    else
        fz = ((k * zr + 16.) / 116)
    end

    Ls = (116 * fy) - 16
    as = 500 * (fx - fy)
    bs = 200 * (fy - fz)

    return vec3(2.55 * Ls + .5, as + .5, bs + .5)
end

local function colorDistance(c1, c2)
    local lab1 = rgb2lab(c1)
    local lab2 = rgb2lab(c2)
    return distance3(lab1, lab2)
end

local function createPalette(indexTable, redPal, greenPal, bluePal, ldFunc)
    ldFunc = ldFunc or function() return 0.1 end
    local colorTable = { }
    local pal = {
        indexTable = indexTable,
        colorTable = colorTable,
        redPal = redPal,
        greenPal = greenPal,
        bluePal = bluePal
    }
    local s = os.clock()
    for r = 0, redPal - 1 do
        colorTable[r] = { }
        for g = 0, greenPal - 1 do
            colorTable[r][g] = { }
            for b = 0, bluePal - 1 do
                local minDist = 9999
                local entry
                for i = 1, #indexTable do
                    local dist = colorDistance(vec3(r/(redPal-1), g/(greenPal-1), b/(bluePal-1)), indexTable[i])
                    if dist < minDist then
                        minDist = dist
                        entry = i
                    end
                end
                colorTable[r][g][b] = entry
            end
        end
        if os.clock() - s > ldFunc((r*greenPal*bluePal)/(redPal*greenPal*bluePal)) then
            sleep(0)
            s = os.clock()
        end
    end
    return pal
end

local function paletteConvert(palette, color)
    local x = color.x < 0 and 0 or color.x > 1 and 1 or color.x
    local y = color.y < 0 and 0 or color.y > 1 and 1 or color.y
    local z = color.z < 0 and 0 or color.z > 1 and 1 or color.z
    local r = floor(x * palette.redPal)
    local g = floor(y * palette.greenPal)
    local b = floor(z * palette.bluePal)
    if r == palette.redPal then r = r - 1 end
    if g == palette.greenPal then g = g - 1 end
    if b == palette.bluePal then b = b - 1 end
    local i = palette.colorTable[r][g][b]
    return i
end

local function createShaderProgram()
    local prog = {
        shaders = { }
    }
    return prog
end

local function addShader(program, shader)
    program.shaders[#program.shaders + 1] = shader
end

local function applyShaderProgram(program, frameBuffer)
    local pixelData = frameBuffer.pixelData
    local w = frameBuffer.size.x
    local h = frameBuffer.size.y
    for i = 1, #program.shaders do
        local shaderFunc = program.shaders[i].shaderFunc
        local inputs = program.shaders[i].inputs
        for x = 1, w do
            for y = 1, h do
                pixelData[x + (y - 1) * w] = shaderFunc(vec4(pixelData[x + (y - 1) * w]), vec2((x - 1), (h - y + 1)*3/2), inputs)
            end
        end
    end
end

local function createShader(shaderFunc)
    local shader = {
        shaderFunc = shaderFunc,
        inputs = { }
    }
    return shader
end

local function setShaderInput(shader, i, v)
    shader.inputs[i] = v
end

local function createFrameBuffer(palette, charTable, size, renderStrategy, blit, scp)
    renderStrategy = renderStrategy or STRATEGY_FULL_CELL
    if renderStrategy == STRATEGY_HEX_CELL and (size.x % 2 ~= 0 or size.y % 3 ~= 0) then
        error("Frame Buffer using STRATEGY_HEX_CELL must have width multiple of 2 and height multiple of 3.")
    end
    if renderStrategy == STRATEGY_TWO_CELL then
        error("STRATEGY_TWO_CELL not yet implemented.")
    end
    local pixelData = { }
    local buf = {
        palette = palette,
        charTable = charTable,
        size = size,
        pixelData = pixelData,
        renderStrategy = renderStrategy,
        blit = blit or term_blit,
        scp = scp or term_setCursorPos
    }
    for i = 1, size.x * size.y do
        pixelData[i] = vec4(0, 0, 0, 0)
    end
    return buf
end

-- Thanks for writing this useful function, oli414.
local function getDrawingCharacter(topLeft, topRight, left, right, bottomLeft, bottomRight)
    local data = 128
    if not bottomRight then
        data = data + (topLeft and 1 or 0)
        data = data + (topRight and 2 or 0)
        data = data + (left and 4 or 0)
        data = data + (right and 8 or 0)
        data = data + (bottomLeft and 16 or 0)
    else
        data = data + (topLeft and 0 or 1)
        data = data + (topRight and 0 or 2)
        data = data + (left and 0 or 4)
        data = data + (right and 0 or 8)
        data = data + (bottomLeft and 0 or 16)
    end
    return string.char(data), bottomRight
end

local function drawFrameBuffer(buf, pos)
    local w, h = buf.size.x, buf.size.y
    local x, y = pos.x, pos.y
    local pixelData = buf.pixelData
    local palette = buf.palette
    local charTable = buf.charTable
    local blit = buf.blit
    local scp = buf.scp
    if buf.renderStrategy == STRATEGY_FULL_CELL then
        for i = y, y + h - 1 do
            local s = { }
            local b = { }
            local f = { }
            for j = (i-1)*w + 1, i*w do
                local colIdx = paletteConvert(palette, pixelData[j])
                local charEntry = charTable[colIdx]
                s[#s + 1] = charEntry[1]
                b[#b + 1] = charEntry[2]
                f[#f + 1] = charEntry[3]
            end
            scp(x, i)
            blit(table_concat(s), table_concat(f), table_concat(b))
        end
    elseif buf.renderStrategy == STRATEGY_HEX_CELL then
        for sy = 0, h-1, 3 do
            local s = { }
            local b = { }
            local f = { }
            for sx = 0, w-1, 2 do
                local uniqueColors = { }
                local colorCounts = { }
                local convColors = { }
                for j = 0, 2 do
                    for i = 0, 1 do
                        local col = charTable[paletteConvert(palette, pixelData[sx+i + (sy+j) * w + 1])][2]
                        convColors[#convColors + 1] = col
                        if colorCounts[col] then
                            colorCounts[col] = colorCounts[col] + 1
                        else
                            colorCounts[col] = 1
                            uniqueColors[#uniqueColors + 1] = col
                        end
                    end
                end
                table.sort(uniqueColors, function(a, b) return colorCounts[a] > colorCounts[b] end)
                local c0 = uniqueColors[1]
                local c1 = uniqueColors[2] or c0
                local char, inv = getDrawingCharacter(
                    convColors[1] ~= c1,
                    convColors[2] ~= c1,
                    convColors[3] ~= c1,
                    convColors[4] ~= c1,
                    convColors[5] ~= c1,
                    convColors[6] ~= c1)
                s[#s + 1] = char
                if inv then
                    b[#b + 1] = c0
                    f[#f + 1] = c1
                else
                    b[#b + 1] = c1
                    f[#f + 1] = c0
                end
            end
            scp(x, y + sy/3)
            blit(table_concat(s), table_concat(f), table_concat(b))
        end
    elseif buf.renderStrategy == STRATEGY_TWO_CELL then
        error("STRATEGY_TWO_CELL not yet implemented.")
    end
end

local function getUnstretchedResolution(buf)
    return vec2(buf.size.x, buf.size.y * 3/2)
end

local function clearFrameBuffer(buf, color)
    local pixelData = buf.pixelData
    for i = 1, #pixelData do
        pixelData[i] = color
    end
end

local g = getfenv()
g.STRATEGY_FULL_CELL = STRATEGY_FULL_CELL
g.STRATEGY_HEX_CELL = STRATEGY_HEX_CELL
g.STRATEGY_TWO_CELL = STRATEGY_TWO_CELL
g.vec2 = vec2
g.vec3 = vec3
g.vec4 = vec4
g.length2 = length2
g.length3 = length3
g.length4 = length4
g.normalize2 = normalize2
g.normalize3 = normalize3
g.normalize4 = normalize4
g.distance2 = distance2
g.distance3 = distance3
g.distance4 = distance4
g.dot2 = dot2
g.dot3 = dot3
g.dot4 = dot4
g.add2 = add2
g.add3 = add3
g.add4 = add4
g.sub2 = sub2
g.sub3 = sub3
g.sub4 = sub4
g.mul2 = mul2
g.mul3 = mul3
g.mul4 = mul4
g.div2 = div2
g.div3 = div3
g.div4 = div4
g.add2s = add2s
g.add3s = add3s
g.add4s = add4s
g.sub2s = sub2s
g.sub3s = sub3s
g.sub4s = sub4s
g.mul2s = mul2s
g.mul3s = mul3s
g.mul4s = mul4s
g.div2s = div2s
g.div3s = div3s
g.div4s = div4s
g.mat2mulvec2 = mat2mulvec2
g.mat3mulvec3 = mat3mulvec3
g.getRotMat2 = getRotMat2
g.getRotMat3x = getRotMat3x
g.getRotMat3y = getRotMat3y
g.getRotMat3z = getRotMat3z
g.fract = fract
g.mix = mix
g.clamp = clamp
g.smoothstep = smoothstep
g.lab2xyz = lab2xyz
g.xyz2rgb = xyz2rgb
g.lab2rgb = lab2rgb
g.rgb2lab = rgb2lab
g.colorDistance = colorDistance
g.createPalette = createPalette
g.paletteConvert = paletteConvert
g.createShaderProgram = createShaderProgram
g.addShader = addShader
g.applyShaderProgram = applyShaderProgram
g.createShader = createShader
g.setShaderInput = setShaderInput
g.createFrameBuffer = createFrameBuffer
g.drawFrameBuffer = drawFrameBuffer
g.getUnstretchedResolution = getUnstretchedResolution
g.clearFrameBuffer = clearFrameBuffer