#!/usr/bin/env python3
"""
LoRa Mesh Attack Detector — Synthetic Data Generator
Generates 30k samples across 5 classes for detecting jamming, replay, and spoofing
on LoRa mesh networks running on Heltec V4 nodes.

Feature vector (32 elements):
  [0-3]   Packet rate features (rx_rate, tx_rate, ack_rate, beacon_rate) — packets/sec
  [4-7]   Error features (crc_error_rate, timeout_rate, dropped_ratio, retry_ratio)
  [8-11]  RSSI channel 0-3 (mean RSSI per channel, 4 LoRa channels)
  [12-15] RSSI statistics (overall_mean, overall_std, min_rssi, max_rssi)
  [16-19] Duplicate features (dup_pkt_id_ratio, dup_timestamp_ratio, dup_payload_hash_ratio, unique_id_rate)
  [20-23] Timing features (iat_mean, iat_std, iat_regularity_score, burst_density)
  [24-27] Node identity features (known_node_ratio, unexpected_rssi_direction, node_id_diversity, id_conflict_count)
  [28-31] Channel features (active_channel_count, channel_hop_rate, noise_floor_mean, snr_mean)

Classes: {0: normal, 1: jam_constant, 2: jam_sweep, 3: replay, 4: spoof}
"""

import numpy as np
import os
import argparse

SEED = 42
NUM_FEATURES = 32
CLASSES = ['normal', 'jam_constant', 'jam_sweep', 'replay', 'spoof']
NUM_CLASSES = len(CLASSES)


def generate_normal(n: int, rng: np.random.Generator) -> np.ndarray:
    X = np.zeros((n, NUM_FEATURES), dtype=np.float32)

    # Packet rates
    X[:, 0] = rng.normal(2.0, 0.6, n).clip(0.5, 4.0)       # rx rate
    X[:, 1] = rng.normal(1.5, 0.5, n).clip(0.3, 3.0)       # tx rate
    X[:, 2] = rng.normal(1.3, 0.4, n).clip(0.3, 2.5)       # ack rate
    X[:, 3] = rng.normal(0.1, 0.03, n).clip(0.03, 0.2)     # beacon rate

    # Error features (low in normal)
    X[:, 4] = rng.normal(0.02, 0.01, n).clip(0, 0.05)       # crc error rate
    X[:, 5] = rng.normal(0.03, 0.015, n).clip(0, 0.07)      # timeout rate
    X[:, 6] = rng.normal(0.01, 0.008, n).clip(0, 0.03)      # dropped ratio
    X[:, 7] = rng.normal(0.05, 0.025, n).clip(0, 0.12)      # retry ratio

    # RSSI per channel (typical LoRa ranges)
    X[:, 8] = rng.normal(-90, 8, n).clip(-120, -60)          # ch0
    X[:, 9] = rng.normal(-88, 8, n).clip(-120, -60)          # ch1
    X[:, 10] = rng.normal(-92, 8, n).clip(-120, -60)         # ch2
    X[:, 11] = rng.normal(-89, 8, n).clip(-120, -60)         # ch3

    # RSSI statistics
    X[:, 12] = rng.normal(-90, 6, n).clip(-115, -65)         # overall mean
    X[:, 13] = rng.normal(5, 2, n).clip(1, 12)               # overall std
    X[:, 14] = rng.normal(-110, 5, n).clip(-125, -95)        # min rssi
    X[:, 15] = rng.normal(-70, 5, n).clip(-85, -55)          # max rssi

    # Duplicate features (very low in normal)
    X[:, 16] = rng.normal(0.005, 0.004, n).clip(0, 0.02)     # dup pkt id ratio
    X[:, 17] = rng.normal(0.003, 0.003, n).clip(0, 0.01)     # dup timestamp ratio
    X[:, 18] = rng.normal(0.002, 0.002, n).clip(0, 0.008)    # dup payload hash ratio
    X[:, 19] = rng.normal(0.95, 0.03, n).clip(0.85, 1.0)     # unique id rate

    # Timing features
    X[:, 20] = rng.normal(0.5, 0.15, n).clip(0.15, 1.0)      # iat mean
    X[:, 21] = rng.normal(0.1, 0.04, n).clip(0.02, 0.25)     # iat std
    X[:, 22] = rng.normal(0.85, 0.06, n).clip(0.65, 0.95)    # regularity score
    X[:, 23] = rng.normal(0.08, 0.04, n).clip(0, 0.2)        # burst density

    # Node identity
    X[:, 24] = rng.normal(0.95, 0.03, n).clip(0.85, 1.0)     # known node ratio
    X[:, 25] = rng.normal(0.05, 0.03, n).clip(0, 0.15)       # unexpected rssi direction
    X[:, 26] = rng.normal(4, 1.5, n).clip(2, 8)              # node id diversity
    X[:, 27] = rng.normal(0, 0.01, n).clip(0, 0.03)          # id conflict count

    # Channel features
    X[:, 28] = rng.normal(3, 0.5, n).clip(2, 4)              # active channels
    X[:, 29] = rng.normal(0.5, 0.2, n).clip(0.1, 1.0)        # channel hop rate
    X[:, 30] = rng.normal(-115, 3, n).clip(-125, -108)        # noise floor
    X[:, 31] = rng.normal(12, 3, n).clip(5, 20)              # SNR

    return X


