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

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"

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

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, name in ipairs(monitors) do
        local monitor = peripheral.wrap(name)
        if monitor then
            local width, height = monitor.getSize()
            print(i .. ". " .. name .. " (" .. 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 selectedMonitorName = monitors[choice]
    local selectedMonitor = peripheral.wrap(selectedMonitorName)
    
    print("Selected monitor: " .. selectedMonitorName)
    
    -- Return monitor details as a table
    return {
        name = selectedMonitorName,
        peripheral = selectedMonitor
    }
end

function assignMonitor(monitorType, monitorData)
    if monitorType == "instrument_panel" then
        Config.Monitors.InstrumentPanelMonitor = monitorData
    elseif monitorType == "autopilot_control" then
        Config.Monitors.AutopilotControlMonitor = monitorData
    end
end

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

Throttles = {
    fore = {
        name = nil,
        side = nil
    },
    aft = {
        name = nil,
        side = nil
    },
    down = {
        name = nil,
        side = nil
    }
}



function ThrottleInit() 

    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()

        for pname, currentState in pairs(currentStates) do
            local initialState = initialStates[pname]
            for k, v in pairs(currentState) do
                if initialState[k] ~= v then
                    Throttles[throttleType].name = pname
                    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(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 = {}
    for vi, vs in pairs(velSensors) do
        local vsAxis = vs.getAxis()
        local vsVelocity = vs.getVelocity()
        VelocityVectors[vsAxis] = vsVelocity
    end
end

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

    -- Altitude Sensor
    if altSensor then
        Altitude = altSensor.getHeight()
        AirPressure = altSensor.getAirPressure()
        VerticalSpeed = altSensor.getVerticalSpeed()
    end
end

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

    -- Navigation Table
    if navTable then
        Heading = navTable.getHeading()
        NavTableHasTarget = navTable.hasTarget()
        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
    else
        Heading = 0
        NavTableHasTarget = false
        BearingToTarget = 0
        TargetClosureRate = 0
        TargetVerticalOffset = 0
        TargetRelativeAngle = 0
    end
end

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

    -- Gimbal Sensor
    if gimbalSensor then
        Angles = gimbalSensor.getAngles()
        AngularRates = gimbalSensor.getAngularRates()
        LinearAcceleration = gimbalSensor.getLinearAcceleration()
    else
        Angles = {0, 0}
        AngularRates = {0, 0, 0}
        LinearAcceleration = {0, 0, 0}
    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()

    -- 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")
    end
    
    if Config.Monitors.AutopilotControlMonitor == nil then
        print("Identifying autopilot control monitor...")
        Config.Monitors.AutopilotControlMonitor = identifyMonitor("autopilot control")
    end


    ThrottleInit()
    
    Thrusters = {}
    
    local configFile = io.open(Config.ConfigPath, "r")
    if configFile then
        local configContent = configFile:read("*all")
        local configChunk, err = load(configContent)
        if configChunk then
            Config = configChunk()
        else
            print("Error loading config string: "..err)
        end
        configFile:close()
    else
        print("Could not open the config file")
    end

    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("----------------------")
    
    if sensorData then
        m.setCursorPos(1, 3)
        m.write("Pitch:   " .. string.format("%.2f", sensorData.gimbal.pitch or 0))
        m.setCursorPos(1, 4)
        m.write("Roll:    " .. string.format("%.2f", sensorData.gimbal.roll 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 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

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 displayAutopilotControls(monitor)
    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.write("Mode: Auto")
end

function Main()
    Init()

    while true do
        -- Update monitor displays
        if Config.Monitors.InstrumentPanelMonitor then
            displayInstrumentPanel(Config.Monitors.InstrumentPanelMonitor)
        end
        
        if Config.Monitors.AutopilotControlMonitor then
            displayAutopilotControls(Config.Monitors.AutopilotControlMonitor)
        end

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

        PollSensors()
        UpdateGlobalThrust()
        UpdateStabilization()
        
        -- Add a small delay to avoid overloading the system
        os.sleep(0.1)
    end
end

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