balloon

-- Autopilot for a Create Aeronautics balloon. On the data cable: the hot air
-- burners (lit by redstone), a gimbal sensor, an altitude sensor if you have
-- one. It keeps the ship level by giving more hot air to the low side, and
-- holds the height you choose. Up/Down: height, Tab: devices, Q: quit.

-- Settings: raise a gain if the ship reacts too slowly, lower it if it swings
local GAIN_LEVEL = 1
local GAIN_HEIGHT = 1
local CLIMB = 1.5  -- blocks per second at most
local FLIP_NS = false  -- true if it leans more instead of less
local FLIP_EW = false

local C = term.colors
local W = term.get_size().w
local H = term.get_size().h
local COLOR = os.tier() >= 4
local LO = 5
local MAPW = W * 37 // 100
local BY = H - 1
local DIRS = {"N", "NE", "E", "SE", "S", "SW", "W", "NW"}
local T = {
  tabs = {"Flight", "Devices"}, height = "Height ", target = "Target ", thrust = "Thrust ", climb = "Climb ",
  tilt = "Tilt ", flat = "level", low = " low ", level = "Level", hold = "Height", man = "MAN",
  noBurner = "No burner: connect them with data cable.", noGyro = "No gimbal sensor: no levelling.",
  noAlt = "No altitude sensor: thrust by hand.", unlit = "no redstone!", noBalloon = "no balloon",
  head = "Burner  Pos.    Amount  RS Balloon", center = "center",
  hint = "Up/Down: height  Tab: devices  Q: quit",
}

local target = nil
local hold = true
local level = true
local burners = {}
local maxOf = {}
local gyro = nil
local alti = nil
local onShip = false
local R = 0
local CX = 0
local CZ = 0
local alt = nil
local vz = 0
local ax = 0  -- > 0: the south side is low
local az = 0  -- > 0: the east side is low
local rx = 0
local rz = 0
local sp = nil  -- set point, gliding to the target
local trim = nil  -- thrust that holds the height
local lift = 0.5
local ix = 0
local iz = 0
local last = os.clock()
local scanned = -100
local lost = false
local page = 1
local shown = {}
local hits = {}
local tock = false

local function clamp(x, a, b)
  return math.max(a, math.min(b, x))
end

local function fit(s, n)
  return string.sub(s .. string.rep(" ", n), 1, n)
end

local function f1(v)
  return string.format("%.1f", v)
end

local function pct(f)
  return math.floor(f * 100 + 0.5) .. "%"
end

local function save()
  fs.write("balloon.cfg", tostring(target) .. " " .. tostring(hold) .. " " .. tostring(level))
end

local function angle(dx, dz)
  return (math.atan2(dx, -dz) * 180 / math.pi + 360) % 360
end

local function compass(dx, dz)
  return DIRS[math.floor((angle(dx, dz) + 22.5) / 45) % 8 + 1]
end

local function place(name)
  local at = string.find(name, "@", 1, true)
  if at then
    local xyz = string.split(string.sub(name, at + 1), ",")
    return {x = tonumber(xyz[1]), z = tonumber(xyz[3])}
  end
  return nil
end

local function addBurner(name)
  local p = peripheral.wrap(name)
  local b = {name = name, p = p, amount = p.amount(), signal = 15, f = 0, dx = 0, dz = 0, order = 0, pos = place(name)}
  if maxOf[name] == nil then
    maxOf[name] = p.set_amount(100000)
    p.set_amount(b.amount)
  end
  b.max = maxOf[name]
  if b.pos == nil then
    rs.set(name, 15)  -- against the computer: lit by it
  end
  table.insert(burners, b)
end

local function byPlace(a, b)
  return a.order < b.order
end

local function first(kind)
  local name = peripheral.list(kind)[1]
  return name and peripheral.wrap(name)
end

