-- @server
-- @include /koptilnya/libs/constants.txt
-- @include ./wire_component.txt
require('/koptilnya/libs/constants.txt')
require('./wire_component.txt')

Differential = class('Differential', WireComponent)

function Differential:initialize(config, order)
    self.power = config.Power or 1
    self.coast = config.Coast or 1
    self.finalDrive = config.FinalDrive or 1
    self.preload = config.Preload or 10
    self.order = order or 1

    self.viscousCoeff = config.ViscousCoeff or 0.9

    self.distributionCoeff = config.DistributionCoeff or 1
    self.maxCorrectingTorque = config.MaxCorrectingTorque or 200

    self.avgRPM = 0

    self.prefix = 'Axle' .. self.order .. '_'

    self.leftWheel = wire.ports[self.prefix .. 'LeftWheel']
    self.rightWheel = wire.ports[self.prefix .. 'RightWheel']

    self.lwPowerCoeff = 0
    self.rwPowerCoeff = 0

    self.lwtq_div_rwtq = 0
    self.rwtq_div_lwtq = 0

    self.lwtq = 0
    self.rwtq = 0

    self.lwav = 0
    self.rwav = 0

    hook.add('input', 'wire_axle_update' .. self.order, function(key, val)
        if key == self.prefix .. 'LeftWheel' then
            self.leftWheel = val
        end

        if key == self.prefix .. 'RightWheel' then
            self.rightWheel = val
        end
    end)
end

function Differential:getInputs()
    local inputs = {}

    inputs[self.prefix .. 'LeftWheel'] = 'entity'
    inputs[self.prefix .. 'RightWheel'] = 'entity'

    return inputs
end

function Differential:getOutputs()
    local outputs = {}

    outputs[self.prefix .. 'LWPowerCoeff'] = 'number'
    outputs[self.prefix .. 'RWPowerCoeff'] = 'number'
    outputs[self.prefix .. 'LWTq'] = 'number'
    outputs[self.prefix .. 'RWTq'] = 'number'
    outputs[self.prefix .. 'LWAV'] = 'number'
    outputs[self.prefix .. 'RWAV'] = 'number'
    outputs[self.prefix .. 'AvgRPM'] = 'number'
    outputs[self.prefix .. 'LW_div_RW'] = 'number'
    outputs[self.prefix .. 'RW_div_LW'] = 'number'
    -- outputs[self.prefix .. 'SGAV'] = 'number'

    return outputs
end

function Differential:updateOutputs()
    wire.ports[self.prefix .. 'LWPowerCoeff'] = self.lwPowerCoeff
    wire.ports[self.prefix .. 'RWPowerCoeff'] = self.rwPowerCoeff
    wire.ports[self.prefix .. 'LWTq'] = self.lwtq
    wire.ports[self.prefix .. 'RWTq'] = self.rwtq
    wire.ports[self.prefix .. 'RWTq'] = self.rwtq
    wire.ports[self.prefix .. 'LWAV'] = self.lwav
    wire.ports[self.prefix .. 'RWAV'] = self.rwav
    wire.ports[self.prefix .. 'AvgRPM'] = self.avgRPM
    wire.ports[self.prefix .. 'LW_div_RW'] = self.lwtq_div_rwtq
    wire.ports[self.prefix .. 'RW_div_LW'] = self.rwtq_div_lwtq
    -- wire.ports[self.prefix .. 'SGAV'] = self.sgav
end

function Differential:linkGearbox(gbox)
    self.gearbox = gbox
end

