#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI-DRX: CNN Model Implementation for Dynamic DRX Parameter Optimization in 5G Networks

This script implements the Convolutional Neural Network (CNN) model used in the AI-DRX
mechanism for predicting optimal sleep times based on traffic patterns and network conditions.

Author: Hammad Lazrek
University Sidi Mohammed Ben Abdellah, L3IA Laboratory of FSDM-Fez, Morocco
"""

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import Dense, Conv1D, MaxPooling1D, Flatten, Input, Dropout, BatchNormalization
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import os
import datetime
import argparse

# Set random seeds for reproducibility
np.random.seed(42)
tf.random.set_seed(42)

class AIDRX_CNN:
    """
    AI-DRX CNN model class for predicting optimal DRX parameters
    """
    
    def __init__(self, input_length=60, feature_dim=5, model_path=None):
        """
        Initialize the AI-DRX CNN model
        
        Parameters:
        -----------
        input_length : int
            Length of the input sequence (time steps)
        feature_dim : int
            Number of features per time step
        model_path : str
            Path to pre-trained model (if loading an existing model)
        """
        self.input_length = input_length
        self.feature_dim = feature_dim
        self.model = None
        
        if model_path and os.path.exists(model_path):
            self.load_model(model_path)
        else:
            self.build_model()
    
    def build_model(self):
        """
        Build the CNN model architecture
        """
        # Define model architecture
        inputs = Input(shape=(self.input_length, self.feature_dim))
        
        # First convolutional block
        x = Conv1D(filters=32, kernel_size=3, activation='relu', padding='same')(inputs)
        x = BatchNormalization()(x)
        x = MaxPooling1D(pool_size=2)(x)
        
        # Second convolutional block
        x = Conv1D(filters=64, kernel_size=3, activation='relu', padding='same')(x)
        x = BatchNormalization()(x)
        x = MaxPooling1D(pool_size=2)(x)
        
        # Third convolutional block
        x = Conv1D(filters=128, kernel_size=3, activation='relu', padding='same')(x)
        x = BatchNormalization()(x)
        x = MaxPooling1D(pool_size=2)(x)
        
        # Flatten and dense layers
        x = Flatten()(x)
        x = Dense(128, activation='relu')(x)
        x = Dropout(0.3)(x)
        x = Dense(64, activation='relu')(x)
        x = Dropout(0.2)(x)
        
        # Output layer - predicting dynamic sleep time (TDY)
        outputs = Dense(1, activation='linear')(x)
        
        # Create model
        self.model = Model(inputs=inputs, outputs=outputs)
        
        # Compile model
        self.model.compile(
            optimizer=Adam(learning_rate=0.001),
            loss='mean_squared_error',
            metrics=['mae']
        )
        
        print(self.model.summary())
    
    def train(self, X_train, y_train, X_val, y_val, batch_size=32, epochs=100, patience=10):
        """
        Train the CNN model
        
        Parameters:
        -----------
        X_train : numpy.ndarray
            Training input data
        y_train : numpy.ndarray
            Training target data
        X_val : numpy.ndarray
            Validation input data
        y_val : numpy.ndarray
            Validation target data
        batch_size : int
            Batch size for training
        epochs : int
            Maximum number of epochs
        patience : int
            Patience for early stopping
        
        Returns:
        --------
        history : tensorflow.keras.callbacks.History
            Training history
        """
        # Define callbacks
        early_stopping = EarlyStopping(
            monitor='val_loss',
            patience=patience,
            restore_best_weights=True
        )
        
        model_checkpoint = ModelCheckpoint(
            filepath='ai_drx_cnn_best_model.h5',
            monitor='val_loss',
            save_best_only=True
        )
        
        # Train model
        history = self.model.fit(
            X_train, y_train,
            validation_data=(X_val, y_val),
            batch_size=batch_size,
            epochs=epochs,
            callbacks=[early_stopping, model_checkpoint],
            verbose=1
        )
        
        return history
    
    def predict(self, X):
        """
        Make predictions using the trained model
        
        Parameters:
        -----------
        X : numpy.ndarray
            Input data
        
        Returns:
        --------
        predictions : numpy.ndarray
            Predicted dynamic sleep times
        """
        if self.model is None:
            raise ValueError("Model not initialized. Call build_model() or load_model() first.")
        
        return self.model.predict(X)
    
    def save_model(self, path='ai_drx_cnn_model.h5'):
        """
        Save the model to disk
        
        Parameters:
        -----------
        path : str
            Path to save the model
        """
        if self.model is None:
            raise ValueError("Model not initialized. Call build_model() first.")
        
        self.model.save(path)
        print(f"Model saved to {path}")
    
    def load_model(self, path):
        """
        Load a pre-trained model from disk
        
        Parameters:
        -----------
        path : str
            Path to the saved model
        """
        self.model = load_model(path)
        print(f"Model loaded from {path}")
    
    def evaluate(self, X_test, y_test):
        """
        Evaluate the model on test data
        
        Parameters:
        -----------
        X_test : numpy.ndarray
            Test input data
        y_test : numpy.ndarray
            Test target data
        
        Returns:
        --------
        metrics : dict
            Evaluation metrics
        """
        if self.model is None:
            raise ValueError("Model not initialized. Call build_model() or load_model() first.")
        
        loss, mae = self.model.evaluate(X_test, y_test, verbose=0)
        
        # Make predictions
        y_pred = self.predict(X_test)
        
        # Calculate additional metrics
        mse = np.mean((y_test - y_pred) ** 2)
        rmse = np.sqrt(mse)
        
        metrics = {
            'loss': loss,
            'mae': mae,
            'mse': mse,
            'rmse': rmse
        }
        
        return metrics


def generate_synthetic_data(num_samples=1000, sequence_length=60, feature_dim=5):
    """
    Generate synthetic data for training and testing the AI-DRX CNN model
    
    Parameters:
    -----------
    num_samples : int
        Number of samples to generate
    sequence_length : int
        Length of each input sequence
    feature_dim : int
        Number of features per time step
    
    Returns:
    --------
    X : numpy.ndarray
        Input data of shape (num_samples, sequence_length, feature_dim)
    y : numpy.ndarray
        Target data of shape (num_samples,)
    feature_names : list
        Names of the features
    """
    # Define feature names
    feature_names = [
        'traffic_load',          # Normalized traffic load
        'signal_strength',       # Normalized signal strength
        'packet_inter_arrival',  # Normalized packet inter-arrival time
        'packet_size',           # Normalized packet size
        'channel_quality'        # Normalized channel quality indicator
    ]
    
    # Initialize data arrays
    X = np.zeros((num_samples, sequence_length, feature_dim))
    y = np.zeros((num_samples,))
    
    # Generate synthetic data
    for i in range(num_samples):
        # Generate base patterns with some randomness
        traffic_pattern = np.clip(np.sin(np.linspace(0, 4*np.pi, sequence_length)) * 0.5 + 0.5 + np.random.normal(0, 0.1, sequence_length), 0, 1)
        signal_pattern = np.clip(np.cos(np.linspace(0, 2*np.pi, sequence_length)) * 0.3 + 0.7 + np.random.normal(0, 0.05, sequence_length), 0, 1)
        
        # Add some random spikes to simulate bursts
        if np.random.random() < 0.3:  # 30% chance of having a traffic burst
            burst_start = np.random.randint(0, sequence_length - 10)
            burst_length = np.random.randint(5, 10)
            traffic_pattern[burst_start:burst_start+burst_length] = np.clip(traffic_pattern[burst_start:burst_start+burst_length] * 2, 0, 1)
        
        # Generate other features based on traffic and signal patterns
        inter_arrival_pattern = 1.0 - traffic_pattern * 0.8 + np.random.normal(0, 0.05, sequence_length)
        packet_size_pattern = traffic_pattern * 0.7 + np.random.normal(0, 0.1, sequence_length)
        channel_quality_pattern = signal_pattern * 0.8 + np.random.normal(0, 0.05, sequence_length)
        
        # Clip values to [0, 1] range
        inter_arrival_pattern = np.clip(inter_arrival_pattern, 0, 1)
        packet_size_pattern = np.clip(packet_size_pattern, 0, 1)
        channel_quality_pattern = np.clip(channel_quality_pattern, 0, 1)
        
        # Combine features
        X[i, :, 0] = traffic_pattern
        X[i, :, 1] = signal_pattern
        X[i, :, 2] = inter_arrival_pattern
        X[i, :, 3] = packet_size_pattern
        X[i, :, 4] = channel_quality_pattern
        
        # Generate target value (dynamic sleep time in ms)
        # Formula based on the paper's description:
        # - Low traffic, high signal strength -> longer sleep time
        # - High traffic, low signal strength -> shorter sleep time
        avg_traffic = np.mean(traffic_pattern[-10:])  # Consider recent traffic
        avg_signal = np.mean(signal_pattern[-10:])    # Consider recent signal strength
        
        if avg_traffic < 0.3:  # Low traffic
            base_sleep_time = 500  # ms
        elif avg_traffic < 0.7:  # Medium traffic
            base_sleep_time = 200  # ms
        else:  # High traffic
            base_sleep_time = 50   # ms
        
        # Adjust based on signal strength
        signal_factor = 1 + (avg_signal - 0.5)
        
        # Final sleep time with some randomness
        sleep_time = base_sleep_time * signal_factor * (0.9 + np.random.random() * 0.2)
        
        # Clip to reasonable range (10-640ms as per paper)
        sleep_time = np.clip(sleep_time, 10, 640)
        
        y[i] = sleep_time
    
    return X, y, feature_names


def plot_training_history(history):
    """
    Plot training history
    
    Parameters:
    -----------
    history : tensorflow.keras.callbacks.History
        Training history
    """
    plt.figure(figsize=(12, 5))
    
    # Plot training & validation loss
    plt.subplot(1, 2, 1)
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('Model Loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Validation'], loc='upper right')
    
    # Plot training & validation mean absolute error
    plt.subplot(1, 2, 2)
    plt.plot(history.history['mae'])
    plt.plot(history.history['val_mae'])
    plt.title('Model Mean Absolute Error')
    plt.ylabel('MAE')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Validation'], loc='upper right')
    
    plt.tight_layout()
    plt.savefig('ai_drx_training_history.png')
    plt.close()


def plot_predictions(y_true, y_pred, title='AI-DRX Predictions'):
    """
    Plot true vs predicted values
    
    Parameters:
    -----------
    y_true : numpy.ndarray
        True values
    y_pred : numpy.ndarray
        Predicted values
    title : str
        Plot title
    """
    plt.figure(figsize=(10, 6))
    
    # Plot scatter of true vs predicted
    plt.scatter(y_true, y_pred, alpha=0.5)
    
    # Plot perfect prediction line
    min_val = min(np.min(y_true), np.min(y_pred))
    max_val = max(np.max(y_true), np.max(y_pred))
    plt.plot([min_val, max_val], [min_val, max_val], 'r--')
    
    plt.title(title)
    plt.xlabel('True Dynamic Sleep Time (ms)')
    plt.ylabel('Predicted Dynamic Sleep Time (ms)')
    plt.grid(True, alpha=0.3)
    
    # Add metrics to plot
    mse = np.mean((y_true - y_pred) ** 2)
    rmse = np.sqrt(mse)
    mae = np.mean(np.abs(y_true - y_pred))
    
    plt.text(0.05, 0.95, f'MSE: {mse:.2f}\nRMSE: {rmse:.2f}\nMAE: {mae:.2f}',
             transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8))
    
    plt.tight_layout()
    plt.savefig('ai_drx_predictions.png')
    plt.close()


def save_sample_data(X, y, feature_names, filename='ai_drx_sample_data.csv', samples=100):
    """
    Save sample data to CSV for reference
    
    Parameters:
    -----------
    X : numpy.ndarray
        Input data
    y : numpy.ndarray
        Target data
    feature_names : list
        Names of the features
    filename : str
        Output filename
    samples : int
        Number of samples to save
    """
    # Take a subset of the data
    X_sample = X[:samples, -1, :]  # Take the last time step of each sample
    y_sample = y[:samples]
    
    # Create DataFrame
    df = pd.DataFrame(X_sample, columns=feature_names)
    df['dynamic_sleep_time_ms'] = y_sample
    
    # Save to CSV
    df.to_csv(filename, index=False)
    print(f"Sample data saved to {filename}")


def main():
    """
    Main function to demonstrate the AI-DRX CNN model
    """
    parser = argparse.ArgumentParser(description='AI-DRX CNN Model Training and Evaluation')
    parser.add_argument('--train', action='store_true', help='Train a new model')
    parser.add_argument('--evaluate', action='store_true', help='Evaluate the model')
    parser.add_argument('--model_path', type=str, default='ai_drx_cnn_model.h5', help='Path to save/load model')
    parser.add_argument('--samples', type=int, default=5000, help='Number of synthetic samples to generate')
    
    args = parser.parse_args()
    
    # Generate synthetic data
    print("Generating synthetic data...")
    X, y, feature_names = generate_synthetic_data(num_samples=args.samples)
    
    # Save sample data for reference
    save_sample_data(X, y, feature_names)
    
    # Split data into train, validation, and test sets
    X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)
    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)
    
    print(f"Data shapes: X_train: {X_train.shape}, X_val: {X_val.shape}, X_test: {X_test.shape}")
    
    # Initialize model
    ai_drx = AIDRX_CNN(input_length=X.shape[1], feature_dim=X.shape[2], model_path=args.model_path if not args.train else None)
    
    if args.train:
        print("Training model...")
        history = ai_drx.train(X_train, y_train, X_val, y_val, epochs=50)
        
        # Plot training history
        plot_training_history(history)
        
        # Save model
        ai_drx.save_model(args.model_path)
    
    if args.evaluate or args.train:
        print("Evaluating model...")
        metrics = ai_drx.evaluate(X_test, y_test)
        print("Evaluation metrics:")
        for metric, value in metrics.items():
            print(f"  {metric}: {value:.4f}")
        
        # Make predictions on test set
        y_pred = ai_drx.predict(X_test).flatten()
        
        # Plot predictions
        plot_predictions(y_test, y_pred)
    
    # Generate example real-time prediction
    print("\nExample real-time prediction:")
    # Take a random sample from test set
    sample_idx = np.random.randint(0, len(X_test))
    sample_input = X_test[sample_idx:sample_idx+1]
    sample_true = y_test[sample_idx]
    
    # Make prediction
    sample_pred = ai_drx.predict(sample_input)[0][0]
    
    print(f"Input features (last time step):")
    for i, name in enumerate(feature_names):
        print(f"  {name}: {sample_input[0, -1, i]:.4f}")
    print(f"True dynamic sleep time: {sample_true:.2f} ms")
    print(f"Predicted dynamic sleep time: {sample_pred:.2f} ms")
    
    # Save a sample prediction to CSV for NS-3 integration
    sample_predictions = []
    for i in range(10):
        idx = np.random.randint(0, len(X_test))
        input_features = X_test[idx:idx+1]
        true_value = y_test[idx]
        pred_value = ai_drx.predict(input_features)[0][0]
        
        # Extract last time step features
        features = {name: input_features[0, -1, i] for i, name in enumerate(feature_names)}
        features['true_sleep_time_ms'] = true_value
        features['predicted_sleep_time_ms'] = pred_value
        
        sample_predictions.append(features)
    
    # Save to CSV
    pd.DataFrame(sample_predictions).to_csv('ai_drx_sample_predictions.csv', index=False)
    print("Sample predictions saved to ai_drx_sample_predictions.csv")


if __name__ == "__main__":
    main()