#!/usr/bin/env python3
"""
MAVLink Anomaly Detector — Training Pipeline
Trains a dense network on 48-feature MAVLink telemetry vectors, exports int8 TFLite.
Target model size: <150 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', 'cmd_injection', 'gps_spoof', 'waypoint_tamper']
NUM_CLASSES = len(CLASSES)
NUM_FEATURES = 48


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(96, activation='relu',
                              kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(48, 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):
    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, target_kb=150):
    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
    status = "OK" if size_kb <= target_kb else "WARNING: exceeds target"
    print(f"\nTFLite saved: {output_path} — {size_kb:.1f} KB [{status}, target {target_kb} KB]")
    return tflite_model


def verify_tflite(tflite_model, X_test, y_test, keras_model, n_verify=10):
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()
    inp = interpreter.get_input_details()
    out = interpreter.get_output_details()
    input_scale, input_zp = inp[0]['quantization']
    output_scale, output_zp = out[0]['quantization']

    print(f"\nTFLite verification ({n_verify} samples):")
    print(f"  Input: {inp[0]['shape']}, scale={input_scale:.6f}, zp={input_zp}")
    print(f"  Output: {out[0]['shape']}, 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_pred = np.argmax(keras_model.predict(sample, verbose=0), axis=1)[0]
        sample_q = (sample / input_scale + input_zp).astype(np.int8)
        interpreter.set_tensor(inp[0]['index'], sample_q)
        interpreter.invoke()
        tflite_pred = np.argmax(interpreter.get_tensor(out[0]['index']), axis=1)[0]
        match = "MATCH" if keras_pred == tflite_pred else "MISMATCH"
        if keras_pred == tflite_pred:
            matches += 1
        print(f"  [{i}] true={CLASSES[y_test[i]]}, keras={CLASSES[keras_pred]}, tflite={CLASSES[tflite_pred]} {match}")
    print(f"  Match rate: {matches}/{n_verify}")


def main():
    print("=" * 60)
    print("MAVLink Anomaly Detector — Training Pipeline")
    print("=" * 60)

    X_train, y_train, X_test, y_test = load_data()
    print(f"Loaded: train={X_train.shape}, test={X_test.shape}")

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

    np.savez(os.path.join(SCRIPT_DIR, 'scaler_params.npz'),
             mean=scaler.mean_.astype(np.float32),
             scale=scaler.scale_.astype(np.float32))

    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]}")

    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),
    ]

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

    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"\n{classification_report(y_test, y_pred, target_names=CLASSES, digits=4)}")

    cm = confusion_matrix(y_test, y_pred)
    header = "".join(f"{c[:10]:>12}" for c in CLASSES)
    print(f"{'':>14}{header}")
    for i, row in enumerate(cm):
        print(f"{CLASSES[i]:>14}{''.join(f'{v:>12}' for v in row)}")

    output_path = os.path.join(SCRIPT_DIR, 'mavlink_anomaly.tflite')
    tflite_model = export_tflite(model, X_tr, output_path, target_kb=150)
    verify_tflite(tflite_model, X_test, y_test, model, n_verify=10)

    model.save(os.path.join(SCRIPT_DIR, 'mavlink_anomaly.keras'))
    print("\nDone.")


if __name__ == '__main__':
    main()
