Thrusters = {
    thruster0 = {
        type = "rotator",
        name = nil,
        thruster = nil,
        transmission = nil,
        -- NOTE:
        -- affectVectors will have values depending on how the thruster's 
        -- thrust vector interacts with the vessel's attitude
        -- possible directions for yaw: "port"|"star"
        -- possible directions for roll: "port"|"star"
        -- possible directions for pitch: "up"|"down"
        -- possible directions for lateral: "port"|"star"|"fore"|"aft"|"up"|"down"
        affectVectors = {
            yaw = nil,
            pitch = nil,
            roll = nil,
            lateral = nil
        },
        power = 0
    },
    thruster1 = {
        type = "thruster",
        name = nil,
        thruster = nil,
        affectVectors = {
            yaw = nil,
            pitch = nil,
            roll = nil,
            lateral = nil
        },
        power = 0
    }
} -- Table of thruster tables

SensorData = {}

ThrusterTypes = {
    "thruster",
    "solid_fuel_thruster",
    "ion_thruster",
    "vector_thruster",
    "liquid_vector_thruster"
}

PropellerTypes = {
    "gyroscopic_propeller_bearing",
    "propeller",
    "propeller_bearing"
}

TransmissionTypes = {
    "analog_transmission",
    "Create_RotationSpeedController"
}

Config = {}
Config.ConfigPath = "/caero-attitude-control/config/config.txt"

-- File serializatino and deserialization
-- taken from https://gist.github.com/yuhanz/6688d474a3c391daa6d6
function tableToString(table)
    return "return"..serializeTable(table)
end

-- loadstring is deprecated, so modified to run without it
function stringToTable(str)
    local f = {}
    local chunk, err = load(str)
    if chunk then
        f = chunk()
    else
        print("Error loading config string: "..err)
    end
    return f
end

function serializeTable(val, name, skipnewlines, depth)
    skipnewlines = skipnewlines or false
    depth = depth or 0

    local tmp = string.rep(" ", depth)
    if name then 
    	if not string.match(name, '^[a-zA-z_][a-zA-Z0-9_]*$') then
    		name = string.gsub(name, "'", "\\'")
    		name = "['".. name .. "']"
    	end
    	tmp = tmp .. name .. " = "
     end

    if type(val) == "table" then
        tmp = tmp .. "{" .. (not skipnewlines and "\n" or "")

        for k, v in pairs(val) do
            tmp =  tmp .. serializeTable(v, k, skipnewlines, depth + 1) .. "," .. (not skipnewlines and "\n" or "")
        end

        tmp = tmp .. string.rep(" ", depth) .. "}"
    elseif type(val) == "number" then
        tmp = tmp .. tostring(val)
    elseif type(val) == "string" then
        tmp = tmp .. string.format("%q", val)
    elseif type(val) == "boolean" then
        tmp = tmp .. (val and "true" or "false")
    else
        tmp = tmp .. "\"[inserializeable datatype:" .. type(val) .. "]\""
    end

    return tmp
end

-- Monitor Configuration  
Config.Monitors = {}
Config.Monitors.InstrumentPanelMonitor = nil
Config.Monitors.AutopilotControlMonitor = nil

local function identifyMonitor(monitorType)
    -- Find all available monitors
    local monitors = { peripheral.find("monitor") }
    
    if not monitors or #monitors == 0 then
        print("No monitors found on the network.")
        return nil
    end
    
    print("Available monitors:")
    for i, monitor in ipairs(monitors) do
        if monitor then
            local width, height = monitor.getSize()
            print(i .. ". " .. peripheral.getName(monitor) .. " (" .. width .. "x" .. height .. ")")
        end
    end
    
    print("\nPlease select the monitor you want to use for " .. monitorType .. " (enter number):")
    
    -- Simple input handling for demonstration
    local input = read()
    local choice = tonumber(input)
    
    if not choice or choice < 1 or choice > #monitors then
        print("Invalid selection.")
        return nil
    end
    
    local selectedMonitor = monitors[choice]
    local selectedMonitorName = peripheral.getName(selectedMonitor)
    
    print("Selected monitor: " .. selectedMonitorName)
    
    -- Return monitor details as a table
    return {
        name = selectedMonitorName,
        peripheral = selectedMonitor
    }
