CC Neural Net

-- Obtain command-line arguments.
local args = { ... }
local training_mode = (args[1] == "training")

-- Configuration
local WIDTH, HEIGHT = 100, 70                   -- Drawing grid dimensions (cells)
local NEURON_WIDTH, NEURON_HEIGHT = 20, 20        -- Normalized grid dimensions
local DISPLAY_SCALE = 1                         -- Not used for scaling now
local LEARNING_RATE = 0.1

-- Grid drawing offsets (for border)
local GRID_OFFSET_X = 2
local GRID_OFFSET_Y = 2

-- Neural network configuration
local INPUT_SIZE = NEURON_WIDTH * NEURON_HEIGHT   -- 400 inputs
local HIDDEN_SIZE = 50                            -- Hidden layer size
local OUTPUT_SIZE = 3                             -- Three classes: square, circle, triangle

-- File for saving network parameters
local NETWORK_FILE = "network_weights.txt"

---------------------------------
-- Neural Network Functions
---------------------------------

local network = {}

local function initialize_network()
    network = { weights1 = {}, bias1 = {},
                weights2 = {}, bias2 = {} }
    for i = 1, HIDDEN_SIZE do
        network.weights1[i] = {}
        for j = 1, INPUT_SIZE do
            network.weights1[i][j] = math.random() * 2 - 1
        end
        network.bias1[i] = math.random() * 2 - 1
    end
    for k = 1, OUTPUT_SIZE do
        network.weights2[k] = {}
        for i = 1, HIDDEN_SIZE do
            network.weights2[k][i] = math.random() * 2 - 1
        end
        network.bias2[k] = math.random() * 2 - 1
    end
end

local function save_network()
    local file = fs.open(NETWORK_FILE, "w")
    file.write(textutils.serialize(network))
    file.close()
end

local function load_network()
    local file = fs.open(NETWORK_FILE, "r")
    if file then
        network = textutils.unserialize(file.readAll())
        file.close()
    else
        initialize_network()
    end
end

-- Activation functions
local function sigmoid(x)
    return 1 / (1 + math.exp(-x))
end

local function dsigmoid(y)
    return y * (1 - y)  -- derivative assuming y = sigmoid(x)
end

local function softmax(vec)
    local max_val = -math.huge
    for i, v in ipairs(vec) do
        if v > max_val then max_val = v end
    end
    local sum_exp = 0
    local exp_vec = {}
    for i, v in ipairs(vec) do
        exp_vec[i] = math.exp(v - max_val)
        sum_exp = sum_exp + exp_vec[i]
    end
    local out = {}
    for i, v in ipairs(exp_vec) do
        out[i] = v / sum_exp
    end
    return out
end

-- Forward pass: returns hidden activations and output probabilities
local function forward(input)
    local hidden = {}
    for i = 1, HIDDEN_SIZE do
        local sum = network.bias1[i]
        for j = 1, INPUT_SIZE do
            sum = sum + network.weights1[i][j] * input[j]
        end
        hidden[i] = sigmoid(sum)
    end
    local output_raw = {}
    for k = 1, OUTPUT_SIZE do
        local sum = network.bias2[k]
        for i = 1, HIDDEN_SIZE do
            sum = sum + network.weights2[k][i] * hidden[i]
        end
        output_raw[k] = sum
    end
    local output = softmax(output_raw)
    return hidden, output
end

-- Train the network using one training sample (input vector and one-hot target)
local function train_network(input, target, learning_rate)
    learning_rate = learning_rate or LEARNING_RATE
    local hidden, output = forward(input)

    -- Compute output error (delta) for softmax with cross-entropy:
    local delta_output = {}
    for k = 1, OUTPUT_SIZE do
        delta_output[k] = output[k] - target[k]
    end

    -- Update weights2 and bias2
    for k = 1, OUTPUT_SIZE do
        for i = 1, HIDDEN_SIZE do
            network.weights2[k][i] = network.weights2[k][i] - learning_rate * delta_output[k] * hidden[i]
        end
        network.bias2[k] = network.bias2[k] - learning_rate * delta_output[k]
    end

    -- Backpropagate to hidden layer
    local delta_hidden = {}
    for i = 1, HIDDEN_SIZE do
        local sum = 0
        for k = 1, OUTPUT_SIZE do
            sum = sum + network.weights2[k][i] * delta_output[k]
        end
        delta_hidden[i] = sum * dsigmoid(hidden[i])
    end
    for i = 1, HIDDEN_SIZE do
        for j = 1, INPUT_SIZE do
            network.weights1[i][j] = network.weights1[i][j] - learning_rate * delta_hidden[i] * input[j]
        end
        network.bias1[i] = network.bias1[i] - learning_rate * delta_hidden[i]
    end
    save_network()