-- The centre is the middle of the burners; those far from it level the ship
local function scan()
  burners = {}
  for _, name in ipairs(peripheral.list("hot_air_burner")) do
    pcall(addBurner, name)
  end
  local n = 0
  CX = 0
  CZ = 0
  for _, b in ipairs(burners) do
    if b.pos then
      n = n + 1
      CX = CX + b.pos.x
      CZ = CZ + b.pos.z
    end
  end
  CX = CX / math.max(n, 1)
  CZ = CZ / math.max(n, 1)
  R = 0
  for _, b in ipairs(burners) do
    if b.pos then
      b.dx = b.pos.x - CX
      b.dz = b.pos.z - CZ
      R = math.max(R, math.sqrt(b.dx * b.dx + b.dz * b.dz))
    end
  end
  for _, b in ipairs(burners) do
    b.outer = R >= 1 and math.sqrt(b.dx * b.dx + b.dz * b.dz) >= R * 0.4
    b.ux = 0
    b.uz = 0
    b.dir = "C"
    if b.outer then
      b.ux = b.dx / R
      b.uz = b.dz / R
      b.dir = compass(b.dx, b.dz)
      b.order = 1 + angle(b.dx, b.dz)
    end
  end
  table.sort(burners, byPlace)
  gyro = first("gimbal_sensor")
  alti = first("altitude_sensor")
  onShip = vehicle ~= nil and vehicle.on_vehicle()
  scanned = os.clock()
  lost = false
end

local function top(b)
  local t = b.max
  if b.cap and b.signal > 0 then
    t = math.min(t, b.cap * 15 / b.signal)
  end
  return math.max(t, LO + 1)
end

local function readBurner(b)
  local r = b.p.read()
  b.amount = r.amount
  b.signal = r.signal
  b.filled = r.filled
  b.cap = r.capacity
end

local function sense(dt)
  for _, b in ipairs(burners) do
    lost = lost or not pcall(readBurner, b).ok
  end
  if gyro then
    local r = pcall(gyro.read)
    lost = lost or not r.ok
    if r.ok then
      local nx = r.value.angle_x
      local nz = r.value.angle_z
      if FLIP_NS then
        nx = -nx
      end
      if FLIP_EW then
        nz = -nz
      end
      rx = (nx - ax) / dt
      rz = (nz - az) / dt
      ax = nx
      az = nz
    end
  end
  local h = nil
  if alti then
    local r = pcall(alti.height)
    lost = lost or not r.ok
    h = r.value
  elseif onShip then
    h = vehicle.position().y
  end
  if onShip then
    vz = vehicle.velocity().y
  elseif h and alt then
    vz = (h - alt) / dt
  end
  alt = h
end

-- Height: thrust = trim + gap to the set point - climb speed. Level: an outer
-- burner gets more when its side is low or going down.
local function control(dt)
  if burners[1] == nil then
    return
  end
  if trim == nil then
    trim = 0
    for _, b in ipairs(burners) do
      trim = trim + clamp((b.amount - LO) / (top(b) - LO), 0, 1) / #burners
    end
    lift = trim
  end
  if hold and alt then
    target = target or math.floor(alt + 0.5)
    sp = sp or alt
    sp = sp + clamp(target - sp, -CLIMB * dt, CLIMB * dt)
    local e = sp - alt
    local want = trim + GAIN_HEIGHT * (0.012 * e - 0.1 * vz)
    if (want < 1 or e < 0) and (want > 0 or e > 0) then
      trim = clamp(trim + GAIN_HEIGHT * 0.003 * e * dt, 0, 1)
    end
    lift = clamp(want, 0, 1)
  else
    sp = nil
    trim = lift
  end
  local cx = 0
  local cz = 0
  if gyro and level then
    ix = clamp(ix + GAIN_LEVEL * 0.005 * ax * dt, -0.4, 0.4)
    iz = clamp(iz + GAIN_LEVEL * 0.005 * az * dt, -0.4, 0.4)
    cx = GAIN_LEVEL * (0.02 * ax + 0.05 * rx) + ix
    cz = GAIN_LEVEL * (0.02 * az + 0.05 * rz) + iz
  end
  local shift = 0
  for _, b in ipairs(burners) do
    b.f = lift + b.ux * cz + b.uz * cx
    if b.f > 1 and b.outer then
      shift = math.min(shift, 1 - b.f)
    end
  end
  for _, b in ipairs(burners) do
    if b.outer then
      b.f = b.f + shift
    end
    if b.f < 0 then
      b.f = 0
    elseif b.f > 1 then
      b.f = 1
    end
    local want = math.floor(LO + b.f * (top(b) - LO) + 0.5)
    if want ~= b.amount then
      pcall(b.p.set_amount, want)
    end
  end
end

local function changed(key, sig)
  if shown[key] == sig then
    return false
  end
  shown[key] = sig
  return true
end

local function put(x, y, s, fg)
  term.set_cursor(x, y)
  term.set_fg(fg)
  term.write(s)
end

local function lines(list, x)
  for y = 2, BY - 1 do
    local l = list[y - 1] or {"", 0}
    if shown[y] ~= l[1] then
      shown[y] = l[1]
      put(x, y, fit(l[1], W - x + 1), l[2])
    end
  end
