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

# -------------------------------
# 3D Holographic HDGL Parameters
# -------------------------------
NUM_NODES = 8           # distributed nodes
NUM_CORES = 2           # cores per node
NUM_STRANDS = 8
SLOTS_PER_STRAND = 4
TOTAL_SLOTS = NUM_STRANDS * SLOTS_PER_STRAND

PHI = 1.6180339887
OMEGA_BASE = 1 / (PHI**np.arange(1, NUM_STRANDS+1))**7

FS = 1_000_000           # sample rate
BLOCK_SIZE = 4096
t = np.arange(BLOCK_SIZE) / FS

# Define holographic target grid (x, y, z)
GRID_RES = 3  # points per axis (3x3x3 = 27 points)
TARGET_GRID = np.array([[x, y, z] 
                        for x in np.linspace(-5,5,GRID_RES)
                        for y in np.linspace(-5,5,GRID_RES)
                        for z in np.linspace(0,10,GRID_RES)])

# Node positions
NODE_POSITIONS = np.array([[random.uniform(-50,50),
                            random.uniform(-50,50),
                            random.uniform(0,20)]
                           for _ in range(NUM_NODES)])

C_LIGHT = 3e8  # speed of light

# -------------------------------
# Node Lattice Initialization
# -------------------------------
def init_node():
    cores = []
    for c in range(NUM_CORES):
        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))
        weights = np.ones((NUM_STRANDS, SLOTS_PER_STRAND))
        cores.append({'lattice': lattice, 'phases': phases, 'weights': weights, 'omega': OMEGA_BASE*(1/(c+1))})
    return cores

node_lattices = [init_node() for _ in range(NUM_NODES)]

# -------------------------------
# Network Simulation (UDP)
# -------------------------------
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_node(node_id, cores, gps_time):
    packet = {'node_id': node_id, 'cores': cores, 'timestamp': gps_time, 'pos': NODE_POSITIONS[node_id]}
    sock.sendto(pickle.dumps(packet), (BROADCAST_IP, UDP_PORT))

def listen_for_peers():
    while True:
        try:
            data, addr = sock.recvfrom(16384)
            packet = pickle.loads(data)
            peers[packet['node_id']] = packet
        except socket.timeout:
            continue
        except Exception as e:
            print(f"Network error: {e}")

threading.Thread(target=listen_for_peers, daemon=True).start()

# -------------------------------
# 3D Phase Compensation
# -------------------------------
def compensate_for_hologram(node_id, cores, target_points):
    node_pos = NODE_POSITIONS[node_id]
    for core in cores:
        for idx, point in enumerate(target_points):
            r = np.linalg.norm(node_pos - point)
            delay = r / C_LIGHT
            phase_shift = 2*np.pi * FS * delay
            core['phases'] = (core['phases'] + phase_shift / len(target_points)) % (2*np.pi)

# -------------------------------
# Generate 3D Holographic RF Block
# -------------------------------
def generate_rf_block(target_points):
    rf_block = np.zeros(BLOCK_SIZE, dtype=np.complex64)
    for node_id, node in enumerate(node_lattices):
        compensate_for_hologram(node_id, node, target_points)
        for core in node:
            lattice = core['lattice']
            phases = core['phases']
            omega = core['omega']
            lattice += 0.01 * omega[:,None]
            for idx in range(TOTAL_SLOTS):
                strand = idx // SLOTS_PER_STRAND
                slot = idx % SLOTS_PER_STRAND
                amp = lattice[strand, slot] / np.max(lattice)
                phi = phases[strand, slot]
                freq_offset = 50e3 * (lattice[strand, slot]-0.5)
                carrier = np.exp(1j*(2*np.pi*(freq_offset)*t + phi))
                rf_block += amp * carrier
    return rf_block / np.max(np.abs(rf_block))

# -------------------------------
# Main Loop
# -------------------------------
try:
    print("3D Holographic HDGL RF Transmission Started.")
    while True:
        gps_time = time.time()
        rf_block = generate_rf_block(TARGET_GRID)
        for node_id, node in enumerate(node_lattices):
            broadcast_node(node_id, node, gps_time)
        # SDR hardware output: sdr.write_samples(rf_block)
        time.sleep(BLOCK_SIZE/FS)

except KeyboardInterrupt:
    print("3D Holographic HDGL RF Transmission Stopped.")
