--@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.mass = config.Mass or 20
    self.radius = config.Radius or 0.27
    self.rollingResistance = config.RollingResistance or 40
    self.squat = config.Squat or 0.1
    self.slipCircleShape = config.SlipCircleShape or 1.05

    self.forwardFriction = Friction:new(config.ForwardFriction)
    self.sideFriction = Friction:new(config.SideFriction)
    self.frictionPreset = FrictionPreset:new(config.FrictionPreset)
    self.satFrictionPreset = FrictionPreset:new({ B = 13, C = 2.4, D = 1, E = 0.48 })

    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.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
        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.frictionPreset: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.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.frictionPreset: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(Vx, Vy, Lc)
    local VxAbs = math.abs(Vx)
    local Sy = 0

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

    Sy = math.clamp(Sy, -90, 90)
    --local phi = (1 - self.E) * Sy + (self.E / self.B) * math.atan(self.B * Sy)
    --local Mz = -slipSign * self.satFrictionPreset:evaluate(math.abs(phi), self.B, C, self.D, 0) * Lc

    return self.satFrictionPreset:evaluate(math.abs(Sy)) * -math.sign(Sy)
end

function Wheel:selfAligningTorque2(Vx, Vy, Fz)

    if Fz == 0 then
        return 0
    end

    local VxAbs = math.abs(Vx)
    local Sy = 0

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

    Sy = math.clamp(Sy, -1, 1)

    local a = {
        -2.72,
        -2.28,
        -1.86,
        -2.73,
        0.110,
        -0.070,
        0.643,
        -4.04,
        0.015,
        -0.066,
        0.945,
        0.030,
        0.070
    }
    local FzSqr = math.pow(Fz, 2)

    local C = 2.4
    local D = a[1] * FzSqr + a[2] * Fz
    local BCD = (a[3] * FzSqr + a[4] * Fz) / math.pow(math.exp(1), a[5] * Fz)
    local B = BCD / (C * D)
    local E = a[6] * FzSqr + a[7] * Fz + a[8]

    local phi = (1 - E) * Sy + (E / B) * math.atan(B * Sy)
    self.phi = phi

    return D * math.sin(C * math.atan(B * phi))
end

function Wheel:selfAligningTorque3(Mx, My, Sx, Sy)
    local M = Mx * math.cos(Sy) + My * math.cos(Sx)

    return self.radius * M
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)

    --self.steerAngle = self.steerAngle + self.mz / self.inertia

    self.forward = self.entity:getForward():rotateAroundAxis(self.entity:getUp(), -self.steerAngle)
    self.right = self.entity:getRight():rotateAroundAxis(self.entity:getUp(), -self.steerAngle)

    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, -- 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.
    )

    --local Mx = Fy * (math.cos(Sx) - 1) + Fx * math.sin(Sx)
    --local My = Fx * (math.cos(Sy) - 1) + Fy * math.sin(Sy)
    --local Mz = self:selfAligningTorque3(Mx, My, Sx, Sy)
    local Mz = self:selfAligningTorque2(forwardSpeed, sideSpeed, self.load)

    self.mz = Mz

    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