import streamlit as st
import matplotlib.pyplot as plt
import numpy as np
import subprocess
import os
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler
import seaborn as sns
from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix, accuracy_score, roc_curve, roc_auc_score
from sklearn.model_selection import train_test_split
from keras.models import Sequential
import pickle
from keras.layers import Dense, Dropout, LSTM
from keras.callbacks import ModelCheckpoint
from keras.utils.np_utils import to_categorical
from tcn import compiled_tcn  # loading TCN model
from sklearn.ensemble import RandomForestClassifier
import shap
import time
import threading
import random
import queue

# Set page configuration
st.set_page_config(
    page_title="Explainable AI for Intrusion Detection Systems",
    page_icon="🔒",
    layout="wide"
)

# Add title and description
st.title("Explainable AI for Intrusion Detection Systems: LIME and SHAP Applicability on Multi-Layer Perceptron")
st.markdown("This application demonstrates the use of explainable AI techniques for intrusion detection systems.")

# Initialize session state variables to store data across reruns
if 'dataset' not in st.session_state:
    st.session_state.dataset = None
if 'filename' not in st.session_state:
    st.session_state.filename = None
if 'labels' not in st.session_state:
    st.session_state.labels = None
if 'X' not in st.session_state:
    st.session_state.X = None
if 'Y' not in st.session_state:
    st.session_state.Y = None
if 'perturb' not in st.session_state:
    st.session_state.perturb = None
if 'features' not in st.session_state:
    st.session_state.features = None
if 'X_train' not in st.session_state:
    st.session_state.X_train = None
if 'X_test' not in st.session_state:
    st.session_state.X_test = None
if 'y_train' not in st.session_state:
    st.session_state.y_train = None
if 'y_test' not in st.session_state:
    st.session_state.y_test = None
if 'label_encoder' not in st.session_state:
    st.session_state.label_encoder = None
if 'scaler' not in st.session_state:
    st.session_state.scaler = None
if 'lstm_model' not in st.session_state:
    st.session_state.lstm_model = None
if 'tcn_model' not in st.session_state:
    st.session_state.tcn_model = None
if 'accuracy' not in st.session_state:
    st.session_state.accuracy = []
if 'precision' not in st.session_state:
    st.session_state.precision = []
if 'recall' not in st.session_state:
    st.session_state.recall = []
if 'fscore' not in st.session_state:
    st.session_state.fscore = []
if 'shap_values' not in st.session_state:
    st.session_state.shap_values = None
# if "collecting_data" not in st.session_state:
#     st.session_state.collecting_data = False
# if "collected_data" not in st.session_state:
#     st.session_state.collected_data = []
# if "data_queue" not in st.session_state:
#     st.session_state.data_queue = queue.Queue()
# if "collection_thread" not in st.session_state:
#     st.session_state.collection_thread = None
# if "collection_type" not in st.session_state:
#     st.session_state.collection_type = None
# if "display_data" not in st.session_state:
#     st.session_state.display_data = pd.DataFrame()


def ensure_session_state():
    """
    Initialize all necessary session state variables
    """
    session_vars = [
        ('collecting_data', False),
        ('collected_data', []),
        ('data_queue', queue.Queue()),
        ('collection_thread', None),
        ('collection_type', None),
        ('display_data', pd.DataFrame()),
        ('csv_data', None)
    ]
    
    for var_name, default_value in session_vars:
        if var_name not in st.session_state:
            st.session_state[var_name] = default_value





# Create sidebar for navigation
st.sidebar.title("Navigation")
page = st.sidebar.radio("Go to", 
    ["Upload Dataset", "Preprocess Dataset", "Train & Test Split", 
     "Model Training", "Evaluation", "Prediction", "Explainability", 
     "Real-time Monitoring"])


def load_csv_data(file_path):
    """
    Load CSV data and store in session state
    """
    try:
        # Read the CSV file
        df = pd.read_csv(file_path)
        
        # Ensure 'Label' column exists
        if 'Label' not in df.columns:
            st.error("CSV file must contain a 'Label' column")
            return None
        
        return df
    except Exception as e:
        st.error(f"Error loading CSV file: {e}")
        return None
def collect_data(data_type, data_queue, csv_data):
    """
    Collect data from CSV based on label type
    """
    try:
        # Filter data based on label
        filtered_data = csv_data[csv_data['Label'] == data_type]
        
        # Iterate through rows with 5-second gap
        for _, row in filtered_data.iterrows():
            # Break if collection is stopped
            if not st.session_state.collecting_data:
                break
            
            # Create a DataFrame for the current row
            current_row_df = pd.DataFrame([row])
            
            # Put the row in the queue
            data_queue.put(current_row_df)
            
            # Wait for 5 seconds
            time.sleep(5)
        
    except Exception as e:
        st.error(f"Error in data collection: {e}")
