local super = require "Regression"

CubicRegression = super:new()

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

function CubicRegression: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 x5 = 0
    local x6 = 0
    local y1 = 0
    local y2 = 0
    local y1x1 = 0
    local y1x2 = 0
    local y1x3 = 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, 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
        x5 = x5 + x ^ 5
        x6 = x6 + x ^ 6
        y1 = y1 + y
        y2 = y2 + y ^ 2
        y1x1 = y1x1 + y * x
        y1x2 = y1x2 + y * x ^ 2
        y1x3 = y1x3 + y * x ^ 3
    end
    
    local det = n * x2 * x4 * x6
        - n * x2 * x5 ^ 2
        - n * x3 ^ 2 * x6
        + 2 * n * x3 * x4 * x5
        - n * x4 ^ 3
        - x1 ^ 2 * x4 * x6
        + x1 ^ 2 * x5 ^ 2
        + 2 * x1 * x2 * x3 * x6
        - 2 * x1 * x3 ^ 2 * x5
        - 2 * x1 * x2 * x4 * x5
        + 2 * x1 * x3 * x4 ^ 2
        - x2 ^ 3 * x6
        + 2 * x2 ^ 2 * x3 * x5
        + x2 ^ 2 * x4 ^ 2
        - 3 * x2 * x3 ^ 2 * x4
        + x3 ^ 4
    if det ~= 0 then
        local adj11 = x2 * x4 * x6 + 2 * x3 * x4 * x5 - x2 * x5 ^ 2 - x3 ^ 2 * x6 - x4 ^ 3
        local adj12 = x1 * x5 ^ 2 + x2 * x3 * x6 + x3 * x4 ^ 2 - x1 * x4 * x6 - x3 ^ 2 * x5 - x2 * x4 * x5
        local adj13 = x1 * x3 * x6 + x2 * x3 * x5 + x2 * x4 ^ 2 - x1 * x4 * x5 - x2 ^ 2 * x6 - x3 ^ 2 * x4
        local adj14 = x1 * x4 ^ 2 + x2 ^ 2 * x5 + x3 ^ 3 - x1 * x3 * x5 - 2 * x2 * x3 * x4
        local adj22 = n * x4 * x6 + 2 * x2 * x3 * x5 - n * x5 ^ 2 - x2 ^ 2 * x6 - x3 ^ 2 * x4
        local adj23 = n * x4 * x5 + x1 * x2 * x6 + x3 ^ 3 - n * x3 * x6 - x1 * x3 * x5 - x2 * x3 * x4
        local adj24 = n * x3 * x5 + x1 * x3 * x4 + x2 ^ 2 * x4 - n * x4 ^ 2 - x1 * x2 * x5 - x2 * x3 ^ 2
        local adj33 = n * x2 * x6 + 2 * x1 * x3 * x4 - n * x4 ^ 2 - x1 ^ 2 * x6 - x2 * x3 ^ 2
        local adj34 = n * x3 * x4 + x1 ^ 2 * x5 + x2 ^ 2 * x3 - n * x2 * x5 - x1 * x3 ^ 2 - x1 * x2 * x4
        local adj44 = n * x2 * x4 + 2 * x1 * x2 * x3 - n * x3 ^ 2 - x1 ^ 2 * x4 - x2 ^ 3
        local coeff0 = (adj11 * y1 + adj12 * y1x1 + adj13 * y1x2 + adj14 * y1x3) / det
        local coeff1 = (adj12 * y1 + adj22 * y1x1 + adj23 * y1x2 + adj24 * y1x3) / det
        local coeff2 = (adj13 * y1 + adj23 * y1x1 + adj33 * y1x2 + adj34 * y1x3) / det
        local coeff3 = (adj14 * y1 + adj24 * y1x1 + adj34 * y1x2 + adj44 * y1x3) / det
        
        local SSreg = coeff0 ^ 2  *  n + 2 * coeff0 * coeff1 * x1 + (2 * coeff0 * coeff2 + coeff1 * coeff1) * x2 + (2 * coeff0 * coeff3 + 2 * coeff1 * coeff2) * x3 + (2 * coeff1 * coeff3 + coeff2 * coeff2) * x4 + (2 * coeff2 * coeff3) * x5 + coeff3 * coeff3 * x6
            - 2 * y1 / n * (coeff0 * n + coeff1 * x1 + coeff2 * x2 + coeff3 * x3)
            + y1 ^ 2 / n
        local SStot = y2 - y1 ^ 2 / n
        local r2 = SSreg / SStot
        
        coeff0 = (coeff0 - yOffset) / yScale
        coeff1 = coeff1 / yScale
        coeff2 = coeff2 / yScale
        coeff3 = coeff3 / yScale
        
        self._tcoeff0 = coeff0
        self._tcoeff1 = coeff1
        self._tcoeff2 = coeff2
        self._tcoeff3 = coeff3
        
        coeff0 = coeff0 + coeff1 * xOffset + coeff2 * xOffset ^ 2 + coeff3 * xOffset ^ 3
        coeff1 = (coeff1 + 2 * coeff2 * xOffset + 3 * coeff3 * xOffset ^ 2) * xScale
        coeff2 = (coeff2 + 3 * coeff3 * xOffset) * xScale ^ 2
        coeff3 = (coeff3) * xScale ^ 3
        
        self._coeff0 = coeff0
        self._coeff1 = coeff1
        self._coeff2 = coeff2
        self._coeff3 = coeff3
        self._r2 = r2
        return true
    end
end

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

return CubicRegression
