--@include ./friction.txt
--@include ./friction_preset.txt
--@include /koptilnya/libs/constants.txt

local Friction = require('./friction.txt')
local FrictionPreset = require('./friction_preset.txt')

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

local Wheel = class('Wheel')

function Wheel:initialize(config)
    config = config or {}

    self.direction = config.Direction or 1
    self.mass = config.Mass or 20
    self.radius = config.Radius or 0.27
    self.rollingResistance = config.RollingResistance or 20
    self.squat = config.Squat or 0.1
    self.slipCircleShape = config.SlipCircleShape or 1.05
    self.casterAngle = math.rad(config.CasterAngle or 0) * self.direction
    self.camberAngle = math.rad(config.CamberAngle or 0)
    self.toeAngle = math.rad(config.ToeAngle or 0) * -self.direction

    self.forwardFriction = Friction:new(config.ForwardFriction)
    self.sideFriction = Friction:new(config.SideFriction)

    self.lateralFrictionPreset = FrictionPreset:new(config.LateralFrictionPreset)
    self.longitudinalFrictionPreset = FrictionPreset:new(config.LongitudinalFrictionPreset)
    self.aligningFrictionPreset = FrictionPreset:new(config.AligningFrictionPreset)

    self.motorTorque = 0
    self.brakeTorque = 0
    self.steerAngle = 0
    self.hasHit = false
    self.counterTorque = 0
    self.inertia = 0
    self.load = 0
    self.angularVelocity = 0
    self.mz = 0

    self.forward = Vector(0)
    self.right = Vector(0)
    self.up = Vector(0)
    self.entity = NULL_ENTITY
    self.physObj = nil
    self.baseInertia = 0.5 * self.mass * math.pow(self.radius, 2)
    self.inertia = self.baseInertia
end

function Wheel:setEntity(entity)
    self.entity = entity
    self.entity:setNoDraw(false)

    self.physObj = isValid(entity) and entity:getPhysicsObject() or nil

    if isValid(self.physObj) then
        self.physObj:setMaterial('friction_00')
        self.physObj:setMass(self.mass)
        self.physObj:enableDrag(false)
        self.physObj:setDragCoefficient(0)
        self.physObj:setAngleDragCoefficient(0)
    end
end

function Wheel:setInertia(inertia)
    if isValid(self.physObj) then
        print(inertia)
        self.physObj:setInertia(Vector(inertia))
        self.inertia = inertia
    end
end

function Wheel:isValid()
    return isValid(self.entity) and isValid(self.physObj)
end

function Wheel:getRPM()
    return self.angularVelocity * RAD_TO_RPM
end

function Wheel:getLongitudinalLoadCoefficient(load)
    return 11000 * (1 - math.exp(-0.00014 * load));
end

function Wheel:getLateralLoadCoefficient(load)
    return 18000 * (1 - math.exp(-0.0001 * load));
end

function Wheel:stepLongitudinal(Tm, Tb, Vx, W, Lc, R, I, kFx, kSx)
    local Winit = W
    local VxAbs = math.abs(Vx)
    local Sx = 0

    if VxAbs >= 0.1 then
        Sx = (Vx - W * R) / VxAbs
    else
        Sx = (Vx - W * R) * 0.6
    end

    Sx = Sx * kSx;
    Sx = math.clamp(Sx, -1, 1)

    W = W + Tm / I * TICK_INTERVAL

    Tb = Tb * (W > 0 and -1 or 1)
    local TbCap = math.abs(W) * I / TICK_INTERVAL
    local Tbr = math.abs(Tb) - math.abs(TbCap)
    Tbr = math.max(Tbr, 0)
    Tb = math.clamp(Tb, -TbCap, TbCap)
    W = W + Tb / I * TICK_INTERVAL

    local maxTorque = self.longitudinalFrictionPreset:evaluate(math.abs(Sx)) * Lc * kFx * R
    local errorTorque = (W - Vx / R) * I / TICK_INTERVAL
    local surfaceTorque = math.clamp(errorTorque, -maxTorque, maxTorque)

    W = W - surfaceTorque / I * TICK_INTERVAL
    local Fx = surfaceTorque / R

    Tbr = Tbr * (W > 0 and -1 or 1)
    local TbCap2 = math.abs(W) * I / TICK_INTERVAL
    local Tb2 = math.clamp(Tbr, -TbCap2, TbCap2)
    W = W + Tb2 / I * TICK_INTERVAL

    local deltaOmegaTorque = (W - Winit) * I / TICK_INTERVAL
    local Tcnt = -surfaceTorque + Tb + Tb2 - deltaOmegaTorque

    if Lc < 0.001 then
    Sx = 0
    end

    return W, Sx, Fx, Tcnt
end

