#!/usr/bin/env python3
"""
Drone Classifier — Synthetic Data Generator
Generates 50k samples across 7 classes for training a 128-bin RF spectrum classifier.

Feature vector (128 elements):
  128-bin power spectral density across 2.4 GHz band (2400-2500 MHz, ~0.78 MHz/bin)
  Each bin = log-scale RSSI proxy (-90 dBm noise floor, 0 = max signal)

Classes:
  0: NOISE        — thermal/environmental noise, flat spectrum
  1: OCUSYNC      — DJI OcuSync FH pattern, 10 MHz hopping channels
  2: FPV_FHSS     — racing FPV, wide sweep, 5.8 GHz mirror artifacts in 2.4 sub-band
  3: SKYDIO       — Skydio AI mesh, proprietary 2.4+5 GHz dual band
  4: UNKNOWN      — unrecognized drone signal, anomalous peaks
  5: WIFI         — standard 802.11 OFDM traffic, channel peaks at 1/6/11
  6: BLUETOOTH    — 1 MHz BLE/BT hops, high spectral density
"""

import numpy as np
import os
import argparse

SEED = 42
NUM_FEATURES = 128
CLASSES = ['NOISE', 'OCUSYNC', 'FPV_FHSS', 'SKYDIO', 'UNKNOWN', 'WIFI', 'BLUETOOTH']
NUM_CLASSES = len(CLASSES)

# Frequency axis: 2400-2500 MHz in 128 bins (~0.78 MHz per bin)
FREQ_MHZ = np.linspace(2400, 2500, NUM_FEATURES)

# WiFi channel center frequencies (MHz) for 2.4 GHz
WIFI_CHANNELS = {1: 2412, 6: 2437, 11: 2462}


def noise_floor(n: int, rng: np.random.Generator) -> np.ndarray:
    """Flat thermal noise floor with slight random variation."""
    # Noise floor: -85 to -75 dBm, normalized to [0,1] where 0=-90, 1=0
    base = rng.uniform(0.05, 0.15, (n, NUM_FEATURES)).astype(np.float32)
    base += rng.normal(0, 0.02, (n, NUM_FEATURES)).astype(np.float32)
    return base.clip(0, 1)


def add_peak(X: np.ndarray, center_mhz: float, width_bins: float,
             amplitude: float, rng: np.random.Generator) -> np.ndarray:
    """Add a Gaussian spectral peak to X (in-place, returns X)."""
    center_bin = (center_mhz - 2400) / (100 / NUM_FEATURES)
    bins = np.arange(NUM_FEATURES, dtype=np.float32)
    gaussian = amplitude * np.exp(-0.5 * ((bins - center_bin) / width_bins) ** 2)
    X += gaussian[np.newaxis, :]
    return X


