import numpy as np
import threading
import time
import struct
# from lora import LoRa  # Uncomment for actual hardware

# -----------------------------
# Config
# -----------------------------
NODE_COUNT = 16
STRANDS, SLOTS = 8, 4
TICK_INTERVAL = 0.1
NEIGHBOR_COUNT = 4  # Phyllotaxis neighbors

LORA_FREQ = 915e6
LORA_SF = 7
LORA_BW = 125e3

# -----------------------------
# Phyllotaxis positions
# -----------------------------
def phyllotaxis_positions(n, c=1.0):
    golden_angle = np.pi * (3 - np.sqrt(5))
    positions = []
    for k in range(n):
        r = c * np.sqrt(k)
        theta = k * golden_angle
        x, y = r * np.cos(theta), r * np.sin(theta)
        positions.append((x, y))
    return np.array(positions)

positions = phyllotaxis_positions(NODE_COUNT)

# -----------------------------
# Node definition
# -----------------------------
class Node:
    def __init__(self, idx, pos):
        self.idx = idx
        self.pos = pos
        self.lattice = np.zeros((STRANDS, SLOTS))
        self.neighbors = []
        self.lock = threading.Lock()
        self.rx_queue = []
        # Initialize LoRa hardware
        # self.lora = LoRa(freq=LORA_FREQ, sf=LORA_SF, bw=LORA_BW)
    
    def update_neighbors(self, all_positions):
        distances = np.linalg.norm(all_positions - self.pos, axis=1)
        nearest = np.argsort(distances)[1:NEIGHBOR_COUNT+1]  # skip self
        self.neighbors = nearest
    
    def propagate(self):
        """Average lattice with neighbor lattices"""
        with self.lock:
            for n_idx in self.neighbors:
                neighbor = nodes[n_idx]
                with neighbor.lock:
                    self.lattice += 0.25 * (neighbor.lattice - self.lattice)
    
    def receive_lattice(self, lattice_bytes):
        """Receive OTA lattice update as bytes"""
        lattice = np.frombuffer(lattice_bytes, dtype=np.float32).reshape((STRANDS, SLOTS))
        self.rx_queue.append(lattice)
    
    def process_rx(self):
        """Process incoming OTA updates"""
        with self.lock:
            for lattice in self.rx_queue:
                self.lattice += 0.1 * (lattice - self.lattice)
            self.rx_queue.clear()
    
    def transmit_lattice(self):
        """Send lattice to neighbors via LoRa"""
        # Serialize lattice
        lattice_bytes = self.lattice.astype(np.float32).tobytes()
        # For real LoRa:
        # self.lora.send(lattice_bytes)
        # For simulation: deliver to neighbor objects
        for n_idx in self.neighbors:
            neighbor = nodes[n_idx]
            neighbor.receive_lattice(lattice_bytes)

# -----------------------------
# Initialize nodes
# -----------------------------
nodes = []
for i, pos in enumerate(positions):
    nodes.append(Node(i, pos))
for node in nodes:
    node.update_neighbors(positions)

# -----------------------------
# Node tick thread
# -----------------------------
def node_tick(node: Node):
    while True:
        node.process_rx()
        node.propagate()
        node.transmit_lattice()
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Start threads
# -----------------------------
for node in nodes:
    threading.Thread(target=node_tick, args=(node,), daemon=True).start()

# -----------------------------
# Global monitoring
# -----------------------------
global_tick = 0
while True:
    avg = np.mean([node.lattice.mean() for node in nodes])
    print(f"[Tick {global_tick}] Global lattice avg: {avg:.3f}")
    global_tick += 1
    time.sleep(0.5)
