-- Automatic Transmission Type-E v1.0
-- Copyright (C) 2023 Vegetable Studio - All Rights Reserved
-- Redistribution, modification and commercial use are not allowed.
-- Please visit "https://www.patreon.com/vegetable_studio" for instructions.

-- Essentials
local raw = ac.accessCarPhysics()
local index = car.index

ac.store('Transmission Type [' .. tostring(index) ..']', 'AT')

local engine = ac.INIConfig.carData(index, 'engine.ini')
local drivetrain = ac.INIConfig.carData(index, 'drivetrain.ini')
local tyre = ac.INIConfig.carData(index, 'tyres.ini')
local compoundDefault = tyre:get('COMPOUND_DEFAULT', 'INDEX', 0)

local gearCount = car.gearCount
local maxGear = gearCount
local limiter = car.rpmLimiter
local torque = ac.DataLUT11.carData(index, engine:get('HEADER', 'POWER_CURVE', 'power.lut'))
local idleRpm = engine:get('ENGINE_DATA', 'MINIMUM', 800)
local clutchMaxTorque = drivetrain:get('CLUTCH', 'MAX_TORQUE', 600)
local gearsFinalRatio = drivetrain:get('GEARS', 'FINAL', 4.4)
local radius = tyre:get(compoundDefault ~= 0 and 'REAR_' .. tostring(compoundDefault) or 'REAR', 'RADIUS', 0.32)
local circumference = radius * 2 * math.pi
local rpmToKmh = circumference / 1000 * 60
local velocityToRpm = 60 / math.pi / 2

local gearRatios = {}
local topSpeeds = {}

for i = 1, gearCount do
    gearRatios[i] = drivetrain:get('GEARS', 'GEAR_' .. tostring(i), 0)
    topSpeeds[i] = limiter * 0.99 / gearRatios[i] / gearsFinalRatio * rpmToKmh
end

gearRatios[-1] = drivetrain:get('GEARS', 'GEAR_R', - gearRatios[1])
gearRatios[0] = gearRatios[1]
topSpeeds[-1] = 0
topSpeeds[0] = 0

local overDriveGear = 100
for i = 1, gearCount do
    if gearRatios[i] < 1 then
        overDriveGear = i
        break
    end
end

local function getBoost(rpm, gear)
    local boost = 1
    if car.turboCount == 0 then return boost end

    for id, sec in engine:iterate('TURBO') do
        local maxBoost = engine:get(sec, 'MAX_BOOST', 0)
        local wasteGate = engine:get(sec, 'WASTEGATE', 0)
        local referenceRpm = engine:get(sec, 'REFERENCE_RPM', 0)
        local gamma = engine:get(sec, 'GAMMA', 0)
        local ctrlTurbo = ac.INIConfig.carData(index, 'ctrl_turbo' .. tostring(id - 1) .. '.ini')
        local boostGamma = boost + math.min(rpm / referenceRpm, 1) ^ gamma * maxBoost

        if ctrlTurbo:tryGetLut('CONTROLLER_0', 'LUT') then
            local boostAdd = 0
            local boostMult = 1
            for index, section in ctrlTurbo:iterate('CONTROLLER') do
                local combinator = ctrlTurbo:get(section, 'COMBINATOR', nil)
                local input = ctrlTurbo:get(section, 'INPUT', nil)
                local lut = ctrlTurbo:tryGetLut(section, 'LUT')
                if not lut or combinator == nil or input == nil then return end

                if combinator[1] == 'ADD' then
                    if input[1] == 'RPMS' then
                        boostAdd = lut:get(rpm)
                    elseif input[1] == 'GEAR' then
                        boostAdd = lut:get(gear)
                    end
                elseif combinator[1] == 'MULT' then
                    if input[1] == 'RPMS' then
                        boostMult = lut:get(rpm)
                    elseif input[1] == 'GEAR' then
                        boostMult = lut:get(gear)
                    end
                end
            end

            boostGamma = boostGamma * boostAdd * boostMult
        end

        boost = boost + math.min(boostGamma, wasteGate)
    end

    return boost
end

