# main.py
import numpy as np
import tensorflow as tf
from sklearn.preprocessing import StandardScaler
from feature_select import feature_selection
import pandas as pd
import os
import time
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Step 1: Load MNIST dataset
print("Loading MNIST dataset...")
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
X_train = X_train.reshape(X_train.shape[0], -1)
X_test = X_test.reshape(X_test.shape[0], -1)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

y_train = y_train.flatten()
y_test = y_test.flatten()


# Define algorithms to test
algorithms = ['BVRSZHTn', 'SAGA', 'BVRSZHT12', 'VRSZHT', 'SARAH']
results = []

# Loop over each algorithm and perform feature selection, training, and evaluation
for algorithm in algorithms:
    print(f"Applying feature selection with {algorithm}...")
    start_time = time.time()
    X_train_selected, X_test_selected = feature_selection(algorithm, X_train, y_train, X_test)
    selection_time = time.time() - start_time

    # Record the number of selected features
    num_features = X_train_selected.shape[1]

    # Train a DNN classifier for the selected features
    print(f"Training DNN classifier for {algorithm}-selected features...")
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=(X_train_selected.shape[1],)),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # Train the model
    model.fit(X_train_selected, y_train, epochs=10, batch_size=32, validation_split=0.1, verbose=2)

    # Evaluate the model
    test_loss, test_acc = model.evaluate(X_test_selected, y_test, verbose=0)
    print(f"Test accuracy for {algorithm}-selected features: {test_acc:.2f}")

    # Record results
    results.append({'Algorithm': algorithm, 'Accuracy': test_acc, 'Num_Features': num_features, 'Selection_Time (s)': selection_time})

# Create a DataFrame to store the results
results_df = pd.DataFrame(results)
print(results_df)

# Save the results to a CSV file
results_df.to_csv('feature_selection_results.csv', index=False)