end

-- Autopilot
Config.Autopilot = {
    AutopilotEngaged = false,
    AutoForeAft = false,
    AutopilotDesiredSpeed = nil,
    AutopilotDesiredHeading = nil
}

-- Correctional values for when some sensor values are offset
Config.SensorCorrection = {
    Heading = 0 -- this should be +/- 180
}

Config.Throttles = nil

function ThrottleInit() 

    if Config.Throttles ~= nil then
        return
    end

    Config.Throttles = {}

    local sides = {
        "top",
        "bottom",
        "left",
        "right",
        "front",
        "back"
    }

    local function getRelayStates()
        local states = {}
        local sides = {
            "top",
            "bottom",
            "left",
            "right",
            "front",
            "back"
        }

        local names = peripheral.getNames()
        for _, name in ipairs(names) do
            if peripheral.getType(name) == "redstone_relay" then
                local wrapper = peripheral.wrap(name)
                for _, side in ipairs(sides) do
                    states[name] = states[name] or {}
                    states[name][side] = wrapper.getAnalogInput(side)
                end
            end
        end

        return states
    end

    local function identifyThrottle(throttleType)
        print("Identifying "..throttleType.." throttle.")
        print("\nPlease change the input signal for the throttle.")

        local initialStates = getRelayStates()

        os.pullEvent("redstone")
        
        local currentStates = getRelayStates()

        Config.Throttles[throttleType] = {}

        for pname, currentState in pairs(currentStates) do
            local initialState = initialStates[pname]
            for k, v in pairs(currentState) do
                if initialState[k] ~= v then
                    Config.Throttles[throttleType].name = pname
                    Config.Throttles[throttleType].side = k
                    return
                end
            end
        end
    end

    identifyThrottle("fore")
    identifyThrottle("aft")
    identifyThrottle("down")
end

function PropellerInit()
    local transmissions = {}
    local propellers = {}

    for _, v in ipairs(TransmissionTypes) do
        table.insert(transmissions, peripheral.find(v))
    end

    for _, v in ipairs(PropellerTypes) do
        table.insert(propellers, peripheral.find(v))
    end

    for pi, pv in ipairs(propellers) do 
        for ti, tv in ipairs(transmissions) do
            if pv.getSubnetworkAnchorId() == tv.getSelfId() then
                if Thrusters[peripheral.getName(pv)] == nil then
                    Thrusters[peripheral.getName(pv)] = {
                        type = "rotator",
                        thruster = pv,
                        transmission = tv,
                        -- NOTE:
                        -- affectVectors will have values depending on how the thruster's 
                        -- thrust vector interacts with the vessel's attitude
                        -- possible directions for yaw: "port"|"star"
                        -- possible directions for roll: "port"|"star"
                        -- possible directions for pitch: "up"|"down"
                        -- possible directions for lateral: "port"|"star"|"fore"|"aft"|"up"|"down"
                        affectVectors = {
                            yaw = nil,
                            pitch = nil,
                            roll = nil,
                            lateral = nil
                        },
                        power = (tv.getSignal())/15
                    }
                end
                if Thrusters[peripheral.getName(pv)] ~= nil then
                    if Thrusters[peripheral.getName(pv)].type == nil then 
                        Thrusters[peripheral.getName(pv)].type = "rotator" 
                    end
                    Thrusters[peripheral.getName(pv)].thruster = pv
                    Thrusters[peripheral.getName(pv)].transmission = tv
                    if Thrusters[peripheral.getName(pv)].affectVectors == nil then 
                        Thrusters[peripheral.getName(pv)].affectVectors = {} 
                    end
                    Thrusters[peripheral.getName(pv)].power = (tv.getSignal())/15
                end
            end
        end
    end
end

