local super = require "Regression"

QuadraticRegression = super:new()

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

function QuadraticRegression:finish()
    local n = self._n
    local xs = self._xs
    local ys = self._ys
    local x1 = 0
    local x2 = 0
    local x3 = 0
    local x4 = 0
    local y1 = 0
    local y2 = 0
    local y1x1 = 0
    local y1x2 = 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
        x3 = x3 + x ^ 3
        x4 = x4 + x ^ 4
        y1 = y1 + y
        y2 = y2 + y ^ 2
        y1x1 = y1x1 + y * x
        y1x2 = y1x2 + y * x ^ 2
    end
    
    local det = n * x2 * x4
        - n * x3 ^ 2
        - x1 ^ 2 * x4
        + 2 * x1 * x2 * x3
        - x2 ^ 3
    if det ~= 0 then
        local adj11 = x2 * x4 - x3 ^ 2
        local adj12 = x2 * x3 - x4 * x1
        local adj13 = x1 * x3 - x2 ^ 2
        local adj22 = n * x4 - x2 ^ 2
        local adj23 = x1 * x2 - n * x3
        local adj33 = n * x2 - x1 ^ 2
        local coeff0 = (adj11 * y1 + adj12 * y1x1 + adj13 * y1x2) / det
        local coeff1 = (adj12 * y1 + adj22 * y1x1 + adj23 * y1x2) / det
        local coeff2 = (adj13 * y1 + adj23 * y1x1 + adj33 * y1x2) / det
        
        local SSreg = coeff0 ^ 2 * n + 2 * coeff0 * coeff1 * x1 + (2 * coeff0 * coeff2 + coeff1 ^ 2) * x2 + 2 * coeff1 * coeff2 * x3 + coeff2 ^ 2 * x4
            - 2 * y1 / n * (coeff0 * n + coeff1 * x1 + coeff2 * x2)
            + y1 ^ 2 / n
        local SStot = y2 - y1 ^ 2 / n
        local r2 = SSreg / SStot
        
        coeff0 = (coeff0 - yOffset) / yScale
        coeff1 = coeff1 / yScale
        coeff2 = coeff2 / yScale
        
        self._tcoeff0 = coeff0
        self._tcoeff1 = coeff1
        self._tcoeff2 = coeff2
        
        coeff0 = coeff0 + coeff1 * xOffset + coeff2 * xOffset ^ 2
        coeff1 = (coeff1 + 2 * coeff2 * xOffset) * xScale
        coeff2 = (coeff2) * xScale ^ 2
        
        self._coeff0 = coeff0
        self._coeff1 = coeff1
        self._coeff2 = coeff2
        self._r2 = r2
        return true
    end
end

function QuadraticRegression:getEquation()
    local equation = 'y = '
    local empty = true
    if self._coeff2 == 1 then
        equation = equation .. 'x^{2}'
        empty = false
    elseif self._coeff2 == -1 then
        equation = equation .. '-x^{2}'
        empty = false
    elseif self._coeff2 ~= 0 then
        equation = equation .. string.format('%gx^{2}', self._coeff2)
        empty = false
    end
    if empty then
        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
    else
        if self._coeff1 == 1 then
            equation = equation .. ' + x'
        elseif self._coeff1 == -1 then
            equation = equation .. ' - x'
        elseif self._coeff1 < 0 then
            equation = equation .. string.format(' - %gx', -self._coeff1)
        elseif self._coeff1 > 0 then
            equation = equation .. string.format(' + %gx', self._coeff1)
        end
    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 QuadraticRegression: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 + self._tcoeff2 * x^2
end

return QuadraticRegression
