import os
import random
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.optimizers import Adam
import json
from tqdm import tqdm

# Set random seed for reproducibility
random.seed(47)
tf.random.set_seed(47)

# Directory to save models
model_dir = "/Users/nihar/Desktop/Desktop - nihar’s MacBook Air/cc/code files/auto-ML/NNbasics/modelllls"
if not os.path.exists(model_dir):
    os.makedirs(model_dir)

# Clear the directory
for f in os.listdir(model_dir):
    os.remove(os.path.join(model_dir, f))

# Load 50% of MNIST data
(x_train_full, y_train_full), (x_test_full, y_test_full) = mnist.load_data()
x_train_full, x_test_full = x_train_full / 255.0, x_test_full / 255.0

# Use 50% of the data
x_train_full, y_train_full = x_train_full[:30000], y_train_full[:30000]
x_test_full, y_test_full = x_test_full[:5000], y_test_full[:5000]

# Split into training, validation, and test sets
train_size = int(0.7 * len(x_train_full))
val_size = int(0.2 * len(x_train_full))
test_size = len(x_train_full) - train_size - val_size

x_train, y_train = x_train_full[:train_size], y_train_full[:train_size]
x_val, y_val = x_train_full[train_size:train_size + val_size], y_train_full[train_size:train_size + val_size]
x_test, y_test = x_train_full[train_size + val_size:], y_train_full[train_size + val_size:]

# Generate more values for learning rates, batch sizes, and epochs
learning_rates = [round(0.00001 * (1.2 ** i), 8) for i in range(20)]
batch_sizes = [8, 16, 32, 64, 128, 256, 512]
epochs_list = [10, 15, 20, 25, 30, 35, 40, 45, 50]

# Generate 999 random combinations
combinations = set()
while len(combinations) < 999:
    lr = random.choice(learning_rates)
    bs = random.choice(batch_sizes)
    epochs = random.choice(epochs_list)
    combinations.add((lr, bs, epochs))

# Train and save models with progress bar
for i, (lr, bs, epochs) in enumerate(tqdm(combinations, desc="Training models")):
    model = Sequential([
        Flatten(input_shape=(28, 28)),
        Dense(128, activation='relu'),
        Dense(10, activation='softmax')
    ])
    
    model.compile(optimizer=Adam(learning_rate=lr),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    
    history = model.fit(x_train, y_train, epochs=epochs, batch_size=bs, validation_data=(x_val, y_val), verbose=0)
    
    model_filename = f"mnist_model_lr{lr}_bs{bs}_epochs{epochs}.keras"
    model.save(os.path.join(model_dir, model_filename))
    
    history_filename = model_filename.replace('.keras', '.json')
    with open(os.path.join(model_dir, history_filename), 'w') as f:
        json.dump(history.history, f)