#!/usr/bin/env python3
"""
RF Attack Classifier — Training Pipeline
Trains a dense network on 64-feature RF traffic vectors, exports int8 TFLite for ESP32-S3.
Target model size: <200 KB
"""

import os
import sys
import numpy as np
import tensorflow as tf
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.preprocessing import StandardScaler

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CLASSES = ['normal', 'deauth', 'beacon_flood', 'evil_twin', 'pmkid', 'arp_spoof']
NUM_CLASSES = len(CLASSES)
NUM_FEATURES = 64


def load_data():
    train_path = os.path.join(SCRIPT_DIR, 'train.npz')
    test_path = os.path.join(SCRIPT_DIR, 'test.npz')
    if not os.path.exists(train_path) or not os.path.exists(test_path):
        print("ERROR: Run generate_synthetic_data.py first.")
        sys.exit(1)
    train = np.load(train_path)
    test = np.load(test_path)
    return train['X'], train['y'], test['X'], test['y']


def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(NUM_FEATURES,)),
        tf.keras.layers.Dense(128, activation='relu',
                              kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(64, activation='relu',
                              kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(NUM_CLASSES, activation='softmax'),
    ])
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'],
    )
    return model


def representative_dataset_gen(X_cal):
    """Generator for int8 quantization calibration."""
    for i in range(min(500, len(X_cal))):
        yield [X_cal[i:i+1].astype(np.float32)]


def export_tflite(model, X_cal, output_path):
    """Export with int8 quantization."""
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = lambda: representative_dataset_gen(X_cal)
    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()
    with open(output_path, 'wb') as f:
        f.write(tflite_model)
    size_kb = os.path.getsize(output_path) / 1024
    print(f"\nTFLite model saved: {output_path}")
    print(f"Model size: {size_kb:.1f} KB", end="")
    if size_kb > 200:
        print(" [WARNING: exceeds 200 KB target]")
    else:
        print(f" [OK, under 200 KB target]")
    return tflite_model


def verify_tflite(tflite_model, X_test, y_test, keras_model, scaler, n_verify=10):
    """Verify TFLite inference matches Keras predictions on n samples."""
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    input_scale, input_zp = input_details[0]['quantization']
    output_scale, output_zp = output_details[0]['quantization']

    print(f"\nTFLite verification on {n_verify} test samples:")
    print(f"  Input: shape={input_details[0]['shape']}, dtype={input_details[0]['dtype']}, scale={input_scale:.6f}, zp={input_zp}")
    print(f"  Output: shape={output_details[0]['shape']}, dtype={output_details[0]['dtype']}, scale={output_scale:.6f}, zp={output_zp}")

    matches = 0
    for i in range(min(n_verify, len(X_test))):
        sample = X_test[i:i+1].astype(np.float32)

        # Keras prediction
        keras_pred = np.argmax(keras_model.predict(sample, verbose=0), axis=1)[0]

        # TFLite prediction — quantize input
        sample_q = (sample / input_scale + input_zp).astype(np.int8)
        interpreter.set_tensor(input_details[0]['index'], sample_q)
        interpreter.invoke()
        tflite_out = interpreter.get_tensor(output_details[0]['index'])
        tflite_pred = np.argmax(tflite_out, axis=1)[0]

        match = "MATCH" if keras_pred == tflite_pred else "MISMATCH"
        if keras_pred == tflite_pred:
            matches += 1
        print(f"  Sample {i}: true={CLASSES[y_test[i]]}, keras={CLASSES[keras_pred]}, tflite={CLASSES[tflite_pred]} [{match}]")

    print(f"  Match rate: {matches}/{n_verify} ({100*matches/n_verify:.0f}%)")


def main():
    print("=" * 60)
    print("RF Attack Classifier — Training Pipeline")
    print("=" * 60)

    # Load data
    X_train, y_train, X_test, y_test = load_data()
    print(f"\nLoaded: train={X_train.shape}, test={X_test.shape}")

    # Normalize
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train).astype(np.float32)
    X_test = scaler.transform(X_test).astype(np.float32)

    # Save scaler params for deployment
    scaler_path = os.path.join(SCRIPT_DIR, 'scaler_params.npz')
    np.savez(scaler_path, mean=scaler.mean_.astype(np.float32),
             scale=scaler.scale_.astype(np.float32))
    print(f"Scaler params saved: {scaler_path}")

    # Split train into train/val (80/20)
    n_val = int(len(X_train) * 0.2)
    idx = np.random.RandomState(42).permutation(len(X_train))
    X_val, y_val = X_train[idx[:n_val]], y_train[idx[:n_val]]
    X_tr, y_tr = X_train[idx[n_val:]], y_train[idx[n_val:]]
    print(f"Split: train={X_tr.shape[0]}, val={X_val.shape[0]}, test={X_test.shape[0]}")

    # Build and train
    model = build_model()
    model.summary()

    callbacks = [
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss', patience=10, restore_best_weights=True, verbose=1
        ),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss', factor=0.5, patience=5, min_lr=1e-6, verbose=1
        ),
    ]

    history = model.fit(
        X_tr, y_tr,
        validation_data=(X_val, y_val),
        epochs=50,
        batch_size=64,
        callbacks=callbacks,
        verbose=1,
    )

    # Evaluate on test set
    print("\n" + "=" * 60)
    print("Test Set Evaluation")
    print("=" * 60)
    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=0)
    print(f"Test Loss: {test_loss:.4f} | Test Accuracy: {test_acc:.4f}")

    y_pred = np.argmax(model.predict(X_test, verbose=0), axis=1)
    print(f"\nClassification Report:\n")
    print(classification_report(y_test, y_pred, target_names=CLASSES, digits=4))
    print(f"Confusion Matrix:\n")
    cm = confusion_matrix(y_test, y_pred)
    # Pretty print
    header = "".join(f"{c[:8]:>10}" for c in CLASSES)
    print(f"{'':>12}{header}")
    for i, row in enumerate(cm):
        row_str = "".join(f"{v:>10}" for v in row)
        print(f"{CLASSES[i]:>12}{row_str}")

    # Export TFLite
    output_path = os.path.join(SCRIPT_DIR, 'rf_attack_classifier.tflite')
    tflite_model = export_tflite(model, X_tr, output_path)

    # Verify
    verify_tflite(tflite_model, X_test, y_test, model, scaler, n_verify=10)

    # Save Keras model too
    keras_path = os.path.join(SCRIPT_DIR, 'rf_attack_classifier.keras')
    model.save(keras_path)
    print(f"\nKeras model saved: {keras_path}")

    print("\nDone.")


if __name__ == '__main__':
    main()
