import numpy as np
import time
import struct
from pylora import LoRa
import rtlsdr  # Or HackRF/LimeSDR library
from scipy.signal import resample

# -----------------------------
# Config
# -----------------------------
STRANDS, SLOTS, NODE_COUNT = 8, 4, 4
ALPHA = 0.3
RIDE_FACTOR = 0.5
SAMPLE_RATE = 2.048e6
IQ_LEN = 2048

# -----------------------------
# LoRa setup
# -----------------------------
lora = LoRa()
lora.set_mode(LoRa.RX)
lora.set_frequency(868e6)
lora.set_sf(7)
lora.set_bw(125e3)
lora.set_cr(5)

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

# -----------------------------
# Decode LoRa packet
# -----------------------------
def decode_packet(packet):
    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)
        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))
    return sig.astype(np.complex64)

# -----------------------------
# Placeholder for environmental RF input
# -----------------------------
def get_environmental_signal(length=IQ_LEN):
    t = np.arange(length) / SAMPLE_RATE
    # Replace with SDR RX samples for live riding
    return 0.3*np.sin(2*np.pi*50e3*t) + 0.2*np.sin(2*np.pi*75e3*t)

# -----------------------------
# Placeholder SDR transmit
# -----------------------------
def transmit_iq(iq_samples):
    # Replace with HackRF/LimeSDR TX code
    pass

# -----------------------------
# Main loop
# -----------------------------
print("[+] OTA HDGL lattice aggregator running...")
while True:
    if lora.packet_available():
        pkt = lora.receive_packet()
        node_id = pkt[0] % NODE_COUNT
        lattice, tick = decode_packet(pkt)
        nodes[node_id] = lattice

    nodes = smith_resonance(nodes)
    agg_lattice = aggregate_lattice(nodes)
    env_sig = get_environmental_signal()
    iq = lattice_to_iq(agg_lattice, carrier=env_sig)

    transmit_iq(iq)
    time.sleep(0.05)
