--@name Custom Wheel
--@author Koptilnya
--@server
--@include /koptilnya/libs/constants.txt

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

Wheel = class('Wheel')

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

    self.name = config.Name or 'Wheel'
    self.mass = config.Mass or 20
    self.radius = config.Radius or 0.25
    self.width = config.Width or 0.25
    self.dragTorque = config.DragTorque or 40

    self.B = 10.86
    self.C = 2.15
    self.D = 0.933
    self.E = 0.992

    self.motorTorque = 0
    self.brakeTorque = 0
    self.steerAngle = 0

    self.forward = Vector(0)
    self.right = Vector(0)
    self.up = Vector(0)
    self.inertia = 0
    self.load = 0
    self.angularVelocity = 0

    self.baseInertia = 0.5 * self.mass * math.pow(self.radius, 2)
    self.inertia = self.baseInertia
    self.entity = NULL_ENTITY
    self.physObject = nil
    self.hasHit = false

    wire.adjustPorts(
        {
            [self.name] = 'entity',
            Torque = 'number',
            Brake = 'number'
        },
        {
            HasHit = 'number'
        }
    )

    hook.add('tick', 'custom_wheel_update', function()
        if isValid(self.physObject) then
            self:update()
        end
    end)

    hook.add('input', 'custom_wheel_input', function(input, value)
        if input == self.name then
            if isValid(value) then
                self.entity = value
                self.entity:enableSphere(true, self.radius * UNITS_PER_METER)

                self.physObject = value:getPhysicsObject()
                self.physObject:setMaterial('friction_00')
                self.physObject:setMass(self.mass)
                self.physObject:setInertia(Vector(self.inertia))
                self.physObject:enableDrag(false)
                self.physObject:setDragCoefficient(0)
                self.physObject:setAngleDragCoefficient(0)
            else
                self.entity = NULL_ENTITY
                self.physObject = nil
            end
        end
    end)
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:pacejka(slip, B, C, D, E)
    local t = math.abs(slip)
    return D * math.sin(C * math.atan(B * t - E * (B * t - math.atan(B * t))))
end

function Wheel:stepLongitudinal(Tm, Tb, Vx, W, Lc, R, I, kFx, kSx)
    if I < 0.0001 then
        I = 0.0001
    end

    local Winit = W
    local VxAbs = math.abs(Vx)
    local WAbs  = math.abs(W)
    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:pacejka(math.abs(Sx), self.B, self.C, self.D, self.E) * 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
    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.3 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:pacejka(math.abs(Sy), self.B, self.C, self.D, self.E) * 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:update()
    local externalStress = self.physObject:getStress()
    local velocity = self.physObject:getVelocity() * UNITS_TO_METERS

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

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

    self.forward = self.entity:getForward()
    self.right = -self.entity:getRight()

    local forwardSpeed = velocity:dot(self.forward)
    local sideSpeed = velocity:dot(self.right)

    self.motorTorque = wire.ports.Torque
    self.brakeTorque = wire.ports.Brake

    local W, Sx, Fx, Tcnt = self:stepLongitudinal(
        self.motorTorque,
        self.brakeTorque + self.dragTorque,
        forwardSpeed,
        self.angularVelocity,
        self.longitudinalLoadCoefficient,
        self.radius,
        self.inertia,
        0.95, -- Force coeff
        0.9 -- Slip coeff
    )

    local Sy, Fy = self:stepLateral(
        forwardSpeed,
        sideSpeed,
        self.lateralLoadCoefficient,
        0.95, -- Force coeff
        0.9 -- Slip coeff
    )

    Sx, Sy, Fx, Fy = self:slipCircle(
        Sx,
        Sy,
        Fx,
        Fy,
        1.05 -- Shape of the slip circle / ellipse.
    )

    if self.hasHit then
        local surfaceForceVector = self.right * Fy + self.forward * Fx

        self.physObject:applyForceOffset(surfaceForceVector, self.entity:getPos())
    end

    wire.ports.HasHit = self.hasHit and 1 or 0
end

local wheel = Wheel:new()