ship_stabilizer

-- Ship PID Stabilizer
-- Reads a gimbal_sensor and outputs analog redstone to control propeller speed.
-- Ship is assumed to face SE/NW, so world X/Z angles are rotated 45 degrees
-- into ship-local pitch and roll before feeding the PID controllers.
--
-- If an altitude_sensor is also attached, a third PID maintains a target height
-- by adjusting all four propellers collectively (↑/↓ arrows change target).

local SAMPLE_RATE = 0.05 -- seconds per update (~20 Hz)

-- How many blocks to move the altitude target per key press
local ALT_STEP = 1

-- Tune these gains for your ship's mass and propeller responsiveness.
-- Tuning order: zero kI, raise kP until it corrects, raise kD until oscillation stops,
-- then add tiny kI to kill steady-state lean. Never raise kP past the oscillation point.
local GAINS = {
    pitch    = { kP = 0.25, kI = 0.0,  kD = 2.0  },
    roll     = { kP = 0.25, kI = 0.0,  kD = 2.0  },
    altitude = { kP = 0.4,  kI = 0.01, kD = 1.5  },
}

-- Low-pass smoothing for the derivative term (0 = no filter, 1 = freeze).
-- Reduces noise-driven output spikes. Raise toward 0.8 if output chatters.
local D_FILTER = 0.4

-- Map logical propeller positions to computer output sides.
local OUT = {
    front = "back",
    back  = "front",
    left  = "left",
    right = "right",
}

-- Propellers cut out above this signal strength — keep output strictly below it.
local MAX  = 14

-- Neutral signal (level flight, no altitude correction).
local BASE = math.floor(MAX / 2)

-- ── PID ──────────────────────────────────────────────────────────────────────

local pidState = {
    pitch    = { integral = 0, prevErr = 0, dSmooth = 0 },
    roll     = { integral = 0, prevErr = 0, dSmooth = 0 },
    altitude = { integral = 0, prevErr = 0, dSmooth = 0 },
}

local function clamp(v, lo, hi)
    return math.max(lo, math.min(hi, v))
end

local function pidStep(axis, err, dt)
    local g = GAINS[axis]
    local s = pidState[axis]

    -- Integral with windup guard (skip accumulation when kI is zeroed out)
    if g.kI > 0 then
        s.integral = s.integral + err * dt
        -- Limit integral so its contribution stays within ±(BASE/2)
        s.integral = clamp(s.integral, -(BASE / 2) / g.kI, (BASE / 2) / g.kI)
    else
        s.integral = 0
    end

    -- Derivative with low-pass filter to reduce sensor noise spikes
    local rawDeriv = (err - s.prevErr) / dt
    s.dSmooth = D_FILTER * s.dSmooth + (1 - D_FILTER) * rawDeriv
    s.prevErr = err

    return g.kP * err + g.kI * s.integral + g.kD * s.dSmooth
end

-- ── Coordinate transform ──────────────────────────────────────────────────────
-- For a SE-facing ship: pitch = xAngle + zAngle, roll = xAngle - zAngle.
-- The √2 normalisation is absorbed into PID gains.

local function worldToShip(xDeg, zDeg)
    return xDeg + zDeg, xDeg - zDeg
end

-- ── Peripherals ───────────────────────────────────────────────────────────────

local gimbal = peripheral.find("gimbal_sensor")
if not gimbal then
    error("gimbal_sensor peripheral not found — attach one to this computer")
end

local altimeter = peripheral.find("altitude_sensor")
local mode3D    = altimeter ~= nil
local targetAlt = 0
if mode3D then
    targetAlt = altimeter.getHeight()
end

-- ── Redstone output ───────────────────────────────────────────────────────────

local function setOutputs(front, back, left, right)
    redstone.setAnalogOutput(OUT.front, math.floor(clamp(front, 0, MAX)))
    redstone.setAnalogOutput(OUT.back,  math.floor(clamp(back,  0, MAX)))
    redstone.setAnalogOutput(OUT.left,  math.floor(clamp(left,  0, MAX)))
    redstone.setAnalogOutput(OUT.right, math.floor(clamp(right, 0, MAX)))
end

local function resetOutputs()
    for _, side in pairs(OUT) do
        redstone.setAnalogOutput(side, 0)
    end
end

-- ── Main ──────────────────────────────────────────────────────────────────────

if mode3D then
    print(string.format("3D mode  —  holding altitude %.1f", targetAlt))
    print("↑/↓ arrows adjust target altitude, Q to stop.")
else
    print("2D mode  —  no altitude_sensor found, tilt-only stabilization.")
    print("Press Q to stop.")
end

local running  = true
local lastTime = os.clock()

parallel.waitForAny(
    -- ── Control loop ──────────────────────────────────────────────────────────
    function()
        while running do
            local now = os.clock()
            local dt  = now - lastTime
            if dt < SAMPLE_RATE then
                os.sleep(SAMPLE_RATE - dt)
                dt = os.clock() - lastTime
            end
            lastTime = os.clock()
            dt = math.max(dt, 0.001)

            local angles     = gimbal.getAngles()
            local xDeg, zDeg = angles[1], angles[2]
            local pitch, roll = worldToShip(xDeg, zDeg)

            local pitchOut = pidStep("pitch", pitch, dt)
            local rollOut  = pidStep("roll",  roll,  dt)

            -- Collective: altitude PID adds the same offset to all four props.
            local collective = 0
            local curAlt     = 0
            if mode3D then
                curAlt     = altimeter.getHeight()
                local altErr = targetAlt - curAlt
                collective = pidStep("altitude", altErr, dt)
            end

            -- Mix collective (altitude) with differential (pitch/roll).
            --   pitch: front/back pair  — positive pitch means nose up → more rear, less front
            --   roll:  left/right pair  — positive roll means right down → more left, less right
            local front = BASE + collective - pitchOut
            local back  = BASE + collective + pitchOut
            local left  = BASE + collective + rollOut
            local right = BASE + collective - rollOut

            setOutputs(front, back, left, right)

            term.clear()
            term.setCursorPos(1, 1)
            print(string.format("Gimbal   X: %+6.2f deg   Z: %+6.2f deg", xDeg, zDeg))
            print(string.format("Ship   Pitch: %+6.2f     Roll: %+6.2f", pitch, roll))
            if mode3D then
                print(string.format("Altitude  cur: %7.2f   target: %7.2f   err: %+.2f",
                    curAlt, targetAlt, targetAlt - curAlt))
            else
                print("Altitude  [no sensor]")
            end
            print(string.format("Output  F:%2d  B:%2d  L:%2d  R:%2d  (max %d)",
                math.floor(clamp(front, 0, MAX)),
                math.floor(clamp(back,  0, MAX)),
                math.floor(clamp(left,  0, MAX)),
                math.floor(clamp(right, 0, MAX)),
                MAX))
            print("")
            if mode3D then
                print("↑/↓ adjust altitude target  |  Q quit")
            else
                print("Q quit")
            end
        end
    end,

    -- ── Key listener ──────────────────────────────────────────────────────────
    function()
        while true do
            local _, key = os.pullEvent("key")
            if key == keys.q then
                running = false
                return
            elseif mode3D and key == keys.up then
                targetAlt = targetAlt + ALT_STEP
            elseif mode3D and key == keys.down then
                targetAlt = targetAlt - ALT_STEP
            end
        end
    end
)

resetOutputs()
print("Stabilizer stopped — all outputs zeroed.")