end

-- Predict function: returns predicted class and probabilities
local function predict(input)
    local _, output = forward(input)
    local max_val = -math.huge
    local max_index = 1
    for k, v in ipairs(output) do
        if v > max_val then
            max_val = v
            max_index = k
        end
    end
    local class = (max_index == 1 and "square") or (max_index == 2 and "circle") or (max_index == 3 and "triangle")
    return class, output
end

-- Convert user shape string to one-hot target vector.
local function target_vector(shape)
    if shape == "square" then
        return {1, 0, 0}
    elseif shape == "circle" then
        return {0, 1, 0}
    elseif shape == "triangle" then
        return {0, 0, 1}
    else
        return {0, 1, 0} -- default to circle if unrecognized
    end
end

---------------------------------
-- Drawing Grid Functions
---------------------------------

local function clear_grid(w, h)
    local grid = {}
    for y = 1, h do
        grid[y] = {}
        for x = 1, w do
            grid[y][x] = 0
        end
    end
    return grid
end

-- Global drawing grid (100x70)
local grid = clear_grid(WIDTH, HEIGHT)
local prev_grid = clear_grid(WIDTH, HEIGHT)

-- Draw a border around the grid.
local function draw_border(offset_x, offset_y, width, height)
    term.setBackgroundColor(colors.gray)
    -- Top border
    term.setCursorPos(offset_x - 1, offset_y - 1)
    term.write(string.rep(" ", width + 2))
    -- Bottom border
    term.setCursorPos(offset_x - 1, offset_y + height)
    term.write(string.rep(" ", width + 2))
    -- Left and right borders
    for row = offset_y, offset_y + height - 1 do
        term.setCursorPos(offset_x - 1, row)
        term.write(" ")
        term.setCursorPos(offset_x + width, row)
        term.write(" ")
    end
end

-- Draw a single cell for input grid.
local function draw_cell_input(x, y)
    local screen_x = GRID_OFFSET_X + x - 1
    local screen_y = GRID_OFFSET_Y + y - 1
    term.setCursorPos(screen_x, screen_y)
    if grid[y][x] == 1 then
        term.setBackgroundColor(colors.lime)
    else
        term.setBackgroundColor(colors.black)
    end
    term.write(" ")
end

-- Fully render the input drawing grid with border.
local function draw_full_grid_input()
    term.setBackgroundColor(colors.black)
    term.clear()
    for y = 1, HEIGHT do
        for x = 1, WIDTH do
            draw_cell_input(x, y)
        end
    end
    draw_border(GRID_OFFSET_X, GRID_OFFSET_Y, WIDTH, HEIGHT)
    term.setCursorPos(1, GRID_OFFSET_Y + HEIGHT + 2)
    term.setBackgroundColor(colors.black)
    term.write("Draw with mouse. Press ENTER when done:")
end

-- Capture drawing input from the user.
local function draw_input()
    draw_full_grid_input()
    while true do
        local event, button, x, y = os.pullEvent()
        if (event == "mouse_click" or event == "mouse_drag") and x and y then
            -- Adjust for grid offset
            local gridX = x - GRID_OFFSET_X + 1
            local gridY = y - GRID_OFFSET_Y + 1
            if gridX >= 1 and gridX <= WIDTH and gridY >= 1 and gridY <= HEIGHT then
                if grid[gridY][gridX] ~= 1 then
                    grid[gridY][gridX] = 1
                    draw_cell_input(gridX, gridY)
                end
            end
        elseif event == "key" and button == keys.enter then
            break
        end
    end
end

-- Detect the bounding box of the drawn shape.
local function get_bounding_box(input_grid)
    local minX, maxX = WIDTH + 1, 0
    local minY, maxY = HEIGHT + 1, 0
    for y = 1, HEIGHT do
        for x = 1, WIDTH do
            if input_grid[y][x] == 1 then
                if x < minX then minX = x end
                if x > maxX then maxX = x end
                if y < minY then minY = y end
                if y > maxY then maxY = y end
            end
        end
    end
    if maxX < minX then
        return 1, WIDTH, 1, HEIGHT  -- No drawing detected; default to full grid.
    end
    return minX, maxX, minY, maxY
end