# Function to simulate data collection (run in a separate thread)

def start_collection(data_type):
    """
    Start data collection for specified type
    """
    # Ensure session state is initialized
    ensure_session_state()
    
    # Check if CSV data is loaded
    if st.session_state.csv_data is None:
        st.warning("Please upload a CSV file first!")
        return
    
    # Check if collection is already in progress
    if st.session_state.collecting_data:
        st.warning("Data collection already in progress!")
        return
    
    # Reset session state variables
    st.session_state.collecting_data = True
    st.session_state.collection_type = data_type
    st.session_state.collected_data = []
    st.session_state.display_data = pd.DataFrame()
    
    # Ensure data queue exists
    if st.session_state.data_queue is None:
        st.session_state.data_queue = queue.Queue()
    
    # Start collection thread
    st.session_state.collection_thread = threading.Thread(
        target=collect_data, 
        args=(data_type, st.session_state.data_queue, st.session_state.csv_data)
    )
    st.session_state.collection_thread.daemon = True
    st.session_state.collection_thread.start()

# Stop data collection
def stop_collection():
    """
    Stop ongoing data collection
    """
    # Ensure session state is initialized
    ensure_session_state()
    
    # Stop collecting
    st.session_state.collecting_data = False
    
    # Wait for thread to finish
    if st.session_state.collection_thread:
        st.session_state.collection_thread.join(timeout=1.0)
    
    # Collect data from queue
    collected_data = []
    while not st.session_state.data_queue.empty():
        collected_data.append(st.session_state.data_queue.get())
    
    # Convert to DataFrame
    if collected_data:
        collected_df = pd.concat(collected_data, ignore_index=True)
        st.write(f"Collected {len(collected_df)} records of {st.session_state.collection_type} traffic")
        return collected_df
    else:
        st.warning("No data was collected!")
        return None

# Main Streamlit app
def real_time_monitoring_1():
    # Ensure session state is initialized
    ensure_session_state()
    
    st.title("Network Traffic Monitoring")
    
    # CSV File Upload
    uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    if uploaded_file is not None:
        # Load CSV data
        st.session_state.csv_data = load_csv_data(uploaded_file)
        
        if st.session_state.csv_data is not None:
            st.success("CSV file loaded successfully!")
    
    # Columns for monitoring
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Normal Traffic Monitoring")
        if not st.session_state.collecting_data:
            if st.button("Start Normal Monitoring"):
                start_collection("Normal")
        else:
            if st.session_state.collection_type == "Normal":
                if st.button("Stop Normal Monitoring"):
                    collected_df = stop_collection()
                    if collected_df is not None:
                        st.success(f"Collected {len(collected_df)} records of normal traffic!")
                        # Display collected data
                        st.dataframe(collected_df)
    
    with col2:
        st.subheader("Attack Traffic Monitoring")
        if not st.session_state.collecting_data:
            if st.button("Start Attack Monitoring"):
                start_collection("Attack")
        else:
            if st.session_state.collection_type == "Attack":
                if st.button("Stop Attack Monitoring"):
                    collected_df = stop_collection()
                    if collected_df is not None:
                        st.success(f"Collected {len(collected_df)} records of attack traffic!")
                        # Display collected data
                        st.dataframe(collected_df)
    
    # Real-time display of current collection
    if st.session_state.collecting_data:
        # Create containers for dynamic content
        status_container = st.empty()
        data_container = st.empty()
        
        # Collect current data from queue
        current_data = []
        temp_queue = queue.Queue()
        
        while not st.session_state.data_queue.empty():
            item = st.session_state.data_queue.get()
            current_data.append(item)
            temp_queue.put(item)
        
        # Restore queue
        while not temp_queue.empty():
            st.session_state.data_queue.put(temp_queue.get())
        
        # Update display if new data exists
        if current_data:
            new_data = pd.concat(current_data, ignore_index=True)
            
            # Update display data
            if st.session_state.display_data.empty:
                st.session_state.display_data = new_data
            else:
                st.session_state.display_data = pd.concat([st.session_state.display_data, new_data], ignore_index=True)
            
            # Keep only latest records
            if len(st.session_state.display_data) > 10:
                st.session_state.display_data = st.session_state.display_data.tail(10)
            
            # Show status
            status_container.write(f"Currently collecting {st.session_state.collection_type} traffic data. Total records: {len(st.session_state.display_data)}")
            
            # Display latest records
            data_container.subheader("Latest Records")
            data_container.dataframe(st.session_state.display_data)

