import struct
import numpy as np
import time
from pylora import LoRa  # LoRa USB stick for PC
from scipy.signal import resample

# -----------------------------
# Configuration
# -----------------------------
STRANDS = 8
SLOTS = 4
PACKET_SIZE = STRANDS*SLOTS + 1  # 32 + 1 tick
SAMPLE_RATE = 2.048e6
CARRIER_FREQ = 100e6
NODE_COUNT = 4

# -----------------------------
# Initialize LoRa
# -----------------------------
lora = LoRa()
lora.set_mode(LoRa.RX)
lora.set_frequency(868e6)
lora.set_sf(7)
lora.set_bw(125e3)
lora.set_cr(5)

# -----------------------------
# Node lattice storage
# -----------------------------
nodes = {i: np.zeros((STRANDS, SLOTS)) for i in range(NODE_COUNT)}
ticks = {i: 0 for i in range(NODE_COUNT)}

# -----------------------------
# Decode packet
# -----------------------------
def decode_packet(packet):
    flat = list(struct.unpack("B"*STRANDS*SLOTS + "B", packet))
    tick = flat[-1]
    lattice = np.array(flat[:-1]).reshape((STRANDS, SLOTS)) / 50.0  # scale back
    return lattice, tick

# -----------------------------
# Aggregate lattices
# -----------------------------
def aggregate_lattice(nodes):
    lattices = np.array(list(nodes.values()))
    agg = np.mean(lattices, axis=0)  # average across nodes
    return agg

# -----------------------------
# Convert lattice to IQ
# -----------------------------
def lattice_to_iq(lattice, length=1024):
    t = np.arange(length) / SAMPLE_RATE
    signal = np.zeros(length)
    for s in range(STRANDS):
        weight = np.mean(lattice[s])
        freq = 1e3*(s+1)
        signal += weight * np.sin(2*np.pi*freq*t)
    # Normalize
    signal /= np.max(np.abs(signal))
    return signal.astype(np.complex64)

# -----------------------------
# Main loop
# -----------------------------
print("[+] Aggregator running...")
while True:
    if lora.packet_available():
        packet = lora.receive_packet()
        node_id = packet[0] % NODE_COUNT  # simple node indexing
        lattice, tick = decode_packet(packet)
        nodes[node_id] = lattice
        ticks[node_id] = tick

    # Aggregate
    agg_lattice = aggregate_lattice(nodes)
    iq_signal = lattice_to_iq(agg_lattice, length=2048)

    # Here, iq_signal is ready for:
    # - Transmission via TX-capable SDR (HackRF/LimeSDR)
    # - Visualization/recording
    # For example: sdr.write_samples(iq_signal)

    # Optional: print hex summary
    binary_agg = (agg_lattice > 1.0).astype(int).flatten()
    hex_agg = int(binary_agg.dot(1 << np.arange(STRANDS*SLOTS)))
    print(f"Aggregated Lattice Hex: 0x{hex_agg:08x}")
    time.sleep(0.1)