-- Normalize the drawn shape (within its bounding box) to a 20x20 grid.
local function normalize_shape_to_20x20(input_grid)
    local minX, maxX, minY, maxY = get_bounding_box(input_grid)
    local norm = clear_grid(NEURON_WIDTH, NEURON_HEIGHT)
    local box_width = maxX - minX + 1
    local box_height = maxY - minY + 1
    for ny = 1, NEURON_HEIGHT do
        for nx = 1, NEURON_WIDTH do
            local srcX = minX + math.floor(((nx - 1) / (NEURON_WIDTH - 1)) * (box_width - 1) + 0.5)
            local srcY = minY + math.floor(((ny - 1) / (NEURON_HEIGHT - 1)) * (box_height - 1) + 0.5)
            norm[ny][nx] = input_grid[srcY][srcX]
        end
    end
    return norm
end

-- Stretch a 20x20 grid back to the original 100x70 drawing area.
local function stretch_norm_to_drawing(norm, targetWidth, targetHeight)
    local stretched = clear_grid(targetWidth, targetHeight)
    for y = 1, targetHeight do
        for x = 1, targetWidth do
            local srcX = math.floor((x - 1) / (targetWidth - 1) * (NEURON_WIDTH - 1)) + 1
            local srcY = math.floor((y - 1) / (targetHeight - 1) * (NEURON_HEIGHT - 1)) + 1
            stretched[y][x] = norm[srcY][srcX]
        end
    end
    return stretched
end

-- Draw a grid (of size targetWidth x targetHeight) on screen with border.
local function draw_stretched_grid_on_screen(aGrid, targetWidth, targetHeight)
    term.setBackgroundColor(colors.black)
    term.clear()
    for y = 1, targetHeight do
        for x = 1, targetWidth do
            local screen_x = GRID_OFFSET_X + x - 1
            local screen_y = GRID_OFFSET_Y + y - 1
            term.setCursorPos(screen_x, screen_y)
            if aGrid[y][x] == 1 then
                term.setBackgroundColor(colors.lime)
            else
                term.setBackgroundColor(colors.black)
            end
            term.write(" ")
        end
    end
    draw_border(GRID_OFFSET_X, GRID_OFFSET_Y, targetWidth, targetHeight)
end

-- Flatten a 2D grid into a 1D vector.
local function flatten_grid(two_d_grid)
    local vector = {}
    for y = 1, #two_d_grid do
        for x = 1, #two_d_grid[y] do
            table.insert(vector, two_d_grid[y][x])
        end
    end
    return vector
end

---------------------------------
-- Main Loop
---------------------------------

-- Initialize or load network parameters.
math.randomseed(os.time())
load_network()

while true do
    -- Reset the drawing grid.
    grid = clear_grid(WIDTH, HEIGHT)
    prev_grid = clear_grid(WIDTH, HEIGHT)
    term.setBackgroundColor(colors.black)
    term.clear()

    -- Capture user drawing input.
    draw_input()

    -- Normalize the drawn shape to a fixed 20x20 grid.
    local normalized_grid = normalize_shape_to_20x20(grid)

    -- Compute the network prediction using the flattened 20x20 input.
    local input_vector = flatten_grid(normalized_grid)
    local predicted_class, output_probs = predict(input_vector)

    -- Stretch the normalized grid back to the original drawing area (100x70).
    local stretched_grid = stretch_norm_to_drawing(normalized_grid, WIDTH, HEIGHT)
    draw_stretched_grid_on_screen(stretched_grid, WIDTH, HEIGHT)

    -- Display the network's prediction and probabilities.
    term.setCursorPos(1, GRID_OFFSET_Y + HEIGHT + 2)
    term.setBackgroundColor(colors.black)
    term.write(("Prediction: %s (Square: %.1f%%, Circle: %.1f%%, Triangle: %.1f%%)"):format(
        predicted_class,
        output_probs[1] * 100,
        output_probs[2] * 100,
        output_probs[3] * 100
    ))

    if training_mode then
        -- In training mode, prompt for correct shape label and train.
        term.setCursorPos(1, GRID_OFFSET_Y + HEIGHT + 3)
        term.write("Enter shape name (square, circle, triangle): ")
        sleep(0.2)  -- brief pause to flush events
        local shape = read()
        local target = target_vector(shape)
        train_network(input_vector, target, LEARNING_RATE)
        term.setCursorPos(1, GRID_OFFSET_Y + HEIGHT + 4)
        term.write("Training complete. Press ENTER to continue...")
        repeat
            local event, key = os.pullEvent("key")
        until key == keys.enter
    else
        -- In inference-only mode, do not train; wait for user to continue.
        term.setCursorPos(1, GRID_OFFSET_Y + HEIGHT + 3)
        term.write("Inference-only mode. Press ENTER to draw again...")
        repeat
            local event, key = os.pullEvent("key")
        until key == keys.enter
    end
end