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

print("=" * 70)
print("🔵 QUICK CNN MODEL - DROWSY vs NON-DROWSY")
print("=" * 70)

def load_small_sample(data_path, samples_per_class=1000):
    """Load a small balanced sample for quick training"""
    images = []
    labels = []
    
    for class_name, class_idx in [('drowsy', 1), ('non_drowsy', 0)]:
        class_path = os.path.join(data_path, 'ddd', class_name)
        if os.path.exists(class_path):
            img_files = [f for f in os.listdir(class_path) if f.endswith(('.jpg', '.png', '.jpeg'))]
            
            # Take random sample
            if len(img_files) > samples_per_class:
                img_files = np.random.choice(img_files, samples_per_class, replace=False)
            
            print(f"  Loading {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 small sample
data_path = 'data/processed'
print("\n[1/4] Loading small balanced sample...")
X, y = load_small_sample(data_path, samples_per_class=1000)

print(f"\n✅ Total images: {len(X)}")
print(f"   Drowsy: {np.sum(y==1)} | Non-Drowsy: {np.sum(y==0)}")

# Preprocess
X = X.astype('float32') / 255.0

# Split
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
y_train_cat = tf.keras.utils.to_categorical(y_train, 2)
y_val_cat = tf.keras.utils.to_categorical(y_val, 2)

print(f"\n[2/4] Data split: Train={len(X_train)}, Val={len(X_val)}")

# Build model
print("\n[3/4] Building model...")
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.Dropout(0.5),
    tf.keras.layers.Dense(2, activation='softmax')
])

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

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

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

print("\n" + "=" * 70)
print("✅ TRAINING COMPLETE!")
print("=" * 70)
print(f"📊 Best Validation Accuracy: {max(history.history['val_accuracy']):.2%}")
print(f"📁 Model saved to: data/models/cnn_ddd_quick.h5")