import os
import numpy as np
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from utils.preprocess import load_images_from_folder  # adjust import if needed

# Parameters
DATASET_PATH = './data/train'
IMG_SIZE = (48, 48)  # Example size (width, height)
NUM_CLASSES = 7      # Number of emotion classes, adjust accordingly
EPOCHS = 30
BATCH_SIZE = 64

# Load data
x_data, y_data, emotions = load_images_from_folder(DATASET_PATH, IMG_SIZE)

# Preprocess labels
y_data_cat = to_categorical(y_data, num_classes=NUM_CLASSES)

# Normalize images
x_data = x_data.astype('float32') / 255.0

# Build a simple CNN model
model = Sequential([
    Conv2D(32, (3,3), activation='relu', input_shape=(IMG_SIZE[0], IMG_SIZE[1], 1)),
    MaxPooling2D(2,2),
    Conv2D(64, (3,3), activation='relu'),
    MaxPooling2D(2,2),
    Conv2D(128, (3,3), activation='relu'),
    MaxPooling2D(2,2),
    Flatten(),
    Dense(128, activation='relu'),
    Dropout(0.5),
    Dense(NUM_CLASSES, activation='softmax')
])

model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])

# Train
model.fit(x_data, y_data_cat, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_split=0.2)

# Save model
os.makedirs('./models', exist_ok=True)
model.save('./models/emotion_model.h5')

print("Training complete and model saved.")