create coasters RCT2 score

--[[
    Records track information using CC: Sable,
    for Create: Coasters Simulated.
]]


MOVING_VELOCITY = .01 -- min velocity to start moving (b/s)
STANDARD_GRAVITY = 9.80665 -- used for calculating Gs
GRAVITY = 9.80665 -- used for adding Gs caused by gravity (only this value changes on the moon)
INVERSION_DOT = -.7 -- 153 degrees away from upwards
MIN_AIRTIME_G = .001 -- min g force pulling up (including gravity) for it to count as airtime
MIN_AIRTIME_DURATION = .01 -- min duration of airtime to count
MIN_DROP_SPEED = .1 -- min speed downwards to count as a drop
MIN_DROP_HEIGHT = .5 -- min height of drop to count (1 slab)


local tracker = {}
tracker.__index = tracker

function tracker.new()
    return setmetatable({
        sum = 0,
        abs_sum = 0,
        values = {}
    }, tracker)
end

function tracker:update(value)
    if self.min == nil or value < self.min then
        self.min = value
    end
    if self.max == nil or value > self.max then
        self.max = value
    end
    local abs = math.abs(value)
    if self.abs_min == nil or abs < self.abs_min then
        self.abs_min = abs
    end
    if self.abs_max == nil or abs > self.abs_max then
        self.abs_max = abs
    end

    self.sum = self.sum + value
    self.abs_sum = self.sum + abs
    table.insert(self.values, value)
end

function tracker:getAverage()
    return self.sum / #self.values
end

function tracker:getAbsAverage()
    return self.abs_sum / #self.values
end