function Wheel:stepLateral(Vx, Vy, Lc, kFy, kSy)
    local VxAbs = math.abs(Vx)
    local Sy = 0

    if VxAbs > 0.1 then
        Sy = math.deg(math.atan(Vy / VxAbs))
        Sy = Sy / 50
    else
        Sy = Vy * (0.003 / TICK_INTERVAL)
    end

    Sy = Sy * kSy * 0.95
    Sy = math.clamp(Sy, -1, 1)
    local slipSign = Sy < 0 and -1 or 1
    local Fy = -slipSign * self.lateralFrictionPreset:evaluate(math.abs(Sy)) * Lc * kFy

    if Lc < 0.0001 then
    Sy = 0
    end

    return Sy, Fy
end

function Wheel:slipCircle(Sx, Sy, Fx, Fy, slipCircleShape)
    local SxAbs = math.abs(Sx)

    if SxAbs > 0.01 then
        local SxClamped = math.clamp(Sx, -1, 1)
        local SyClamped = math.clamp(Sy, -1, 1)
        local combinedSlip = Vector2(
            SxClamped * slipCircleShape,
            SyClamped
        )
        local slipDir = combinedSlip:getNormalized()

        local F = math.sqrt(Fx * Fx + Fy * Fy)
        local absSlipDirY = math.abs(slipDir.y)

        Fy = F * absSlipDirY * (Fy < 0 and -1 or 1)
    end

    return Sx, Sy, Fx, Fy
end

function Wheel:selfAligningTorque(Sy, load)
    if math.abs(Sy) < 0.001 or load < 0.001 then
        return 0
    end

    local B = self.aligningFrictionPreset.B
    local C = self.aligningFrictionPreset.C
    local D = self.aligningFrictionPreset.D
    local E = self.aligningFrictionPreset.E

    local loadScale = load * 1000
    local mechanicalTrail = 0.15
    local casterEffect = math.tan(self.casterAngle)
    local effectiveTrail = mechanicalTrail + casterEffect * self.radius
    local D_scaled = D * loadScale * effectiveTrail

    local camberEffect = 1 + 0.5 * math.abs(self.camberAngle)
    local toeEffect = 1 + 0.3 * math.abs(self.toeAngle)

    local term = B * Sy - E * (B * Sy - math.atan(B * Sy))
    local Mz = D_scaled * camberEffect * toeEffect * math.sin(C * math.atan(term))

    return Mz
end

function Wheel:rotateVector(vector)
    local ang = self.entity:getAngles()
    local baseForward = ang:getForward()
    local baseUp = ang:getUp()
    local baseRight = -ang:getRight()

    local steerRotated = vector:rotateAroundAxis(baseUp, nil, math.rad(-self.steerAngle) + self.toeAngle)
    local camberRotated = steerRotated:rotateAroundAxis(baseForward, nil, -self.camberAngle)
    local casterRotated = camberRotated:rotateAroundAxis(baseRight, nil, -self.casterAngle)

    return casterRotated:getNormalized()
end

function Wheel:update()
    if not isValid(self.entity) or not isValid(self.physObj) then
        return
    end

    local externalStress = self.physObj:getStress()
    local velocity = self.physObj:getVelocity() * UNITS_TO_METERS

    self.hasHit = externalStress ~= 0
    self.load = externalStress * KG_TO_KN

    self.longitudinalLoadCoefficient = self:getLongitudinalLoadCoefficient(self.load * 1000)
    self.lateralLoadCoefficient = self:getLateralLoadCoefficient(self.load * 1000)

    local ang = self.entity:getAngles()
    local baseForward = ang:getForward()
    local baseUp = ang:getUp()
    local baseRight = -ang:getRight()
    
    self.forward = self:rotateVector(baseForward)
    self.right = self:rotateVector(baseRight)
    self.up = self:rotateVector(baseUp)

    local forwardSpeed = 0
    local sideSpeed = 0

    if self.hasHit then
        forwardSpeed = velocity:dot(self.forward)
        sideSpeed = velocity:dot(self.right)
    end

    local W, Sx, Fx, CounterTq = self:stepLongitudinal(
        self.motorTorque,
        self.brakeTorque + self.rollingResistance,
        forwardSpeed,
        self.angularVelocity,
        self.longitudinalLoadCoefficient,
        self.radius,
        self.inertia,
        0.95,
        0.9
    )

    local Sy, Fy = self:stepLateral(
        forwardSpeed,
        sideSpeed,
        self.lateralLoadCoefficient,
        0.95,
        0.9

    )

    Sx, Sy, Fx, Fy = self:slipCircle(
        Sx,
        Sy,
        Fx,
        Fy,
        1.05
    )

    self.mz = self:selfAligningTorque(Sy, self.load)

    self.angularVelocity = W
    self.counterTorque = CounterTq
    self.forwardFriction.slip = Sx
    self.forwardFriction.force = Fx
    self.forwardFriction.speed = forwardSpeed
    self.sideFriction.slip = Sy
    self.sideFriction.force = Fy
    self.sideFriction.speed = sideSpeed
end

return Wheel