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


local PowertrainComponent = require('./powertrain_component.txt')
local WheelComponent = require('./wheel.txt')

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

local Differential = class('Differential', PowertrainComponent)

function Differential:initialize(vehicle, name, config)
    PowertrainComponent.initialize(self, vehicle, name, config)

    self.wireOutputs = {
        Diff_Torque = 'number'
    }

    self.finalDrive = config.FinalDrive or 4
    self.inertia = config.Inertia or 0.2
    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.stiffness = config.Stiffness or 0.1
    self.slipTorque = config.SlipTorque or 1000
end

function Differential:linkComponent(component)
    if not component:isInstanceOf(PowertrainComponent) then
        return
    end

    if self.outputA == nil then
        self.outputA = component
        component.input = self
    elseif self.outputB == nil then
        self.outputB = component
        component.input = self
    end
end

function Differential:updateWireOutputs()
    wire.ports.Diff_Torque = self.torque

    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()) / math.pow(self.finalDrive, 2)
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()

    self.torque = torque * self.finalDrive

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

    tqA = self.outputA:forwardStep(tqA, inertia * 0.5 * math.pow(self.finalDrive, 2) + aI) / self.finalDrive
    tqB = self.outputB:forwardStep(tqB, inertia * 0.5 * math.pow(self.finalDrive, 2) + bI) / self.finalDrive

    return (tqA + tqB) / self.finalDrive
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 = math.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

return Differential