import numpy as np
import tensorflow as tf
import cv2
import os

# Tamaño deseado de las imágenes
IMG_SIZE = (48, 48)

# Rutas a las carpetas de entrenamiento y testeo
train_dir = 'archive/train'
test_dir = 'archive/test'  

# Función para cargar imágenes y etiquetas
def load_data(directory):
    X = []
    y = []
    emotions = ['fear', 'neutral', 'happy']  
    for idx, emotion in enumerate(emotions):
        emotion_dir = os.path.join(directory, emotion)  
        for image_name in os.listdir(emotion_dir):
            image_path = os.path.join(emotion_dir, image_name)  
            image = cv2.imread(image_path)
            if image is not None:  
                # Redimensionar la imagen
                resized_image = cv2.resize(image, IMG_SIZE)
                X.append(resized_image)
                y.append(idx)  
            else:
                print(f"Error al cargar la imagen: {image_path}")
    return np.array(X), np.array(y)

# Cargar datos de entrenamiento y prueba
X_train, y_train = load_data(train_dir)
X_test, y_test = load_data(test_dir)

# Normalizar las imágenes
X_train = X_train.astype('float32') / 255.0
X_test = X_test.astype('float32') / 255.0

# Definir la arquitectura del modelo
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(*IMG_SIZE, 3)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')  
])

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

# Entrenar el modelo
model.fit(X_train, y_train, epochs=17, batch_size=32, validation_split=0.2)

# Evaluar el modelo con los datos de prueba
test_loss, test_acc = model.evaluate(X_test, y_test)  
print('Precisión del modelo en los datos de prueba:', test_acc)

# Guardar el modelo
model.save('modelo_emociones.h5')