function ThrusterInit()
    local thrusters = {}

    for _, v in ipairs(ThrusterTypes) do
        table.insert(thrusters, peripheral.find(v))
    end

    for ti, tv in ipairs(thrusters) do
        if Thrusters[peripheral.getName(tv)] == nil then
            Thrusters[peripheral.getName(tv)] = {
                thruster = tv,
                transmission = nil,
                type = "thruster",
                -- NOTE:
                -- affectVectors will have values depending on how the thruster's 
                -- thrust vector interacts with the vessel's attitude
                -- possible directions for yaw: "port"|"star"
                -- possible directions for roll: "port"|"star"
                -- possible directions for pitch: "up"|"down"
                -- possible directions for lateral: "port"|"star"|"fore"|"aft"|"up"|"down"
                affectVectors = {
                    yaw = nil,
                    pitch = nil,
                    roll = nil,
                    lateral = nil
                },
                power = tv.getPower()
            }
        end
        if Thrusters[peripheral.getName(tv)] ~= nil then
            if Thrusters[peripheral.getName(tv)].type == nil then 
                Thrusters[peripheral.getName(tv)].type = "thruster" 
            end
            Thrusters[peripheral.getName(tv)].thruster = tv
            Thrusters[peripheral.getName(tv)].transmission = nil
            if Thrusters[peripheral.getName(tv)].affectVectors == nil then 
                Thrusters[peripheral.getName(tv)].affectVectors = {} 
            end
            Thrusters[peripheral.getName(tv)].power = tv.getPower()
        end
    end
end

function Update()
    PropellerInit()
    ThrusterInit()
end

local function checkIfThrusterIsIndexed()
    local thrusters = {}
    local unindexedThrusters = {}

    for _, v in ipairs(ThrusterTypes) do
        local thrustersList = peripheral.find(v)
        if thrustersList then
            table.insert(thrusters, thrustersList)
        end
    end

    for _, v in ipairs(PropellerTypes) do
        local thrustersList = peripheral.find(v)
        if thrustersList then
            table.insert(thrusters, thrustersList)
        end
    end

    for _, v in ipairs(thrusters) do
        if Thrusters[peripheral.getName(v)] == nil then
            table.insert(unindexedThrusters, v)
        end
    end

    if #unindexedThrusters == 0 then
        return nil
    end
    return unindexedThrusters
end

local function partiallyUpdateThrusters(thrusterList)
    local thrusterTypes = {
        "thruster",
        "ion_thruster",
        "vector_thruster",
        "liquid_vector_thruster"
    }

    local propellerTypes = {
        "gyroscopic_propeller_bearing",
        "propeller",
        "propeller_bearing"
    }

    -- for every thruster in the input array
    -- check if they are one of thrusterTypes
    for _, tv in ipairs(thrusterList) do
        for _, tt in ipairs(ThrusterTypes) do
            if peripheral.getType(tv) == tt then
                Thrusters[peripheral.getName(tv)] = {
                    thruster = tv,
                    transmission = nil,
                    type = "thruster",
                    affectVectors = {
                        yaw = nil,
                        roll = nil,
                        pitch = nil,
                        lateral = nil
                    },
                    power = tv.getPower()
                }
            end
        end

        for _, pt in ipairs(PropellerTypes) do
            if peripheral.getType(tv) == pt then
                local thisThruster = {
                    thruster = tv,
                    transmission = nil,
                    type = "rotator",
                    affectVectors = {
                        yaw = nil,
                        roll = nil,
                        pitch = nil,
                        lateral = nil
                    },
                    power = tv.getPower()
                }

                for _, tt in ipairs(TransmissionTypes) do 
                    local transmissions = peripheral.find(tt)
                    for _, t in ipairs(transmissions) do
                        if tv.getSubnetworkAnchorId() == t.getSelfId() then
                            thisThruster.transmission = t
                        end
                    end
                end

                Thrusters[peripheral.getName(tv)] = thisThruster
            end
        end
    end
end

function UpdateStabilization()
    -- This function needs full implementation according to requirements in docs.md
    -- It should adjust thrust based on sensor data to achieve target angles/velocities
    
    -- For now, implement placeholder functionality that calls the existing thrust methods
    print("UpdateStabilization called")
end

function PollThrottle()
    for _, v in pairs(Config.Throttles) do
        if v.name and v.side then
            v.input = peripheral.wrap(v.name).getAnalogInput(v.side)
        end
    end
end

function PollSensors()
    PollVelocity()
    PollAltitude()
    PollNavTable()
    PollGimbal()
end

