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.InstrumentPanelMonitor = nil
Config.Monitors.AutopilotControlMonitor = nil

-- Autopilot
Config.Autopilot.AutopilotEngaged = false -- Is auto-yaw enabled?
Config.Autopilot.AutoForeAft = false -- Is auto-thrust enabled?
Config.Autopilot.AutopilotDesiredSpeed = 0 -- Signed thrust value to compare against VelocityVectors["x"]
Config.Autopilot.AutopilotDesiredHeading = 0 -- If NavTableHasTarget == true this gets ignored in favor of TargetRelativeAngle

-- Gimbal Sensor Values
Angles = {
    0, -- xAngle: rotation about body-X (pitch); 0 is aligned with world-up
    0  -- zAngle: rotation about body-Z (roll); 0 is aligned with world-up
}
AngularRates = {
    0, -- wx: pitch rate about body-X
    0, -- wy: yaw rate about body-Y
    0  -- wz: roll rate about body-Z
}
LinearAcceleration = {
    0, -- Acceleration along body-X
    0, -- Acceleration along body-Y
    0  -- Acceleration along body-Z
}

-- Nav Table Values
Heading = 0

-- Nav Table Values (targets)
NavTableHasTarget = false
BearingToTarget = 0
TargetClosureRate = 0
TargetVerticalOffset = 0
TargetRelativeAngle = 0

-- Altitude Sensor Values
Altitude = 0
AirPressure = 0
VerticalSpeed = 0

-- Velocity Sensors
VelocityVectors = {
    x = 0, -- +X is east (port), +Y is up, +Z is south (aft)
    y = 0, -- -X is west (star), -Y is down, -Z is north (fore)
    z = 0
}

-- Throttle Inputs
Throttles = {
    ForeThrottle = {
        name = "",
        side = "",
        input = 0
    },
    DownThrottle = {
        name = "",
        side = "",
        input = 0
    },
    AftThrottle = {
        name = "",
        side = "",
        input = 0
    }
}

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][side] = wrapper.getAnalogInput(side)
            end
        end
    end

    return states
end

local function identifyRelay(label)
    print("\nIdentifying "..label.." throttle")
    print("\nPlease change the input signal for the throttle")

    local initialStates = getRelayStates

    while true do
        os.pullEvent("redstone_relay")

        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
                    return {
                        name = pname,
                        side = k
                    }
                end
            end
        end
    end
end

function ThrottleInit() 
    Throttles.ForeThrottle = identifyRelay("fore")
    Throttles.DownThrottle = identifyRelay("down")
    Throttles.AftThrottle = identifyRelay("aft")
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 = nil end
                    Thrusters[peripheral.getName(pv)].power = (tv.getSingnal())/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 = nil 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
        table.insert(thrusters, peripheral.find(v))
    end

    for _, v in ipairs(PropellerTypes) do
        table.insert(thrusters, peripheral.find(v))
    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()

end

function PollThrottle()
    for _, v in ipairs(Throttles) do
        v.input = peripheral.call(v.name, "getAnalogInput("..v.side..")")
    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
    Altitude = altSensor.getHeight()
    AirPressure = altSensor.getAirPressure()
    VerticalSpeed = altSensor.getVerticalSpeed()
end

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

    -- Navigation Table
    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
end

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

    -- Gimbal Sensor
    Angles = gimbalSensor.getAngles()
    AngularRates = gimbalSensor.getAngularRates()
    LinearAcceleration = gimbalSensor.getLinearAcceleration()
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 ipairs(VelocityVectors) do
        desiredLateralThrustVectors[f] = CustomSigmoid(v)
    end

    local thrustDirections = {}
    for f, v in 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 ipairs(thrustDirections) do
        for _, t in ipairs(Thrusters) do
            if t.affectVectors.lateral == d then
                SetThrusterPower(t, d)
            end
        end
    end

    -- angular thrust
    local desiredAngluarThrustVectors = {}

    for f, v in ipairs(AngularRates) do
        desiredAngluarThrustVectors[f] = CustomSigmoid(v)
    end
end

function Init()
    ThrottleInit()
    Thrusters = {}
    
    local configFile = io.open(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
    else
        print("Could not open the config file")
    end
    configFile.close()

    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

        else
            print("Could not open the thruster config file")
        end
        thrusterConfigFile.close()
    end

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

function Main()
    Init()

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

        PollSensors()
        UpdateGlobalThrust()
    end
end