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

# -----------------------------
# Configuration
# -----------------------------
NODE_COUNT = 16           # Total lattice nodes
STRANDS, SLOTS = 8, 4
TICK_INTERVAL = 0.05
DECAY = 0.9
TX_FREQUENCY = 915e6

# -----------------------------
# Node State
# -----------------------------
nodes = {}
node_lock = threading.Lock()
packet_queue = Queue()
global_tick = 0

# -----------------------------
# Phyllotaxis Node Positioning
# -----------------------------
def phyllotaxis_positions(n, c=1.0):
    """Returns n (x,y) positions in a sunflower spiral"""
    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)

# Initialize nodes with lattice and positions
positions = phyllotaxis_positions(NODE_COUNT, c=1.0)
for i, pos in enumerate(positions):
    nodes[i] = {'lattice': np.zeros((STRANDS, SLOTS)),
                'last_tick': -1,
                'weight': 1.0,
                'position': pos}

# -----------------------------
# Mock OTA Packet Reception
# -----------------------------
class LoRaRAK:
    def receive(self):
        import random
        if random.random() < 0.2:
            node_id = random.randint(0, NODE_COUNT-1)
            lattice = np.random.rand(STRANDS*SLOTS).astype(np.float32)
            return node_id, lattice
        return None, None

def lora_rx_thread(lora, queue):
    global global_tick
    while True:
        node_id, lattice_flat = lora.receive()
        if node_id is not None:
            lattice = lattice_flat.reshape((STRANDS, SLOTS))
            queue.put((node_id, lattice, global_tick))
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Queue Processor
# -----------------------------
def process_queue():
    while True:
        try:
            node_id, lattice, tick = packet_queue.get()
            with node_lock:
                nodes[node_id]['lattice'] = lattice
                nodes[node_id]['last_tick'] = tick
                nodes[node_id]['weight'] = 1.0
        except Exception as e:
            print("[Queue Error]", e)

# -----------------------------
# Phyllotaxis-weighted Aggregation
# -----------------------------
def aggregate_lattice():
    with node_lock:
        lattices, weights = [], []
        for node in nodes.values():
            age = global_tick - node['last_tick']
            w = node['weight'] * (DECAY ** max(age,0))
            lattices.append(node['lattice'] * w)
            weights.append(w)
        lattices = np.array(lattices)
        weights = np.array(weights)
        # Optional: apply spatial weighting based on phyllotaxis distance
        spatial_weights = np.linalg.norm(positions, axis=1)
        combined_weights = weights / (1 + spatial_weights)
        weighted_avg = np.sum(lattices * combined_weights[:, None, None], axis=0) / (np.sum(combined_weights)+1e-12)
        return weighted_avg

# -----------------------------
# Tick Loop
# -----------------------------
def tick_loop():
    global global_tick
    while True:
        lattice = aggregate_lattice()
        print(f"[Tick {global_tick}] Lattice mean: {lattice.mean():.3f}")
        global_tick += 1
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Main
# -----------------------------
def main():
    lora = LoRaRAK()
    threading.Thread(target=lora_rx_thread, args=(lora, packet_queue), daemon=True).start()
    threading.Thread(target=process_queue, daemon=True).start()
    threading.Thread(target=tick_loop, daemon=True).start()
    
    while True:
        time.sleep(1)

if __name__ == "__main__":
    main()