function tracker:__tostring()
    return ("min %.3g max %.3g (abs min %.3g abs max %.3g) avg ~%.3g (%.3g total over %d values)"):format(self.min, self.max, self.abs_min, self.abs_max, self:getAverage(), self.sum, #self.values)
end

local function getScoreText(intensity)
    if intensity <= 2.55 then
        return "Low"
    elseif intensity <= 5.11 then
        return "Medium"
    elseif intensity <= 7.67 then
        return "High"
    elseif intensity <= 10.23 then
        return "Very high"
    elseif intensity <= 12.79 then
        return "Extreme"
    elseif intensity <= 50.01 then
        return "Ultra-Extreme"
    else
        return "Uber-Extreme"
    end
end

local function split(text)
    local words = {}
    for word in (text or ""):gmatch("%S+") do
        table.insert(words, word:lower())
    end
    return words
end


-- Setup

assert(sublevel, "Sublevel API is missing, is CC: Sable installed?")
assert(quaternion, "Quaternion API is missing, is CC: Sable installed?")
assert(sublevel.isInPlotGrid(), "Computer is not on a sublevel.")


print("\nReady to dispatch...")

local forward = vector.new(0, 0, -1)
local up = vector.new(0, 1, 0) -- seats also always face up in their sublevel
do
    local last_time = os.epoch("utc") / 1000
    local last_p = sublevel.getLogicalPose().position
    while true do
        local time = os.epoch("utc") / 1000
        local delta = time - last_time

        local p = sublevel.getLogicalPose().position
        local v = p:sub(last_p):mul(delta)
        if v:length() > MOVING_VELOCITY then
            local q = sublevel.getLogicalPose().orientation
            forward = q:inverse():mul(v):round(MOVING_VELOCITY):normalize()
            break
        end

        last_time = time
        last_p = p
        os.sleep(0)
    end
end
local left = up:cross(forward):normalize()


-- Recording

print(("\nCart started moving in %s, recording started..."):format(forward:round(.01)))
print("Hold t to end the recording.")
local start = os.epoch("utc") / 1000

local tracker_times = tracker.new()
local total_g = tracker.new()
local lateral_g = tracker.new()
local vertical_g = tracker.new()
local accel_g = tracker.new()
local speed = tracker.new()
local airtime = tracker.new()
local drops = tracker.new()
local inversions = 0
local length = 0

local end_time = start
do
    local airtime_start = nil
    local drop_total = nil
    local was_inverted = false

    local last_time = start
    local last_p = sublevel.getLogicalPose().position
    local last_v = vector.new(0, 0, 0)

    local timer = os.startTimer(0)
    while true do
        local event, p1 = os.pullEvent()

        if event == "key" and p1 == keys.t then
            break
        elseif event == "timer" and p1 == timer then

            local now = os.epoch("utc") / 1000
            tracker_times:update(now)
            local delta = now - last_time

            local p = sublevel.getLogicalPose().position
            local p_delta = p:sub(last_p)
            length = length + p_delta:length()

            local v = p_delta:div(delta)
            speed:update(v:length())

            local q = sublevel.getLogicalPose().orientation

            local global_g = v:sub(last_v):div(delta):add(vector.new(0, GRAVITY, 0)):div(STANDARD_GRAVITY)
            local global_g_length = global_g:length()
            total_g:update(global_g_length)
            local g = q:inverse():mul(global_g):mul(global_g_length) -- quaternion:mul() normalizes the vector for some reason

            local lat = g:dot(left)
            lateral_g:update(lat)

            local vert = g:dot(up)
            vertical_g:update(vert)

            local accel = g:dot(forward)
            accel_g:update(accel)

            local global_up = q:inverse():mul(up)
            if global_up:dot(up) < INVERSION_DOT then
                if not was_inverted then
                    print("> inversion")
                    was_inverted = true
                    inversions = inversions + 1
                end
            else
                was_inverted = false
            end

            if g.y < -MIN_AIRTIME_G then
                if airtime_start == nil then
                    airtime_start = now
                end
            elseif airtime_start ~= nil then
                local time = last_time - airtime_start
                if time > MIN_AIRTIME_DURATION then
                    airtime:update(time)
                    print(("> airtime ended of %.3g seconds"):format(last_time - airtime_start))
                end
                airtime_start = nil
            end

            if v.y < -MIN_DROP_SPEED then
                -- TODO don't count loops as drops somehow
                if drop_total == nil then
                    drop_total = 0
                end
                drop_total = drop_total - v.y * delta
            elseif drop_total ~= nil then
                if drop_total > MIN_DROP_HEIGHT then
                    drops:update(drop_total)
                    print(("> drop ended of %.3g blocks"):format(drop_total))
                end
                drop_total = nil
            end

            last_time = now
            -- don't let user being slow to stop script inflate duration
            if v:length() > MOVING_VELOCITY then
                end_time = now
            end
            last_v = v
            last_p = p
            timer = os.startTimer(0)
        end
    end
end


-- Data compilation

local duration = end_time - start

local trackers = {
    time = tracker_times,
    total_g = total_g,
    lateral_g = lateral_g,
    vertical_g = vertical_g,
    accel_g = accel_g,
    speed = speed
}
local events = {
    airtime = airtime,
    drops = drops
}
local misc = {
    inversions = inversions,
    length = length,
    duration = duration
}

-- https://github-wiki-see.page/m/OpenRCT2/OpenRCT2/wiki/Ride-rating-calculation
-- https://github.com/OpenRCT2/OpenRCT2/blob/a3f7b5d3b09f470dee3e1b28ba157e40e3ff95c4/src/openrct2/ride/RideRatings.cpp
local excitement =
    (.8 * vertical_g.max) -- g-force
    + (.24 * math.max(-2.5, vertical_g.min))
    + (.4 * math.min(1.5, lateral_g.abs_max))
    + (.11 * math.min(9, #drops.values)) -- drops
    + (.0049 * drops.max)
    + (.27 * math.min(6, inversions)) -- inversions
    + (.125 * airtime.sum) -- airtime

local intensity =
    (.8 * vertical_g.max) -- g-force
    + (.8 * (1 + vertical_g.min))
    + lateral_g.abs_max
    + (.14 * #drops.values) -- drops
    + (.0098 * drops.max)
    + (.5 * inversions) -- inversions

local nausea =
    (.26 * vertical_g.max) -- g-force
    + (.22 * (1 + vertical_g.min))
    + (.33 * lateral_g.abs_max)
    + (.1 * #drops.values) -- drops
    + (.0016 * drops.max)
    + (.22 * inversions) -- inversions
    + (.0625 * airtime.sum) -- airtime

-- https://github.com/OpenRCT2/OpenRCT2/blob/5aa3f85c56c4bd68500380190d9deac357686423/src/openrct2/ride/rtd/coaster/TwisterRollerCoaster.h
excitement = excitement
    + (.0116577 * math.min(6000, length)) -- track length
    + (.4 * math.min(150, duration)) -- duration -- https://github-wiki-see.page/m/Sadret/openrct2-plugin-wiki/wiki/Date-and-Time
    + (.01736 * speed:getAverage()) -- speed avg -- https://gall.dcinside.com/mgallery/board/view/?id=rct&no=8507
    + (.00264 * speed.max) -- speed max

intensity = intensity
    + (.026 * speed:getAverage()) -- speed avg
    + (.00528 * speed.max) -- speed max

nausea = nausea
    + (.0021 * speed.max) -- speed max

if lateral_g.abs_max > 2.8 then
    intensity = intensity + 3.75
    nausea = nausea + 2
end
if lateral_g.abs_max > 3.1 then
    excitement = excitement * .5
    intensity = intensity + 8.5
    nausea = nausea + 4
end

for _,threshold in pairs({10, 11, 12, 13.2, 14.5}) do
    if intensity > threshold then
        excitement = excitement * .75
    end
end

local rct2_scores = {Excitement = excitement, Intensity = intensity, Nausea = nausea}
local function printGeneralInfo()
    local lines = {
        "",
        ("Speed: max %.2fkb/h avg ~%.2fkb/h"):format(speed.max * 3.6, speed:getAverage() * 3.6),
        ("Duration: %.2fs"):format(duration),
        ("%.1f blocks long"):format(length),
        "",
        ("G-force: max %.2fg avg ~%.2fg"):format(total_g.max, total_g:getAverage()),
        ("Lateral g-force: max %.2fg avg ~%.2fg"):format(lateral_g.abs_max, lateral_g:getAbsAverage()),
        ("Vertical g-force: down max %.2fg up max %.2fg avg ~%.2fg"):format(vertical_g.min, vertical_g.max, vertical_g:getAverage()),
        ("Acceleration g-force: brake max %.2fg speed max %.2fg avg ~%.2fg"):format(accel_g.min, accel_g.max, accel_g:getAbsAverage()),
        "",
        ("%d drops max %.1fb total %.1fb avg ~%.1fb"):format(#drops.values, drops.max, drops.sum, drops:getAverage()),
        ("%d airtimes max %.2fs total %.2fs avg ~%.2fs"):format(#airtime.values, airtime.max, airtime.sum, airtime:getAverage()),
        ("%d inversions"):format(inversions),
        ""
    }

    for name, val in pairs(rct2_scores) do
        table.insert(lines, ("%s: %.2f (%s)"):format(name, val, getScoreText(val)))
    end

    local _, height = term.getCursorPos()
    textutils.pagedPrint(table.concat(lines, "\n"), height - 2)
end
printGeneralInfo()

local function writeMisc(path)
    local full_path = ("%s/misc.lua"):format(path)
    local file = fs.open(full_path, "w")

    file.writeLine(textutils.serialize(misc, {
        compact = true,
        allow_repetitions = true
    }))
    file.flush()
    file.close()

    return full_path
end

local function writeTrackerSummaries(path)
    local full_path = ("%s/summaries.csv"):format(path)
    local file = fs.open(full_path, "w")

    file.writeLine("name,min,max,abs_min,abs_max,avg,sum,count")
    for name, tracker in pairs(trackers) do
        file.writeLine(("%s,%s,%s,%s,%s,%s,%s,%d"):format(
            name,
            tracker.min, tracker.max,
            tracker.abs_min, tracker.abs_max,
            tracker:getAverage(),
            tracker.sum,
            #tracker.values
        ))
    end
    file.flush()
    file.close()

    return full_path
end

local function writeTrackers(path)
    local full_path = ("%s/trackers.csv"):format(path)
    local file = fs.open(full_path, "w")

    local columns = {}
    for name, tracker in pairs(trackers) do
        table.insert(columns, name)
    end
    file.writeLine(table.concat(columns, ","))

    local i = 1
    while true do
        local time = trackers.time.values[i]
        if time == nil or time > end_time then
            break
        end

        local values = {}
        for _, name in ipairs(columns) do
            local tracker = trackers[name]
            if i > #tracker.values then
                table.insert(values, "")
            else
                table.insert(values, tracker.values[i])
            end
        end
        file.writeLine(table.concat(values, ","))

        i = i + 1
    end

    file.flush()
    file.close()

    return full_path
end


-- CLI

print("\nWelcome to the interactive CLI,\ntype 'help' for usage information.")

local completion = require("cc.completion")
local function complete(text)
    local words = split(text)

    if #words <= 1 then
        return completion.choice(text, {
            "help",
            "exit",
            "general",
            "graph",
            "export"
        })
    end
    if #words == 2 then
        if words[1] == "graph" then
            local keys = {}
            for name, _ in pairs(trackers) do
                table.insert(keys, name)
            end
            return completion.choice(words[2], keys)

        elseif words[1] == "export" then
            return fs.complete(words[2], shell.dir(), false, true)
        end
    end
end

local cmd_history = {}
while true do
    write("\n? ")
    local cmd = read(nil, cmd_history, complete)
    table.insert(cmd_history, cmd)
    local words = split(cmd)

    if words[1] == "help" then
        local _, height = term.getCursorPos()
        textutils.pagedPrint([[Usage:
* help - this information list
* exit - exit this program
* general - reprints the previously printed general data, like RCT2 scores
* graph <metric> - prints graph for the specified metric
* export <path> - exports all recorded data to multiple files in specified path]],
        height - 2)

    elseif words[1] == "exit" then
        break
    elseif words[1] == "general" or words[1] == "rct2" then
        printGeneralInfo()
    elseif words[1] == "export" then
        if #words ~= 2 or words[2] == nil then
            print("Missing path to export to.")
        else
            local success, result = pcall(function()
                local path = words[2]
                fs.makeDir(path)

                local funcs = {
                    writeMisc,
                    writeTrackerSummaries,
                    writeTrackers
                }
                for i, func in pairs(funcs) do
                    print(("Exported to '%s'... (%d/%d)"):format(func(path), i, #funcs))
                end
            end)

            if not success then
                print(("Export failed: %s"):format(result))
            end
        end
    else
        print("Unknown command, use 'help' for usage information.")
    end
end