function PollVelocity()
    local velSensors = { peripheral.find("velocity_sensor") }

    -- Velocity Sensors
    VelocityVectors = {}
    local next = next
    if next(velSensors) == nil then

        SensorData.Velocity = {}

        for vsindex, velsensor in pairs(velSensors) do
            
            local vsAxis = velsensor.getAxis()
            local vsVelocity = velsensor.getVelocity()
            VelocityVectors[vsAxis] = vsVelocity
            
            SensorData.Velocity.Raw = VelocityVectors
        end
    end
    
end

function PollAltitude()
    local altSensor = peripheral.find("altitude_sensor") -- enforce one alt sensor

    -- Altitude Sensor
    if altSensor then

        SensorData.Altitude = {}

        Altitude = altSensor.getHeight()
        AirPressure = altSensor.getAirPressure()
        VerticalSpeed = altSensor.getVerticalSpeed()

        SensorData.Altitude.Altitude = Altitude
        SensorData.Altitude.AirPressure = AirPressure
        SensorData.Altitude.VerticalSpeed = VerticalSpeed
    end
end

function PollNavTable()
    local navTable = peripheral.find("navigation_table") -- enforce one nav table
    
    -- Navigation Table
    if navTable then

        SensorData.NavTable = {}

        Heading = navTable.getHeading()

        if Config.SensorCorrection.Heading ~= 0 then
            Heading = Heading + Config.SensorCorrection.Heading
            if Heading > 360 then
                Heading = Heading - 360
            elseif Heading < 0 then
                Heading = Heading + 360
            else
                Heading = Heading + Config.SensorCorrection.Heading
            end
        end

        NavTableHasTarget = navTable.hasTarget()

        SensorData.NavTable.Heading = Heading
        SensorData.NavTable.HasTarget = NavTableHasTarget
        if NavTableHasTarget then
            BearingToTarget = navTable.getBearing()
            TargetClosureRate = navTable.getClosureRate()
            TargetVerticalOffset = navTable.getVerticalOffsetToTarget()
            TargetRelativeAngle = navTable.getRelativeAngle()
        end
        if not NavTableHasTarget then
            BearingToTarget = 0
            TargetClosureRate = 0
            TargetVerticalOffset = 0
            TargetRelativeAngle = 0
        end

        SensorData.NavTable.TargetBearing = BearingToTarget
        SensorData.NavTable.TargetClosureRate = TargetClosureRate
        SensorData.NavTable.TargetVerticalOffset = TargetVerticalOffset
        SensorData.NavTable.TargetRelativeAngle = TargetRelativeAngle
    end
end

function PollGimbal()
    local gimbalSensor = peripheral.find("gimbal_sensor") -- enforce one gimbal sensor

    -- Gimbal Sensor
    if gimbalSensor then

        SensorData.Gimbal = {}

        Angles = gimbalSensor.getAngles()
        AngularRates = gimbalSensor.getAngularRates()
        LinearAcceleration = gimbalSensor.getLinearAcceleration()

        SensorData.Gimbal.Angles = Angles
        SensorData.Gimbal.AngularRates = AngularRates
        SensorData.Gimbal.LinearAcceleration = LinearAcceleration
    end
end

function Sigmoid(x)
    local e = 2.718281828459045
    return 1 / (1 + e^(-x))
end

function CustomSigmoid(x)
    return ((2 * Sigmoid(x)) - 1) * -1
end

-- thruster is a thruster type, power is a vector from 0.0 to 1.0
function SetThrusterPower(thruster, power)
    if thruster.type == "rotator" then
        thruster.power = power
        local actualPower = 0
        if peripheral.getType(thruster.transmission) == "Create_RotationSpeedController" then
            if (math.ceil(power*256) - power*256 <= 0.5) then
                actualPower = math.ceil(power*256)
            end
            if (math.ceil(power*256) - power*256) > 0.5 then
                actualPower = math.floor(power*256)
            end

            thruster.transmission.setTargetSpeed(actualPower)
        end
        if peripheral.getType(thruster.transmission) == "analog_transmission" then
            if (math.ceil(power*15) - power*15) <= 0.5 then
                actualPower = math.ceil(power*15)
            end
            if (math.ceil(power*15) - power*15) > 0.5 then
                actualPower = math.floor(power*15)
            end

            if actualPower < 1 then
                actualPower = 15 -- signal 15 will decouple the speed and stop the motion
            end
            if actualPower <= 15 then
                actualPower = actualPower - 1
            end

            thruster.transmission.setSignal(actualPower)
        end
        
    end
    if thruster.type == "thruster" then
        thruster.power = power
        thruster.thruster.setPowerNormalized(power)
    end
