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

# -------------------------------
# Multi-Core HDGL Parameters
# -------------------------------
NUM_CORES = 2
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_BASE = 1 / (PHI**N_BASE)**7

# Initialize cores
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))})

# -------------------------------
# Networking: Multi-Hop Mesh
# -------------------------------
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 = {}  # address -> last seen timestamp

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

def listen_for_peers():
    global cores
    while True:
        try:
            data, addr = sock.recvfrom(8192)
            packet = pickle.loads(data)
            peer_cores = packet['cores']

            peers[addr[0]] = time.time()

            for c, peer_core in enumerate(peer_cores):
                core = cores[c]
                total_weight = core['weights'] + peer_core['weights']
                core['lattice'][:] = (core['lattice']*core['weights'] + peer_core['lattice']*peer_core['weights'])/total_weight
                core['phases'][:] = (core['phases']*core['weights'] + peer_core['phases']*peer_core['weights'])/total_weight
                core['weights'][:] = np.maximum(core['weights'], peer_core['weights'])
        except socket.timeout:
            continue
        except Exception as e:
            print(f"Network error: {e}")

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

# -------------------------------
# Frequency Hopping
# -------------------------------
HOP_INTERVAL = 0.1
FREQ_HOP_RANGES = [(100e3, 200e3), (200e3, 300e3), (300e3, 400e3)]

def get_hopped_frequencies():
    freqs = []
    for idx in range(TOTAL_SLOTS * NUM_CORES):
        hop_range = random.choice(FREQ_HOP_RANGES)
        freqs.append(random.uniform(*hop_range))
    return np.array(freqs)

# -------------------------------
# Generate RF block with inter-core resonance
# -------------------------------
def generate_rf_block(freqs):
    rf_block = np.zeros(BLOCK_SIZE, dtype=np.complex64)

    # Update cores
    for c, core in enumerate(cores):
        lattice = core['lattice']
        phases = core['phases']
        weights = core['weights']
        omega = core['omega']

        lattice_new = np.copy(lattice)

        # Intra-strand resonance
        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)
                intra_resonance = np.sum([omega[j]*lattice[j,i] for j in range(NUM_STRANDS) if j!=s])
                lattice_new[s,i] += 0.01*intra_resonance
                lattice_new[s,i] = lattice_new[s,i] if lattice_new[s,i] > SQRT_PHI else lattice_new[s,i]*0.8

        # -------------------------------
        # Inter-Core Resonance
        # -------------------------------
        for other_c, other_core in enumerate(cores):
            if other_c != c:
                influence = 0.005 * (other_core['lattice'] - lattice)
                lattice_new += influence

        core['lattice'][:] = lattice_new
        weights[:] = 0.9*weights + 0.1*lattice_new
        phases[:] += 0.05*lattice

        # Generate RF for this core
        for idx in range(TOTAL_SLOTS):
            strand = idx // SLOTS_PER_STRAND
            slot = idx % SLOTS_PER_STRAND
            carrier = np.exp(1j*(2*np.pi*freqs[c*TOTAL_SLOTS + idx]*t + phases[strand,slot]))
            amp = lattice[strand,slot] / np.max(lattice)
            rf_block += amp * carrier

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

# -------------------------------
# Multi-Hop Cleanup
# -------------------------------
def prune_peers(timeout=5):
    while True:
        now = time.time()
        stale = [addr for addr,last_seen in peers.items() if now-last_seen>timeout]
        for addr in stale:
            del peers[addr]
        time.sleep(1)

prune_thread = threading.Thread(target=prune_peers, daemon=True)
prune_thread.start()

# -------------------------------
# Main Streaming Loop
# -------------------------------
FS = 1_000_000
BLOCK_SIZE = 4096
t = np.arange(BLOCK_SIZE) / FS

try:
    print("HDGL Multi-Core Lattice with Inter-Core Resonance streaming. Ctrl+C to stop.")
    last_hop = time.time()
    freqs = get_hopped_frequencies()

    while True:
        if time.time() - last_hop > HOP_INTERVAL:
            freqs = get_hopped_frequencies()
            last_hop = time.time()

        block = generate_rf_block(freqs)
        broadcast_lattice()
        # sdr.write_samples(block)  # connect to SDR/DAC
        time.sleep(BLOCK_SIZE/FS)

except KeyboardInterrupt:
    print("Streaming stopped.")
