#!/usr/bin/env python3
"""
MAVLink Anomaly Detector — Synthetic Data Generator
Generates 40k samples across 4 classes for detecting command injection,
GPS spoofing, and waypoint tampering in MAVLink telemetry streams.

Feature vector (48 elements):
  [0-7]   Message type histogram (heartbeat, gps_raw, attitude, command_long,
           mission_item, mission_count, param_value, statustext) — normalized
  [8-11]  Timing stats (mean_iat, std_iat, heartbeat_regularity, msg_burst_ratio)
  [12-15] Heartbeat features (interval_mean, interval_std, missed_count, type_consistency)
  [16-19] System ID features (unique_sysid, unexpected_sysid_ratio, compid_diversity, sysid_switches)
  [20-23] Sequence features (seq_gap_count, seq_dup_count, seq_reset_count, seq_monotonic_ratio)
  [24-27] GPS position features (lat_delta_mean, lon_delta_mean, alt_delta_mean, position_jump_count)
  [28-31] GPS velocity features (speed_mean, speed_std, speed_max, impossible_accel_count)
  [32-35] GPS quality features (fix_type_mode, hdop_mean, satellites_mean, gps_time_consistency)
  [36-39] Attitude features (roll_rate_mean, pitch_rate_mean, yaw_rate_mean, attitude_discontinuity)
  [40-43] Command features (cmd_rate, cmd_diversity, cmd_without_ack_ratio, restricted_cmd_count)
  [44-47] Mission features (mission_item_rate, mission_without_count, item_sequence_violation, mission_change_freq)

Classes: {0: normal, 1: cmd_injection, 2: gps_spoof, 3: waypoint_tamper}
"""

import numpy as np
import os
import argparse

SEED = 42
NUM_FEATURES = 48
CLASSES = ['normal', 'cmd_injection', 'gps_spoof', 'waypoint_tamper']
NUM_CLASSES = len(CLASSES)