end

function UpdateGlobalThrust() 
    -- lateral thrust
    local desiredLateralThrustVectors = {}

    for f, v in pairs(VelocityVectors) do
        desiredLateralThrustVectors[f] = CustomSigmoid(v)
    end

    local thrustDirections = {}
    for f, v in pairs(desiredLateralThrustVectors) do
        if f == "x" then
            if v > 0 then table.insert(thrustDirections, { port = v }) else table.insert(thrustDirections, { star = math.abs(v) }) end
        end
        if f == "y" then
            if v > 0 then table.insert(thrustDirections, { down = v }) else table.insert(thrustDirections, { up = math.abs(v) }) end
        end
        if f == "z" then
            if v > 0 then table.insert(thrustDirections, { fore = v }) else table.insert(thrustDirections, { aft = math.abs(v) }) end
        end
    end

    for _, d in pairs(thrustDirections) do
        for _, t in pairs(Thrusters) do
            if t.affectVectors.lateral == d then
                SetThrusterPower(t, d)
            end
        end
    end

    -- angular thrust
    local desiredAngluarThrustVectors = {}

    for f, v in pairs(AngularRates) do
        desiredAngluarThrustVectors[f] = CustomSigmoid(v)
    end
    
    -- Need to complete this function implementation based on docs.md requirements
end

function Init()
    
    local configFile = io.open(Config.ConfigPath, "r")
    if configFile then
        local configContent = configFile:read("*all")
        Config = stringToTable(configContent)
        configFile:close()
    else
        print("Could not open the config file")
    end

    -- Initialize monitors - identify which monitor is for which type  
    if Config.Monitors.InstrumentPanelMonitor == nil then
        print("Identifying instrument panel monitor...")
        Config.Monitors.InstrumentPanelMonitor = identifyMonitor("instrument panel")
    else
        Config.Monitors.InstrumentPanelMonitor.peripheral = peripheral.wrap(Config.Monitors.InstrumentPanelMonitor.name)
    end
    
    if Config.Monitors.AutopilotControlMonitor == nil then
        print("Identifying autopilot control monitor...")
        Config.Monitors.AutopilotControlMonitor = identifyMonitor("autopilot control")
    else
        Config.Monitors.AutopilotControlMonitor.peripheral = peripheral.wrap(Config.Monitors.AutopilotControlMonitor.name)
    end


    ThrottleInit()
    
    Thrusters = {}

    if Config.thrusterConfigPath ~= nil then
        local thrusterConfigFile = io.open(Config.thrusterConfigPath, "r")
        if thrusterConfigFile then
            local configContent = thrusterConfigFile:read("*all")
            local configChunk, err = load(configContent)
            if configChunk then
                Thrusters = configChunk()
            else
                print("Error loading thruster config string: "..err)
            end
            thrusterConfigFile:close()
        else
            print("Could not open the thruster config file")
        end
    end

    -- If we're here, the config loaded successfully and the initialization is done.
end

