#!/usr/bin/env python3
"""
build_all_128.py — Build all 5 TFLite-Micro models with 128-feature inputs.
Run from models/ directory on DGX (or any machine with TF 2.x).

Models:
  drone_classifier    128→7   (NOISE/OCUSYNC/FPV_FHSS/SKYDIO/UNKNOWN/WIFI/BLUETOOTH)
  rf_attack_classifier 128→8  (none/deauth/evil_twin/jamming/replay/port_scan/gps_spoof/c2_beacon)
  mavlink_anomaly     128→4   (normal/cmd_injection/gps_spoof/waypoint_tamper)
  lora_mesh_attack    128→5   (normal/jam_constant/jam_sweep/replay/spoof)
  protocol_anomaly    128→6   (normal/dns_tunnel/c2_beacon/port_scan/mitm/data_exfil)
"""

import os
import sys
import numpy as np

# ─────────────────────────────────────────────────────────────────────────────
# Synthetic data generators — one per model
# ─────────────────────────────────────────────────────────────────────────────

def gen_drone_classifier(n_per_class: int, rng: np.random.Generator):
    """128-bin RF spectrum → 7 drone/signal classes."""
    classes = ['NOISE','OCUSYNC','FPV_FHSS','SKYDIO','UNKNOWN','WIFI','BLUETOOTH']
    NUM_F = 128
    all_X, all_y = [], []

    def noise_floor(n):
        return rng.uniform(0.03, 0.12, (n, NUM_F)).astype(np.float32)

    def gaussian_peak(center_bin, width_bins, amp, shape):
        bins = np.arange(shape[1], dtype=np.float32)
        return amp * np.exp(-0.5 * ((bins - center_bin) / width_bins) ** 2)

    # 0: NOISE
    X = noise_floor(n_per_class)
    all_X.append(X); all_y.append(np.zeros(n_per_class, dtype=np.int32))

    # 1: OCUSYNC — 3-5 broad hopping peaks
    X = noise_floor(n_per_class)
    for i in range(n_per_class):
        for _ in range(rng.integers(3,6)):
            cb = rng.uniform(5, 123)
            X[i] += gaussian_peak(cb, 6.4, rng.uniform(0.3,0.7), (1,NUM_F))
    all_X.append(X.clip(0,1)); all_y.append(np.ones(n_per_class, dtype=np.int32))

    # 2: FPV_FHSS — elevated baseline + 5-10 random peaks
    X = noise_floor(n_per_class) + rng.uniform(0.12, 0.30, (n_per_class, NUM_F)).astype(np.float32)
    for i in range(n_per_class):
        for cb in rng.uniform(3, 125, rng.integers(5,11)):
            X[i] += gaussian_peak(cb, 2.0, rng.uniform(0.2,0.5), (1,NUM_F))
    all_X.append(X.clip(0,1)); all_y.append(np.full(n_per_class, 2, dtype=np.int32))

    # 3: SKYDIO — 2-3 structured 20 MHz channels
    X = noise_floor(n_per_class)
    skydio_bins = [28, 66, 92]
    for i in range(n_per_class):
        chosen = rng.choice(skydio_bins, rng.integers(2,4), replace=False)
        for cb in chosen:
            X[i] += gaussian_peak(float(cb), 12.8, rng.uniform(0.4,0.75), (1,NUM_F))
        X[i] += rng.uniform(0.02, 0.08, NUM_F).astype(np.float32)
    all_X.append(X.clip(0,1)); all_y.append(np.full(n_per_class, 3, dtype=np.int32))

    # 4: UNKNOWN — asymmetric/irregular peaks
    X = noise_floor(n_per_class)
    for i in range(n_per_class):
        for cb in rng.uniform(3, 125, rng.integers(2,7)):
            w = rng.uniform(1.5, 10)
            X[i] += gaussian_peak(cb, w, rng.uniform(0.2,0.6), (1,NUM_F))
        tilt = np.arange(NUM_F) * rng.uniform(-0.001, 0.001)
        X[i] += tilt.astype(np.float32)
    all_X.append(X.clip(0,1)); all_y.append(np.full(n_per_class, 4, dtype=np.int32))

    # 5: WIFI — peaks at bins ~16 (ch1/2412), ~48 (ch6/2437), ~80 (ch11/2462)
    wifi_bins = [16, 48, 80]
    X = noise_floor(n_per_class)
    for i in range(n_per_class):
        chosen = rng.choice(wifi_bins, rng.integers(1,4), replace=False)
        for cb in chosen:
            X[i] += gaussian_peak(float(cb), 12.8, rng.uniform(0.4,0.9), (1,NUM_F))
    all_X.append(X.clip(0,1)); all_y.append(np.full(n_per_class, 5, dtype=np.int32))

    # 6: BLUETOOTH — uniform elevation + BLE adv peaks at bins ~3, ~34, ~103
    X = noise_floor(n_per_class) + rng.uniform(0.2, 0.45, (n_per_class, NUM_F)).astype(np.float32)
    for i in range(n_per_class):
        for cb in [3.0, 34.0, 103.0]:
            X[i] += gaussian_peak(cb, 1.3, rng.uniform(0.15,0.35), (1,NUM_F))
    all_X.append(X.clip(0,1)); all_y.append(np.full(n_per_class, 6, dtype=np.int32))

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)
    # small noise
    X = (X + rng.normal(0, 0.005, X.shape).astype(np.float32)).clip(0, 1)
    return X, y, classes