def collect_data_1(data_type, queue_obj):
    try:
        # Hard-coded path to the dataset
        dataset_path = r"C:\Users\mani3\Downloads\ExplainableIDS\ExplainableIDS\Dataset\CIC-IDS2017.csv"
        
        # Read the CSV file
        csv_data = pd.read_csv(dataset_path)
        
        # Filter based on data type
        if data_type == "normal":
            # Filter rows where Label indicates normal traffic
            # Assuming 0 or "BENIGN" or "Normal" could be used for normal traffic
            if "Label" in csv_data.columns:
                if csv_data["Label"].dtype == 'object':
                    subset = csv_data[csv_data["Label"].str.contains("BENIGN|Normal", case=False, na=False)]
                else:
                    subset = csv_data[csv_data["Label"] == 0]
            else:
                # If Label column doesn't exist, use the first 100 rows as a fallback
                subset = csv_data.head(100)
        else:  # attack
            # Filter rows where Label indicates attack traffic
            if "Label" in csv_data.columns:
                if csv_data["Label"].dtype == 'object':
                    subset = csv_data[~csv_data["Label"].str.contains("BENIGN|Normal", case=False, na=False)]
                else:
                    subset = csv_data[csv_data["Label"] != 0]
            else:
                # If Label column doesn't exist, use rows 100-200 as a fallback
                subset = csv_data.iloc[100:200]
        
        # If subset is empty, use a fallback
        if len(subset) == 0:
            subset = csv_data.head(100)
        
        # Take only the first 100 rows to keep it manageable
        subset = subset.head(100)
        
        # Start collecting
        rows_per_second = 2  # Adjust this to control how fast data appears
        row_index = 0
        
        while st.session_state.collecting_data and row_index < len(subset):
            # Get the next row
            row = subset.iloc[row_index:row_index+1].copy()
            row_index += 1
            
            # Add timestamp
            row['timestamp'] = time.strftime("%H:%M:%S")
            
            # Add to queue
            queue_obj.put(row)
            
            # Sleep to simulate real-time collection
            time.sleep(1/rows_per_second)
            
            # If we've gone through all rows, loop back to the beginning
            if row_index >= len(subset):
                row_index = 0
        
    except Exception as e:
        print(f"Error in data collection: {str(e)}")
        # Put error message in queue
        error_df = pd.DataFrame({"Error": [str(e)]})
        queue_obj.put(error_df)

# Function to start data collection
def start_collection_1(data_type):
    if st.session_state.collecting_data:
        st.warning("Data collection already in progress!")
        return
    
    st.session_state.collecting_data = True
    st.session_state.collection_type = data_type
    st.session_state.collected_data = []
    st.session_state.display_data = pd.DataFrame()
    
    # Start collection in a separate thread
    st.session_state.collection_thread = threading.Thread(
        target=collect_data, 
        args=(data_type, st.session_state.data_queue)
    )
    st.session_state.collection_thread.daemon = True
    st.session_state.collection_thread.start()

# Function to stop data collection
def stop_collectio_1():
    if not st.session_state.collecting_data:
        st.warning("No data collection in progress!")
        return
    
    st.session_state.collecting_data = False
    
    # Wait for thread to finish
    if st.session_state.collection_thread:
        st.session_state.collection_thread.join(timeout=1.0)
    
    # Get all data from queue
    while not st.session_state.data_queue.empty():
        st.session_state.collected_data.append(st.session_state.data_queue.get())
    
    # Convert to DataFrame
    if st.session_state.collected_data:
        collected_df = pd.concat(st.session_state.collected_data, ignore_index=True)
        return collected_df
    else:
        st.warning("No data was collected!")
        return None

# Function to predict on real-time data
def predict_real_time_data(collected_df, data_type):
    # For hard-coded predictions based on the collection type
    results = []
    
    # Create timestamps
    if 'timestamp' in collected_df.columns:
        timestamps = collected_df['timestamp'].tolist()
    else:
        timestamps = [time.strftime("%H:%M:%S") for _ in range(len(collected_df))]
    
    # Generate prediction results based on the requested data type
    for i in range(len(collected_df)):
        # Hard-code the prediction to match the collection type
        if data_type == "normal":
            prediction = "BENIGN"
            confidence = random.uniform(85.0, 99.9)  # High confidence for correct predictions
        else:  # attack
            prediction = "ATTACK"
            confidence = random.uniform(85.0, 99.9)  # High confidence for correct predictions
        
        results.append({
            "Timestamp": timestamps[i],
            "Confidence": f"{confidence:.2f}%",
            "Prediction": prediction
        })
    
    return pd.DataFrame(results)

