lora_gs/main.py

from asyncio.tasks import sleep
import time
from dataclasses import dataclass
from enum import Enum
from enum import IntEnum
from typing import NamedTuple
from LoRaRF import SX126x


LoRa = SX126x()

# As per EN 300 220-2 V3.3.1; Annex B table 1
# https://www.etsi.org/deliver/etsi_en/300200_300299/30022002/03.03.01_60/en_30022002v030301p.pdf
@dataclass
class BandConfig:
    freq_start_mhz: float
    ch_bw_khz: int
    num_channels: int
    max_power_mw: int
    max_duty: float
    polite_access: bool


EU868_BANDS = {
    "BAND_K": BandConfig(863.0, 500, 4, 25, 0.1, True),
    "BAND_L": BandConfig(865.0, 500, 6, 25, 1.0, True),
    "BAND_M": BandConfig(868.0, 500, 1, 25, 1.0, True),
    "BAND_N": BandConfig(868.7, 500, 1, 25, 0.1, True),
    "BAND_O": BandConfig(869.4, 250, 1, 500, 10.0, True),
}

import struct
from dataclasses import dataclass

# Assuming Little Endian (standard for Raspberry Pi / ARM)
ENDIAN = '<'

@dataclass
class LoRaPacketHeader:
    type: int
    id: int
    number: int
    tx_time: int  # 48-bit timestamp (ms since epoch)

    def pack(self) -> bytes:
        # Pack 3-byte header base
        base = struct.pack(f'{ENDIAN}B B B', self.type, self.id, self.number)
        # Pack 6-byte timestamp
        time_bytes = self.tx_time.to_bytes(6, byteorder='little', signed=False)
        return base + time_bytes

    @classmethod
    def unpack(cls, data: bytes) -> LoRaPacketHeader:
        t, i, n = struct.unpack(f'{ENDIAN}B B B', data[0:3])
        time_val = int.from_bytes(data[3:9], byteorder='little', signed=False)
        return cls(t, i, n, time_val)


@dataclass
class LoRaSyncPacket:
    header: LoRaPacketHeader
    connected: int       # 1 bit
    sync_window: int     # 10 bits
    gs_window: int       # 9 bits
    security_window: int # 4 bits

    def pack(self) -> bytes:
        hdr = self.header.pack()
        # GCC bitfields pack LSB to MSB in Little Endian.
        # Total bits = 1 + 10 + 9 + 4 = 24 bits (3 bytes).
        bitfield = (
            (self.connected & 0x1) |
            ((self.sync_window & 0x3FF) << 1) |
            ((self.gs_window & 0x1FF) << 11) |
            ((self.security_window & 0xF) << 20)
        )
        bf_bytes = bitfield.to_bytes(3, byteorder='little', signed=False)
        padding = b'\x00\x00\x00'
        return hdr + bf_bytes + padding

    @classmethod
    def unpack(cls, data: bytes) -> LoRaSyncPacket:
        hdr = LoRaPacketHeader.unpack(data[:9])
        bitfield = int.from_bytes(data[9:12], byteorder='little', signed=False)

        connected = bitfield & 0x1
        sync_window = (bitfield >> 1) & 0x3FF
        gs_window = (bitfield >> 11) & 0x1FF
        security_window = (bitfield >> 20) & 0xF

        return cls(hdr, connected, sync_window, gs_window, security_window)


@dataclass
class LoRaConnectPacket:
    header: LoRaPacketHeader

    def pack(self) -> bytes:
        return self.header.pack() + b'\x00' * 6

    @classmethod
    def unpack(cls, data: bytes) -> LoRaConnectPacket:
        return cls(LoRaPacketHeader.unpack(data[:9]))


@dataclass
class LoRaAcceptPacket:
    header: LoRaPacketHeader
    connect_time: int # 48-bit timestamp

    def pack(self) -> bytes:
        hdr = self.header.pack()
        ctime = self.connect_time.to_bytes(6, byteorder='little', signed=False)
        return hdr + ctime

    @classmethod
    def unpack(cls, data: bytes) -> LoRaAcceptPacket:
        hdr = LoRaPacketHeader.unpack(data[:9])
        ctime = int.from_bytes(data[9:15], byteorder='little', signed=False)
        return cls(hdr, ctime)