-- Display functions for monitors
function displayInstrumentPanel(monitor, sensorData)
    if not monitor or not monitor.peripheral then return end
    
    local m = monitor.peripheral
    m.clear()
    m.setCursorPos(1, 1)
    
    -- Simple instrument panel display
    m.write("Caero Attitude Control")
    m.setCursorPos(1, 2)
    m.write("----------------------")

    local next = next
    if next(sensorData) ~= nil then
        if sensorData then
            m.setCursorPos(1, 3)
            m.write("Pitch:   " .. string.format("%.2f", SensorData.Gimbal.Angles.xAngle or 0))
            m.setCursorPos(1, 4)
            m.write("Roll:    " .. string.format("%.2f", SensorData.Gimbal.Angles.zAngle or 0))
            m.setCursorPos(1, 5)
            m.write("Yaw:     " .. string.format("%.2f", SensorData.NavTable.Heading or 0))
            m.setCursorPos(1, 6)
            m.write("Velocity: " .. string.format("%.2f", sensorData.velocity or 0))
            m.setCursorPos(1, 7)
            m.write("Altitude: " .. string.format("%.2f", SensorData.Altitude.Altitude or 0))
        else
            m.setCursorPos(1, 3)
            m.write("Pitch:   0.00")
            m.setCursorPos(1, 4)
            m.write("Roll:    0.00")
            m.setCursorPos(1, 5)
            m.write("Yaw:     0.00")
            m.setCursorPos(1, 6)
            m.write("Velocity: 0.00")
            m.setCursorPos(1, 7)
            m.write("Altitude: 0.00")
        end
        
        -- Display autopilot status
        local autopilotStatus = Config.Autopilot.AutopilotEngaged and "ON" or "OFF"
        m.setCursorPos(1, 9)
        m.write("Autopilot: " .. autopilotStatus)
    end
    
    
end

function displayAutopilotControls(monitor, sensorData)
    if not monitor or not monitor.peripheral then return end
    
    local m = monitor.peripheral
    m.clear()
    m.setCursorPos(1, 1)
    
    -- Simple autopilot control panel  
    m.write("Autopilot Controls")
    m.setCursorPos(1, 2)
    m.write("------------------")
    m.setCursorPos(1, 3)
    m.write("Engage: [ ]")
    m.setCursorPos(1, 4)
    m.write("Speed:  [0.0]")
    m.setCursorPos(1, 5)
    m.write("Heading: [0.0]")
    m.setCursorPos(1, 6)
    m.write("Manual: [ ]")
    m.setCursorPos(1, 7)
    m.setCursorPos(1, 8)
    m.write("Mode: Auto")
end

-- Sensor data collection - returns a table with all relevant sensor readings
function collectSensorData()
    local data = {}
    
    -- Collect gimbal data
    data.gimbal = {
        pitch = 0,
        roll = 0
    }
    
    -- Collect navigation table data
    data.navTable = {
        heading = 0,
        hasTarget = false
    }
    
    -- Collect altitude data
    data.altitude = 0
    
    -- Collect velocity data
    data.velocity = 0
    
    return data
end

function WriteConfigFiles()
    print("Writing to thruster config file.")
    local thrusterConfigFile = fs.open(Config.thrusterConfigPath, "w+")
    thrusterConfigFile.write(tableToString(Thrusters))
    thrusterConfigFile.close()

    print("Writing to config file")
    local configFile = fs.open(Config.ConfigPath, "w+")
    configFile.write(tableToString(Config))
    configFile.close()
end

function Main()
    Init()

    local sentinel = true
    print("Mainloop starting. Press 'q' to stop the loop and save configuration changes.")
    while sentinel do

        -- Update monitor displays
        if Config.Monitors.InstrumentPanelMonitor then
            displayInstrumentPanel(Config.Monitors.InstrumentPanelMonitor, SensorData)
        end
        
        if Config.Monitors.AutopilotControlMonitor then
            displayAutopilotControls(Config.Monitors.AutopilotControlMonitor, SensorData)
        end

        local unindexedThrusters = checkIfThrusterIsIndexed()
        if unindexedThrusters ~= nil then
            partiallyUpdateThrusters(unindexedThrusters)
        end

        PollSensors()
        UpdateGlobalThrust()
        UpdateStabilization()

        local function getTerminateEvent()
            local event, key, is_held = os.pullEventRaw("key")
            if keys.getName(key) == "q" then 
                print("Quit Input Received") 
                
                print("Writing to thruster config file.")
                local thrusterConfigFile = fs.open(Config.thrusterConfigPath, "w+")
                thrusterConfigFile.write(tableToString(Thrusters))
                thrusterConfigFile.close()

                print("Writing to config file")
                local configFile = fs.open(Config.ConfigPath, "w+")
                configFile.write(tableToString(Config))
                configFile.close()

                sentinel = false
            else
                return
            end
        end

        local function do_sleep() os.sleep(0.1) end

        parallel.waitForAny(do_sleep, getTerminateEvent)
    end
end

-- Run the main function when the script starts
Main()