def generate_normal(n: int, rng: np.random.Generator) -> np.ndarray:
    """Normal MAVLink stream — regular heartbeats, smooth GPS, expected commands."""
    X = np.zeros((n, NUM_FEATURES), dtype=np.float32)

    # Message type histogram
    X[:, 0] = rng.normal(0.15, 0.03, n).clip(0.05, 0.25)   # heartbeat
    X[:, 1] = rng.normal(0.20, 0.04, n).clip(0.08, 0.35)   # gps_raw
    X[:, 2] = rng.normal(0.25, 0.05, n).clip(0.10, 0.40)   # attitude
    X[:, 3] = rng.normal(0.03, 0.015, n).clip(0, 0.07)     # command_long
    X[:, 4] = rng.normal(0.01, 0.008, n).clip(0, 0.03)     # mission_item
    X[:, 5] = rng.normal(0.005, 0.003, n).clip(0, 0.015)   # mission_count
    X[:, 6] = rng.normal(0.08, 0.03, n).clip(0.02, 0.15)   # param_value
    X[:, 7] = rng.normal(0.02, 0.01, n).clip(0, 0.05)      # statustext

    # Timing stats
    X[:, 8] = rng.normal(0.05, 0.015, n).clip(0.02, 0.1)    # mean IAT
    X[:, 9] = rng.normal(0.015, 0.006, n).clip(0.003, 0.04) # std IAT
    X[:, 10] = rng.normal(0.92, 0.04, n).clip(0.8, 0.99)    # heartbeat regularity
    X[:, 11] = rng.normal(0.08, 0.04, n).clip(0, 0.2)       # burst ratio

    # Heartbeat features
    X[:, 12] = rng.normal(1.0, 0.05, n).clip(0.8, 1.2)      # interval mean (1 Hz)
    X[:, 13] = rng.normal(0.02, 0.01, n).clip(0.005, 0.05)   # interval std
    X[:, 14] = rng.normal(0.5, 0.4, n).clip(0, 2)            # missed count
    X[:, 15] = rng.normal(0.98, 0.015, n).clip(0.93, 1.0)    # type consistency

    # System ID features
    X[:, 16] = rng.normal(2, 0.5, n).clip(1, 4)              # unique sysid
    X[:, 17] = rng.normal(0.01, 0.008, n).clip(0, 0.04)      # unexpected sysid ratio
    X[:, 18] = rng.normal(2, 0.8, n).clip(1, 4)              # compid diversity
    X[:, 19] = rng.normal(0.5, 0.4, n).clip(0, 2)            # sysid switches

    # Sequence features
    X[:, 20] = rng.normal(1, 0.8, n).clip(0, 4)              # gap count
    X[:, 21] = rng.normal(0.2, 0.2, n).clip(0, 1)            # dup count
    X[:, 22] = rng.normal(0.3, 0.25, n).clip(0, 1)           # reset count
    X[:, 23] = rng.normal(0.97, 0.02, n).clip(0.9, 1.0)      # monotonic ratio

    # GPS position features — smooth
    X[:, 24] = rng.normal(0.0001, 0.00005, n).clip(0, 0.0003)  # lat delta mean (degrees)
    X[:, 25] = rng.normal(0.0001, 0.00005, n).clip(0, 0.0003)  # lon delta mean
    X[:, 26] = rng.normal(0.5, 0.3, n).clip(0, 1.5)            # alt delta mean (meters)
    X[:, 27] = rng.normal(0.1, 0.1, n).clip(0, 0.5)            # position jump count

    # GPS velocity
    X[:, 28] = rng.normal(5, 3, n).clip(0, 15)               # speed mean (m/s)
    X[:, 29] = rng.normal(2, 1, n).clip(0.3, 5)              # speed std
    X[:, 30] = rng.normal(12, 4, n).clip(3, 25)              # speed max
    X[:, 31] = rng.normal(0.1, 0.1, n).clip(0, 0.5)          # impossible accel count

    # GPS quality
    X[:, 32] = rng.normal(3, 0.3, n).clip(2, 4)              # fix type (3=3D fix)
    X[:, 33] = rng.normal(1.2, 0.3, n).clip(0.6, 2.5)        # hdop
    X[:, 34] = rng.normal(10, 2, n).clip(5, 16)              # satellites
    X[:, 35] = rng.normal(0.98, 0.015, n).clip(0.93, 1.0)    # time consistency

    # Attitude
    X[:, 36] = rng.normal(3, 1.5, n).clip(0.5, 8)            # roll rate mean (deg/s)
    X[:, 37] = rng.normal(2, 1, n).clip(0.3, 6)              # pitch rate mean
    X[:, 38] = rng.normal(1.5, 0.8, n).clip(0.2, 4)          # yaw rate mean
    X[:, 39] = rng.normal(0.02, 0.015, n).clip(0, 0.06)      # attitude discontinuity

    # Command features
    X[:, 40] = rng.normal(0.1, 0.05, n).clip(0, 0.25)        # cmd rate
    X[:, 41] = rng.normal(3, 1.5, n).clip(1, 8)              # cmd diversity
    X[:, 42] = rng.normal(0.05, 0.03, n).clip(0, 0.15)       # cmd without ack ratio
    X[:, 43] = rng.normal(0.02, 0.015, n).clip(0, 0.06)      # restricted cmd count

    # Mission features
    X[:, 44] = rng.normal(0.01, 0.008, n).clip(0, 0.03)      # mission item rate
    X[:, 45] = rng.normal(0.02, 0.015, n).clip(0, 0.06)      # mission without count
    X[:, 46] = rng.normal(0.01, 0.01, n).clip(0, 0.04)       # sequence violation
    X[:, 47] = rng.normal(0.05, 0.03, n).clip(0, 0.15)       # mission change freq

    return X


def generate_cmd_injection(n: int, rng: np.random.Generator) -> np.ndarray:
    """Command injection — unexpected COMMAND_LONG packets, wrong sysid/compid."""
    X = generate_normal(n, rng)

    # Command_long ratio spikes
    X[:, 3] = rng.normal(0.20, 0.06, n).clip(0.10, 0.40)
    X[:, 40] = rng.normal(1.5, 0.5, n).clip(0.5, 3.0)       # cmd rate very high
    X[:, 41] = rng.normal(12, 4, n).clip(5, 25)              # high cmd diversity
    X[:, 42] = rng.normal(0.45, 0.12, n).clip(0.2, 0.75)     # many cmds without ack
    X[:, 43] = rng.normal(0.35, 0.1, n).clip(0.15, 0.6)      # restricted cmds HIGH

    # System ID anomalies
    X[:, 16] = rng.normal(5, 2, n).clip(3, 12)               # more unique sysids
    X[:, 17] = rng.normal(0.3, 0.1, n).clip(0.12, 0.55)      # unexpected sysid ratio HIGH
    X[:, 18] = rng.normal(6, 2, n).clip(3, 12)               # compid diversity HIGH
    X[:, 19] = rng.normal(8, 3, n).clip(3, 18)               # sysid switches HIGH

    # Sequence anomalies
    X[:, 20] = rng.normal(8, 3, n).clip(3, 18)               # gap count high
    X[:, 21] = rng.normal(3, 1.5, n).clip(1, 8)              # dup count
    X[:, 23] = rng.normal(0.75, 0.08, n).clip(0.55, 0.9)     # monotonic ratio drops

    # Timing — irregular
    X[:, 9] = rng.normal(0.05, 0.02, n).clip(0.02, 0.1)      # std IAT higher
    X[:, 11] = rng.normal(0.35, 0.1, n).clip(0.15, 0.6)      # burst ratio higher

    return X


