Location System

-- Location System
-- by MysticT

local tLocations = {}
local tKnownLocations = {}
local bDebug = false

-- Helper functions

local function clear()
    term.clear()
    term.setCursorPos(1, 1)
end

local function table_size(t)
    local i = 0
    for _,_ in pairs(t) do
        i = i + 1
    end
    return i
end

local function connect()
    print("Starting location system...")
    print("Connecting to rednet...")
    for _,s in ipairs(rs.getSides()) do
        if peripheral.isPresent(s) and peripheral.getType(s) == "modem" then
            rednet.open(s)
            return true
        end
    end
    print("No modem found")
    print("Aborting")
    return false
end

local function _print(...)
    if bDebug then
        print(...)
    end
end

local function vector_equal(v1, v2)
    return v1.x == v2.x and v1.y == v2.y and v1.z == v2.z
end

-- Load/Save functions

local function save()
    _print("Saving locations...")
    local file = fs.open("locations.dat", "w")
    if file then
        for id, loc in pairs(tKnownLocations) do
            file.writeLine(tostring(id)..": "..tostring(loc.x)..", "..tostring(loc.y)..", "..tostring(loc.z))
        end
        file.close()
        _print("Locations saved")
        return true
    end
    _print("Error saving locations")
    return false
end

local function load()
    _print("Loading locations")
    local file = fs.open("locations.dat", "r")
    if file then
        local i = 0
        local sLine = file.readLine()
        while sLine do
            local sId, sX, sY, sZ = string.match(sLine, "(%d+): (%-?%d+), (%-?%d+), (%-?%d+)")
            local id, x, y, z = tonumber(sId), tonumber(sX), tonumber(sY), tonumber(sZ)
            if id and x and y and z then
                tKnownLocations[id] = vector.new(x, y, z)
                i = i + 1
            end
            sLine = file.readLine()
        end
        file.close()
        _print(n, " locations loaded")
        return true, i
    end
    _print("Error loading locations")
    return false
end

-- Location functions

local function trilaterate(A, B, C)
    local a2b = B.position - A.position
    local a2c = C.position - A.position
    if math.abs( a2b:normalize():dot( a2c:normalize() ) ) > 0.999 then
        return nil
    end
    local d = a2b:length()
    local ex = a2b:normalize( )
    local i = ex:dot( a2c )
    local ey = (a2c - (ex * i)):normalize()
    local j = ey:dot( a2c )
    local ez = ex:cross( ey )
    local r1 = A.distance
    local r2 = B.distance
    local r3 = C.distance
    local x = (r1*r1 - r2*r2 + d*d) / (2*d)
    local y = (r1*r1 - r3*r3 - x*x + (x-i)*(x-i) + j*j) / (2*j)
    local result = A.position + (ex * x) + (ey * y)
    local zSquared = r1*r1 - x*x - y*y
    if zSquared > 0 then
        local z = math.sqrt(zSquared)
        local result1 = result + (ez * z)
        local result2 = result - (ez * z)
        local rounded1, rounded2 = result1:round(), result2:round()
        if rounded1.x ~= rounded2.x or rounded1.y ~= rounded2.y or rounded1.z ~= rounded2.z then
            return result1, result2
        else
            return rounded1
        end
    end
    return result:round()
end

local function narrow(p1, p2, fix)
    local dist1 = math.abs((p1 - fix.position):length() - fix.distance)
    local dist2 = math.abs((p2 - fix.position):length() - fix.distance)
    if math.abs(dist1 - dist2) < 0.05 then
        return p1, p2
    elseif dist1 < dist2 then
        return p1:round()
    else
        return p2:round()
    end
end

function locate(id)
    local tFixes = {}
    local p1, p2
    for _,t in pairs(tLocations[id]) do
        table.insert(tFixes, { ["position"] = vector.new(t.x, t.y, t.z), ["distance"] = t.d })
    end
    local i = 3
    while (p1 == nil or p2 ~= nil) and i <= #tFixes do
        if not p1 then
            p1, p2 = trilaterate(tFixes[1], tFixes[2], tFixes[i])
        else
            p1, p2 = narrow(p1, p2, tFixes[i])
        end
        i = i + 1
    end
    if p1 and p2 then
        _print("Ambiguous position")
        _print("Could be "..p1.x..","..p1.y..","..p1.z.." or "..p2.x..","..p2.y..","..p2.z)
    elseif p1 then
        if not tKnownLocations[id] or not vector_equal(tKnownLocations[id], p1) then
            print("Location added:")
            print(id, ": ", p1)
            tKnownLocations[id] = p1
            save()
        end
    else
        _print("Couldn't determine location for ", id)
    end
end

local function addLocation(id, t)
    if not tLocations[t.id] then
        tLocations[t.id] = {}
    end
    local tLoc = {}
    tLoc.x = t.x
    tLoc.y = t.y
    tLoc.z = t.z
    tLoc.d = t.d
    tLocations[t.id][id] = tLoc
    if table_size(tLocations[t.id]) >= 3 then
        _print("Trying to locate #", t.id)
        locate(t.id)
    end
end

-- Event handlers

local function handle_message(id, msg)
    local st = string.match(msg, "<LOCATION> (.+)")
    if st then
        local t = textutils.unserialize(st)
        if t and type(t) == "table" then
            _print("Location received from ", id, ":")
            _print(t.id, ": ", "(", t.x, ", ", t.y, ", ", t.z, ", ", t.d, ")")
            addLocation(id, t)
        end
    end
end

local function handle_char(c)
    c = string.lower(c)
    if c == "l" then
        print("Locations:")
        for id, loc in pairs(tKnownLocations) do
            print(id, ": ", loc)
        end
    elseif c == "c" then
        clear()
    end
end

-- Start program

local function printUsage()
    local sName = fs.getName(shell.getRunningProgram())
    print("Usage:")
    print(sName, " receive [debug]")
    print(sName, " host id x y z")
end

local tArgs = {...}
local sAction = tArgs[1]
if not sAction then
    printUsage()
    return
end

if sAction == "receive" then
    if #tArgs > 1 then
        if string.lower(tArgs[2]) == "debug" then
            bDebug = true
        else
            printUsage()
            return
        end
    end
    clear()
    if not connect() then
        return
    end
    local ok, n = load()
    if ok then
        print("Loaded ", n, " locations")
    end
    print("Waiting for messages...")
    print("Press l to list located computers")
    print("Press c to clear the screen")
    while true do
        local evt, arg1, arg2 = os.pullEvent()
        if evt == "rednet_message" then
            handle_message(arg1, arg2)
        elseif evt == "char" then
            handle_char(arg1)
        end
    end
elseif sAction == "host" then
    local nMasterID = tonumber(tArgs[2])
    local tLoc = {}
    tLoc.x = tonumber(tArgs[3])
    tLoc.y = tonumber(tArgs[4])
    tLoc.z = tonumber(tArgs[5])
    if not nMasterID or not tLoc.x or not tLoc.y or not tLoc.z then
        printUsage()
        return
    end
    clear()
    if not connect() then
        return
    end
    print("Waiting for messages...")
    while true do
        local id, msg, dist = rednet.receive()
        print("Message received from ", id, " at ", dist, " meters.")
        tLoc.id = id
        tLoc.d = dist
        rednet.send(nMasterID, "<LOCATION> "..textutils.serialize(tLoc))
    end
else
    printUsage()
end