def generate_noise(n: int, rng: np.random.Generator) -> np.ndarray:
    """Pure thermal/environmental noise — flat spectrum."""
    X = noise_floor(n, rng)
    # Occasional narrowband interference spikes (not drone-related)
    for _ in range(3):
        center = rng.uniform(2401, 2499, n)
        amp = rng.uniform(0.05, 0.15, n)
        for i in range(n):
            c_bin = (center[i] - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amp[i] * np.exp(-0.5 * ((bins - c_bin) / 1.5) ** 2)
    return X.clip(0, 1).astype(np.float32)


def generate_ocusync(n: int, rng: np.random.Generator) -> np.ndarray:
    """DJI OcuSync — frequency-hopping, wide peaks across full band."""
    X = noise_floor(n, rng)
    # OcuSync uses 10 MHz channels, hops pseudo-randomly
    # We see time-averaged smear: 3-5 dominant 10 MHz channels active
    num_channels = rng.integers(3, 6, n)
    for i in range(n):
        positions = rng.uniform(2410, 2490, num_channels[i])
        amps = rng.uniform(0.3, 0.7, num_channels[i])
        for j in range(num_channels[i]):
            c_bin = (positions[j] - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amps[j] * np.exp(-0.5 * ((bins - c_bin) / 6.4) ** 2)
        # Characteristic inter-hop energy smear
        X[i] += rng.uniform(0.05, 0.15, NUM_FEATURES).astype(np.float32)
    return X.clip(0, 1).astype(np.float32)


def generate_fpv_fhss(n: int, rng: np.random.Generator) -> np.ndarray:
    """FPV racing drone FHSS — fast sweeping, chaotic peaks."""
    X = noise_floor(n, rng)
    # FHSS sweeps whole band — time-averaged shows near-uniform elevation
    # with rapid hop residuals
    for i in range(n):
        # Elevated baseline from fast hopping
        X[i] += rng.uniform(0.15, 0.35, NUM_FEATURES).astype(np.float32)
        # Residual dwell peaks (5-10 random positions)
        num_peaks = rng.integers(5, 11)
        positions = rng.uniform(2402, 2498, num_peaks)
        amps = rng.uniform(0.2, 0.5, num_peaks)
        for j in range(num_peaks):
            c_bin = (positions[j] - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amps[j] * np.exp(-0.5 * ((bins - c_bin) / 2.0) ** 2)
    return X.clip(0, 1).astype(np.float32)


def generate_skydio(n: int, rng: np.random.Generator) -> np.ndarray:
    """Skydio — proprietary AI mesh, multiple 20 MHz channels, structured pattern."""
    X = noise_floor(n, rng)
    # Skydio uses structured multi-channel: typically 2-3 wide 20 MHz channels
    # with a characteristic low-power mesh backhaul
    skydio_centers = [2422, 2452, 2472]
    for i in range(n):
        num_active = rng.integers(2, 4)
        chosen = rng.choice(skydio_centers, num_active, replace=False)
        for c in chosen:
            amp = rng.uniform(0.4, 0.75)
            c_bin = (c - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amp * np.exp(-0.5 * ((bins - c_bin) / 12.8) ** 2)
        # Mesh backhaul — low-level activity across band
        X[i] += rng.uniform(0.02, 0.08, NUM_FEATURES).astype(np.float32)
        # Characteristic control channel at 2.4 GHz center
        c_bin = (2450 - 2400) / (100 / NUM_FEATURES)
        bins = np.arange(NUM_FEATURES, dtype=np.float32)
        X[i] += 0.15 * np.exp(-0.5 * ((bins - c_bin) / 3.2) ** 2)
    return X.clip(0, 1).astype(np.float32)


def generate_unknown(n: int, rng: np.random.Generator) -> np.ndarray:
    """Unknown drone — anomalous patterns, doesn't fit other categories."""
    X = noise_floor(n, rng)
    # Mix of unusual signatures: asymmetric peaks, irregular spacing
    for i in range(n):
        # Irregular number of peaks at unusual positions
        num_peaks = rng.integers(2, 7)
        positions = rng.choice(np.linspace(2405, 2495, 50), num_peaks, replace=False)
        widths = rng.uniform(1.5, 10, num_peaks)
        amps = rng.uniform(0.2, 0.6, num_peaks)
        for j in range(num_peaks):
            c_bin = (positions[j] - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amps[j] * np.exp(-0.5 * ((bins - c_bin) / widths[j]) ** 2)
        # Asymmetric spectrum tilt
        tilt = rng.uniform(-0.001, 0.001)
        X[i] += (np.arange(NUM_FEATURES) * tilt).astype(np.float32)
    return X.clip(0, 1).astype(np.float32)


def generate_wifi(n: int, rng: np.random.Generator) -> np.ndarray:
    """Standard 802.11 WiFi — dominant peaks at channels 1, 6, 11."""
    X = noise_floor(n, rng)
    for i in range(n):
        # Determine which channels are active (at least 1, usually 2-3)
        num_active = rng.integers(1, 4)
        active_channels = rng.choice(list(WIFI_CHANNELS.values()), num_active, replace=False)
        for ch_mhz in active_channels:
            amp = rng.uniform(0.4, 0.9)
            # WiFi 20 MHz channel = ~25.6 bins
            c_bin = (ch_mhz - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amp * np.exp(-0.5 * ((bins - c_bin) / 12.8) ** 2)
        # Secondary channels sometimes visible
        if rng.random() > 0.6:
            other = rng.choice([2422, 2432, 2442, 2447], 1)[0]
            c_bin = (other - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += rng.uniform(0.1, 0.3) * np.exp(-0.5 * ((bins - c_bin) / 6.4) ** 2)
    return X.clip(0, 1).astype(np.float32)


def generate_bluetooth(n: int, rng: np.random.Generator) -> np.ndarray:
    """Bluetooth/BLE — 1 MHz hops, high spectral density, near-uniform elevation."""
    X = noise_floor(n, rng)
    for i in range(n):
        # BT hops 79 × 1 MHz channels — time-averaged = elevated baseline
        X[i] += rng.uniform(0.2, 0.45, NUM_FEATURES).astype(np.float32)
        # BLE advertising channels at 2402, 2426, 2480 are slightly stronger
        for ble_adv in [2402, 2426, 2480]:
            amp = rng.uniform(0.15, 0.35)
            c_bin = (ble_adv - 2400) / (100 / NUM_FEATURES)
            bins = np.arange(NUM_FEATURES, dtype=np.float32)
            X[i] += amp * np.exp(-0.5 * ((bins - c_bin) / 1.28) ** 2)
        # Fine-grained hop texture
        X[i] += rng.normal(0, 0.04, NUM_FEATURES).astype(np.float32)
    return X.clip(0, 1).astype(np.float32)


GENERATORS = [
    generate_noise,
    generate_ocusync,
    generate_fpv_fhss,
    generate_skydio,
    generate_unknown,
    generate_wifi,
    generate_bluetooth,
]


def main():
    parser = argparse.ArgumentParser(description='Generate synthetic drone classifier data')
    parser.add_argument('--total', type=int, default=50000, help='Total samples')
    parser.add_argument('--test-frac', type=float, default=0.2, help='Test fraction')
    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 = samples_per_class + (1 if cls_idx < remainder else 0)
        X_cls = gen_fn(n, rng)
        y_cls = np.full(n, cls_idx, dtype=np.int32)
        all_X.append(X_cls)
        all_y.append(y_cls)
        print(f"  Class {cls_idx} ({CLASSES[cls_idx]}): {n} samples")

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)

    # Small noise for augmentation
    X += rng.normal(0, 0.005, X.shape).astype(np.float32)
    X = X.clip(0, 1)

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

    # Split
    n_test = int(len(X) * args.test_frac)
    X_train, X_test = X[n_test:], X[:n_test]
    y_train, y_test = y[n_test:], y[:n_test]

    os.makedirs(args.output_dir, exist_ok=True)
    np.savez_compressed(os.path.join(args.output_dir, 'train.npz'), X=X_train, y=y_train)
    np.savez_compressed(os.path.join(args.output_dir, 'test.npz'), X=X_test, y=y_test)

    print(f"\nDataset saved to: {args.output_dir}")
    print(f"  Total: {len(X)}  Train: {len(X_train)}  Test: {len(X_test)}")
    print(f"  Features: {NUM_FEATURES}  Classes: {NUM_CLASSES} — {CLASSES}")


if __name__ == '__main__':
    main()
