#!/usr/bin/env python3
"""
Drone Classifier — Training Pipeline
Trains a compact MLP on 128-bin RF spectra, exports int8 TFLite for ESP32-S3.
Target model size: <=32 KB (MODEL_ARENA_SIZE_PER_SLOT).
"""

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 = ['NOISE', 'OCUSYNC', 'FPV_FHSS', 'SKYDIO', 'UNKNOWN', 'WIFI', 'BLUETOOTH']
NUM_CLASSES = len(CLASSES)
NUM_FEATURES = 128


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():
    """Small MLP that fits inside 32 KB int8 TFLite budget."""
    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(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):
    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 > 32:
        print(" [WARNING: exceeds 32 KB target]")
    else:
        print(" [OK, within 32 KB budget]")
    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()
    in_d  = interpreter.get_input_details()[0]
    out_d = interpreter.get_output_details()[0]
    in_scale, in_zp = in_d['quantization']
    out_scale, out_zp = out_d['quantization']
    print(f"\nTFLite verify — {n_verify} samples:")
    print(f"  Input dtype={in_d['dtype']}, scale={in_scale:.6f}, zp={in_zp}")
    print(f"  Output dtype={out_d['dtype']}, scale={out_scale:.6f}, zp={out_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))
        sample_q = (sample / in_scale + in_zp).astype(np.int8)
        interpreter.set_tensor(in_d['index'], sample_q)
        interpreter.invoke()
        tflite_out  = interpreter.get_tensor(out_d['index'])
        tflite_pred = np.argmax(tflite_out)
        ok = keras_pred == tflite_pred
        matches += int(ok)
        print(f"  [{i}] true={CLASSES[y_test[i]]}, keras={CLASSES[keras_pred]}, "
              f"tflite={CLASSES[tflite_pred]} {'MATCH' if ok else 'MISMATCH'}")
    print(f"  Match rate: {matches}/{n_verify}")


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

    X_train, y_train, X_test, y_test = load_data()
    print(f"\nLoaded: 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)

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

    # 80/20 train/val split
    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=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)

    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=0)
    print(f"\nTest Loss: {test_loss:.4f} | Test Accuracy: {test_acc:.4f}")

    y_pred = np.argmax(model.predict(X_test, verbose=0), axis=1)
    print(classification_report(y_test, y_pred, target_names=CLASSES, digits=4))

    output_path = os.path.join(SCRIPT_DIR, 'drone_classifier.tflite')
    tflite_model = export_tflite(model, X_tr, output_path)
    verify_tflite(tflite_model, X_test, y_test, model)

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


if __name__ == '__main__':
    main()