end

local function buttons(list)
  hits = {}
  for y = BY, H do
    term.set_cursor(1, y)
    term.clear_line()
  end
  local x = 1
  local y = BY
  for _, b in ipairs(list) do
    local label = " " .. b[1] .. " "
    if x + #label > W + 1 then
      x = 1
      y = y + 1
    end
    table.insert(hits, {x1 = x, x2 = x + #label - 1, y = y, a = b[2]})
    if COLOR then
      term.set_bg(b[3])
      put(x, y, label, C.black)
      term.set_bg(C.black)
    else
      put(x, y, "[" .. b[1] .. "]", C.white)
    end
    x = x + #label + 1
  end
  if y == BY and H >= 15 then
    put(1, H, T.hint, C.gray)
  end
end

-- Seen from above, north up: a disc per burner, red when it works hard; the
-- line points to the low side, longer when the ship leans more.
local function map()
  local pw = MAPW * 6
  local ph = (BY - 2) * 9
  local mx = pw // 2
  local my = 10 + ph // 2
  local rad = math.min(pw, ph) // 2 - 1
  local r = math.max(3, rad // 7)
  local sig = math.floor(ax * 4) .. "," .. math.floor(az * 4)
  for _, b in ipairs(burners) do
    sig = sig .. "|" .. math.floor(b.f * 20) .. b.signal
  end
  if not changed("map", sig) then
    return
  end
  gfx.rect(1, 10, pw, ph, "none", true)
  gfx.circle(mx, my, rad, C.gray)
  gfx.text(mx - 1, my - rad + 4, DIRS[1], C.light_gray)
  gfx.text(mx + rad - 6, my - 2, DIRS[3], C.light_gray)
  gfx.text(mx - 1, my + rad - 8, DIRS[5], C.light_gray)
  gfx.text(mx - rad + 4, my - 2, DIRS[7], C.light_gray)
  local s = (rad - r - 17) / math.max(R, 1)
  for _, b in ipairs(burners) do
    local x = mx + math.floor(b.dx * s + 0.5)
    local y = my + math.floor(b.dz * s + 0.5)
    local color = C.yellow
    if b.signal == 0 then
      color = C.gray
    elseif b.f > 0.66 then
      color = C.red
    elseif b.f > 0.33 then
      color = C.orange
    end
    gfx.circle(x, y, r, color, true)
    gfx.text(x - #pct(b.f) * 2 + 1, y + r + 2, pct(b.f), C.white)
  end
  local tilt = math.sqrt(ax * ax + az * az)
  if tilt >= 0.3 then
    local len = math.min(rad, tilt * rad / 10) / tilt
    local color = C.red
    if tilt < 1 then
      color = C.lime
    elseif tilt < 3 then
      color = C.yellow
    end
    gfx.line(mx, my, mx + math.floor(az * len), my + math.floor(ax * len), color, 2)
  end
end

local function flight()
  map()
  local l = {{T.height .. "--", C.gray}}
  if alt then
    l[1] = {T.height .. f1(alt), C.white}
  end
  local auto = hold and alt ~= nil
  if auto then
    table.insert(l, {T.target .. target .. "  " .. T.thrust .. pct(lift), C.yellow})
  else
    table.insert(l, {T.thrust .. pct(lift) .. " " .. T.man, C.orange})
  end
  table.insert(l, {T.climb .. f1(vz), C.white})
  local tilt = math.sqrt(ax * ax + az * az)
  if gyro == nil then
    table.insert(l, {T.tilt .. "--", C.gray})
  elseif tilt < 0.5 then
    table.insert(l, {T.tilt .. T.flat, C.lime})
  else
    table.insert(l, {T.tilt .. f1(tilt) .. T.low .. compass(az, ax), C.orange})
  end
  if H >= 15 then
    table.insert(l, {"", 0})
  end
  for _, b in ipairs(burners) do
    local s = fit(b.dir, 3) .. fit(pct(b.f), 5)
    if b.signal == 0 then
      table.insert(l, {s .. T.unlit, C.red})
    elseif b.cap then
      table.insert(l, {s .. math.floor(b.filled or 0) .. "/" .. math.floor(b.cap), C.white})
    else
      table.insert(l, {s .. T.noBalloon, C.orange})
    end
  end
  if burners[1] == nil then
    table.insert(l, {T.noBurner, C.red})
  end
  lines(l, MAPW + 2)
  if changed("buttons", tostring(level) .. tostring(auto)) then
    local mode = T.man
    if auto then
      mode = "AUTO"
    end
    buttons({{T.level .. (level and " ON" or " OFF"), "a", level and C.lime or C.gray},
      {T.hold .. " " .. mode, "h", auto and C.lime or C.gray},
      {"-10", -10, C.light_blue}, {"-1", -1, C.light_blue}, {"+1", 1, C.light_blue}, {"+10", 10, C.light_blue}})
  end
end

-- Every device: where it is from the centre, and what it reads
local function devices()
  local l = {{T.head, C.light_gray}}
  for _, b in ipairs(burners) do
    local s = fit(b.outer and b.dir or T.center, 8) .. fit(math.floor(b.dx + 0.5) .. "," .. math.floor(b.dz + 0.5), 8)
    if b.pos == nil then
      s = fit(T.center, 8) .. fit(b.name, 8)
    end
    s = s .. fit(b.amount .. "/" .. b.max, 8) .. fit(b.signal .. "", 3)
    if b.cap then
      table.insert(l, {s .. math.floor(b.filled or 0) .. "/" .. math.floor(b.cap), b.signal > 0 and C.white or C.red})
    else
      table.insert(l, {s .. T.noBalloon, C.orange})
    end
  end
  for _, name in ipairs(peripheral.list()) do
    local p = peripheral.wrap(name)
    if p and p.type() ~= "hot_air_burner" then
      local at = place(name)
      local s = fit(p.type(), 16) .. fit(name, 8)
      if at then
        s = fit(p.type(), 16) .. fit(math.floor(at.x - CX + 0.5) .. "," .. math.floor(at.z - CZ + 0.5), 8)
      end
      if p.read then
        for k, v in pairs(p.read()) do
          if type(v) == "number" then
            s = s .. k .. " " .. f1(v) .. " "
          end
        end
      end
      table.insert(l, {s, C.lime})
    end
  end
  if gyro == nil then
    table.insert(l, {T.noGyro, C.red})
  end
  if alti == nil and not onShip then
    table.insert(l, {T.noAlt, C.orange})
  end
  if burners[1] == nil then
    table.insert(l, {T.noBurner, C.red})
  end
  lines(l, 1)
  if changed("buttons", "scan") then
    buttons({{"Scan", "s", C.light_blue}})
  end
end

local function draw()
  if changed("page", page) then
    term.clear()
    gfx.clear()
    if COLOR then
      gfx.rect(1, 1, W * 6, 9, "blue", true)
    end
    put(1, 1, " " .. T.tabs[page] .. "  (Tab: " .. T.tabs[3 - page] .. ")", C.white)
  end
  if page == 1 then
    flight()
  else
    devices()
  end
end

local function act(a)
  if a == "a" then
    level = not level
  elseif a == "h" then
    hold = not hold
    target = alt and math.floor(alt + 0.5)
  elseif a == "s" then
    scan()
  elseif a == "tab" then
    page = 3 - page
    shown = {}
  elseif hold and alt then
    target = target + a
  else
    lift = clamp(lift + a / 100, 0, 1)
  end
  save()
  draw()
end

local s = fs.read("balloon.cfg")
if s then
  local v = string.split(s, " ")
  target = tonumber(v[1])
  hold = v[2] == "true"
  level = v[3] == "true"
end
local KEYS = {up = 1, down = -1, right = 10, left = -10, tab = "tab"}
scan()
draw()
local timer = os.start_timer(0.5)
while true do
  local e = os.pull_event()
  if e.name == "timer" and e.id == timer then
    timer = os.start_timer(0.5)
    local now = os.clock()
    local dt = math.max(0.05, now - last)
    last = now
    if lost or now - scanned > 30 then
      scan()
    end
    sense(dt)
    control(dt)
    tock = not tock
    if tock then
      draw()
    end
  elseif e.name == "click" then
    if e.y == 1 then
      act("tab")
    end
    for _, h in ipairs(hits) do
      if e.y == h.y and e.x >= h.x1 and e.x <= h.x2 then
        act(h.a)
      end
    end
  elseif e.name == "key" and KEYS[e.key] then
    act(KEYS[e.key])
  elseif e.name == "char" and string.find("ahs", string.lower(e.char), 1, true) then
    act(string.lower(e.char))
  elseif e.name == "char" and string.lower(e.char) == "q" then
    break
  end
end
term.clear()
gfx.clear()
term.set_cursor(1, 1)