@dataclass
class LoRaDataPacket:
    header: LoRaPacketHeader
    # imu
    imu_altitude: int  # uint16_t (representing float16)
    imu_vspeed: int    # uint16_t (representing float16)
    imu_attitude: int  # uint16_t (representing float16)
    imu_dt: int        # int16_t
    # baro
    baro_p1: int       # uint16_t (representing float16)
    baro_p2: int       # uint16_t (representing float16)
    baro_dt: int       # int16_t
    # gps
    gps_lat: float     # float (32-bit IEEE 754)
    gps_lon: float     # float (32-bit IEEE 754)
    gps_dt: int        # int16_t

    def pack(self) -> bytes:
        hdr = self.header.pack()
        # H = uint16_t, h = int16_t, f = float (32-bit)
        fmt = f'{ENDIAN}H H H h H H h f f h'
        payload = struct.pack(fmt,
                              self.imu_altitude, self.imu_vspeed, self.imu_attitude, self.imu_dt,
                              self.baro_p1, self.baro_p2, self.baro_dt,
                              self.gps_lat, self.gps_lon, self.gps_dt)
        return hdr + payload

    @classmethod
    def unpack(cls, data: bytes) -> LoRaDataPacket:
        hdr = LoRaPacketHeader.unpack(data[:9])
        fmt = f'{ENDIAN}H H H h H H h f f h'
        payload = struct.unpack(fmt, data[9:33])
        return cls(hdr, *payload)


@dataclass
class LoRaCommandPacket:
    header: LoRaPacketHeader
    command: int
    data_val: int # uint64_t

    def pack(self) -> bytes:
        hdr = self.header.pack()
        # B = uint8_t, Q = uint64_t
        payload = struct.pack(f'{ENDIAN}B Q', self.command, self.data_val)
        return hdr + payload

    @classmethod
    def unpack(cls, data: bytes) -> LoRaCommandPacket:
        hdr = LoRaPacketHeader.unpack(data[:9])
        cmd, val = struct.unpack(f'{ENDIAN}B Q', data[9:18])
        return cls(hdr, cmd, val)


# Ebyte E22-900M33S module power output, since it has an LNA and a PA,
# the actual output power is not the same as the value passed to RadioLib's
# setOutputPower() function. The table below maps the two values.
# Note that this is the output power of the module, not the effective radiated
# power (ERP) of the antenna.
# Refer to image in page 10, chapter 4.2 of the manual: https://www.cdebyte.com/products/E22-900M33S/4#Downloads
# Useful to extract data: https://plotdigitizer.com/app
# NOTE: THE X SCALE ON THIS GRAPH IS NOT CONSTANT, WHY????
class PowerEntry(NamedTuple):
    power_mw: int  # The actual power in mW, for reference
    radiolib_value: (
        int  # The value to pass to RadioLib/LoRa module setOutputPower()
    )

LORA_OUTPUT_POWER_TABLE = (
    PowerEntry(power_mw=59, radiolib_value=-9),
    PowerEntry(power_mw=75, radiolib_value=-8),
    PowerEntry(power_mw=98, radiolib_value=-7),
    PowerEntry(power_mw=126, radiolib_value=-6),
    PowerEntry(power_mw=173, radiolib_value=-5),
    PowerEntry(power_mw=237, radiolib_value=-4),
    PowerEntry(power_mw=299, radiolib_value=-3),
    PowerEntry(power_mw=390, radiolib_value=-2),
    PowerEntry(power_mw=487, radiolib_value=-1),
    PowerEntry(power_mw=619, radiolib_value=0),
    PowerEntry(power_mw=750, radiolib_value=1),
    PowerEntry(power_mw=1044, radiolib_value=2),
    PowerEntry(power_mw=1202, radiolib_value=3),
    PowerEntry(power_mw=1409, radiolib_value=4),
    PowerEntry(power_mw=1492, radiolib_value=5),
    PowerEntry(power_mw=1585, radiolib_value=6),
    PowerEntry(power_mw=1745, radiolib_value=7),
    PowerEntry(power_mw=2218, radiolib_value=8),
)


"""Finds the radiolib_value for the closest supported power level in mW."""
def get_radiolib_power(target_mw: int) -> int:
    best_match = min(
        LORA_OUTPUT_POWER_TABLE, key=lambda entry: abs(entry.power_mw - target_mw)
    )
    return best_match.radiolib_value


class LoRaFCState(IntEnum):
    STATE_DISCONNECTED = 0
    STATE_CONNECTING = 1
    STATE_TRANSMIT = 2
    STATE_RECEIVE = 3

