require "table2"
local fmt = require "fmt"

local function polynomial_regression(x, y)
    local xm   = x:mean()
    local ym   = y:mean()
    local x2m  = x:mapped(|e| -> e * e):mean()
    local x3m  = x:mapped(|e| -> e * e * e):mean()
    local x4m  = x:mapped(|e| -> e * e * e * e):mean()
    local z    = table.zip(x, y)
    local xym  = z:mapped(|p| -> p[1] * p[2]):mean()
    local x2ym = z:mapped(|p| -> p[1] * p[1] * p[2]):mean()

    local sxx   = x2m - xm * xm
    local sxy   = xym - xm * ym
    local sxx2  = x3m - xm * x2m
    local sx2x2 = x4m - x2m * x2m
    local sx2y  = x2ym - x2m * ym

    local b = (sxy * sx2x2 - sx2y * sxx2) / (sxx * sx2x2 - sxx2 * sxx2)
    local c = (sx2y * sxx - sxy * sxx2) / (sxx * sx2x2 - sxx2 * sxx2)
    local a = ym - b * xm - c * x2m

    local function abc(xx) return a + b * xx + c * xx * xx end

    fmt.print("y = %g + %gx + %gx%s\n", a, b, c, fmt.super(2))
    print(" Input  Approximation")
    print(" x   y     y1")
    for z as p do fmt.print("%2d %3d  %5.1f", p[1], p[2], abc(p[1])) end
end

local x = {}
for i = 1, 11 do x[i] = i - 1 end
local y = {1, 6, 17, 34, 57, 86, 121, 162, 209, 262, 321}
polynomial_regression(x, y)