import numpy as np
import time
import struct
from rtlsdr import RtlSdr
from hackrf import HackRF  # pip install pyhackrf
from scipy.signal import resample

# -----------------------------
# Config
# -----------------------------
STRANDS, SLOTS, NODE_COUNT = 8, 4, 4
ALPHA = 0.3        # Smith-graph resonance factor
RIDE_FACTOR = 0.5  # Environmental riding
SAMPLE_RATE = 2.048e6
IQ_LEN = 2048
TX_FREQ = 915e6

# -----------------------------
# Node lattices
# -----------------------------
nodes = {i: np.zeros((STRANDS, SLOTS)) for i in range(NODE_COUNT)}

# -----------------------------
# Decode Node Packet (LoRa/RX)
# -----------------------------
def decode_packet(packet):
    # packet format: STRANDS*SLOTS bytes + tick byte
    flat = list(struct.unpack("B"*STRANDS*SLOTS + "B", packet))
    tick = flat[-1]
    lattice = np.array(flat[:-1]).reshape((STRANDS, SLOTS)) / 50.0
    return lattice, tick

# -----------------------------
# Smith-graph resonance
# -----------------------------
def smith_resonance(nodes):
    node_list = list(nodes.values())
    blended_nodes = {}
    for i, lattice in enumerate(node_list):
        resonance = np.zeros_like(lattice)
        for j, other in enumerate(node_list):
            if i == j: continue
            resonance += other
        resonance /= (len(node_list)-1)
        blended_nodes[i] = (1-ALPHA)*lattice + ALPHA*resonance
    return blended_nodes

# -----------------------------
# Lattice aggregation
# -----------------------------
def aggregate_lattice(nodes):
    lattices = np.array(list(nodes.values()))
    return np.mean(lattices, axis=0)

# -----------------------------
# Lattice -> IQ with environmental riding
# -----------------------------
def lattice_to_iq(lattice, carrier=None, length=IQ_LEN):
    t = np.arange(length) / SAMPLE_RATE
    sig = np.zeros(length)
    for s in range(STRANDS):
        weight = np.mean(lattice[s])
        freq = 1e3*(s+1)  # each strand is its own carrier harmonic
        sig += weight * np.sin(2*np.pi*freq*t)

    if carrier is not None:
        if len(carrier) != length:
            carrier = resample(carrier, length)
        sig = (1-RIDE_FACTOR)*carrier + RIDE_FACTOR*sig

    sig /= np.max(np.abs(sig)) + 1e-12
    return sig.astype(np.complex64)

# -----------------------------
# Environmental signal sampling
# -----------------------------
def get_environmental_signal(length=IQ_LEN):
    sdr = RtlSdr()
    sdr.sample_rate = SAMPLE_RATE
    sdr.center_freq = TX_FREQ
    sdr.gain = 'auto'
    samples = sdr.read_samples(length)
    sdr.close()
    return samples

# -----------------------------
# HackRF TX
# -----------------------------
def transmit_iq(iq_samples):
    hackrf = HackRF()
    hackrf.setup()
    hackrf.sample_rate = SAMPLE_RATE
    hackrf.center_freq = TX_FREQ
    hackrf.tx(iq_samples.tobytes())  # transmit buffer
    hackrf.close()

# -----------------------------
# Main OTA Loop
# -----------------------------
print("[+] OTA HDGL lattice aggregator starting...")
while True:
    # -- Receive Node Packets Here --
    # TODO: Replace with LoRa RX code to fill `nodes` dict

    nodes = smith_resonance(nodes)
    agg_lattice = aggregate_lattice(nodes)

    # Sample environmental RF
    try:
        env_sig = get_environmental_signal()
    except:
        env_sig = None

    # Convert to IQ
    iq = lattice_to_iq(agg_lattice, carrier=env_sig)

    # Transmit OTA
    transmit_iq(iq)

    time.sleep(0.05)