# Function for real-time monitoring page
def real_time_monitoring_1():
    st.title("Real-time Network Monitoring")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Normal Traffic Monitoring")
        if not st.session_state.collecting_data:
            if st.button("Start Normal Monitoring"):
                start_collection("normal")
        else:
            if st.session_state.collection_type == "normal":
                if st.button("Stop Normal Monitoring"):
                    collected_df = stop_collection()
                    if collected_df is not None:
                        st.success(f"Collected {len(collected_df)} records of normal traffic!")
                        
                        # Show collected data
                        st.subheader("Collected Normal Traffic Data")
                        display_df = collected_df.drop('timestamp', axis=1, errors='ignore')
                        # Display only first 5 columns to keep the display manageable
                        display_cols = display_df.columns[:5].tolist()
                        st.write(display_df[display_cols])
                        
                        # Predict on collected data
                        st.subheader("Prediction Results")
                        results_df = predict_real_time_data(collected_df, "normal")
                        if results_df is not None:
                            st.write(results_df)
                            
                            # Count predictions
                            pred_counts = results_df['Prediction'].value_counts()
                            
                            # Create pie chart
                            fig, ax = plt.subplots()
                            ax.pie(pred_counts, labels=pred_counts.index, autopct='%1.1f%%')
                            ax.set_title("Prediction Distribution")
                            st.pyplot(fig)
    
    with col2:
        st.subheader("Attack Traffic Monitoring")
        if not st.session_state.collecting_data:
            if st.button("Start Attack Monitoring"):
                start_collection("attack")
        else:
            if st.session_state.collection_type == "attack":
                if st.button("Stop Attack Monitoring"):
                    collected_df = stop_collection()
                    if collected_df is not None:
                        st.success(f"Collected {len(collected_df)} records of attack traffic!")
                        
                        # Show collected data
                        st.subheader("Collected Attack Traffic Data")
                        display_df = collected_df.drop('timestamp', axis=1, errors='ignore')
                        # Display only first 5 columns to keep the display manageable
                        display_cols = display_df.columns[:5].tolist()
                        st.write(display_df[display_cols])
                        
                        # Predict on collected data
                        st.subheader("Prediction Results")
                        results_df = predict_real_time_data(collected_df, "attack")
                        if results_df is not None:
                            st.write(results_df)
                            
                            # Count predictions
                            pred_counts = results_df['Prediction'].value_counts()
                            
                            # Create pie chart
                            fig, ax = plt.subplots()
                            ax.pie(pred_counts, labels=pred_counts.index, autopct='%1.1f%%')
                            ax.set_title("Prediction Distribution")
                            st.pyplot(fig)
    
    # Display status and live data
    if st.session_state.collecting_data:
        # Create containers for dynamic content
        status_container = st.empty()
        data_container = st.empty()
        prediction_container = st.empty()
        
        # Get all current data from queue without removing
        current_data = []
        temp_queue = queue.Queue()
        
        while not st.session_state.data_queue.empty():
            item = st.session_state.data_queue.get()
            current_data.append(item)
            temp_queue.put(item)
        
        # Restore queue
        while not temp_queue.empty():
            st.session_state.data_queue.put(temp_queue.get())
        
        # Update display data with new records
        if current_data:
            new_data = pd.concat(current_data, ignore_index=True)
            if st.session_state.display_data.empty:
                st.session_state.display_data = new_data
            else:
                st.session_state.display_data = pd.concat([st.session_state.display_data, new_data], ignore_index=True)
            
            # Keep only the latest 10 records to avoid overloading the display
            if len(st.session_state.display_data) > 10:
                st.session_state.display_data = st.session_state.display_data.tail(10)
            
            # Show status
            status_container.write(f"Currently collecting {st.session_state.collection_type} traffic data. Total records captured: {len(st.session_state.display_data)}")
            
            # Display the latest records
            data_container.subheader("Latest Records (Live)")
            display_df = st.session_state.display_data.drop('timestamp', axis=1, errors='ignore')
            # Display only first 5 columns to keep the display manageable
            display_cols = display_df.columns[:5].tolist()
            data_container.write(display_df[display_cols])
            
            # Show live predictions
            prediction_container.subheader("Live Predictions")
            results_df = predict_real_time_data(st.session_state.display_data, st.session_state.collection_type)
            prediction_container.write(results_df)
        else:
            status_container.write(f"Currently collecting {st.session_state.collection_type} traffic data. Waiting for records...")



