import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import hilbert

# -------------------------------
# Simulation Parameters
# -------------------------------
NUM_NODES = 4
NUM_CORES = 2
NUM_STRANDS = 4
SLOTS_PER_STRAND = 4

FS = 2000           # Sampling rate (Hz)
T = 2.0             # Simulation duration (s)
t = np.arange(0, T, 1/FS)

CARRIER_FREQ = 5    # Environmental carrier (Hz)
CARRIER_AMP = 1.0

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

SIDE_AMPLITUDE = 0.05  # Sideband fraction (non-interfering)

# -------------------------------
# Generate Carrier Signal
# -------------------------------
carrier = CARRIER_AMP * np.sin(2*np.pi*CARRIER_FREQ*t)

# -------------------------------
# Initialize HDGL Lattice Nodes
# -------------------------------
def init_lattice():
    cores = []
    for _ 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))
        cores.append({'lattice': lattice, 'phases': phases})
    return cores

nodes = [init_lattice() for _ in range(NUM_NODES)]

# -------------------------------
# Continuous Lattice Evolution
# -------------------------------
def evolve_lattice(node_cores, carrier_signal, dt=0.001):
    """Phase-lock each lattice slot to carrier and apply small sideband."""
    sideband = np.zeros_like(carrier_signal)
    analytic_carrier = np.angle(hilbert(carrier_signal))  # Instantaneous carrier phase
    for core in node_cores:
        for s in range(NUM_STRANDS):
            for slot in range(SLOTS_PER_STRAND):
                # Target phase: carrier + lattice encoding offset
                phase_target = analytic_carrier + 0.2 * core['lattice'][s, slot]
                # Smooth update (continuous analog evolution)
                core['phases'][s, slot] += 0.1 * (phase_target - core['phases'][s, slot])
                # Generate sideband contribution
                sideband += np.sin(2*np.pi*CARRIER_FREQ*t + core['phases'][s, slot])
    sideband /= (NUM_CORES * NUM_STRANDS * SLOTS_PER_STRAND)
    return SIDE_AMPLITUDE * sideband

# -------------------------------
# Aggregate Multi-Node Sideband
# -------------------------------
composite_signal = carrier.copy()
for node_id, node_cores in enumerate(nodes):
    composite_signal += evolve_lattice(node_cores, carrier)

# Normalize to preserve carrier amplitude
composite_signal /= (1 + NUM_NODES * SIDE_AMPLITUDE)

# -------------------------------
# Extract HDGL Sideband (Receiver Simulation)
# -------------------------------
def extract_sideband(signal, carrier_freq):
    analytic = hilbert(signal)
    inst_phase = np.unwrap(np.angle(analytic))
    inst_amplitude = np.abs(analytic)
    return inst_amplitude, inst_phase

amplitude, phase = extract_sideband(composite_signal, CARRIER_FREQ)

# -------------------------------
# Plot Results
# -------------------------------
plt.figure(figsize=(14,6))
plt.plot(t, carrier, label='Original Carrier', alpha=0.6)
plt.plot(t, composite_signal, label='Carrier + HDGL Sideband', alpha=0.9)
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.title('HDGL Lattice Riding Live Carrier (Analog Sideband)')
plt.legend()
plt.show()