def generate_gps_spoof(n: int, rng: np.random.Generator) -> np.ndarray:
    """GPS spoofing — sudden position jumps, altitude discontinuities, impossible velocities."""
    X = generate_normal(n, rng)

    # GPS position jumps
    X[:, 24] = rng.normal(0.005, 0.003, n).clip(0.001, 0.015)  # lat delta LARGE (>100m)
    X[:, 25] = rng.normal(0.005, 0.003, n).clip(0.001, 0.015)  # lon delta LARGE
    X[:, 26] = rng.normal(15, 8, n).clip(3, 40)                # alt delta LARGE
    X[:, 27] = rng.normal(6, 3, n).clip(2, 15)                 # many position jumps

    # Impossible velocities
    X[:, 28] = rng.normal(40, 15, n).clip(15, 80)            # speed mean ridiculous
    X[:, 29] = rng.normal(20, 8, n).clip(5, 45)              # speed std high
    X[:, 30] = rng.normal(80, 25, n).clip(35, 150)           # speed max impossible
    X[:, 31] = rng.normal(5, 2, n).clip(2, 12)               # impossible accel HIGH

    # GPS quality may look good (sophisticated spoof) or degrade
    # Mix: some spoofers maintain quality, others don't
    spoof_quality = rng.random(n) > 0.5
    X[spoof_quality, 32] = rng.normal(3, 0.1, spoof_quality.sum()).clip(2.8, 3.2)
    X[~spoof_quality, 33] = rng.normal(3.5, 1.0, (~spoof_quality).sum()).clip(2, 6)  # hdop degrades
    X[:, 35] = rng.normal(0.8, 0.08, n).clip(0.6, 0.95)      # time consistency drops

    # Attitude may conflict with GPS
    X[:, 39] = rng.normal(0.15, 0.06, n).clip(0.05, 0.3)     # attitude discontinuity

    return X


def generate_waypoint_tamper(n: int, rng: np.random.Generator) -> np.ndarray:
    """Waypoint tampering — MISSION_ITEM changes without proper MISSION_COUNT."""
    X = generate_normal(n, rng)

    # Mission features spike
    X[:, 4] = rng.normal(0.08, 0.03, n).clip(0.03, 0.18)    # mission_item ratio high
    X[:, 44] = rng.normal(0.15, 0.05, n).clip(0.05, 0.3)    # mission item rate HIGH
    X[:, 45] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)      # mission without count HIGH
    X[:, 46] = rng.normal(0.4, 0.12, n).clip(0.15, 0.7)     # sequence violation HIGH
    X[:, 47] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)      # mission change freq HIGH

    # mission_count stays low (attacker skips the protocol)
    X[:, 5] = rng.normal(0.002, 0.001, n).clip(0, 0.005)

    # System ID may be unexpected
    X[:, 17] = rng.normal(0.15, 0.06, n).clip(0.05, 0.3)
    X[:, 19] = rng.normal(3, 1.5, n).clip(1, 8)

    # Sequence anomalies
    X[:, 20] = rng.normal(5, 2, n).clip(2, 12)
    X[:, 23] = rng.normal(0.82, 0.06, n).clip(0.65, 0.93)

    return X


GENERATORS = [generate_normal, generate_cmd_injection, generate_gps_spoof, generate_waypoint_tamper]


def main():
    parser = argparse.ArgumentParser(description='Generate synthetic MAVLink anomaly data')
    parser.add_argument('--total', type=int, default=40000)
    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 = samples_per_class + (1 if cls_idx < remainder else 0)
        X_cls = gen_fn(n, rng)
        all_X.append(X_cls)
        all_y.append(np.full(n, cls_idx, dtype=np.int32))
        print(f"  Class {cls_idx} ({CLASSES[cls_idx]}): {n} 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)
    X_train, X_test = X[n_test:], X[:n_test]
    y_train, y_test = y[n_test:], y[:n_test]

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

    print(f"\nTotal: {len(X)} | Train: {len(X_train)} | Test: {len(X_test)}")
    print(f"Features: {NUM_FEATURES} | Classes: {NUM_CLASSES} — {CLASSES}")
    for name, sy in [('Train', y_train), ('Test', y_test)]:
        u, c = np.unique(sy, return_counts=True)
        print(f"  {name}: { {CLASSES[i]: v for i, v in zip(u, c)} }")


if __name__ == '__main__':
    main()
