local super = require "Regression"

LinearRegression = super:new()

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

function LinearRegression: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
    
    local xScale = 1 / (self._xmax - self._xmin)
    local xOffset = 1 - self._xmin / (self._xmax - self._xmin)
    local yScale = 1 / (self._ymax - self._ymin)
    local yOffset = 1 - self._ymin / (self._ymax - self._ymin)
    
    for index = 1, self._n do
        local x = xs[index] * xScale + xOffset
        local y = ys[index] * yScale + yOffset
        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 = 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))
        
        coeff0 = (coeff0 - yOffset) / yScale
        coeff1 = coeff1 / yScale
        
        self._tcoeff0 = coeff0
        self._tcoeff1 = coeff1
        
        coeff0 = coeff0 + coeff1 * xOffset
        coeff1 = (coeff1) * xScale
        
        self._coeff0 = coeff0
        self._coeff1 = coeff1
        self._r2 = r2
        return true
    end
end

function LinearRegression:getEquation()
    local equation = 'y = '
    local empty = true
    if self._coeff1 == 1 then
        equation = equation .. 'x'
        empty = false
    elseif self._coeff1 == -1 then
        equation = equation .. '-x'
        empty = false
    elseif self._coeff1 ~= 0 then
        equation = equation .. string.format('%gx', self._coeff1)
        empty = false
    end
    if empty then
        equation = equation .. string.format('%g', self._coeff0)
    else
        if self._coeff0 < 0 then
            equation = equation .. string.format(' - %g', -self._coeff0)
        elseif self._coeff0 > 0 then
            equation = equation .. string.format(' + %g', self._coeff0)
        end
    end
    return equation
end

function LinearRegression:getValue(x)
    local xScale = 1 / (self._xmax - self._xmin)
    local xOffset = 1 - self._xmin / (self._xmax - self._xmin)
    x = x * xScale + xOffset
    return self._tcoeff0 + self._tcoeff1 * x
end

return LinearRegression