def gen_rf_attack_classifier(n_per_class: int, rng: np.random.Generator):
    """128 features → 8 RF attack classes."""
    classes = ['none','deauth','evil_twin','jamming','replay','port_scan','gps_spoof','c2_beacon']
    NUM_F = 128

    def base():
        return rng.normal(0, 1, (n_per_class, NUM_F)).astype(np.float32)

    all_X, all_y = [], []

    # 0: none — random normal traffic
    X = base()
    all_X.append(X); all_y.append(np.zeros(n_per_class, dtype=np.int32))

    # 1: deauth — spikes in mgmt/ctrl features, burst timing
    X = base()
    X[:, 0:8] += rng.normal(3.0, 0.5, (n_per_class, 8))   # mgmt heavy
    X[:, 32:40] += rng.normal(-2.0, 0.4, (n_per_class, 8)) # fast IAT
    X[:, 96:104] += rng.normal(4.0, 0.6, (n_per_class, 8)) # deauth confidence
    all_X.append(X); all_y.append(np.full(n_per_class, 1, dtype=np.int32))

    # 2: evil_twin — RSSI variance spikes
    X = base()
    X[:, 16:24] += rng.normal(4.0, 0.5, (n_per_class, 8))  # high RSSI variance
    X[:, 48:56] += rng.normal(3.0, 0.5, (n_per_class, 8))  # beacon elevated
    X[:, 104:112] += rng.normal(3.5, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 2, dtype=np.int32))

    # 3: jamming — broad spectrum elevation, high power
    X = rng.normal(2.5, 0.6, (n_per_class, NUM_F)).astype(np.float32)
    X[:, 60:80] += rng.normal(3.0, 0.5, (n_per_class, 20))
    all_X.append(X); all_y.append(np.full(n_per_class, 3, dtype=np.int32))

    # 4: replay — timing anomalies, seq number issues
    X = base()
    X[:, 40:48] += rng.normal(3.5, 0.5, (n_per_class, 8))  # seq gaps
    X[:, 56:64] += rng.normal(-1.5, 0.4, (n_per_class, 8)) # timing regularity
    X[:, 112:120] += rng.normal(3.0, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 4, dtype=np.int32))

    # 5: port_scan — high connection rate, diverse dst
    X = base()
    X[:, 64:72] += rng.normal(4.0, 0.5, (n_per_class, 8))
    X[:, 80:88] += rng.normal(3.5, 0.5, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 5, dtype=np.int32))

    # 6: gps_spoof — GPS feature anomalies
    X = base()
    X[:, 88:96] += rng.normal(4.5, 0.5, (n_per_class, 8))
    X[:, 24:32] += rng.normal(-2.0, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 6, dtype=np.int32))

    # 7: c2_beacon — periodic low-bandwidth, hidden channel
    X = base()
    X[:, 8:16] += rng.normal(3.0, 0.5, (n_per_class, 8))
    X[:, 120:128] += rng.normal(4.0, 0.5, (n_per_class, 8))
    X[:, 72:80] += rng.normal(3.5, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 7, dtype=np.int32))

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)
    X += rng.normal(0, 0.05, X.shape).astype(np.float32)
    return X, y, classes


