--@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.getSplitStrategy(type)
    return SPLIT_STRATEGIES[type] or SPLIT_STRATEGIES[Differential.TYPES.Open]
end

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
    self.splitStrategy = Differential.getSplitStrategy(config.Type or Differential.TYPES.Open)

    self.steerLock = config.SteerLock or 0
    self.steerAngle = 0
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:splitStrategy(
            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

    -- // REFACTOR
    if self.steerLock ~= 0 then
        local steerInertia = (aI + bI) / 2
        local inputForce = 228.0
        local maxSteerSpeed = math.rad(1337)

        local inputTorque = self.vehicle.steer * inputForce

        local avgSteerAngle = (self.outputA.customWheel.steerAngle + self.outputB.customWheel.steerAngle) / 2
        local mz = self.outputA.customWheel.mz + self.outputB.customWheel.mz
        local avgMz = (self.outputA.customWheel.mz + self.outputB.customWheel.mz) / 2
        local maxMz = math.max(self.outputA.customWheel.mz + self.outputB.customWheel.mz)
        local steerTorque = mz * -1 + inputTorque

        local steerAngularAccel = steerTorque / steerInertia

        self.steerAngle = math.clamp(
            self.steerAngle + steerAngularAccel * TICK_INTERVAL,
            -self.steerLock,
            self.steerLock
        )

        local wheelbase = 2.05
        local trackWidth = 1.124
        local radius = wheelbase / math.tan(math.rad(self.steerAngle))

        local innerAngle = math.deg(math.atan(wheelbase / (radius - (trackWidth / 2))))
        local outerAngle = math.deg(math.atan(wheelbase / (radius + (trackWidth / 2))))

        self.outputA.customWheel.steerAngle = outerAngle
        self.outputB.customWheel.steerAngle = innerAngle
    end

    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

Differential.TYPES = {
    Open = 'Open',
    Locking = 'Locking',
    VLSD = 'VLSD',
    HLSD = 'HLSD',
}

SPLIT_STRATEGIES = {
    [Differential.TYPES.Open] = Differential._openDiffTorqueSplit,
    [Differential.TYPES.Locking] = Differential._lockingDiffTorqueSplit,
    [Differential.TYPES.VLSD] = Differential._VLSDTorqueSplit,
    [Differential.TYPES.HLSD] = Differential._HLSDTorqueSplit,
}

return Differential