# hdgl_enviro_lora.py
import threading
import time
import struct
import numpy as np
from queue import Queue
from rtlsdr import RtlSdr
from hackrf import HackRF
from scipy.signal import resample

# LoRa library for RAK4630
from lora import LoRa  # Replace with your RAK4630 Python binding

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

# -----------------------------
# Shared State
# -----------------------------
nodes = {i: np.zeros((STRANDS, SLOTS)) for i in range(NODE_COUNT)}
node_lock = threading.Lock()
packet_queue = Queue()
env_signal = np.zeros(IQ_LEN)
env_lock = threading.Lock()
tick_count = 0

# -----------------------------
# LoRa Node RX Thread
# -----------------------------
def lora_rx_thread(lora: LoRa, queue: Queue):
    while True:
        try:
            pkt = lora.receive()
            if not pkt:
                time.sleep(0.01)
                continue

            # Expect 1 byte node_id + 32 floats
            if len(pkt) != 1 + 32*4:
                print("[WARN] Packet length mismatch")
                continue

            node_id = pkt[0] % NODE_COUNT
            lattice_values = struct.unpack('<32f', pkt[1:])
            lattice = np.array(lattice_values).reshape((STRANDS, SLOTS))
            lattice /= np.max(np.abs(lattice)) + 1e-12  # normalize

            queue.put((node_id, lattice))
        except Exception as e:
            print("[LoRa RX Error]", e)

# -----------------------------
# Queue Processor
# -----------------------------
def process_queue(queue: Queue):
    global nodes
    while True:
        try:
            node_id, lattice = queue.get()
            with node_lock:
                nodes[node_id] = lattice
        except Exception as e:
            print("[Queue Processor Error]", e)

# -----------------------------
# Environmental RX Thread
# -----------------------------
def env_rx():
    global env_signal
    sdr = RtlSdr()
    sdr.sample_rate = SAMPLE_RATE
    sdr.center_freq = TX_FREQ
    sdr.gain = 'auto'
    while True:
        try:
            samples = sdr.read_samples(IQ_LEN)
            with env_lock:
                env_signal = samples.astype(np.complex64)
        except Exception as e:
            print("[Env RX Error]", e)
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Smith-graph resonance
# -----------------------------
def smith_resonance():
    with node_lock:
        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 /= max(len(node_list)-1, 1)
            blended_nodes[i] = (1-ALPHA)*lattice + ALPHA*resonance
        # Update nodes
        for i in blended_nodes:
            nodes[i] = blended_nodes[i]

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

# -----------------------------
# Lattice -> IQ
# -----------------------------
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)) + 1e-12
    return sig.astype(np.complex64)

# -----------------------------
# HackRF TX Thread
# -----------------------------
def tx_loop():
    global tick_count
    hackrf = HackRF()
    hackrf.setup()
    hackrf.sample_rate = SAMPLE_RATE
    hackrf.center_freq = TX_FREQ
    print("[+] HackRF TX initialized...")
    while True:
        smith_resonance()
        agg = aggregate_lattice()
        with env_lock:
            carrier = env_signal.copy()
        iq = lattice_to_iq(agg, carrier)
        hackrf.tx(iq.tobytes())
        tick_count += 1
        print(f"[Tick {tick_count}] Aggregated Lattice mean: {agg.mean():.3f}")
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Main
# -----------------------------
def main():
    lora = LoRa()
    lora.setup(frequency=915e6, sf=7, bw=125e3)

    threading.Thread(target=lora_rx_thread, args=(lora, packet_queue), daemon=True).start()
    threading.Thread(target=process_queue, args=(packet_queue,), daemon=True).start()
    threading.Thread(target=env_rx, daemon=True).start()
    threading.Thread(target=tx_loop, daemon=True).start()

    while True:
        time.sleep(1)

if __name__ == "__main__":
    main()