def gen_mavlink_anomaly(n_per_class: int, rng: np.random.Generator):
    """128 MAVLink protocol features → 4 classes."""
    classes = ['normal','cmd_injection','gps_spoof','waypoint_tamper']
    NUM_F = 128

    def base():
        return rng.normal(0, 1, (n_per_class, NUM_F)).astype(np.float32)

    all_X, all_y = [], []

    # 0: normal
    X = base()
    all_X.append(X); all_y.append(np.zeros(n_per_class, dtype=np.int32))

    # 1: cmd_injection — unexpected command IDs, timing anomalies
    X = base()
    X[:, 0:16] += rng.normal(4.0, 0.6, (n_per_class, 16))
    X[:, 64:72] += rng.normal(3.0, 0.5, (n_per_class, 8))
    X[:, 112:128] += rng.normal(3.5, 0.5, (n_per_class, 16))
    all_X.append(X); all_y.append(np.full(n_per_class, 1, dtype=np.int32))

    # 2: gps_spoof — position delta anomalies, HDOP spikes
    X = base()
    X[:, 32:48] += rng.normal(5.0, 0.5, (n_per_class, 16))
    X[:, 80:88] += rng.normal(-3.0, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 2, dtype=np.int32))

    # 3: waypoint_tamper — waypoint sequence anomalies
    X = base()
    X[:, 48:64] += rng.normal(4.0, 0.6, (n_per_class, 16))
    X[:, 96:112] += rng.normal(3.0, 0.5, (n_per_class, 16))
    all_X.append(X); all_y.append(np.full(n_per_class, 3, dtype=np.int32))

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)
    X += rng.normal(0, 0.05, X.shape).astype(np.float32)
    return X, y, classes


def gen_lora_mesh_attack(n_per_class: int, rng: np.random.Generator):
    """128 LoRa mesh features → 5 classes."""
    classes = ['normal','jam_constant','jam_sweep','replay','spoof']
    NUM_F = 128

    def base():
        return rng.normal(0, 1, (n_per_class, NUM_F)).astype(np.float32)

    all_X, all_y = [], []

    # 0: normal
    X = base()
    all_X.append(X); all_y.append(np.zeros(n_per_class, dtype=np.int32))

    # 1: jam_constant — constant high-power broadband
    X = rng.normal(3.5, 0.5, (n_per_class, NUM_F)).astype(np.float32)
    X[:, 40:88] += rng.normal(2.0, 0.4, (n_per_class, 48))
    all_X.append(X); all_y.append(np.full(n_per_class, 1, dtype=np.int32))

    # 2: jam_sweep — sweeping interference
    X = base()
    for i in range(n_per_class):
        sweep_start = rng.integers(0, 64)
        X[i, sweep_start:sweep_start+40] += rng.normal(4.0, 0.5, 40).astype(np.float32)
    all_X.append(X); all_y.append(np.full(n_per_class, 2, dtype=np.int32))

    # 3: replay — duplicate packets, timing signatures
    X = base()
    X[:, 0:32] += rng.normal(3.0, 0.5, (n_per_class, 32))
    X[:, 96:128] += rng.normal(-2.5, 0.4, (n_per_class, 32))
    all_X.append(X); all_y.append(np.full(n_per_class, 3, dtype=np.int32))

    # 4: spoof — false node identity
    X = base()
    X[:, 32:64] += rng.normal(4.5, 0.5, (n_per_class, 32))
    X[:, 64:96] += rng.normal(-2.0, 0.4, (n_per_class, 32))
    all_X.append(X); all_y.append(np.full(n_per_class, 4, dtype=np.int32))

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)
    X += rng.normal(0, 0.05, X.shape).astype(np.float32)
    return X, y, classes