def generate_jam_constant(n: int, rng: np.random.Generator) -> np.ndarray:
    """Constant carrier jamming — RSSI spike everywhere, zero valid packets."""
    X = generate_normal(n, rng)

    # RSSI spikes across all channels (jammer overwhelms)
    X[:, 8] = rng.normal(-40, 5, n).clip(-55, -25)
    X[:, 9] = rng.normal(-42, 5, n).clip(-55, -25)
    X[:, 10] = rng.normal(-41, 5, n).clip(-55, -25)
    X[:, 11] = rng.normal(-40, 5, n).clip(-55, -25)

    X[:, 12] = rng.normal(-41, 4, n).clip(-52, -28)          # overall mean spikes
    X[:, 13] = rng.normal(2, 1, n).clip(0.5, 5)              # std LOW (constant)
    X[:, 15] = rng.normal(-30, 5, n).clip(-45, -20)          # max rssi very high

    # Packet rates collapse
    X[:, 0] = rng.normal(0.05, 0.04, n).clip(0, 0.15)       # rx near zero
    X[:, 1] = rng.normal(1.5, 0.5, n).clip(0.5, 3.0)        # tx unchanged (our side)
    X[:, 2] = rng.normal(0.02, 0.02, n).clip(0, 0.08)       # ack near zero
    X[:, 3] = rng.normal(0.0, 0.005, n).clip(0, 0.02)       # beacon gone

    # Error rates spike
    X[:, 4] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)       # crc errors
    X[:, 5] = rng.normal(0.7, 0.12, n).clip(0.4, 0.95)      # timeouts
    X[:, 6] = rng.normal(0.8, 0.1, n).clip(0.55, 0.98)      # dropped
    X[:, 7] = rng.normal(0.9, 0.05, n).clip(0.75, 0.99)     # retries

    # Noise floor rises
    X[:, 30] = rng.normal(-60, 5, n).clip(-75, -45)          # noise floor way up
    X[:, 31] = rng.normal(-2, 2, n).clip(-8, 3)              # SNR negative or near zero

    return X


def generate_jam_sweep(n: int, rng: np.random.Generator) -> np.ndarray:
    """Sweep jamming — RSSI spikes rotate across channels."""
    X = generate_normal(n, rng)

    # One channel at a time has spike, others partially affected
    # We simulate the aggregate: high variance across channels
    base = rng.normal(-85, 5, (n, 4))
    # Random channel gets jammed per sample
    jammed_ch = rng.integers(0, 4, n)
    for i in range(n):
        base[i, jammed_ch[i]] = rng.normal(-35, 5)
    X[:, 8:12] = base.clip(-120, -20)

    X[:, 13] = rng.normal(18, 5, n).clip(8, 35)             # RSSI std HIGH (rotating)
    X[:, 12] = rng.normal(-70, 8, n).clip(-95, -45)         # overall mean elevated

    # Partial packet loss (not total like constant jam)
    X[:, 0] = rng.normal(0.6, 0.3, n).clip(0.1, 1.5)
    X[:, 2] = rng.normal(0.3, 0.2, n).clip(0.05, 0.8)

    # Error rates moderate-high
    X[:, 4] = rng.normal(0.3, 0.1, n).clip(0.1, 0.55)
    X[:, 5] = rng.normal(0.35, 0.12, n).clip(0.12, 0.6)
    X[:, 6] = rng.normal(0.4, 0.12, n).clip(0.15, 0.65)

    # Channel hopping disrupted
    X[:, 29] = rng.normal(1.5, 0.5, n).clip(0.5, 3.0)       # hop rate elevated
    X[:, 30] = rng.normal(-80, 8, n).clip(-100, -60)         # noise floor elevated
    X[:, 31] = rng.normal(4, 2, n).clip(-1, 9)               # SNR degraded

    return X