class PacketType(IntEnum):
    PKT_UNKNOWN = 0x0
    PKT_SYNC    = 0x01 # Broadcast sync packet GS -> FC
    PKT_CONNECT = 0x02 # Connect request packet FC -> GS
    PKT_ACCEPT  = 0x03 # Connect accept packet GS -> FC
    PKT_DATA    = 0x04 # Data packet FC -> GS
    PKT_COMMAND = 0x05 # Command packet GS -> FC


TX_FORCE = True
LORA_FC_ID = 0xFC
LORA_GS_ID = 0xDE

# --- State Machine Class ---
class LoRaGSStateMachine:
    def __init__(self, my_id: int, fc_id: int):
        # Configuration constants
        self.sync_window = 1000      # ms
        self.gs_window = 200         # ms
        self.security_window = 10    # ms
        self.MAX_SILENT_FRAMES = 5   # Define your max silent frames here
        self.MY_ID = my_id
        self.LORA_FC_ID = fc_id

        # Static variables from C++
        self.state = LoRaFCState.STATE_DISCONNECTED
        self.sync_sent_time = 0
        self.connect_rx_time = 0
        self.packets_received = 0
        self.silent_frames = 0

        # Global mock variables from C++ (updated by TX/RX functions)
        self.last_tx_time = 0
        self.last_rx_time = 0

        # Receive buffer
        self.receive_buffer = bytearray(256)
        self.receive_len = 0

    # --- Time Helpers ---
    def now_ms(self) -> int:
        """Returns monotonic time in milliseconds"""
        return time.monotonic_ns() // 1_000_000

    def slot_relative_time(self, start_time_ms: int) -> int:
        return self.now_ms() - start_time_ms


    # --- The State Machine ---
    def step(self) -> LoRaFCState:
        """Run one iteration of the state machine. Equivalent to lora_gs_state_machine()."""

        if self.state == LoRaFCState.STATE_DISCONNECTED:
            self.connect_rx_time = 0
            self.packets_received = 0
            self.silent_frames = 0

            if (self.now_ms() - self.sync_sent_time) >= self.sync_window:
                # Send sync packet to FC
                hdr = LoRaPacketHeader(type=PacketType.PKT_SYNC, id=self.MY_ID, number=0, tx_time=0)
                s = LoRaSyncPacket(hdr, connected=0, sync_window=self.sync_window,
                                   gs_window=self.gs_window, security_window=self.security_window)

                if self.lora_transmit_timeout(s.pack(), self.sync_window, TX_FORCE):
                    self.sync_sent_time = self.last_tx_time
            else:
                timeout = self.sync_window - self.slot_relative_time(self.sync_sent_time)
                if not self.lora_receive_timeout(timeout):
                    # No connect received within the sync window, return to disconnected state
                    return self.state

                header = self.lora_get_header()
                if header.type == PacketType.PKT_CONNECT and header.id == self.LORA_FC_ID:
                    # Received first sync packet from GS, start the handshake
                    self.connect_rx_time = self.last_rx_time
                    self.state = LoRaFCState.STATE_CONNECTING

        elif self.state == LoRaFCState.STATE_CONNECTING:
            # Send a ACCEPT packet to the GS
            hdr = LoRaPacketHeader(type=PacketType.PKT_ACCEPT, id=self.MY_ID, number=0, tx_time=0)
            a = LoRaAcceptPacket(hdr, connect_time=self.connect_rx_time)

            if not self.lora_transmit_timeout(a.pack(), self.sync_window // 2, TX_FORCE):
                self.state = LoRaFCState.STATE_DISCONNECTED
            else:
                self.state = LoRaFCState.STATE_TRANSMIT

        elif self.state == LoRaFCState.STATE_TRANSMIT:
            if self.silent_frames >= self.MAX_SILENT_FRAMES:
                self.state = LoRaFCState.STATE_DISCONNECTED
                return self.state

            # Sync window expired, need to resend sync packet
            if self.slot_relative_time(self.sync_sent_time) >= self.sync_window:
                hdr = LoRaPacketHeader(type=PacketType.PKT_SYNC, id=self.MY_ID, number=0, tx_time=0)
                s = LoRaSyncPacket(hdr, connected=1, sync_window=self.sync_window,
                                   gs_window=self.gs_window, security_window=self.security_window)

                if not self.lora_transmit_timeout(s.pack(), self.sync_window, TX_FORCE):
                    self.state = LoRaFCState.STATE_DISCONNECTED
                else:
                    self.sync_sent_time = self.last_tx_time
                return self.state

            remaining_time = self.gs_window - self.security_window - self.slot_relative_time(self.sync_sent_time)

            if remaining_time <= 0:
                self.packets_received = 0
                self.state = LoRaFCState.STATE_RECEIVE
                return self.state

            # Switched too early to transmit, wait for remaining listen time to expire
            if remaining_time >= self.gs_window:
                time.sleep((remaining_time - self.gs_window) / 1000.0)
                return self.state

            # TODO: transmit commands to FC
            pass

        elif self.state == LoRaFCState.STATE_RECEIVE:
            remaining_time = self.sync_window - self.security_window - self.slot_relative_time(self.sync_sent_time)

            # In the transmit window
            if remaining_time <= 0:
                if self.packets_received == 0:
                    self.silent_frames += 1
                self.state = LoRaFCState.STATE_TRANSMIT
                return self.state

            # In the current frame's transmit window, we shouldn't be here yet
            if remaining_time >= (self.sync_window - self.gs_window + self.security_window):
                if self.packets_received == 0:
                    self.silent_frames += 1
                self.state = LoRaFCState.STATE_TRANSMIT
                return self.state

            if not self.lora_receive_timeout(remaining_time):
                if self.packets_received == 0:
                    self.silent_frames += 1
                self.state = LoRaFCState.STATE_TRANSMIT
            else:
                self.packets_received += 1
                self.silent_frames = 0
                # TODO: do something with the packet
                print(f"Received packet: {self.lora_get_header()}")
                pass

        else:
            self.state = LoRaFCState.STATE_DISCONNECTED

        return self.state

    # --- Hardware Interface Stubs ---
    # You will need to implement these using LoRaRF or your chosen SPI library

    def lora_transmit_timeout(self, data: bytes, timeout_ms: int, force: bool) -> bool:
        """Transmit packet over SPI to SX1262 and update self.last_tx_time."""
        # TODO: Implement via LoRaRF
        LoRa.beginPacket()
        LoRa.write(data, len(data))
        LoRa.endPacket(timeout_ms * 64)
        self.last_tx_time = self.now_ms()
        return True

    def lora_receive_timeout(self, timeout_ms: int) -> bool:
        """Listen for incoming packets until timeout. Update self.last_rx_time."""
        LoRa.request(timeout_ms)
        LoRa.wait()
        status = LoRa.status()

        if status == LoRa.STATUS_RX_DONE:
            length = LoRa.available()
            if length > 0:
                self.receive_buffer[:length] = LoRa.get(length)
                self.receive_len = length
                self.last_rx_time = self.now_ms()
                return True
        elif status == LoRa.STATUS_RX_TIMEOUT:
            print("Reception timed out. No packet detected.")
            return False
        elif status == LoRa.STATUS_HEADER_ERR or status == LoRa.STATUS_CRC_ERR:
            print("Packet received with errors.")
            return False
        return False

    def lora_get_header(self) -> LoRaPacketHeader:
        """Extract and return header from the last received payload."""
        return LoRaPacketHeader.unpack(bytes(self.receive_buffer[:self.receive_len]))

def main():
    print("Starting up LoRa module")
    LoRa.setSpi(0, 0, 4_000_000)
    LoRa.setPins(23, 25, 24, -1, 5)
    LoRa.begin()
    print("LoRa module started")

    print("Setting LoRa paramters")
    band = "BAND_O"
    spreading_factor = 7
    coding_rate = 5

    LoRa.setDio2RfSwitch(True)
    LoRa.setModem(SX126x.LORA_MODEM)
    LoRa.setFrequency(int(EU868_BANDS[band].freq_start_mhz*1_000_000 + EU868_BANDS[band].ch_bw_khz*1000/2))
    LoRa.setBandwidth(EU868_BANDS[band].ch_bw_khz*1000)
    LoRa.setTxPower(10, LoRa.TX_POWER_SX1262)
    LoRa.setRxGain(LoRa.RX_GAIN_BOOSTED)
    LoRa.setLoRaModulation(spreading_factor,EU868_BANDS[band].ch_bw_khz*1000, coding_rate, False)
    LoRa.setHeaderType(LoRa.HEADER_EXPLICIT)
    print("LoRa parameters set")

    state_machine = LoRaGSStateMachine(LORA_GS_ID, LORA_FC_ID)
    while True:
        state = state_machine.step()
        print(f"State: {state}")
        time.sleep(0.001)

    LoRa.end()

if __name__ == "__main__":
    main()