def gen_protocol_anomaly(n_per_class: int, rng: np.random.Generator):
    """128 protocol features → 6 classes."""
    classes = ['normal','dns_tunnel','c2_beacon','port_scan','mitm','data_exfil']
    NUM_F = 128

    def base():
        return rng.normal(0, 1, (n_per_class, NUM_F)).astype(np.float32)

    all_X, all_y = [], []

    # 0: normal
    X = base()
    all_X.append(X); all_y.append(np.zeros(n_per_class, dtype=np.int32))

    # 1: dns_tunnel — high DNS volume, large TXT records
    X = base()
    X[:, 0:20] += rng.normal(4.5, 0.5, (n_per_class, 20))
    X[:, 100:112] += rng.normal(3.0, 0.5, (n_per_class, 12))
    all_X.append(X); all_y.append(np.full(n_per_class, 1, dtype=np.int32))

    # 2: c2_beacon — periodic low-BW, distinctive jitter
    X = base()
    X[:, 20:36] += rng.normal(4.0, 0.5, (n_per_class, 16))
    X[:, 112:128] += rng.normal(3.5, 0.5, (n_per_class, 16))
    all_X.append(X); all_y.append(np.full(n_per_class, 2, dtype=np.int32))

    # 3: port_scan — many destination ports, low payload
    X = base()
    X[:, 36:52] += rng.normal(5.0, 0.5, (n_per_class, 16))
    X[:, 72:80] += rng.normal(-2.0, 0.4, (n_per_class, 8))
    all_X.append(X); all_y.append(np.full(n_per_class, 3, dtype=np.int32))

    # 4: mitm — ARP anomalies, session hijack patterns
    X = base()
    X[:, 52:68] += rng.normal(4.0, 0.6, (n_per_class, 16))
    X[:, 80:96] += rng.normal(3.0, 0.5, (n_per_class, 16))
    all_X.append(X); all_y.append(np.full(n_per_class, 4, dtype=np.int32))

    # 5: data_exfil — large outbound, unusual timing
    X = base()
    X[:, 68:84] += rng.normal(4.5, 0.5, (n_per_class, 16))
    X[:, 96:108] += rng.normal(3.5, 0.4, (n_per_class, 12))
    all_X.append(X); all_y.append(np.full(n_per_class, 5, dtype=np.int32))

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)
    X += rng.normal(0, 0.05, X.shape).astype(np.float32)
    return X, y, classes


# ─────────────────────────────────────────────────────────────────────────────
# TFLite training & export (shared for all models)
# ─────────────────────────────────────────────────────────────────────────────

def build_model_mlp(num_features: int, num_classes: int):
    """Compact MLP targeting <32 KB int8 TFLite."""
    import tensorflow as tf
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(num_features,)),
        tf.keras.layers.Dense(64, activation='relu',
                              kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.25),
        tf.keras.layers.Dense(32, activation='relu',
                              kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
        tf.keras.layers.Dense(num_classes, activation='softmax'),
    ])
    model.compile(
        optimizer=tf.keras.optimizers.Adam(1e-3),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'],
    )
    return model


def train_and_export(name: str, X: np.ndarray, y: np.ndarray, classes: list,
                     out_dir: str, num_features: int = 128):
    """Full pipeline: preprocess → train → int8-quantize → save .tflite"""
    import tensorflow as tf
    from sklearn.preprocessing import StandardScaler

    print(f"\n{'='*60}")
    print(f"  {name}  ({X.shape[0]} samples, {num_features} features, {len(classes)} classes)")
    print(f"{'='*60}")

    os.makedirs(out_dir, exist_ok=True)

    # Normalize
    scaler = StandardScaler()
    X = scaler.fit_transform(X).astype(np.float32)
    np.savez(os.path.join(out_dir, 'scaler_params.npz'),
             mean=scaler.mean_.astype(np.float32),
             scale=scaler.scale_.astype(np.float32))

    # Shuffle + split
    rng_state = np.random.RandomState(42)
    idx = rng_state.permutation(len(X))
    X, y = X[idx], y[idx]
    n_test = int(len(X) * 0.2)
    X_test, y_test = X[:n_test], y[:n_test]
    X_train, y_train = X[n_test:], y[n_test:]
    n_val = int(len(X_train) * 0.2)
    idx2 = rng_state.permutation(len(X_train))
    X_val, y_val = X_train[idx2[:n_val]], y_train[idx2[:n_val]]
    X_tr, y_tr   = X_train[idx2[n_val:]], y_train[idx2[n_val:]]

    model = build_model_mlp(num_features, len(classes))
    model.summary()

    callbacks = [
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss', patience=8, restore_best_weights=True, verbose=1),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss', factor=0.5, patience=4, min_lr=1e-6, verbose=1),
    ]
    model.fit(X_tr, y_tr,
              validation_data=(X_val, y_val),
              epochs=40, batch_size=64,
              callbacks=callbacks, verbose=1)

    loss, acc = model.evaluate(X_test, y_test, verbose=0)
    print(f"\nTest accuracy: {acc:.4f} | Loss: {loss:.4f}")

    # Int8 quantization
    def rep_dataset():
        for i in range(min(500, len(X_tr))):
            yield [X_tr[i:i+1]]

    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = rep_dataset
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.inference_input_type  = tf.int8
    converter.inference_output_type = tf.int8
    tflite_model = converter.convert()

    tflite_path = os.path.join(out_dir, f'{name}.tflite')
    with open(tflite_path, 'wb') as f:
        f.write(tflite_model)
    size_kb = os.path.getsize(tflite_path) / 1024
    status = "OK" if size_kb <= 32 else f"OVER BUDGET ({size_kb:.1f} KB)"
    print(f"Saved: {tflite_path}  ({size_kb:.1f} KB) [{status}]")
    return size_kb


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

