local Stabilitron = {}
Stabilitron.__index = Stabilitron

local pid = require ".pid"
local ship = _G.sublevel

local function clamp(value, min, max)
    return math.min(max, math.max(value, min))
end

local function normalizeAngle(angle)
    return ((angle + 180) % 360) - 180
end

local function yawFromPose(pose)
    local w = pose.orientation.a
    local x = pose.orientation.v.x
    local y = pose.orientation.v.y
    local z = pose.orientation.v.z
    local sinyCosp = 2 * (w * y + x * z)
    local cosyCosp = 1 - 2 * (y * y + z * z)
    return math.deg(math.atan(sinyCosp, cosyCosp))
end

function Stabilitron:new(autoAngleRight, autoAngleLeft, gimbal)
    local self = setmetatable({}, Stabilitron)
    self.autoAngleRight = autoAngleRight
    self.autoAngleLeft = autoAngleLeft
    self.gimbal = gimbal

    self.pidPitch = pid:new(0.9, 0.0, 0.25, -10, 10)
    self.pidRoll = pid:new(0.55, 0.0, 0.12, -15, 15)

    self.targetPitch = 0
    self.baseTrim = 12
    self.rollGain = 1
    self.targetX = 0
    self.targetZ = 0
    self.arrivalRadius = 8
    self.maxBank = 5
    self.headingToRoll = 0.04
    self.velocityLookahead = 0.6
    self.bankLiftComp = 0.12
    self.minCollective = 10
    self.pitchCmdLimit = 6
    self.rollRecoverAngle = 15
    self.maxDifferential = 12
    self.airframeDifferentialTrim = 1.5
    self.lastX = nil
    self.lastZ = nil
    self.lastTime = nil
    self.debugDistance = 0
    self.debugHeadingError = 0
    self.debugTargetRoll = 0
    self.debugYaw = 0

    return self
end

function Stabilitron:drawGraph(y, label, val, minV, maxV, color)
    local w = term.getSize()
    local graphW = w - 15
    local range = maxV - minV
    local normalized = (val - minV) / range
    local dotX = math.floor(12 + normalized * graphW)
    dotX = clamp(dotX, 12, w - 4)

    term.setCursorPos(1, y)
    term.setTextColor(colors.white)
    term.write(string.format("%-6s", label))
    term.setCursorPos(12, y)
    term.setTextColor(colors.gray)
    term.write(string.rep("-", graphW))
    term.setCursorPos(dotX, y)
    term.setTextColor(color)
    term.write("O")
    term.setCursorPos(w - 4, y)
    term.setTextColor(colors.white)
    term.write(string.format("%4.0f", val))
end

function Stabilitron:update()
    while true do
        local pose = ship and ship.getLogicalPose and ship.getLogicalPose()
        if pose then
            local now = os.clock()
            local curX = pose.position.x
            local curZ = pose.position.z
            local angles = self.gimbal.getAngles()
            local currentPitch = angles[1]
            local currentRoll = angles[2]
            local yaw = yawFromPose(pose)
            local velX = 0
            local velZ = 0

            if self.lastX and self.lastZ and self.lastTime then
                local dt = now - self.lastTime
                if dt > 0 then
                    velX = (curX - self.lastX) / dt
                    velZ = (curZ - self.lastZ) / dt
                end
            end

            self.lastX = curX
            self.lastZ = curZ
            self.lastTime = now

            local dx = self.targetX - curX
            local dz = self.targetZ - curZ
            local distance = math.sqrt(dx * dx + dz * dz)
            local targetRoll = 0
            local headingError = 0

            if distance > self.arrivalRadius then
                local aimX = self.targetX - velX * self.velocityLookahead
                local aimZ = self.targetZ - velZ * self.velocityLookahead
                local aimDx = aimX - curX
                local aimDz = aimZ - curZ
                local aimDistance = math.sqrt(aimDx * aimDx + aimDz * aimDz)

                if aimDistance > 0.001 then
                    local desiredHeading = math.deg(math.atan(aimDx, aimDz))
                    headingError = normalizeAngle(desiredHeading - yaw)
                    targetRoll = clamp(-headingError * self.headingToRoll, -self.maxBank, self.maxBank)
                end
            end

            if math.abs(currentRoll) > self.rollRecoverAngle then
                targetRoll = 0
            end

            local pitchCmd = self.pidPitch:compute(self.targetPitch, currentPitch, 0.05)
            local rollCmd = self.pidRoll:compute(targetRoll, currentRoll, 0.05)

            local bankComp = math.abs(currentRoll) * self.bankLiftComp
            local pitchCorrection = clamp(pitchCmd, -self.pitchCmdLimit, self.pitchCmdLimit)
            local collective = clamp(
                self.baseTrim + bankComp + pitchCorrection,
                self.minCollective,
                40
            )
            local differential = clamp(
                -rollCmd * self.rollGain + self.airframeDifferentialTrim,
                -self.maxDifferential,
                self.maxDifferential
            )

            local lOut = clamp(collective - differential, -40, 40)
            local rOut = clamp(collective + differential, -40, 40)

            self.debugDistance = distance
            self.debugHeadingError = headingError
            self.debugTargetRoll = targetRoll
            self.debugYaw = yaw

            self.autoAngleLeft:setAngle(lOut)
            self.autoAngleRight:setAngle(rOut)

            term.clear()
            term.setCursorPos(1, 1)
            term.setTextColor(colors.yellow)
            print("--- STABILITRON ---")

            self:drawGraph(3, "PITCH", currentPitch, -90, 90, colors.red)
            self:drawGraph(4, "ROLL ", currentRoll, -45, 45, colors.blue)
            self:drawGraph(5, "TRIM ", self.baseTrim, -40, 40, colors.orange)
            self:drawGraph(7, "OUT_L", lOut, -40, 40, colors.green)
            self:drawGraph(8, "OUT_R", rOut, -40, 40, colors.green)

            term.setCursorPos(1, 10)
            term.setTextColor(colors.lightGray)
            print(string.format("ERR_P : %.2f", self.targetPitch - currentPitch))
            print(string.format("P_CMD : %.2f", pitchCmd))
            print(string.format("R_CMD : %.2f", rollCmd))
            print(string.format("COLL  : %.2f", collective))
            print(string.format("DIST  : %.1f", self.debugDistance))
            print(string.format("YAW_E : %.2f", self.debugHeadingError))
            print(string.format("T_ROL : %.2f", self.debugTargetRoll))
        end

        sleep(0.05)
    end
end

return Stabilitron