# hdgl_enviro_ota.py
import threading, time, struct, numpy as np
from queue import Queue

# Optional SDR interfaces
try:
    from pyrtlsdr import RtlSdr
except ImportError:
    RtlSdr = None

try:
    import hackrf
except ImportError:
    hackrf = None

# -----------------------------
# Configuration
# -----------------------------
STRANDS, SLOTS, NODE_COUNT = 8, 4, 4
TICK_INTERVAL = 0.05       # 20 Hz lattice ticks
ALPHA = 0.3                # Smith-graph resonance coefficient
DECAY = 0.9                # Stale node decay
TX_FREQUENCY = 915e6       # Hz

# -----------------------------
# Node State
# -----------------------------
nodes = {i: {
    'lattice': np.zeros((STRANDS,SLOTS)),
    'last_tick': -1,
    'weight': 1.0
} for i in range(NODE_COUNT)}
node_lock = threading.Lock()
packet_queue = Queue()
global_tick = 0

# -----------------------------
# LoRa RX (RAK4630)
# -----------------------------
class LoRaMock:
    """Replace with actual RAK4630 LoRa interface"""
    def receive(self):
        # Mock: return None or random packet
        import random
        if random.random() < 0.2:
            node_id = random.randint(0,NODE_COUNT-1)
            lattice = np.random.rand(STRANDS*SLOTS).astype(np.float32)
            pkt = struct.pack('<B32f', node_id, *lattice)
            return pkt
        return None

def lora_rx_thread(lora, queue):
    while True:
        pkt = lora.receive()
        if not pkt: continue
        if len(pkt) != 1 + 32*4: continue
        node_id = pkt[0] % NODE_COUNT
        lattice_values = struct.unpack('<32f', pkt[1:])
        lattice = np.array(lattice_values).reshape((STRANDS,SLOTS))
        queue.put((node_id, lattice, global_tick))

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

# -----------------------------
# Smith-Graph + Weighted Failover
# -----------------------------
def smith_resonance():
    with node_lock:
        lattices = []
        weights = []
        for i, node in nodes.items():
            age = global_tick - node['last_tick']
            weight = node['weight'] * (DECAY ** max(age,0))
            lattices.append(node['lattice'] * weight)
            weights.append(weight)
        lattices = np.array(lattices)
        weights = np.array(weights)
        weighted_avg = np.sum(lattices, axis=0) / (np.sum(weights)+1e-12)
        return weighted_avg

# -----------------------------
# Lattice -> IQ Conversion
# -----------------------------
def lattice_to_iq(lattice, carrier=1.0):
    # Convert 8x4 lattice to I/Q signal for HackRF
    iq_len = 256
    i_signal = np.zeros(iq_len, dtype=np.float32)
    q_signal = np.zeros(iq_len, dtype=np.float32)
    flat = lattice.flatten()
    for n in range(min(iq_len, flat.size)):
        i_signal[n] = np.sin(2*np.pi*flat[n])
        q_signal[n] = np.cos(2*np.pi*flat[n])
    iq = i_signal + 1j*q_signal
    return iq * carrier

# -----------------------------
# HackRF TX (optional)
# -----------------------------
def hackrf_tx_thread():
    if hackrf is None:
        print("[HackRF] Module not found, skipping TX")
        return
    radio = hackrf.HackRF()
    radio.setup_tx(TX_FREQUENCY)
    while True:
        lattice = smith_resonance()
        iq = lattice_to_iq(lattice)
        radio.transmit(iq)
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Environmental Riding via RTL-SDR
# -----------------------------
def sdr_riding_thread():
    if RtlSdr is None:
        print("[RTL-SDR] Module not found, skipping environmental riding")
        return
    sdr = RtlSdr()
    sdr.sample_rate = 2.048e6
    sdr.center_freq = 915e6
    sdr.gain = 'auto'
    while True:
        samples = sdr.read_samples(256)
        # Mix with lattice for riding
        lattice = smith_resonance()
        iq = lattice_to_iq(lattice)
        modulated = iq + 0.1*samples[:len(iq)]
        # Could feed to HackRF or store
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Tick Loop
# -----------------------------
def tick_loop():
    global global_tick
    while True:
        lattice = smith_resonance()
        print(f"[Tick {global_tick}] Lattice mean: {lattice.mean():.3f}")
        global_tick += 1
        time.sleep(TICK_INTERVAL)

# -----------------------------
# Main
# -----------------------------
def main():
    lora = LoRaMock()
    threading.Thread(target=lora_rx_thread, args=(lora, packet_queue), daemon=True).start()
    threading.Thread(target=process_queue, daemon=True).start()
    threading.Thread(target=tick_loop, daemon=True).start()
    threading.Thread(target=hackrf_tx_thread, daemon=True).start()
    threading.Thread(target=sdr_riding_thread, daemon=True).start()
    
    while True:
        time.sleep(1)

if __name__ == "__main__":
    main()