# Function to upload dataset
def upload_dataset():
    uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    if uploaded_file is not None:
        st.session_state.filename = uploaded_file.name
        st.session_state.dataset = pd.read_csv(uploaded_file)
        st.write(f"Dataset {st.session_state.filename} loaded")
        st.write(st.session_state.dataset)
        
        # Display label distribution
        st.session_state.labels = np.unique(st.session_state.dataset["Label"])
        st.subheader("Normal & Attack Distribution")
        fig, ax = plt.subplots(figsize=(10, 5))
        label_counts = st.session_state.dataset.groupby('Label').size()
        label_counts.plot(kind="bar", ax=ax)
        ax.set_xlabel("Data Type")
        ax.set_ylabel("Count")
        ax.set_title("Normal & Attack Graph")
        st.pyplot(fig)

# Function to preprocess dataset
def preprocess_dataset():
    if st.session_state.dataset is None:
        st.error("Please upload dataset first!")
        return
    
    st.subheader("Preprocessing Dataset")
    
    with st.spinner("Preprocessing in progress..."):
        st.session_state.label_encoder = []
        columns = st.session_state.dataset.columns
        types = st.session_state.dataset.dtypes.values
        
        for j in range(len(types)):
            name = types[j]
            if name == 'object':  # finding column with object type
                le = LabelEncoder()
                st.session_state.dataset[columns[j]] = pd.Series(
                    le.fit_transform(st.session_state.dataset[columns[j]].astype(str)))  # encode all str columns to numeric
                st.session_state.label_encoder.append([columns[j], le])
        
        st.session_state.dataset.fillna(0, inplace=True)
        st.session_state.dataset = st.session_state.dataset[np.isfinite(st.session_state.dataset).all(1)]
        st.session_state.perturb = st.session_state.dataset['FlowDuration'].ravel()
        st.session_state.Y = st.session_state.dataset['Label'].ravel()
        st.session_state.dataset.drop(['Label', 'FlowDuration'], axis=1, inplace=True)
        st.session_state.features = st.session_state.dataset.columns
        st.session_state.X = st.session_state.dataset.values
        
        st.session_state.scaler = StandardScaler()
        st.session_state.X = st.session_state.scaler.fit_transform(st.session_state.X)
        st.session_state.Y = to_categorical(st.session_state.Y)
    
    st.success("Preprocessing completed!")
    st.write("Processed & Normalized Dataset Features:")
    st.write(pd.DataFrame(st.session_state.X, columns=st.session_state.features))

# Function to split dataset into train and test sets
def train_test_split_dataset():
    if st.session_state.X is None or st.session_state.Y is None:
        st.error("Please preprocess dataset first!")
        return
    
    st.subheader("Train Test Split")
    
    with st.spinner("Splitting dataset..."):
        # Reshape X for LSTM input (samples, time steps, features)
        X_reshaped = np.reshape(st.session_state.X, (st.session_state.X.shape[0], st.session_state.X.shape[1], 1))
        st.session_state.X_train, st.session_state.X_test, st.session_state.y_train, st.session_state.y_test = train_test_split(
            X_reshaped, st.session_state.Y, test_size=0.2)
        
        # Set perturb values for some test samples
        for i in range(0, min(10, st.session_state.X_test.shape[0])):
            st.session_state.X_test[i, 0] = st.session_state.perturb[i]
    
    st.success("Dataset split completed!")
    st.write(f"Total records found in dataset: {st.session_state.X.shape[0]}")
    st.write(f"80% dataset records used to train Algorithms: {st.session_state.X_train.shape[0]}")
    st.write(f"20% dataset records used to test Algorithms: {st.session_state.X_test.shape[0]}")

# Function to calculate metrics and display results
def calculate_metrics(algorithm, testY, predict):
    p = precision_score(testY, predict, average='macro') * 100
    r = recall_score(testY, predict, average='macro') * 100
    f = f1_score(testY, predict, average='macro') * 100
    a = accuracy_score(testY, predict) * 100
    
    # Store metrics in session state
    st.session_state.accuracy.append(a)
    st.session_state.precision.append(p)
    st.session_state.recall.append(r)
    st.session_state.fscore.append(f)
    
    # Display metrics
    st.write(f"{algorithm} Accuracy: {a:.2f}%")
    st.write(f"{algorithm} Precision: {p:.2f}%")
    st.write(f"{algorithm} Recall: {r:.2f}%")
    st.write(f"{algorithm} F1 Score: {f:.2f}%")
    
    # Create confusion matrix and ROC curve
    fig, axs = plt.subplots(1, 2, figsize=(15, 6))
    
    # Confusion Matrix
    conf_matrix = confusion_matrix(testY, predict)
    sns.heatmap(conf_matrix, xticklabels=st.session_state.labels, yticklabels=st.session_state.labels, 
                annot=True, cmap="viridis", fmt="g", ax=axs[0])
    axs[0].set_ylim([0, len(st.session_state.labels)])
    axs[0].set_title(f"{algorithm} Confusion Matrix")
    
    # ROC Curve
    random_probs = [0 for i in range(len(testY))]
    p_fpr, p_tpr, _ = roc_curve(testY, random_probs, pos_label=1)
    axs[1].plot(p_fpr, p_tpr, linestyle='--', color='orange', label="True classes")
    ns_fpr, ns_tpr, _ = roc_curve(testY, predict, pos_label=1)
    axs[1].plot(ns_fpr, ns_tpr, linestyle='--', label='Predicted Classes')
    axs[1].set_title(f"{algorithm} ROC AUC Curve")
    axs[1].set_xlabel('False Positive Rate')
    axs[1].set_ylabel('True Positive Rate')
    axs[1].legend()
    
    st.pyplot(fig)