-- Values for torque converter simulation
local stallRpm = (car.isEngineDiesel and 1500 or 2000) * math.round(idleRpm / 1000)
local stallTorque = torque:get(stallRpm)

if car.turboCount > 0 then
    local stallBoost = getBoost(stallRpm, 1)
    local loadBoost = getBoost(stallRpm * 0.75, 1)
    if stallBoost * 0.75 ^ 2 > loadBoost then stallBoost = loadBoost end

    stallTorque = stallTorque * stallBoost
end

-- Torque capacity of a converter is proportional to r * N ^ 2 * D ^ 5,
-- where r is the mass density of the fluid (kg/m3), N is the impeller speed (rpm), and D is the diameter (m).
local kFactor = math.round(stallRpm / math.sqrt(stallTorque), -1) * (1 + math.round(idleRpm / 750) / 3)
local torqueMult = 1
local torqueConverter = 0
local lockUp = false
local lockUpClutch = 0
local torqueRatio = 2
local kCoupling = (1 / torqueRatio ^ 0.5) / kFactor
local kStator = (1 - kCoupling * kFactor) / kFactor
local couplingRatio = 1.2 / torqueRatio ^ 0.5
-- local diameter = (kCoupling / 844 / math.pi * 2) ^ 0.2 * 2

-- Debug
-- ac.debug("K Factor (Min)", kFactor)
-- ac.debug('Diameter', math.round(diameter * 1000) .. ' mm')
-- ac.debug('Engine inertia to add', math.round(kCoupling * 25 - 0.025, 3) .. ' kg*m^2')

-- Values for shifting operation
local targetGear = 0
local vGear = 0
local gearRangeUp = false
local gearRangeDown = false
local shiftThreshold = 0
local shifting = false
local unlockSpeed = 1 / 0.1
local lockSpeed = 1 / 0.4
local delay = 0.5
local pGear = 0
local shiftSpeed = 1 / (2 / gearCount)
local gearClutch = 0

local maxTorque = {}
local maxPower, maxPowerRpm, maxPowerTorque = {}, {}, {}

for g = 1, gearCount do
    maxTorque[g], maxPower[g] = 0, 0

    for i = idleRpm, limiter + 1000, 100 do
        local boost = getBoost(i, g)

        if torque:get(i) * boost > maxTorque[g] then
            maxTorque[g] = torque:get(i) * boost
        end

        if torque:get(i) * boost * i / 5250 > maxPower[g] then
            maxPower[g] = torque:get(i) * boost * i / 5250
            maxPowerRpm[g] = i
            maxPowerTorque[g] = torque:get(i) * boost
        end
    end
end

local minSplippage = maxPowerTorque[overDriveGear ~= 100 and overDriveGear - 1 or gearCount] / ((maxPowerRpm[overDriveGear ~= 100 and overDriveGear - 1 or gearCount] / velocityToRpm) ^ 2 * kCoupling / (1 - couplingRatio))

-- f(V) = f0 + f01 * V + f02 * V ^ 2
-- where V [kph] is vehicle speed and f0, f01 and f02 are coefficients function of tire construction. For example, for a radial tire, the coefficients can be set as:
-- f0 = 1 * 10 ^ -2
-- f01 = 5 * 10 ^ -7
-- f02 = 2 * 10 ^ -7
local targetSpeedMap = {}
for i = 1, gearCount do
    local rollResistCoef = 1 * 10 ^ -2 + 5 * 10 ^ -7 * topSpeeds[i] + 2 * 10 ^ -7 * topSpeeds[i] ^ 2
    local rollResistForce = car.mass * 9.8 * rollResistCoef
    local airResistForce = 0.5 * 1.2 * (topSpeeds[i] / 3.6) ^ 2
    local accelForce = maxPowerTorque[i] * gearRatios[i] * gearsFinalRatio / radius
    local speedDiff = (accelForce - rollResistForce - airResistForce) / car.mass * delay * 3.6 -- (delay + 1 / unlockSpeed)
    local minUpRpm = stallRpm
    local maxUpRpm = maxPowerRpm[i]
    if overDriveGear - i < 2 then
        minUpRpm = minUpRpm + 500
        maxUpRpm = maxUpRpm + 500
    end
    targetSpeedMap[i] = ac.DataLUT11.parse('(|0=' .. tostring(minUpRpm / limiter * topSpeeds[i]) .. '|1=' .. tostring(topSpeeds[i] * math.min(maxUpRpm / limiter, 0.95 - minSplippage) - speedDiff) .. '|')