MODELS_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_MODELS_DIR = os.path.join(os.path.dirname(MODELS_DIR), 'data', 'models')

CONFIGS = [
    ('drone_classifier',     gen_drone_classifier,     7142,   # ~50k / 7 classes
     os.path.join(MODELS_DIR, 'drone_classifier')),
    ('rf_attack_classifier', gen_rf_attack_classifier, 6250,   # ~50k / 8 classes
     os.path.join(MODELS_DIR, 'rf_attack_classifier')),
    ('mavlink_anomaly',      gen_mavlink_anomaly,      12500,  # ~50k / 4 classes
     os.path.join(MODELS_DIR, 'mavlink_anomaly')),
    ('lora_mesh_attack',     gen_lora_mesh_attack,     10000,  # ~50k / 5 classes
     os.path.join(MODELS_DIR, 'lora_mesh_attack')),
    ('protocol_anomaly',     gen_protocol_anomaly,     8333,   # ~50k / 6 classes
     os.path.join(MODELS_DIR, 'protocol_anomaly')),
]


def main():
    import tensorflow as tf
    print(f"TensorFlow {tf.__version__}")
    print(f"MODELS_DIR: {MODELS_DIR}")
    print(f"DATA_MODELS_DIR: {DATA_MODELS_DIR}")
    os.makedirs(DATA_MODELS_DIR, exist_ok=True)

    rng = np.random.default_rng(42)
    results = []

    for name, gen_fn, n_per_class, out_dir in CONFIGS:
        print(f"\nGenerating data for {name}...")
        X, y, classes = gen_fn(n_per_class, rng)
        print(f"  shape={X.shape}  classes={classes}")

        size_kb = train_and_export(name, X, y, classes, out_dir)
        results.append((name, size_kb))

    # Copy .tflite files to data/models/
    print(f"\n{'='*60}")
    print("  Copying models to data/models/")
    print(f"{'='*60}")
    import shutil
    final_results = []
    for name, size_kb in results:
        src = os.path.join(MODELS_DIR, name, f'{name}.tflite')
        dst = os.path.join(DATA_MODELS_DIR, f'{name}.tflite')
        if os.path.exists(src):
            shutil.copy2(src, dst)
            final_size = os.path.getsize(dst) / 1024
            status = "OK" if final_size <= 32 else "OVER BUDGET"
            print(f"  [{status}] {name}.tflite — {final_size:.1f} KB")
            final_results.append((name, final_size, status))
        else:
            print(f"  [MISS] {name}.tflite — not found at {src}")

    print(f"\n{'='*60}")
    print("  FINAL SUMMARY")
    print(f"{'='*60}")
    all_ok = True
    for name, size_kb, status in final_results:
        print(f"  {status:4s}  {name}.tflite  {size_kb:.1f} KB")
        if status != "OK":
            all_ok = False
    print()
    if all_ok and len(final_results) == 5:
        print("All 5 models built and within 32 KB budget.")
        sys.exit(0)
    else:
        print("WARNING: Some models missing or over budget.")
        sys.exit(1)


if __name__ == '__main__':
    main()