# Function to train and evaluate LSTM model
def run_lstm():
    if st.session_state.X_train is None or st.session_state.y_train is None:
        st.error("Please split dataset first!")
        return
    
    st.subheader("LSTM Model Training and Evaluation")
    
    with st.spinner("Training LSTM model..."):
        # Clear previous metrics
        st.session_state.accuracy = []
        st.session_state.precision = []
        st.session_state.recall = []
        st.session_state.fscore = []
        
        # Define LSTM model
        st.session_state.lstm_model = Sequential()
        st.session_state.lstm_model.add(LSTM(32, input_shape=(st.session_state.X_train.shape[1], st.session_state.X_train.shape[2])))
        st.session_state.lstm_model.add(Dropout(0.3))
        st.session_state.lstm_model.add(Dense(32, activation='relu'))
        st.session_state.lstm_model.add(Dense(st.session_state.y_train.shape[1], activation='softmax'))
        st.session_state.lstm_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
        
        # Create directory for model if it doesn't exist
        os.makedirs("model", exist_ok=True)
        
        # Train or load model
        if not os.path.exists("model/lstm_weights.hdf5"):
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            # Custom callback to update Streamlit progress bar
            class StreamlitCallback(ModelCheckpoint):
                def on_epoch_end(self, epoch, logs=None):
                    progress = (epoch + 1) / 20  # Assuming 20 epochs
                    progress_bar.progress(progress)
                    status_text.text(f"Epoch {epoch+1}/20 - loss: {logs['loss']:.4f} - accuracy: {logs['accuracy']:.4f}")
                    super().on_epoch_end(epoch, logs)
            
            model_check_point = StreamlitCallback(filepath='model/lstm_weights.hdf5', verbose=1, save_best_only=True)
            hist = st.session_state.lstm_model.fit(
                st.session_state.X_train, st.session_state.y_train, 
                batch_size=32, epochs=20, 
                validation_data=(st.session_state.X_test, st.session_state.y_test),
                callbacks=[model_check_point], verbose=0
            )
            
            with open('model/lstm_history.pckl', 'wb') as f:
                pickle.dump(hist.history, f)
        else:
            st.info("Loading pre-trained LSTM model...")
            st.session_state.lstm_model.load_weights("model/lstm_weights.hdf5")
        
        # Evaluate model
        predict = st.session_state.lstm_model.predict(st.session_state.X_test)
        predict = np.argmax(predict, axis=1)
        y_test1 = np.argmax(st.session_state.y_test, axis=1)
        
        st.success("LSTM model evaluation completed!")
        calculate_metrics("LSTM", y_test1, predict)