def generate_replay(n: int, rng: np.random.Generator) -> np.ndarray:
    """Replay attack — duplicate packet IDs, identical timestamps from different directions."""
    X = generate_normal(n, rng)

    # Duplicate features spike
    X[:, 16] = rng.normal(0.35, 0.1, n).clip(0.15, 0.6)     # dup pkt id ratio HIGH
    X[:, 17] = rng.normal(0.4, 0.12, n).clip(0.18, 0.7)     # dup timestamp ratio HIGH
    X[:, 18] = rng.normal(0.45, 0.12, n).clip(0.2, 0.75)    # dup payload hash HIGH
    X[:, 19] = rng.normal(0.55, 0.1, n).clip(0.3, 0.78)     # unique id rate drops

    # Packet rate slightly elevated (replays add traffic)
    X[:, 0] = rng.normal(3.5, 1.0, n).clip(1.5, 6.0)

    # Timing — replays may cluster
    X[:, 21] = rng.normal(0.2, 0.06, n).clip(0.08, 0.35)    # iat std higher
    X[:, 22] = rng.normal(0.6, 0.1, n).clip(0.35, 0.8)      # regularity drops

    # RSSI from unexpected direction (replayer is elsewhere)
    X[:, 25] = rng.normal(0.4, 0.12, n).clip(0.15, 0.65)    # unexpected direction HIGH

    return X


def generate_spoof(n: int, rng: np.random.Generator) -> np.ndarray:
    """Node spoofing — known node ID but unexpected RSSI/direction."""
    X = generate_normal(n, rng)

    # Node identity anomalies
    X[:, 24] = rng.normal(0.85, 0.05, n).clip(0.7, 0.95)    # known ratio still high (spoofed IDs)
    X[:, 25] = rng.normal(0.55, 0.12, n).clip(0.3, 0.8)     # unexpected direction HIGH
    X[:, 27] = rng.normal(3, 1.5, n).clip(1, 8)              # id conflict count HIGH

    # RSSI inconsistency for known nodes
    X[:, 13] = rng.normal(12, 4, n).clip(5, 25)              # overall RSSI std higher
    X[:, 12] = rng.normal(-80, 10, n).clip(-110, -55)        # mean shifted

    # Some duplicates (but not as extreme as replay)
    X[:, 16] = rng.normal(0.08, 0.04, n).clip(0.02, 0.2)
    X[:, 19] = rng.normal(0.88, 0.05, n).clip(0.75, 0.97)

    # Timing may be slightly off
    X[:, 22] = rng.normal(0.7, 0.08, n).clip(0.5, 0.88)

    return X


GENERATORS = [generate_normal, generate_jam_constant, generate_jam_sweep, generate_replay, generate_spoof]


def main():
    parser = argparse.ArgumentParser(description='Generate synthetic LoRa mesh attack data')
    parser.add_argument('--total', type=int, default=30000)
    parser.add_argument('--test-frac', type=float, default=0.2)
    parser.add_argument('--seed', type=int, default=SEED)
    parser.add_argument('--output-dir', type=str, default=os.path.dirname(os.path.abspath(__file__)))
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    samples_per_class = args.total // NUM_CLASSES
    remainder = args.total - samples_per_class * NUM_CLASSES

    all_X, all_y = [], []
    for cls_idx, gen_fn in enumerate(GENERATORS):
        n_cls = samples_per_class + (1 if cls_idx < remainder else 0)
        all_X.append(gen_fn(n_cls, rng))
        all_y.append(np.full(n_cls, cls_idx, dtype=np.int32))
        print(f"  Class {cls_idx} ({CLASSES[cls_idx]}): {n_cls} samples")

    X = np.concatenate(all_X)
    y = np.concatenate(all_y)
    X += rng.normal(0, 0.01, X.shape).astype(np.float32)

    idx = rng.permutation(len(X))
    X, y = X[idx], y[idx]

    n_test = int(len(X) * args.test_frac)
    np.savez_compressed(os.path.join(args.output_dir, 'train.npz'), X=X[n_test:], y=y[n_test:])
    np.savez_compressed(os.path.join(args.output_dir, 'test.npz'), X=X[:n_test], y=y[:n_test])

    print(f"\nTotal: {len(X)} | Train: {len(X)-n_test} | Test: {n_test}")
    print(f"Features: {NUM_FEATURES} | Classes: {NUM_CLASSES} — {CLASSES}")


if __name__ == '__main__':
    main()
