--@include /koptilnya/libs/constants.txt
--@include ./powertrain_component.txt
--@include ./wheel.txt

require('/koptilnya/libs/constants.txt')
require('./powertrain_component.txt')
require('./wheel.txt')

Differential = class('Differential', PowertrainComponent)

function Differential:initialize(config, order)
    PowertrainComponent.initialize(self, 'Axle')

    local prefix = 'Axle' .. (order or '') .. '_'

    self.finalDrive = config.FinalDrive or 4
    self.biasAB = config.Bias or 0.5
    self.coastRamp = config.CoastRamp or 0.5
    self.powerRamp = config.PowerRamp or 1
    self.preload = config.Preload or 10
    self.slipTorque = config.SlipTorque or 1000

    local wheelConfig = {
        CalculateInertia = true
    }
    self.outputA = Wheel:new(wheelConfig, prefix .. 'Left')
    self.outputB = Wheel:new(wheelConfig, prefix .. 'Right')

    table.merge(self.wireInputs, self.outputA.wireInputs)
    table.merge(self.wireInputs, self.outputB.wireInputs)
    table.merge(self.wireOutputs, self.outputA.wireOutputs)
    table.merge(self.wireOutputs, self.outputB.wireOutputs)
end

function Differential:updateWireOutputs()
    if self.outputA ~= nil then
        self.outputA:updateWireOutputs()
    end

    if self.outputB ~= nil then
        self.outputB:updateWireOutputs()
    end
end

function Differential:queryInertia()
    if self.outputA == nil or self.outputB == nil then
        return self.inertia
    end

    return self.inertia + self.outputA:queryInertia() + self.outputB:queryInertia()
end

function Differential:queryAngularVelocity(angularVelocity)
    self.angularVelocity = angularVelocity

    if self.outputA == nil or self.outputB == nil then
        return angularVelocity
    end

    local aW = self.outputA:queryAngularVelocity(angularVelocity)
    local bW = self.outputB:queryAngularVelocity(angularVelocity)

    return (aW - bW) * self.finalDrive * 0.5
end

function Differential:forwardStep(torque, inertia)
    if self.outputA == nil or self.outputB == nil then
        return torque
    end

    local aW = self.outputA:queryAngularVelocity(self.angularVelocity)
    local bW = -self.outputB:queryAngularVelocity(self.angularVelocity)
    local aI = self.outputA:queryInertia()
    local bI = self.outputB:queryInertia()

    local tqA, tqB = self:_openDiffTorqueSplit(
            torque * self.finalDrive,
            aW,
            bW,
            aI,
            bI,
            self.biasAB,
            self.preload,
            self.stiffness,
            self.powerRamp,
            self.coastRamp,
            self.slipTorque
    )

    --tqA = tqA * aI * TICK_INTERVAL
    --tqB = tqB * bI * TICK_INTERVAL

    tqA = self.outputA:forwardStep(tqA, inertia * 0.5 + aI)
    tqB = self.outputB:forwardStep(tqB, inertia * 0.5 + bI)

    return tqA + tqB
end

function Differential:_openDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
    return tq * (1 - biasAB), tq * biasAB;
end

function Differential:_lockingDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
    local sumI = aI + bI
    local w = aI / sumI * aW + bI / sumI * bW
    local aTqCorr = (w - aW) * aI -- / dt
    aTqCorr = aTqCorr * stiffness

    local bTqCorr = (w - bW) * bI -- / dt
    bTqCorr = bTqCorr * stiffness

    local biasA = clamp(0.5 + (bW - aW) * 0.1 * stiffness, 0, 1)

    return tq * biasA + aTqCorr, tq * (1 - biasA) * bTqCorr
end

function Differential:_VLSDTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
    if aW < 0 or bW < 0 then
        return tq * (1 - biasAB), tq * biasAB
    end

    local c = tq > 0 and powerRamp or coastRamp
    local totalW = math.abs(aW) + math.abs(bW)
    local slip = 0
    if aW > 0 or bW > 0 then
        slip = (aW - bW) / totalW
    end
    local dTq = slip * stiffness * c * slipTorque

    return tq * (1 - biasAB) - dTq, tq * biasAB + dTq
end

function Differential:_HLSDTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
    if aW < 0 or bW < 0 then
        return tq * (1 - biasAB), tq * biasAB
    end

    local c = tq > 0 and powerRamp or coastRamp
    local tqFactor = math.clamp(math.abs(tq) / slipTorque, -1, 1)
    local totalW = math.abs(aW) + math.abs(bW)
    local slip = 0
    if aW > 0 or bW > 0 then
        slip = (aW - bW) / totalW
    end
    local dTq = slip * stiffness * c * slipTorque * tqFactor

    return tq * (1 - biasAB) - dTq, tq * biasAB + dTq
end