# Function to train and evaluate TCN model
def run_tcn():
    if st.session_state.X_train is None or st.session_state.y_train is None:
        st.error("Please split dataset first!")
        return
    
    st.subheader("TCN Model Training and Evaluation")
    
    with st.spinner("Training TCN model..."):
        # Define TCN model
        st.session_state.tcn_model = compiled_tcn(
            return_sequences=False, 
            num_feat=1, 
            num_classes=2, 
            nb_filters=20, 
            kernel_size=6, 
            dilations=[2 ** i for i in range(9)], 
            nb_stacks=1,
            max_len=st.session_state.X_train[0:1].shape[1], 
            use_skip_connections=True
        )
        
        # Create directory for model if it doesn't exist
        os.makedirs("model", exist_ok=True)
        
        # Train or load model
        if not os.path.exists("model/tcn_weights.hdf5"):
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            # Custom callback to update Streamlit progress bar
            class StreamlitCallback(ModelCheckpoint):
                def on_epoch_end(self, epoch, logs=None):
                    progress = (epoch + 1) / 20  # Assuming 20 epochs
                    progress_bar.progress(progress)
                    status_text.text(f"Epoch {epoch+1}/20 - loss: {logs['loss']:.4f} - accuracy: {logs['accuracy']:.4f}")
                    super().on_epoch_end(epoch, logs)
            
            model_check_point = StreamlitCallback(filepath='model/tcn_weights.hdf5', verbose=1, save_best_only=True)
            hist = st.session_state.tcn_model.fit(
                st.session_state.X_train, 
                st.session_state.y_train.squeeze().argmax(axis=1), 
                epochs=20, 
                validation_data=(st.session_state.X_test, st.session_state.y_test.squeeze().argmax(axis=1)),
                callbacks=[model_check_point], 
                verbose=0
            )
            
            with open('model/tcn_history.pckl', 'wb') as f:
                pickle.dump(hist.history, f)
        else:
            st.info("Loading pre-trained TCN model...")
            st.session_state.tcn_model.load_weights("model/tcn_weights.hdf5")
        
        # Evaluate model
        predict = st.session_state.tcn_model.predict(st.session_state.X_test)
        predict = np.argmax(predict, axis=1)
        y_test1 = np.argmax(st.session_state.y_test, axis=1)
        
        st.success("TCN model evaluation completed!")
        calculate_metrics("TCN", y_test1, predict)

# Function to display comparison graph
def show_comparison_graph():
    if len(st.session_state.accuracy) < 2:
        st.error("Please train both LSTM and TCN models first!")
        return
    
    st.subheader("Model Comparison")
    
    # Create DataFrame for comparison
    df = pd.DataFrame([
        ['LSTM', 'Precision', st.session_state.precision[0]],
        ['LSTM', 'Recall', st.session_state.recall[0]],
        ['LSTM', 'F1 Score', st.session_state.fscore[0]],
        ['LSTM', 'Accuracy', st.session_state.accuracy[0]],
        ['TCN', 'Precision', st.session_state.precision[1]],
        ['TCN', 'Recall', st.session_state.recall[1]],
        ['TCN', 'F1 Score', st.session_state.fscore[1]],
        ['TCN', 'Accuracy', st.session_state.accuracy[1]],
    ], columns=['Parameters', 'Algorithms', 'Value'])
    
    # Plot comparison
    fig, ax = plt.subplots(figsize=(12, 6))
    df.pivot("Parameters", "Algorithms", "Value").plot(kind='bar', ax=ax)
    ax.set_ylabel("Percentage (%)")
    ax.set_title("Model Performance Comparison")
    st.pyplot(fig)

# Function to make predictions
# Function to make predictions
def predict_attack():
    if st.session_state.lstm_model is None:
        st.error("Please train LSTM model first!")
        return
    
    st.subheader("Predict Attack from Test Data")
    
    uploaded_file = st.file_uploader("Upload test data (CSV)", type="csv")
    if uploaded_file is not None:
        testdata = pd.read_csv(uploaded_file)
        temp = testdata.values
        
        try:
            # Preprocess test data
            for i in range(len(st.session_state.label_encoder)):
                le = st.session_state.label_encoder[i]
                if le[0] in testdata.columns:
                    testdata[le[0]] = pd.Series(le[1].transform(testdata[le[0]].astype(str)))
            
            testdata.fillna(0, inplace=True)
            testdata = testdata[np.isfinite(testdata).all(1)]
            
            # Make sure we drop FlowDuration if it exists
            if 'FlowDuration' in testdata.columns:
                testdata.drop(['FlowDuration'], axis=1, inplace=True)
            
            # Check if columns match what the model expects
            expected_features = st.session_state.features
            missing_cols = set(expected_features) - set(testdata.columns)
            extra_cols = set(testdata.columns) - set(expected_features)
            
            if missing_cols:
                st.warning(f"Missing columns in test data: {missing_cols}")
                for col in missing_cols:
                    testdata[col] = 0  # Add missing columns with default values
            
            if extra_cols:
                st.warning(f"Extra columns in test data, removing: {extra_cols}")
                testdata = testdata.drop(columns=extra_cols)
            
            # Reorder columns to match training data
            testdata = testdata[expected_features]
            
            # Apply scaling
            testdata_values = st.session_state.scaler.transform(testdata)
            
            # Reshape for LSTM input (samples, time steps, features)
            testdata_reshaped = np.reshape(testdata_values, (testdata_values.shape[0], testdata_values.shape[1], 1))
            
            # Make predictions with explicit graph and session handling
            import tensorflow as tf
            with tf.Graph().as_default():
                # Create a new instance of the model with the same architecture
                model = Sequential()
                model.add(LSTM(32, input_shape=(testdata_reshaped.shape[1], testdata_reshaped.shape[2])))
                model.add(Dropout(0.3))
                model.add(Dense(32, activation='relu'))
                model.add(Dense(st.session_state.y_train.shape[1], activation='softmax'))
                model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
                
                # Load weights from the trained model
                model.load_weights("model/lstm_weights.hdf5")
                
                # Make predictions
                predict = model.predict(testdata_reshaped)
            
            predict = np.argmax(predict, axis=1)
            
            # Display results
            results = []
            for i in range(len(predict)):
                results.append({
                    "Test Data": str(temp[i][0:5]),  # Show first 5 values for readability
                    "Prediction": st.session_state.labels[predict[i]]
                })
            
            st.write(pd.DataFrame(results))
            
        except Exception as e:
            st.error(f"Error during prediction: {str(e)}")
            st.text("Detailed error information:")
            st.exception(e)
            
            # Debugging information
            st.subheader("Debugging Information")
            st.write("Test data shape:", testdata.shape if 'testdata' in locals() else "Not available")
            st.write("Expected features:", expected_features if 'expected_features' in locals() else "Not available")
            if 'testdata' in locals():
                st.write("Test data columns:", testdata.columns.tolist())
            if 'testdata_reshaped' in locals():
                st.write("Reshaped test data shape:", testdata_reshaped.shape)
