local super = require "Regression"

PowerRegression = super:new()

local _exp = math.exp
local _log = math.log

function PowerRegression:init()
    super.init(self)
    self._coeff0 = nil
    self._coeff1 = nil
end

function PowerRegression:finish()
    local n = self._n
    local xs = self._xs
    local ys = self._ys
    local x1 = 0
    local x2 = 0
    local y1 = 0
    local y2 = 0
    local y1x1 = 0
    
    for index = 1, self._n do
        local x = xs[index]
        local y = ys[index]
        x = _log(x)
        y = _log(y)
        x1 = x1 + x
        x2 = x2 + x ^ 2
        y1 = y1 + y
        y2 = y2 + y ^ 2
        y1x1 = y1x1 + y * x
    end
    
    local det = n * x2 - x1 ^ 2
    if det ~= 0 then
        local inv11 = x2 / det
        local inv12 = -x1 / det
        local inv22 = n / det
        local coeff0 = _exp(inv11 * y1 + inv12 * y1x1)
        local coeff1 = inv12 * y1 + inv22 * y1x1
        local r2 = ((n * y1x1 - x1 * y1) ^ 2) / ((n * x2 - x1 ^ 2) * (n * y2 - y1 ^ 2))
        
        self._coeff0 = coeff0
        self._coeff1 = coeff1
        self._r2 = r2
        return true
    end
end

function PowerRegression:getEquation()
    local equation = ''
    local empty = true
    if self._coeff1 == 1 then
        equation = 'x'
        empty = false
    elseif self._coeff1 ~= 0 then
        equation = string.format('x^{%g}', self._coeff1)
        empty = false
    end
    if empty then
        equation = string.format('%g', self._coeff0)
    else
        if self._coeff0 == -1 then
            equation = '-' .. equation
        elseif self._coeff0 == 0 then
            equation = '0'
        elseif self._coeff0 ~= 1 then
            equation = string.format('%g', self._coeff0) .. equation
        end
    end
    equation = 'y = ' .. equation
    return equation
end

function PowerRegression:getValue(x)
    return self._coeff0 * x ^ self._coeff1
end

return PowerRegression
