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

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

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

---@alias DifferentialSplitStrategyFn fun(tq: number, aW: number, bW: number, aI: number, bI: number, biasAB: number, preload: number, stiffness: number, powerRamp: number, coastRamp: number, slipTorque: number): number, number

---@class DifferentialConfig: PowertrainComponentConfig
---@field Type? string
---@field FinalDrive? number
---@field Bias? number
---@field CoastRamp? number
---@field PowerRamp? number
---@field Preload? number
---@field Stiffness? number
---@field SlipTorque? number
---@field SteerLock? number
---@field ToeAngle? number

---@class Differential: PowertrainComponent
---@field finalDrive number
---@field biasAB number
---@field coastRamp number
---@field powerRamp number
---@field preload number
---@field stiffness number
---@field slipTorque number
---@field private splitStrategy DifferentialSplitStrategyFn
---@field steerLock number
---@field toeAngle number
---@field steerAngle number
---@field private mzDiff number
local Differential = class('Differential', PowertrainComponent)

---@private
---@param vehicle KPTLVehicle
---@param name string
---@param config DifferentialConfig
function Differential:initialize(vehicle, name, config)
	PowertrainComponent.initialize(self, vehicle, name, config)

	if CLIENT then
		return
	end

	self.wireOutputs = {
		[string.format('%s_Torque', self.name)] = 'number',
		[string.format('%s_W', self.name)] = 'number',
		[string.format('%s_MzDiff', self.name)] = '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.toeAngle = config.ToeAngle or 0
	self.steerAngle = 0

	self.mzDiff = 0
end

---@param component PowertrainComponent
---@return nil
function Differential:linkComponent(component)
	---@diagnostic disable-next-line: undefined-field
	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

---@return nil
function Differential:updateWireOutputs()
	PowertrainComponent.updateWireOutputs(self)

	wire.ports[string.format('%s_Torque', self.name)] = self.torque
	wire.ports[string.format('%s_W', self.name)] = self.angularVelocity
	wire.ports[string.format('%s_MzDiff', self.name)] = self.mzDiff

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

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

---@return number
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

---@return number
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

---@param torque number
---@param inertia number
---@return number
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)
	tqB = self.outputB:forwardStep(tqB, inertia * 0.5 * math.pow(self.finalDrive, 2) + bI)

	-- // REFACTOR
	if self.steerLock ~= 0 then
		local outputA = self.outputA --[[@as Wheel]]
		local outputB = self.outputB --[[@as Wheel]]

		local steerInertia = (aI + bI) / 2
		local driverInput = 15

		local localVelocityLength = chip():getVelocity():getLength()
		local MPH = localVelocityLength * (15 / 352)
		local KPH = MPH * 1.609

		local assist = math.clamp(10.0 / math.sqrt(KPH / 3), 2.0, 10.0)

		local inputForce = driverInput * assist

		local inputTorque = self.vehicle.steer * inputForce

		local mzDiff = outputA.customWheel.mz - outputB.customWheel.mz

		self.mzDiff = mzDiff

		local steerTorque = mzDiff * -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))))

		outputA.customWheel.steerAngle = outerAngle + self.toeAngle * outputA.customWheel.direction
		outputB.customWheel.steerAngle = innerAngle + self.toeAngle * outputB.customWheel.direction
	else
		local outputA = self.outputA --[[@as Wheel]]
		local outputB = self.outputB --[[@as Wheel]]

		outputA.customWheel.steerAngle = self.toeAngle * outputA.customWheel.direction
		outputB.customWheel.steerAngle = self.toeAngle * outputB.customWheel.direction
	end

	return tqA + tqB
end

---@return number, number
local function _openDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	return tq * (1 - biasAB), tq * biasAB
end

---@return number, number
local function _lockingDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	aTq = tq * 0.5
	bTq = tq * 0.5

	local syncTorque = (aW - bW) * stiffness * (aI + bI) * 0.5 / TICK_INTERVAL

	aTq = aTq - syncTorque
	bTq = bTq + syncTorque

	return aTq, bTq
	-- 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

---@return number, number
local function _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

---@return number, number
local function _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

---@enum DIFFERENTIAL_TYPE
DIFFERENTIAL_TYPES = {
	Open = 'Open',
	Locking = 'Locking',
	VLSD = 'VLSD',
	HLSD = 'HLSD',
}

---@type { [DIFFERENTIAL_TYPE]: DifferentialSplitStrategyFn }
local SPLIT_STRATEGIES = {
	[DIFFERENTIAL_TYPES.Open] = _openDiffTorqueSplit,
	[DIFFERENTIAL_TYPES.Locking] = _lockingDiffTorqueSplit,
	[DIFFERENTIAL_TYPES.VLSD] = _VLSDTorqueSplit,
	[DIFFERENTIAL_TYPES.HLSD] = _HLSDTorqueSplit,
}

---@param type DIFFERENTIAL_TYPE
---@return DifferentialSplitStrategyFn
Differential.getSplitStrategy = function(type)
	return SPLIT_STRATEGIES[type] or SPLIT_STRATEGIES[DIFFERENTIAL_TYPES.Open]
end

Differential.TYPES = DIFFERENTIAL_TYPES

return Differential