# Function for SHAP explanation
def explanation():
    if st.session_state.X_train is None or st.session_state.y_train is None:
        st.error("Please split dataset first!")
        return
    
    st.subheader("SHAP Model Explanation")
    
    with st.spinner("Generating SHAP explanation..."):
        # Reshape data for RandomForest
        X_train_reshaped = np.reshape(st.session_state.X_train, (st.session_state.X_train.shape[0], st.session_state.X_train.shape[1] * st.session_state.X_train.shape[2]))
        X_test_reshaped = np.reshape(st.session_state.X_test, (st.session_state.X_test.shape[0], st.session_state.X_test.shape[1] * st.session_state.X_test.shape[2]))
        y_test_reshaped = np.argmax(st.session_state.y_test, axis=1)
        y_train_reshaped = np.argmax(st.session_state.y_train, axis=1)
        
        # Train RandomForest
        rf = RandomForestClassifier()
        rf.fit(X_train_reshaped, y_train_reshaped)
        
        # Generate SHAP values
        explainer = shap.TreeExplainer(rf, X_test_reshaped)
        st.session_state.shap_values = explainer.shap_values(X_test_reshaped, check_additivity=False)
        
        # Plot SHAP summary
        st.write("SHAP Summary Plot")
        fig, ax = plt.subplots(figsize=(12, 8))
        shap.summary_plot(st.session_state.shap_values[0], feature_names=st.session_state.features, show=False)
        st.pyplot(fig)

# Function for Violin explanation
def violin_explanation():
    if st.session_state.shap_values is None:
        st.error("Please generate SHAP explanation first!")
        return
    
    st.subheader("SHAP Violin Plot Explanation")
    
    with st.spinner("Generating violin plot..."):
        # Plot SHAP violin plot
        fig, ax = plt.subplots(figsize=(12, 8))
        shap.plots.violin(st.session_state.shap_values[0], feature_names=st.session_state.features, show=False)
        st.pyplot(fig)

def real_time_monitoring():
    """Open a1.py when Real-time Monitoring is selected"""
    script_path = os.path.join(os.path.dirname(__file__), r"C:\Users\mani3\Downloads\ExplainableIDS\ExplainableIDS\a1.py")
    
    # Run 'a1.py' as a separate Streamlit app
    subprocess.Popen(["streamlit", "run", script_path])
    st.stop()  # Stop execution of the current app
# Render the selected page
if page == "Upload Dataset":
    upload_dataset()
elif page == "Preprocess Dataset":
    preprocess_dataset()
elif page == "Train & Test Split":
    train_test_split_dataset()
elif page == "Model Training":
    model_tab = st.radio("Select Model", ["LSTM", "TCN"])
    if model_tab == "LSTM":
        run_lstm()
    else:
        run_tcn()
elif page == "Evaluation":
    show_comparison_graph()
elif page == "Prediction":
    predict_attack()
elif page == "Explainability":
    explain_tab = st.radio("Select Explanation Type", ["Summary Plot", "Violin Plot"])
    if explain_tab == "Summary Plot":
        explanation()
    else:
        violin_explanation()
elif page == "Real-time Monitoring":
    real_time_monitoring()

# Add footer
st.markdown("---")
st.markdown("Explainable AI for Intrusion Detection Systems - Streamlit Version")


I want to add another model here which is hybrid so can u add the hybrid model which is the combonation of lstm and tcn and shows all the metrics like it shows for all the other model in the code