# ==============================
# RAK4630 Lattice Mesh Node
# ==============================
# MicroPython / CircuitPython compatible
# Hardware: RAK4630 with LoRa
# ==============================

import time
import math
import struct
import urandom as random  # MicroPython random
from machine import Timer
from network import LoRa

# -----------------------------
# Config
# -----------------------------
NODE_ID = 0
STRANDS, SLOTS = 8, 4
TICK_INTERVAL = 0.05  # 50 ms tick
MOD_INDEX = 0.05      # small FM modulation
LORA_FREQ = 915e6

# -----------------------------
# Initialize LoRa
# -----------------------------
lora = LoRa(mode=LoRa.LORA, freq=LORA_FREQ, sf=7, bw=125e3)
lora_sock = lora.socket()
lora_sock.setblocking(False)

# -----------------------------
# Lattice initialization
# -----------------------------
lattice = [[0.5 for _ in range(SLOTS)] for _ in range(STRANDS)]

# -----------------------------
# Phyllotaxis neighbors
# -----------------------------
NODE_COUNT = 16
NEIGHBORS = 4  # connect to 4 nearest nodes

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

positions = phyllotaxis_positions(NODE_COUNT)

def compute_neighbors(node_idx):
    px, py = positions[node_idx]
    distances = []
    for i, (x, y) in enumerate(positions):
        if i == node_idx:
            continue
        d = math.sqrt((px - x)**2 + (py - y)**2)
        distances.append((d, i))
    distances.sort()
    return [idx for _, idx in distances[:NEIGHBORS]]

neighbors = compute_neighbors(NODE_ID)

# -----------------------------
# Lattice → FM modulation
# -----------------------------
def lattice_to_waveform(lattice):
    # Flatten lattice
    flat = [val for strand in lattice for val in strand]
    # Simple FM: sum sine deviations
    waveform = []
    for i, val in enumerate(flat):
        phase = 2 * math.pi * (val * MOD_INDEX)
        waveform.append(math.sin(phase))
    return waveform

# -----------------------------
# Transmit lattice
# -----------------------------
def transmit_lattice():
    # Serialize lattice as bytes (float32)
    flat = [val for strand in lattice for val in strand]
    data = b''.join([struct.pack('<f', v) for v in flat])
    for n_idx in neighbors:
        # In real network: send with LoRa to neighbor
        lora_sock.send(data)  # broadcast, neighbors receive

# -----------------------------
# Receive lattice
# -----------------------------
def receive_lattice():
    try:
        data = lora_sock.recv(256)
        if data:
            # Deserialize
            new_vals = [struct.unpack('<f', data[i*4:i*4+4])[0] for i in range(STRANDS*SLOTS)]
            # Mix with current lattice (superposition)
            for s in range(STRANDS):
                for slot in range(SLOTS):
                    idx = s*SLOTS + slot
                    lattice[s][slot] += 0.1 * (new_vals[idx] - lattice[s][slot])
    except:
        pass

# -----------------------------
# Node tick
# -----------------------------
def node_tick(timer):
    # Random small fluctuation
    for s in range(STRANDS):
        for slot in range(SLOTS):
            lattice[s][slot] += 0.01 * (random.getrandbits(1)*2 - 1)
            lattice[s][slot] = max(0.0, min(1.0, lattice[s][slot]))
    # Transmit to neighbors
    transmit_lattice()
    # Receive updates
    receive_lattice()
    # Optional: waveform generation (DAC) for actual analog output
    waveform = lattice_to_waveform(lattice)
    # DAC output can go here (hardware-specific)
    print(f"Lattice avg: {sum([sum(strand) for strand in lattice])/ (STRANDS*SLOTS):.3f}")

# -----------------------------
# Start periodic timer
# -----------------------------
timer = Timer.Alarm(node_tick, TICK_INTERVAL, periodic=True)

# -----------------------------
# Keep alive
# -----------------------------
while True:
    time.sleep(1)