end

local checked = false
-- local valueToGear = {
--     [-5] = 'S',
--     [-4] = 'M',
--     [-3] = 'L',
--     [-2] = 'D',
--     [-1] = 'R',
--     [0] = 'N'
-- }

function script.update(dt)
    -- Disable H-shifter to stop overriding engaged gear
    raw.isShifterSupported = false

    -- Essentials
    local areControlsLocked = raw.areControlsLocked
    local gear = raw.gear - 1
    local rpmOutputWheel = raw.drivetrainDriveVelocity * velocityToRpm
    local throttle = raw.gas ^ 2
    local targetFinal = gearsFinalRatio

    if raw.inputMethod ~= ac.InputMethod.AI then
        -- Torque converter simulation
        local rpmInput = math.max(raw.rpm, 1)
        local rpmOutput = raw.drivetrainRootVelocity * velocityToRpm / torqueMult
        local speedRatio = math.max(rpmOutput / rpmInput, 0)

        local loadReduction = 1 - math.min(speedRatio / couplingRatio, 1)
        torqueMult = 1 + loadReduction * (torqueRatio - 1)
        targetFinal = targetFinal * torqueMult

        local velocityInput = rpmInput / velocityToRpm
        local torqueCoupling = math.max(velocityInput, raw.drivetrainRootVelocity) ^ 2 * kCoupling * math.min(math.abs(speedRatio - 1) / (1 - couplingRatio), 1)
        local torqueStator = velocityInput ^ 2 * kStator * (1 - math.min(speedRatio / couplingRatio, 1)) ^ 0.5 -- (1 - math.min(speedRatio, 1) ^ (15 / torqueRatio))

        if speedRatio > 1 then torqueCoupling = torqueCoupling * 0.075 end -- reverse mode
        torqueConverter = math.applyLag(torqueConverter, torqueCoupling + torqueStator, 0.5, dt)

        -- Clutch lock-up feature
        if not (vGear < overDriveGear) then
            lockUp = true
        else
            if throttle > 0.5 or gear == 0 then
                lockUp = false
            else
                rpmOutput = math.abs(rpmOutput)
                if rpmOutput > stallRpm + 250 and throttle < 0.25 then lockUp = true end
                if rpmOutput < stallRpm then lockUp = false end
            end
        end

        if lockUp then
            if lockUpClutch < 1 then lockUpClutch = math.min(lockUpClutch + dt * lockSpeed, 1) end
        else
            if lockUpClutch > 0 then lockUpClutch = math.max(lockUpClutch - dt * unlockSpeed, 0) end
        end

        -- Manuel shifting detection
        if raw.gearUp then
            gearRangeDown = true
        end

        if raw.gearDown then
            gearRangeUp = true
        end

        -- Debug
        -- ac.debug("Torque Multiplication", math.round(torqueMult, 2))
        -- ac.debug("Clutch Lock-Up", math.round(lockUpClutch * 100) .. ' %')
        -- ac.debug("Torque Converter", math.round(math.min(torqueConverter, clutchMaxTorque)) .. ' Nm')
        -- ac.debug("Speed Ratio", math.round(speedRatio, 2))
    else
        -- AIs are only allowed to shift between 'R', 'N', 'D' gears.
        if targetGear < 1 and raw.gearUp then
            gearRangeDown = true
        end

        if targetGear < 2 and raw.gearDown then
            gearRangeUp = true
        end
    end

    -- Calculate target gear by wheel speed & throttle input
    local speed = math.round(rpmOutputWheel * rpmToKmh, 2)
    local targetSpeedUp = 0
    local targetSpeedDown = 0

    if targetGear > 0 then
        targetSpeedUp = targetSpeedMap[targetGear]:get(throttle)
        targetSpeedDown = targetSpeedMap[math.max(targetGear - 1, 1)]:get(throttle) - 5

        if shiftThreshold < delay then
            local inSafeRange = speed < targetSpeedMap[math.max(gear, 1)]:get(1) and speed > targetSpeedMap[math.max(gear - 1, 1)]:get(0) - 5
            if speed > targetSpeedUp and targetGear < maxGear then
                targetGear = targetGear + 1
                if inSafeRange then shiftThreshold = 0 end
            elseif targetGear > 1 and (speed < targetSpeedDown or (targetGear > maxGear and speed < topSpeeds[maxGear] - 5)) then -- speed < topSpeeds[math.max(targetGear - 1, maxGear)]
                targetGear = targetGear - 1
                if inSafeRange then shiftThreshold = 0 end
            end
        end
    end

    -- Shifting mechanism
    if targetGear ~= vGear then
        shiftThreshold = shiftThreshold + dt
    else
        shiftThreshold = 0
    end

    if gearRangeUp and not raw.gearDown then
        if maxGear == gearCount then
            targetGear = math.clampN(targetGear - 1, -1, 0)
            shiftThreshold = delay - 0.2
        else
            maxGear = math.min(gearCount, maxGear + 1)
            if speed > targetSpeedUp and targetGear == maxGear - 1 then
                targetGear = targetGear + 1
                shiftThreshold = delay - 0.2
            end
        end
        gearRangeUp = false
    end

    if gearRangeDown and not raw.gearUp then
        if targetGear < 1 then
            targetGear = targetGear + 1
            shiftThreshold = delay - 0.2
        elseif maxGear > 1 then
            maxGear = maxGear - 1
            if targetGear > maxGear and speed < topSpeeds[maxGear] - 5 then
                targetGear = targetGear - 1
                shiftThreshold = delay - 0.2
            end
        end
        gearRangeDown = false
    end

    -- Disable shifting in forward gears to stop interfering new algorithm
    raw.gearUp = false
    raw.gearDown = false

    if shiftThreshold > delay then
        pGear = vGear
        vGear = targetGear
        if pGear ~= 0 and vGear ~= 0 then shifting = true end
        shiftThreshold = 0
    end

    if shifting then
        gearClutch = math.min(gearClutch + dt * shiftSpeed, 1)
        targetFinal = targetFinal * math.lerp(gearRatios[pGear] / gearRatios[vGear], 1, gearClutch)
    end

    if gearClutch == 1 then
        shifting = false
        gearClutch = 0
    end

    -- Actual physics override
    if not areControlsLocked then
        ac.setGearsFinalRatio(targetFinal)
        raw.clutch = math.min(torqueConverter / clutchMaxTorque + lockUpClutch, 1)
        checked = false
    else
        if not checked then
            ac.setGearsFinalRatio(gearsFinalRatio)
            maxGear = gearCount
            targetGear = 0
            vGear = 0
            pGear = 0
            checked = true
            return
        end

        -- User final gear ratio changed detection
        local actualFinal = raw.gearsFinalRatio

        if actualFinal ~= gearsFinalRatio then
            gearsFinalRatio = actualFinal
        end
    end

    raw.requestedGearIndex = vGear + 1

    local gearDisplay = targetGear

    if targetGear > 0 then
        if maxGear == gearCount then
            gearDisplay = -2
        elseif maxGear < gearCount then
            gearDisplay = maxGear
        end
    end

    ac.store('Gear Select [' .. tostring(index) ..']', gearDisplay)

    --Debug
    -- ac.debug('Gear Display', gearDisplay < 1 and valueToGear[gearDisplay] or tostring(gearDisplay))
    -- ac.debug('Gear Range', maxGear)
    -- ac.debug("Target Gear", targetGear)
    -- ac.debug("Target Speed (Up Shift)", math.round(targetSpeedUp) .. ' km/h')
    -- ac.debug("Target Speed (Down Shift)", math.round(targetSpeedDown) .. ' km/h')
    -- ac.debug('Shift Threshold', shiftThreshold)
    -- ac.debug("Final Ratio (Original)", gearsFinalRatio)
    -- ac.debug("Final Ratio (Current)", math.round(raw.gearsFinalRatio, 3))
end