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

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

CARRIER_FREQ = 50   # Hz, demo carrier
CARRIER_AMP = 1.0

NUM_NODES = 4
NUM_CORES = 2
NUM_STRANDS = 4
SLOTS_PER_STRAND = 4
SIDE_AMPLITUDE = 0.05  # Sideband fraction

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

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

# -------------------------------
# Initialize 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 are now just scalar offsets per slot
        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 & Modulation
# -------------------------------
def evolve_lattice(node_cores, carrier_signal):
    """
    Phase-lock lattice slots to carrier with small sideband.
    Fixed: do not update scalar phase with array.
    """
    sideband = np.zeros_like(carrier_signal)
    for core in node_cores:
        for s in range(NUM_STRANDS):
            for slot in range(SLOTS_PER_STRAND):
                phase_offset = core['phases'][s, slot] + 0.2 * core['lattice'][s, slot]
                # Each slot contributes a continuous sine wave with phase offset
                sideband += np.sin(2*np.pi*CARRIER_FREQ*t + phase_offset)
    sideband /= (NUM_CORES * NUM_STRANDS * SLOTS_PER_STRAND)
    return SIDE_AMPLITUDE * sideband

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

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

# -------------------------------
# Receiver: Extract HDGL Sideband
# -------------------------------
def decode_lattice(signal):
    """Extract instantaneous amplitude and phase (Hilbert transform)."""
    analytic = hilbert(signal)
    amplitude = np.abs(analytic)
    phase = np.unwrap(np.angle(analytic))
    return amplitude, phase

amplitude, phase = decode_lattice(composite_signal)

# -------------------------------
# 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()