-- this is a test function that can spin wheels in different directions
function Differential:testUpdate()
    if isValid(self.leftWheel) and isValid(self.rightWheel) then
        local lwav = math.rad(self.leftWheel:getAngleVelocity()[2])
        local rwav = math.rad(self.rightWheel:getAngleVelocity()[2])

        self.avgRPM = (lwav - rwav) / 2 * RAD_TO_RPM * self.finalDrive

        self.lwav = lwav
        self.rwav = rwav

        -- spider gear angular velocity
        local sgav = (lwav + rwav) / 2

        self.sgav = sgav

        -- RAV = relative angular velocity
        local lwrav = lwav - sgav
        local rwrav = rwav - sgav

        local lwInertia = self.leftWheel:getInertia():getLength()
        local rwInertia = self.rightWheel:getInertia():getLength()

        local lwtq = lwrav * lwInertia * 2
        local rwtq = rwrav * rwInertia * 2

        self.lwtq = lwtq
        self.rwtq = rwtq

        local lwDampingTorque = self.viscousCoeff * lwtq
        local rwDampingTorque = self.viscousCoeff * rwtq

        local gboxTorque = self.gearbox and self.gearbox.torque / 10 or 0

        local usedCoeff = gboxTorque > self.preload and self.power or self.coast

        local lwCorrectingTorque = math.clamp(lwDampingTorque * usedCoeff, -self.maxCorrectingTorque,
            self.maxCorrectingTorque)
        local rwCorrectingTorque = math.clamp(rwDampingTorque * usedCoeff, -self.maxCorrectingTorque,
            self.maxCorrectingTorque)

        local driveTorque = gboxTorque * self.distributionCoeff * 52.49 * self.finalDrive

        local lwCorrectedTorque = lwCorrectingTorque * 600 - driveTorque
        local rwCorrectedTorque = rwCorrectingTorque * 600 + driveTorque

        self.leftWheel:applyTorque(lwCorrectedTorque * self.leftWheel:getRight() * TICK_INTERVAL)
        self.rightWheel:applyTorque(rwCorrectedTorque * self.rightWheel:getRight() * TICK_INTERVAL)
    end
end

function Differential:update()
    if isValid(self.leftWheel) and isValid(self.rightWheel) then
        local lwav = math.rad(self.leftWheel:getAngleVelocity()[2])
        local rwav = math.rad(self.rightWheel:getAngleVelocity()[2]) * -1

        self.lwav = lwav
        self.rwav = rwav

        -- RAV = relative anglar velocity
        local lwrav = rwav - lwav
        local rwrav = lwav - rwav

        local lwInertia = self.leftWheel:getInertia()[2]
        local rwInertia = self.rightWheel:getInertia()[2]

        local lwtq = lwrav * lwInertia
        local rwtq = rwrav * rwInertia

        self.lwtq = lwtq
        self.rwtq = rwtq

        -- Whether use diff power or diff coast
        local usePower = self.gearbox and self.gearbox.torque > self.preload and false
        local usedCoeff = usePower and self.power or self.coast

        -- Calculating damping torque  here
        local lwDampingTorque = self.viscousCoeff * lwtq * usedCoeff
        local rwDampingTorque = self.viscousCoeff * rwtq * usedCoeff

        local lwav_abs = math.abs(lwav)
        local rwav_abs = math.abs(rwav)

        local avg_abs = lwav_abs + rwav_abs

        local lwav_to_avg = avg_abs ~= 0 and 2 * lwav_abs / avg_abs or 0
        local rwav_to_avg = avg_abs ~= 0 and 2 * rwav_abs / avg_abs or 0

        self.lwtq_div_rwtq = lwav_to_avg
        self.rwtq_div_lwtq = rwav_to_avg

        local lwPowerCoeff = math.clamp(lwav_to_avg, usedCoeff, 1 + (1 - usedCoeff))
        local rwPowerCoeff = math.clamp(rwav_to_avg, usedCoeff, 1 + (1 - usedCoeff))

        self.lwPowerCoeff = lwPowerCoeff
        self.rwPowerCoeff = rwPowerCoeff

        local lwCorrectingTorque = math.clamp(lwDampingTorque, -self.maxCorrectingTorque, self.maxCorrectingTorque)
        local rwCorrectingTorque = math.clamp(rwDampingTorque, -self.maxCorrectingTorque, self.maxCorrectingTorque)

        local gboxTorque = self.gearbox and self.gearbox.torque or 0

        local driveTorque = gboxTorque / 2 * self.distributionCoeff * self.finalDrive

        local lwCorrectedTorque = lwCorrectingTorque * 20 + driveTorque * lwPowerCoeff -- * lwPowerCoeff
        local rwCorrectedTorque = rwCorrectingTorque * 20 + driveTorque * rwPowerCoeff -- * rwPowerCoeff

        self.leftWheel:applyTorque(lwCorrectedTorque * -self.leftWheel:getRight())
        self.rightWheel:applyTorque(rwCorrectedTorque * self.rightWheel:getRight())

        self.avgRPM = (lwav + rwav) / 2 * RAD_TO_RPM * self.finalDrive
    end
end