"""
LSTM Autoencoder for Unsupervised Anomaly Detection on Sequence Data

This script defines, trains, and evaluates an LSTM Autoencoder using Keras.
It saves the best model based on validation loss and plots reconstruction loss trends.
"""

import matplotlib.pyplot as plt
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.layers import Input, LSTM, RepeatVector, TimeDistributed, Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

# Ensure these variables are defined elsewhere in your code
# timesteps, n_features, X_train_normal, X_val


def build_lstm_autoencoder(timesteps, n_features):
    """Builds and returns an LSTM Autoencoder model."""
    inputs = Input(shape=(timesteps, n_features), name="input_layer")
    encoded = LSTM(128, return_sequences=True, name="encoder_lstm1")(inputs)
    encoded = LSTM(64, return_sequences=False, name="encoder_lstm2")(encoded)
    bottleneck = RepeatVector(timesteps, name="bottleneck")(encoded)
    decoded = LSTM(64, return_sequences=True, name="decoder_lstm1")(bottleneck)
    decoded = LSTM(128, return_sequences=True, name="decoder_lstm2")(decoded)
    outputs = TimeDistributed(Dense(n_features), name="output_layer")(decoded)
    autoencoder = Model(inputs, outputs, name="LSTM_Autoencoder")
    return autoencoder


def plot_loss(history, save_path="reconstruction_loss_trends.pdf"):
    """Plots and saves training/validation loss."""
    plt.figure(figsize=(10, 6))
    plt.plot(history.history['loss'], label='Training Loss')
    plt.plot(history.history['val_loss'], label='Validation Loss')
    plt.title('Reconstruction Loss Trends', fontsize=16)
    plt.xlabel('Epoch', fontsize=12)
    plt.ylabel('Loss (MAE)', fontsize=12)
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.savefig(save_path, dpi=600, format='pdf')
    plt.show()


def main():
    # Hyperparameters and data (should be defined/imported as needed)
    # timesteps, n_features, X_train_normal, X_val

    # Build and compile the autoencoder
    autoencoder = build_lstm_autoencoder(timesteps, n_features)
    autoencoder.compile(optimizer=Adam(1e-3), loss='mae')

    # Callbacks
    early_stop = EarlyStopping(
        monitor='val_loss', patience=10, restore_best_weights=True)
    checkpoint_cb = ModelCheckpoint(
        filepath="best_model_Unsupervised.keras",
        monitor="val_loss",
        save_best_only=True,
        verbose=1
    )

    # Train the model
    history = autoencoder.fit(
        X_train_normal,
        X_train_normal,
        epochs=300,
        batch_size=32,
        validation_data=(X_val, X_val),
        callbacks=[early_stop, checkpoint_cb],
        verbose=1
    )

    # Load the best saved model
    best_autoencoder = load_model('best_model_Unsupervised.keras')

    # Plot results
    plot_loss(history)


if __name__ == "__main__":
    main()