import tensorflow as tf
import numpy as np
import os
import cv2
from sklearn.model_selection import train_test_split

print("=" * 70)
print("👁️ TRAINING EYE MODEL - OPEN vs CLOSED")
print("=" * 70)

def load_eye_data(data_path):
    images, labels = [], []
    
    for split in ['train', 'test']:
        for class_name, class_idx in [('open', 0), ('closed', 1)]:
            class_path = os.path.join(data_path, 'yawn_eye', split, class_name)
            if os.path.exists(class_path):
                img_files = [f for f in os.listdir(class_path) if f.endswith(('.jpg', '.png', '.jpeg'))]
                print(f"  Loading {split}/{class_name}: {len(img_files)} images")
                
                for img_name in img_files:
                    img_path = os.path.join(class_path, img_name)
                    img = cv2.imread(img_path)
                    if img is not None:
                        img = cv2.resize(img, (224, 224))
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                        images.append(img)
                        labels.append(class_idx)
    
    return np.array(images), np.array(labels)

# Load data
data_path = 'data/processed'
print("\n[1/4] Loading Eye dataset...")
X, y = load_eye_data(data_path)

if len(X) == 0:
    print("\n⚠️ No eye images found. Creating synthetic data for testing...")
    X = np.random.random((200, 224, 224, 3)) * 255
    y = np.random.randint(0, 2, 200)
else:
    print(f"\n✅ Total eye images: {len(X)}")
    print(f"   Open: {np.sum(y==0)} | Closed: {np.sum(y==1)}")

X = X.astype('float32') / 255.0
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y if len(y)>0 else None)
y_train = tf.keras.utils.to_categorical(y_train, 2)
y_val = tf.keras.utils.to_categorical(y_val, 2)

# Build model
print("\n[2/4] Building Eye Model (MobileNetV2)...")
base_model = tf.keras.applications.MobileNetV2(
    weights='imagenet', include_top=False, input_shape=(224, 224, 3)
)
base_model.trainable = False

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.4),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(2, activation='softmax')
])

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

# Train
print("\n[3/4] Training Eye Model...")
history = model.fit(X_train, y_train, validation_data=(X_val, y_val), epochs=10, batch_size=16, verbose=1)

os.makedirs('data/models', exist_ok=True)
model.save('data/models/eye_model.h5')

print("\n[4/4] ✅ Eye Model saved to: data/models/eye_model.h5")
print(f"   Validation Accuracy: {history.history['val_accuracy'][-1]:.2%}")