import numpy as np
import threading
import time
from queue import Queue

# -----------------------------
# Config
# -----------------------------
NODE_COUNT = 16
STRANDS, SLOTS = 8, 4
TICK_INTERVAL = 0.05
NEIGHBOR_COUNT = 4  # Phyllotaxis neighbors
TX_ENABLED = False  # OTA on/off for testing

# -----------------------------
# 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 = Queue()

    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):
        """Receive OTA lattice update"""
        self.rx_queue.put(lattice)

    def process_rx(self):
        """Process incoming OTA updates"""
        while not self.rx_queue.empty():
            lattice = self.rx_queue.get()
            with self.lock:
                self.lattice += 0.1 * (lattice - self.lattice)  # weighted superposition

    def transmit_lattice(self):
        """Send lattice to neighbors OTA (stub for real RAK/HackRF)"""
        if not TX_ENABLED:
            return
        for n_idx in self.neighbors:
            neighbor = nodes[n_idx]
            neighbor.receive_lattice(self.lattice.copy())

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