import numpy as np
import time
import socket
import pickle
import threading

# -------------------------------
# HDGL Lattice Parameters
# -------------------------------
NUM_STRANDS = 8
SLOTS_PER_STRAND = 4
TOTAL_SLOTS = NUM_STRANDS * SLOTS_PER_STRAND

PHI = 1.6180339887
SQRT_PHI = np.sqrt(PHI)

N_BASE = np.arange(1, NUM_STRANDS+1)
OMEGA = 1 / (PHI**N_BASE)**7

# Initialize lattice and phase matrices
lattice = np.random.uniform(0.5, 1.0, (NUM_STRANDS, SLOTS_PER_STRAND))
phases = np.random.uniform(0, 2*np.pi, (NUM_STRANDS, SLOTS_PER_STRAND))

# -------------------------------
# RF Output Parameters
# -------------------------------
FS = 1_000_000          # Sampling rate (Hz)
BLOCK_SIZE = 4096
t = np.arange(BLOCK_SIZE) / FS
SLOT_FREQS = np.linspace(100e3, 400e3, TOTAL_SLOTS)

# -------------------------------
# Networking: Multi-Device Sync
# -------------------------------
UDP_PORT = 5005
BROADCAST_IP = "255.255.255.255"
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
sock.bind(("", UDP_PORT))
sock.settimeout(0.001)

peers = []

def broadcast_lattice():
    data = pickle.dumps({'lattice': lattice, 'phases': phases})
    sock.sendto(data, (BROADCAST_IP, UDP_PORT))

def listen_for_peers():
    global lattice, phases
    while True:
        try:
            data, addr = sock.recvfrom(8192)
            packet = pickle.loads(data)
            peer_lattice = packet['lattice']
            peer_phases = packet['phases']
            # Simple lattice fusion: average with peer
            lattice[:] = (lattice + peer_lattice)/2
            phases[:] = (phases + peer_phases)/2
            if addr[0] not in peers:
                peers.append(addr[0])
        except socket.timeout:
            continue
        except Exception as e:
            print(f"Network error: {e}")

# Start listener in background thread
listener_thread = threading.Thread(target=listen_for_peers, daemon=True)
listener_thread.start()

# -------------------------------
# Generate lattice RF block
# -------------------------------
def generate_rf_block():
    global lattice, phases
    # Update lattice
    lattice_new = np.copy(lattice)
    for s in range(NUM_STRANDS):
        for i in range(SLOTS_PER_STRAND):
            lattice_new[s,i] += 0.02*OMEGA[s]*(1+0.05*i)
            resonance = np.sum([OMEGA[j]*lattice[j,i] for j in range(NUM_STRANDS) if j!=s])
            lattice_new[s,i] += 0.01*resonance
            lattice_new[s,i] = lattice_new[s,i] if lattice_new[s,i] > SQRT_PHI else lattice_new[s,i]*0.8
            phases[s,i] += 0.05*lattice[s,i]
    lattice[:] = lattice_new

    # Generate multi-slot RF
    rf_block = np.zeros(BLOCK_SIZE, dtype=np.complex64)
    for idx in range(TOTAL_SLOTS):
        strand = idx // SLOTS_PER_STRAND
        slot = idx % SLOTS_PER_STRAND
        carrier = np.exp(1j*(2*np.pi*SLOT_FREQS[idx]*t + phases[strand,slot]))
        amp = lattice[strand,slot] / np.max(lattice)
        rf_block += amp * carrier

    # Normalize
    rf_block /= np.max(np.abs(rf_block))
    return rf_block

# -------------------------------
# Main Streaming Loop
# -------------------------------
try:
    print("HDGL Multi-Device RF Lattice streaming. Ctrl+C to stop.")
    while True:
        block = generate_rf_block()
        broadcast_lattice()  # share current lattice with peers
        # sdr.write_samples(block)  # Replace with SDR/DAC API
        time.sleep(BLOCK_SIZE/FS)

except KeyboardInterrupt:
    print("Streaming stopped.")
