require "matrix"
require "table2"
local fmt = require "fmt"

local function multiple_regression(y, x)
    local cy = y:transpose()
    local cx = x:transpose()
    return x:matmul(cx):inverse():matmul(x):matmul(cy):transpose():getRow(1)
end

local ys = {
    matrix.from({ {1, 2, 3, 4, 5} }),
    matrix.from({ {3, 4, 5} }),
    matrix.from({ {52.21, 53.12, 54.48, 55.84, 57.20, 58.57, 59.93, 61.29,
                   63.11, 64.47, 66.28, 68.10, 69.92, 72.19, 74.46} })
}

local a = {1.47, 1.50, 1.52, 1.55, 1.57, 1.60, 1.63, 1.65, 1.68, 1.70, 1.73, 1.75, 1.78, 1.80, 1.83}

local xs = {
    matrix.from({ {2, 1, 3, 4, 5} }),
    matrix.from({ {1, 2, 1}, {1, 1, 2} }),
    matrix.from({ table.rep(#a, 1), a, a:mapped(|e| -> e * e) })
}

for i = 1, #ys do
    local v = multiple_regression(ys[i], xs[i])
    fmt.lprint(v)
    if i < #ys then print() end
end