# @title
# =============================================================================
# GOOGLE COLAB: MOUNT DRIVE & PROJECT PATH CONFIGURATION
# =============================================================================

from google.colab import drive
import os

# Mount Google Drive
drive.mount('/content/drive')

# @title
# =============================================================================
# PROJECT PATH CONFIGURATION
# =============================================================================

DRIVE_ROOT = "/content/drive/MyDrive"
COURSE_DIR = "AAIDS-Course"
PROJECT_DIR = "3-Capstone Project"
SUBJECT_DIR = "Deep Learning"
TOPIC_DIR = "Facial Emotion"

# Construct full path
BASE_PATH = os.path.join(DRIVE_ROOT, COURSE_DIR, PROJECT_DIR, SUBJECT_DIR, TOPIC_DIR)

# Verify and change to project directory
if not os.path.exists(BASE_PATH):
    print(f'❌ ERROR: Project path not found: {BASE_PATH}')
    print(f'   Please verify your Google Drive folder structure.')
else:
    os.chdir(BASE_PATH)
    print(f'✅ Google Drive mounted')
    print(f'✅ Working directory: {os.getcwd()}')

    # List available datasets (commented out - slow on startup)
    print(f'\n📁 Available datasets:')
    for item in sorted(os.listdir('.')):
        if os.path.isdir(item) and not item.startswith('.') and 'emotion' in item.lower():
            try:
                count = sum(1 for root, dirs, files in os.walk(item) for f in files if f.lower().endswith(('.jpg', '.jpeg', '.png')))
                print(f'   {item}/ ({count:,} images)')
            except:
                print(f'   {item}/')

# @title
# =============================================================================
# IMPORTS
# =============================================================================

import os
import sys
import time
import pickle
import hashlib
import warnings
from datetime import datetime
from collections import defaultdict, Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np
import pandas as pd
from PIL import Image
from tqdm.notebook import tqdm

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import (
    Input, Conv2D, MaxPooling2D, Dense, Dropout, Flatten,
    BatchNormalization, Activation, RandomFlip, RandomRotation, RandomZoom, RandomContrast
)
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, ModelCheckpoint
from tensorflow.keras.regularizers import l2
from tensorflow.keras.utils import to_categorical

from sklearn.utils.class_weight import compute_class_weight
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px

warnings.filterwarnings('ignore')

print(f'TensorFlow version: {tf.__version__}')
print(f'GPU available: {tf.config.list_physical_devices("GPU")}')

# @title
# =============================================================================
# REPRODUCIBILITY & CONFIGURATION
# =============================================================================

import numpy as np
import tensorflow as tf

SEED = 42
np.random.seed(SEED)
tf.random.set_seed(SEED)

# Model constants
IMG_SIZE = 48
INPUT_SHAPE = (IMG_SIZE, IMG_SIZE, 1)  # Grayscale images
BATCH_SIZE = 64
NUM_CLASSES = 4
CLASS_NAMES = ['happy', 'neutral', 'sad', 'surprise']
SPLITS = ['train', 'validation', 'test']
MAX_EPOCHS = 75
INITIAL_LR = 0.0005  # Initial learning rate for cosine decay
LABEL_SMOOTHING = 0.1  # Label smoothing for B+ and B++ models

print(f'✅ Configuration set:')
print(f'   Random seed: {SEED}')
print(f'   Image size: {IMG_SIZE}x{IMG_SIZE}')
print(f'   Input shape: {INPUT_SHAPE}')
print(f'   Batch size: {BATCH_SIZE}')
print(f'   Classes: {CLASS_NAMES}')

# @title
# =============================================================================
# EXECUTION TIMING TRACKER
# =============================================================================
# Tracks execution times for all key operations to enable performance analysis
# and comparison across notebook runs.
# =============================================================================

import time
from datetime import datetime

# Initialize timing tracker
TIMING_DATA = {
    'notebook_start': time.time(),
    'notebook_start_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
    'data_loading': {},
    'model_training': {},
    'model_parameters': {},
    'system_info': {}
}

def start_timer(operation_name):
    """Start timing an operation."""
    TIMING_DATA[f'_start_{operation_name}'] = time.time()
    return time.time()

def stop_timer(operation_name, category='misc'):
    """Stop timing and record the duration."""
    start_key = f'_start_{operation_name}'
    if start_key in TIMING_DATA:
        duration = time.time() - TIMING_DATA[start_key]
        if category not in TIMING_DATA:
            TIMING_DATA[category] = {}
        TIMING_DATA[category][operation_name] = duration
        del TIMING_DATA[start_key]
        return duration
    return 0

def format_time(seconds):
    """Format seconds into human-readable string."""
    if seconds < 60:
        return f"{seconds:.1f}s"
    elif seconds < 3600:
        mins = seconds / 60
        return f"{mins:.1f}m"
    else:
        hours = seconds / 3600
        return f"{hours:.2f}h"

# Capture system info
try:
    import subprocess
    # Check for GPU
    try:
        gpu_info = subprocess.check_output(['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader'],
                                           stderr=subprocess.DEVNULL).decode().strip()
        TIMING_DATA['system_info']['gpu'] = gpu_info
        TIMING_DATA['system_info']['accelerator'] = 'GPU'
    except:
        # Check for TPU
        try:
            import tensorflow as tf
            tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
            TIMING_DATA['system_info']['accelerator'] = 'TPU'
            TIMING_DATA['system_info']['tpu'] = str(tpu.cluster_spec())
        except:
            TIMING_DATA['system_info']['accelerator'] = 'CPU'

    # Get memory info
    with open('/proc/meminfo', 'r') as f:
        meminfo = f.read()
        for line in meminfo.split('\n'):
            if 'MemTotal' in line:
                mem_kb = int(line.split()[1])
                TIMING_DATA['system_info']['ram_gb'] = round(mem_kb / 1024 / 1024, 1)
                break

    # Get CPU info
    with open('/proc/cpuinfo', 'r') as f:
        cpuinfo = f.read()
        cpu_count = cpuinfo.count('processor')
        TIMING_DATA['system_info']['cpu_cores'] = cpu_count
        for line in cpuinfo.split('\n'):
            if 'model name' in line:
                TIMING_DATA['system_info']['cpu_model'] = line.split(':')[1].strip()
                break
except Exception as e:
    TIMING_DATA['system_info']['error'] = str(e)

print('✅ Execution Timing Tracker initialized')
print(f'   Notebook started: {TIMING_DATA["notebook_start_datetime"]}')
print(f'   Accelerator: {TIMING_DATA["system_info"].get("accelerator", "Unknown")}')
if 'gpu' in TIMING_DATA['system_info']:
    print(f'   GPU: {TIMING_DATA["system_info"]["gpu"]}')
print(f'   RAM: {TIMING_DATA["system_info"].get("ram_gb", "Unknown")} GB')
print(f'   CPU Cores: {TIMING_DATA["system_info"].get("cpu_cores", "Unknown")}')

# @title
# =============================================================================
# TRAINING VISUALIZATION FUNCTION
# =============================================================================
# Comprehensive visualization of training progress including:
# - Accuracy curves (train vs validation)
# - Loss curves (train vs validation)
# - Overfitting analysis (accuracy gap)
# - Learning progression summary
# =============================================================================

def plot_training_history(history, model_name="Model", best_epoch=None):
    """
    Create comprehensive training visualization using Plotly.

    Args:
        history: Keras History object from model.fit()
        model_name: Name for chart titles
        best_epoch: Best epoch number (optional, will calculate if not provided)
    """
    hist = history.history
    epochs = list(range(1, len(hist['accuracy']) + 1))

    # Calculate best epoch if not provided
    if best_epoch is None:
        best_epoch = hist['val_accuracy'].index(max(hist['val_accuracy'])) + 1

    best_val = max(hist['val_accuracy'])

    # Create subplot figure
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=(
            'Training & Validation Accuracy',
            'Training & Validation Loss',
            'Accuracy Gap (Overfitting Analysis)',
            'Learning Progression'
        ),
        vertical_spacing=0.15,
        horizontal_spacing=0.1
    )

    # Colors
    train_color = '#3498db'
    val_color = '#e74c3c'

    # ===== Plot 1: Accuracy =====
    fig.add_trace(
        go.Scatter(x=epochs, y=hist['accuracy'], mode='lines+markers',
                  name='Train Accuracy', line=dict(color=train_color),
                  marker=dict(size=4)),
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(x=epochs, y=hist['val_accuracy'], mode='lines+markers',
                  name='Val Accuracy', line=dict(color=val_color),
                  marker=dict(size=4)),
        row=1, col=1
    )
    # Add best epoch marker
    fig.add_vline(x=best_epoch, line_dash='dash', line_color='green', row=1, col=1)
    fig.add_annotation(x=best_epoch, y=best_val, text=f'Best: {best_val*100:.1f}%',
                      showarrow=True, arrowhead=2, row=1, col=1)

    # ===== Plot 2: Loss =====
    fig.add_trace(
        go.Scatter(x=epochs, y=hist['loss'], mode='lines+markers',
                  name='Train Loss', line=dict(color=train_color),
                  marker=dict(size=4), showlegend=False),
        row=1, col=2
    )
    fig.add_trace(
        go.Scatter(x=epochs, y=hist['val_loss'], mode='lines+markers',
                  name='Val Loss', line=dict(color=val_color),
                  marker=dict(size=4), showlegend=False),
        row=1, col=2
    )
    fig.add_vline(x=best_epoch, line_dash='dash', line_color='green', row=1, col=2)

    # ===== Plot 3: Overfitting Analysis (accuracy gap) =====
    acc_gap = [(t - v) * 100 for t, v in zip(hist['accuracy'], hist['val_accuracy'])]

    # Color based on gap (positive = overfitting, negative = unusual)
    gap_colors = ['#e74c3c' if g > 10 else '#f39c12' if g > 5 else '#2ecc71' if g >= 0 else '#9b59b6' for g in acc_gap]

    fig.add_trace(
        go.Bar(x=epochs, y=acc_gap, name='Accuracy Gap %',
               marker_color=gap_colors),
        row=2, col=1
    )
    # Add reference lines
    fig.add_hline(y=0, line_dash="solid", line_color="black", line_width=1, row=2, col=1)
    fig.add_hline(y=10, line_dash="dash", line_color="orange",
                  annotation_text="High overfitting", row=2, col=1)
    fig.add_hline(y=-5, line_dash="dash", line_color="purple",
                  annotation_text="Negative gap (unusual)", row=2, col=1)

    # ===== Plot 4: Learning Progression =====
    progression_epochs = [1, len(epochs)//4, len(epochs)//2, 3*len(epochs)//4, len(epochs)]
    progression_epochs = [e for e in progression_epochs if e <= len(epochs)]
    progression_vals = [hist['val_accuracy'][e-1] * 100 for e in progression_epochs]
    progression_labels = [f'Ep {e}' for e in progression_epochs]

    fig.add_trace(
        go.Bar(x=progression_labels, y=progression_vals,
               marker_color=['#95a5a6', '#3498db', '#3498db', '#3498db', '#27ae60'],
               text=[f'{v:.1f}%' for v in progression_vals],
               textposition='auto',
               name='Val Accuracy'),
        row=2, col=2
    )

    # Update layout
    fig.update_layout(
        title=dict(text=f'<b>{model_name} Training Performance</b>', x=0.5, font=dict(size=18)),
        height=700,
        template='plotly_white',
        showlegend=True,
        legend=dict(orientation='h', yanchor='bottom', y=1.02, xanchor='center', x=0.3)
    )

    # Axis labels
    fig.update_xaxes(title_text='Epoch', row=1, col=1)
    fig.update_xaxes(title_text='Epoch', row=1, col=2)
    fig.update_xaxes(title_text='Epoch', row=2, col=1)
    fig.update_yaxes(title_text='Accuracy', row=1, col=1)
    fig.update_yaxes(title_text='Loss', row=1, col=2)
    fig.update_yaxes(title_text='Gap (%)', row=2, col=1)
    fig.update_yaxes(title_text='Validation Accuracy (%)', row=2, col=2)

    fig.show()

    # Print summary statistics
    final_gap = acc_gap[-1]
    print("\n" + "=" * 70)
    print(f"📊 {model_name.upper()} TRAINING SUMMARY")
    print("=" * 70)
    print(f"  Total epochs trained: {len(hist['accuracy'])}")
    print(f"  Best epoch: {best_epoch}")
    print(f"  Best validation accuracy: {best_val*100:.2f}%")
    print(f"  Best validation loss: {min(hist['val_loss']):.4f}")
    print(f"  Final accuracy gap: {final_gap:+.2f}%")

    if final_gap > 15:
        print("  🔴 SEVERE overfitting - consider stronger regularization")
    elif final_gap > 10:
        print("  🟠 HIGH overfitting - add regularization")
    elif final_gap > 5:
        print("  🟡 MODERATE overfitting - regularization helping")
    elif final_gap >= 0:
        print("  🟢 GOOD generalization!")
    else:
        print("  🟣 NEGATIVE gap - unusual, check for data issues")
    print("=" * 70)

print("✅ plot_training_history() function defined")

# @title
# =============================================================================
# DATASET CONFIGURATION
# =============================================================================
#
# Dataset Evolution:
# 1. Facial_emotion_images              → Original MIT dataset (messy splits)
# 2. facial_emotion_stratified_preaffect → After 80/10/10 stratification (~19K)
# 3. facial_emotion_stratified          → After AffectNet merge (~22K balanced)
#
# Note: affectnet_emotion_images is the RAW AffectNet source - not for training!
#
# =============================================================================

# All paths relative to BASE_PATH (set in drive mount cell)
DATASETS = {
    'original': {
        'path': './Facial_emotion_images',
        'cache': './cache_original.pkl',
        'description': 'Original MIT FER dataset (messy splits, 0.6% test)',
        'models': ['model_0'],
        'expected_accuracy': '~76% (inflated due to data leakage)',
        'image_count': 20214,
        'phase': 1
    },
    'stratified_preaffect': {
        'path': './facial_emotion_stratified_preaffect',
        'cache': './cache_stratified_preaffect.pkl',
        'description': 'After 80/10/10 stratification, before AffectNet merge (~19K images)',
        'models': ['model_a', 'model_b', 'model_c'],
        'expected_accuracy': '82-84%',
        'image_count': 18981,
        'phase': 2
    },
    'stratified_with_affectnet': {
        'path': './facial_emotion_stratified',
        'cache': './cache_stratified_affectnet.pkl',
        'description': 'Final dataset with AffectNet images merged for class balance (~22K images)',
        'models': ['model_b_plus', 'model_b_plus_plus'],
        'expected_accuracy': '85-86%',
        'image_count': 21938,
        'phase': 3
    }
}

# Output paths
MODELS_PATH = './models'
EVALUATION_PATH = './evaluation'
OUTPUTS_PATH = './outputs'

# Create output directories
for path in [MODELS_PATH, EVALUATION_PATH, OUTPUTS_PATH]:
    os.makedirs(path, exist_ok=True)

# =============================================================================
# PLOTLY VISUALIZATION: Dataset Evolution
# =============================================================================

# Prepare data for visualization
phases = ['Phase 1:\nOriginal', 'Phase 2:\nStratified', 'Phase 3:\nAffectNet Merged']
image_counts = [20214, 18981, 21938]
accuracies_low = [76, 82, 85]
accuracies_high = [76, 84, 86]
models_list = ['Model 0', 'Models A, B, C', 'Models B+, B++']
colors = ['#e74c3c', '#f39c12', '#27ae60']
issues = ['❌ 0.6% test split\n❌ Data leakage', '✅ 80/10/10 split\n✅ Cleaned', '✅ Balanced classes\n✅ +3K images']

# Create subplot figure
fig = make_subplots(
    rows=2, cols=2,
    subplot_titles=(
        'Dataset Size Evolution',
        'Expected Accuracy Range',
        'Dataset Evolution Timeline',
        ''
    ),
    specs=[
        [{'type': 'bar'}, {'type': 'bar'}],
        [{'type': 'table', 'colspan': 2}, None]
    ],
    row_heights=[0.6, 0.4],
    vertical_spacing=0.15,
    horizontal_spacing=0.1
)

# Chart 1: Image counts bar chart
fig.add_trace(
    go.Bar(
        x=phases,
        y=image_counts,
        marker_color=colors,
        text=[f'{c:,}' for c in image_counts],
        textposition='outside',
        name='Images',
        hovertemplate='%{x}<br>Images: %{y:,}<extra></extra>'
    ),
    row=1, col=1
)

# Chart 2: Accuracy range (using bar with error bars style)
fig.add_trace(
    go.Bar(
        x=phases,
        y=[(l+h)/2 for l, h in zip(accuracies_low, accuracies_high)],
        marker_color=colors,
        text=[f'{l}-{h}%' if l != h else f'~{l}%' for l, h in zip(accuracies_low, accuracies_high)],
        textposition='outside',
        name='Accuracy',
        hovertemplate='%{x}<br>Expected: %{text}<extra></extra>',
        error_y=dict(
            type='data',
            symmetric=False,
            array=[(h-l)/2 for l, h in zip(accuracies_low, accuracies_high)],
            arrayminus=[(h-l)/2 for l, h in zip(accuracies_low, accuracies_high)],
            color='rgba(0,0,0,0.3)',
            thickness=2,
            width=10
        )
    ),
    row=1, col=2
)

# Chart 3: Summary table
fig.add_trace(
    go.Table(
        header=dict(
            values=['<b>Phase</b>', '<b>Dataset</b>', '<b>Images</b>', '<b>Models</b>', '<b>Key Changes</b>'],
            fill_color='#34495e',
            font=dict(color='white', size=12),
            align='center',
            height=30
        ),
        cells=dict(
            values=[
                ['Phase 1', 'Phase 2', 'Phase 3'],
                ['Original MIT FER', 'Stratified (Pre-AffectNet)', 'Stratified + AffectNet'],
                ['20,214', '18,981', '21,938'],
                models_list,
                issues
            ],
            fill_color=[['#fadbd8', '#fdebd0', '#d5f5e3'] * 5],
            font=dict(size=11),
            align='center',
            height=50
        )
    ),
    row=2, col=1
)

# Update layout
fig.update_layout(
    title=dict(
        text='📊 FER Capstone: Dataset Evolution Overview',
        font=dict(size=18)
    ),
    showlegend=False,
    height=650,
    template='plotly_white'
)

# Update axes
fig.update_yaxes(title_text='Image Count', row=1, col=1, range=[0, 25000])
fig.update_yaxes(title_text='Val Accuracy (%)', row=1, col=2, range=[70, 92])

fig.show()

# Print text summary
print('\n📁 Dataset Configuration:')
for name, config in DATASETS.items():
    print(f'\n   {name}:')
    print(f'      Path: {config["path"]}')
    print(f'      Cache: {config["cache"]}')
    print(f'      Models: {config["models"]}')

print(f'\n📂 Output directories created: {MODELS_PATH}, {EVALUATION_PATH}, {OUTPUTS_PATH}')

# @title
# =============================================================================
# DATA LOADING WITH CACHING
# =============================================================================
#
# Each dataset gets its own cache file. This means:
# - Switching between phases is instant (just change phase)
# - No need to reload 20,000+ images each time
# - Cache automatically rebuilds if dataset changes or has old format
# - Automatically creates validation split if missing (10% of training)
#
# =============================================================================

class ImageRecord:
    """Container for image data and metadata."""
    __slots__ = ['filepath', 'filename', 'split', 'label', 'label_idx', 'image_data']

    def __init__(self, filepath, filename, split, label, label_idx, image_data):
        self.filepath = filepath
        self.filename = filename
        self.split = split  # Normalized to: train, val, test
        self.label = label
        self.label_idx = label_idx
        self.image_data = image_data


def normalize_split_name(split_name):
    """
    Normalize split folder names to standard names.
    Handles variations like 'validation' vs 'val', 'valid', etc.
    Also handles FER2013-style names like 'PublicTest', 'PrivateTest'.
    """
    split_lower = split_name.lower().replace('_', '').replace('-', '')

    # Training variations
    if split_lower in ['train', 'training']:
        return 'train'
    # Validation variations
    elif split_lower in ['val', 'valid', 'validation', 'dev', 'eval',
                         'publictest', 'public']:
        return 'val'
    # Test variations
    elif split_lower in ['test', 'testing', 'privatetest', 'private']:
        return 'test'
    else:
        # Return as-is but warn
        print(f'   ⚠️ Unknown split name: {split_name} (keeping as {split_lower})')
        return split_lower


def load_single_image(args):
    """Load a single image (for parallel processing)."""
    filepath, filename, split, label, label_idx = args
    try:
        with Image.open(filepath) as img:
            img = img.convert('L').resize((IMG_SIZE, IMG_SIZE))
            image_data = np.array(img, dtype=np.uint8)
        # Normalize split name when creating record
        normalized_split = normalize_split_name(split)
        return ImageRecord(filepath, filename, normalized_split, label, label_idx, image_data)
    except Exception as e:
        return None


def load_dataset_parallel(data_dir, num_workers=8):
    """
    Load all images in parallel with progress bar.
    Automatically detects split folder names (train/validation/val/test).
    """
    tasks = []

    # Auto-detect split folders
    available_splits = [d for d in os.listdir(data_dir)
                        if os.path.isdir(os.path.join(data_dir, d))
                        and not d.startswith('.')]

    print(f'   Detected split folders: {available_splits}')

    for split in available_splits:
        split_path = os.path.join(data_dir, split)
        for label_idx, label in enumerate(CLASS_NAMES):
            folder = os.path.join(split_path, label)
            if os.path.exists(folder):
                for fname in os.listdir(folder):
                    if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
                        filepath = os.path.join(folder, fname)
                        tasks.append((filepath, fname, split, label, label_idx))

    records = []
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(load_single_image, task) for task in tasks]
        for future in tqdm(as_completed(futures), total=len(futures), desc='Loading images'):
            result = future.result()
            if result:
                records.append(result)

    return records


def create_validation_split(all_records, val_fraction=0.1, random_seed=42):
    """
    Create a validation split from training data if one doesn't exist.

    Args:
        all_records: List of ImageRecord objects
        val_fraction: Fraction of training data to use for validation (default 10%)
        random_seed: Random seed for reproducibility

    Returns:
        Updated list of ImageRecord objects with val split created
    """
    splits = set(r.split for r in all_records)

    if 'val' in splits:
        print('   ✅ Validation split already exists')
        return all_records

    print(f'   ⚠️ No validation split found. Creating from {val_fraction*100:.0f}% of training data...')

    # Separate train and other splits
    train_records = [r for r in all_records if r.split == 'train']
    other_records = [r for r in all_records if r.split != 'train']

    # Stratified split by label
    np.random.seed(random_seed)
    new_train = []
    new_val = []

    for label in CLASS_NAMES:
        label_records = [r for r in train_records if r.label == label]
        np.random.shuffle(label_records)

        n_val = int(len(label_records) * val_fraction)

        # Create new records with updated split
        for r in label_records[:n_val]:
            new_val.append(ImageRecord(
                r.filepath, r.filename, 'val', r.label, r.label_idx, r.image_data
            ))
        for r in label_records[n_val:]:
            new_train.append(r)  # Keep as train

    print(f'   Created validation split: {len(new_val):,} images')
    print(f'   Remaining training: {len(new_train):,} images')

    return new_train + new_val + other_records


def load_dataset_with_cache(phase_name):
    """
    Load dataset for a specific phase, using cache if available.
    Automatically rebuilds cache if it has old-style split names.
    Creates validation split from training data if missing.

    Args:
        phase_name: One of 'original', 'stratified', 'affectnet_merged'

    Returns:
        all_records: List of ImageRecord objects
    """
    config = DATASETS[phase_name]
    data_dir = config['path']
    cache_file = config['cache']

    print('=' * 70)
    print(f'📂 Loading Dataset: {phase_name.upper()}')
    print('=' * 70)
    print(f'Path: {data_dir}')
    print(f'Cache: {cache_file}')
    print(f'Description: {config["description"]}')

    if not os.path.exists(data_dir):
        raise FileNotFoundError(f'Dataset not found: {data_dir}')

    need_rebuild = False
    all_records = None

    # Try loading from cache
    if os.path.exists(cache_file):
        print(f'\n📦 Loading from cache: {cache_file}')
        with open(cache_file, 'rb') as f:
            all_records = pickle.load(f)
        print(f'   Loaded {len(all_records):,} images from cache')

        # Validate cache has correct split names (train, val, test)
        cache_splits = set(r.split for r in all_records)
        expected_splits = {'train', 'val', 'test'}

        # Check for old naming (validation instead of val)
        if 'validation' in cache_splits:
            print(f'   ⚠️ Cache has old split names: {cache_splits}')
            print(f'   🔄 Rebuilding cache with normalized names...')
            need_rebuild = True
            os.remove(cache_file)
        # Check if validation is missing entirely
        elif 'val' not in cache_splits and 'train' in cache_splits:
            print(f'   ⚠️ Cache missing validation split: {cache_splits}')
            all_records = create_validation_split(all_records)
            # Re-save cache with validation split
            with open(cache_file, 'wb') as f:
                pickle.dump(all_records, f)
            print(f'   💾 Updated cache with validation split')
    else:
        print(f'\n🔄 Cache not found. Loading from disk...')
        need_rebuild = True

    # Rebuild cache if needed
    if need_rebuild:
        all_records = load_dataset_parallel(data_dir)

        # Create validation split if missing
        all_records = create_validation_split(all_records)

        # Save to cache
        with open(cache_file, 'wb') as f:
            pickle.dump(all_records, f)
        print(f'\n💾 Saved cache: {cache_file}')

    # Show split distribution
    split_counts = Counter(r.split for r in all_records)
    print(f'\n   Split distribution: {dict(split_counts)}')

    return all_records


def prepare_data_arrays(all_records):
    """
    Convert ImageRecord list to train/val/test numpy arrays.

    Returns:
        Dictionary with X_train, y_train, X_val, y_val, X_test, y_test
    """
    # Standard split names (already normalized in ImageRecord)
    standard_splits = ['train', 'val', 'test']

    # Check what splits we actually have
    actual_splits = set(r.split for r in all_records)
    print(f'\n   Found splits in data: {actual_splits}')

    # Verify we have the expected splits
    missing = set(standard_splits) - actual_splits
    if missing:
        print(f'   ⚠️ Missing splits: {missing}')

    # Split records by set
    splits_data = {s: [r for r in all_records if r.split == s] for s in standard_splits}

    data = {}
    for split_name in standard_splits:
        records = splits_data[split_name]

        if len(records) == 0:
            print(f'   ⚠️ Warning: No images found for split: {split_name}')
            # Create empty arrays with correct shape
            data[f'X_{split_name}'] = np.zeros((0, IMG_SIZE, IMG_SIZE, 1), dtype='float32')
            data[f'y_{split_name}'] = np.array([], dtype='int32')
            data[f'y_{split_name}_cat'] = np.zeros((0, NUM_CLASSES), dtype='float32')
            continue

        # Stack images and labels
        X = np.stack([r.image_data for r in records], axis=0)
        y = np.array([r.label_idx for r in records])

        # Normalize and reshape
        X = X.reshape(-1, IMG_SIZE, IMG_SIZE, 1).astype('float32') / 255.0

        # One-hot encode
        y_cat = to_categorical(y, NUM_CLASSES)

        data[f'X_{split_name}'] = X
        data[f'y_{split_name}'] = y
        data[f'y_{split_name}_cat'] = y_cat

    # Print summary
    print('\n📊 Dataset Summary:')
    for split_name in standard_splits:
        X = data[f'X_{split_name}']
        label = {'train': 'Train', 'val': 'Validation', 'test': 'Test'}[split_name]
        print(f'   {label:12}: {X.shape[0]:>6,} images')

    total = sum(data[f'X_{s}'].shape[0] for s in standard_splits)
    print(f'   {"─"*30}')
    print(f'   {"Total":12}: {total:>6,} images')

    return data


def compute_class_weights(y_train):
    """Compute class weights for imbalanced data."""
    weights = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train)
    class_weight_dict = dict(enumerate(weights))

    print('\n⚖️ Class Weights (for imbalanced classes):')
    for idx, name in enumerate(CLASS_NAMES):
        print(f'   {name}: {class_weight_dict[idx]:.3f}')

    return class_weight_dict


print('✅ Data loading functions defined')
print('   • normalize_split_name(): Handles val/validation/valid variations')
print('   • create_validation_split(): Creates val from train if missing')
print('   • load_dataset_with_cache(): Loads with automatic caching + validation')
print('   • prepare_data_arrays(): Creates X_train, X_val, X_test arrays')

# @title
# =============================================================================
# PHASE 1: LOAD ORIGINAL DATASET
# =============================================================================

start_timer('phase1_load')
CURRENT_PHASE = 'original'

# ⚠️ Set to True to force rebuild cache (use if you get unexpected results)
# If model accuracy is unexpectedly low (~36% instead of ~70%), DELETE CACHE!
FORCE_REBUILD_CACHE = False

if FORCE_REBUILD_CACHE:
    cache_file = DATASETS[CURRENT_PHASE]['cache']
    if os.path.exists(cache_file):
        os.remove(cache_file)
        print(f'🗑️ Deleted cache: {cache_file}')

# Load data with caching
records_original = load_dataset_with_cache(CURRENT_PHASE)

# Prepare arrays
data_original = prepare_data_arrays(records_original)

# =============================================================================
# DATA VERIFICATION (Critical for debugging)
# =============================================================================
print('\n' + '=' * 70)
print('🔍 DATA VERIFICATION')
print('=' * 70)

X_train = data_original['X_train']
y_train = data_original['y_train']

# Check pixel value range
print(f'\n📊 Pixel Value Range:')
print(f'   Min: {X_train.min():.4f}')
print(f'   Max: {X_train.max():.4f}')
print(f'   Mean: {X_train.mean():.4f}')
if X_train.max() > 1.0:
    print('   ❌ ERROR: Images not normalized! Should be 0-1 range.')
elif X_train.max() < 0.01:
    print('   ❌ ERROR: Images appear to be all zeros!')
else:
    print('   ✅ Images properly normalized (0-1 range)')

# Check label distribution
print(f'\n📊 Label Distribution (Training):')
unique, counts = np.unique(y_train, return_counts=True)
for idx, count in zip(unique, counts):
    pct = count / len(y_train) * 100
    print(f'   {CLASS_NAMES[idx]:<12}: {count:>5,} ({pct:>5.1f}%)')

# Verify labels match images (spot check)
print(f'\n📊 Sample Verification (first 5 images):')
for i in range(min(5, len(y_train))):
    label_idx = y_train[i]
    label_name = CLASS_NAMES[label_idx]
    pixel_sum = X_train[i].sum()
    print(f'   Image {i}: label={label_idx} ({label_name}), pixel_sum={pixel_sum:.2f}')

# Verify image data is not corrupted (check variance)
print(f'\n📊 Image Data Quality:')
sample_variances = [X_train[i].var() for i in range(min(100, len(X_train)))]
avg_var = np.mean(sample_variances)
print(f'   Average variance (first 100 images): {avg_var:.6f}')
if avg_var < 0.001:
    print('   ❌ ERROR: Images have very low variance - may be corrupted or all same!')
else:
    print('   ✅ Image variance looks normal')

print('=' * 70)

# Record timing
load_time_1 = stop_timer('phase1_load', 'data_loading')
TIMING_DATA['data_loading']['phase1_details'] = {
    'name': 'Original Dataset',
    'images': len(records_original),
    'cached': os.path.exists(DATASETS['original']['cache']),
    'time_seconds': load_time_1
}
print(f'\n⏱️ Phase 1 load time: {format_time(load_time_1)}')

# @title
# =============================================================================
# ORIGINAL DATASET SPLIT ANALYSIS
# =============================================================================

# Count by split (note: splits are normalized to 'train', 'val', 'test')
split_counts = Counter(r.split for r in records_original)
total = len(records_original)

print('=' * 70)
print('📊 ORIGINAL DATASET SPLIT DISTRIBUTION')
print('=' * 70)

# Create visualization
fig = make_subplots(
    rows=1, cols=2,
    specs=[[{'type': 'pie'}, {'type': 'bar'}]],
    subplot_titles=('Split Distribution', 'Expected vs Actual')
)

# Pie chart of actual distribution
labels = ['Train', 'Validation', 'Test']
# Use 'val' key (normalized from 'validation')
values = [split_counts.get('train', 0),
          split_counts.get('val', 0),
          split_counts.get('test', 0)]
colors = ['#2ecc71', '#3498db', '#e74c3c']

fig.add_trace(
    go.Pie(
        labels=labels,
        values=values,
        marker_colors=colors,
        textinfo='label+percent',
        hole=0.3
    ),
    row=1, col=1
)

# Bar chart comparing expected vs actual
expected = [0.80, 0.10, 0.10]
actual = [v/total for v in values]

fig.add_trace(
    go.Bar(name='Expected', x=labels, y=expected, marker_color='lightgray'),
    row=1, col=2
)
fig.add_trace(
    go.Bar(name='Actual', x=labels, y=actual, marker_color=colors),
    row=1, col=2
)

fig.update_layout(
    title_text='⚠️ Original Dataset: Severe Split Imbalance',
    height=400,
    showlegend=True
)
fig.show()

# Print details
print(f'\n{"Split":<12} {"Count":>8} {"Actual":>10} {"Expected":>10} {"Issue":>15}')
print('-' * 60)
for split_display, split_key, expected_pct in [('train', 'train', 0.80),
                                                 ('validation', 'val', 0.10),
                                                 ('test', 'test', 0.10)]:
    count = split_counts.get(split_key, 0)
    actual_pct = count / total if total > 0 else 0
    diff = actual_pct - expected_pct
    issue = '⚠️ CRITICAL' if abs(diff) > 0.05 else '✅ OK'
    print(f'{split_display:<12} {count:>8,} {actual_pct*100:>9.1f}% {expected_pct*100:>9.0f}% {issue:>15}')

print(f'\n{"─"*60}')
print(f'{"TOTAL":<12} {total:>8,}')

# @title
# =============================================================================
# CLASS DISTRIBUTION VISUALIZATION
# =============================================================================
# Create interactive Plotly visualizations showing:
# 1. Class distribution within each split
# 2. Overall class imbalance
# 3. Split size comparison
# =============================================================================

def visualize_class_distribution(records, title="Dataset Distribution"):
    """
    Create Plotly visualization of class distribution from ImageRecord list.

    Args:
        records: List of ImageRecord objects
        title: Chart title
    """
    # Build DataFrame from records
    from collections import defaultdict

    # Count images per (split, class)
    counts = defaultdict(int)
    for r in records:
        counts[(r.split, r.label)] += 1

    # Convert to lists for DataFrame
    data = []
    for (split, label), count in counts.items():
        data.append({'split': split, 'class': label, 'count': count})

    df = pd.DataFrame(data)

    if df.empty:
        print("⚠️ No data to visualize")
        return

    # Create subplot figure
    fig = make_subplots(
        rows=1, cols=2,
        subplot_titles=('Images per Class by Split', 'Overall Class Distribution'),
        specs=[[{'type': 'bar'}, {'type': 'pie'}]]
    )

    # Color scheme
    colors = {'happy': '#2ecc71', 'neutral': '#3498db',
              'sad': '#9b59b6', 'surprise': '#e74c3c'}

    # Get unique classes from data
    classes_in_data = sorted(df['class'].unique())

    # Bar chart - grouped by split
    for class_name in classes_in_data:
        class_data = df[df['class'] == class_name]
        fig.add_trace(
            go.Bar(
                name=class_name.title(),
                x=class_data['split'],
                y=class_data['count'],
                marker_color=colors.get(class_name, '#95a5a6'),
                text=class_data['count'],
                textposition='auto',
            ),
            row=1, col=1
        )

    # Pie chart - overall distribution
    total_by_class = df.groupby('class')['count'].sum()
    fig.add_trace(
        go.Pie(
            labels=[c.title() for c in total_by_class.index],
            values=total_by_class.values,
            marker_colors=[colors.get(c, '#95a5a6') for c in total_by_class.index],
            textinfo='label+percent',
            hole=0.3
        ),
        row=1, col=2
    )

    fig.update_layout(
        title=dict(text=f'<b>{title}</b>', x=0.5, font=dict(size=18)),
        height=450,
        showlegend=True,
        template='plotly_white',
        barmode='group'
    )
    fig.show()

    # Print statistics
    print("\n" + "=" * 70)
    print("CLASS DISTRIBUTION STATISTICS")
    print("=" * 70)

    total = df['count'].sum()
    num_classes = len(classes_in_data)
    ideal_pct = 100.0 / num_classes  # Perfect balance

    print(f"{'Class':<12} {'Count':>8} {'Percent':>10} {'vs Ideal':>12}")
    print("-" * 45)

    for class_name in classes_in_data:
        count = df[df['class'] == class_name]['count'].sum()
        pct = (count / total) * 100
        diff = pct - ideal_pct
        status = "✅" if abs(diff) < 5 else "⚠️"
        print(f"{class_name.title():<12} {count:>8,} {pct:>9.1f}% {diff:+10.1f}% {status}")

    print(f"{'─'*45}")
    print(f"{'TOTAL':<12} {total:>8,}")
    print(f"\nIdeal distribution: {ideal_pct:.1f}% per class")

# Visualize original dataset class distribution
print("\n" + "=" * 70)
print("📊 ORIGINAL DATASET CLASS DISTRIBUTION")
print("=" * 70)
visualize_class_distribution(records_original, "Original MIT/FER+ Dataset - Class Distribution")

# @title
# =============================================================================
# SAMPLE IMAGE VISUALIZATION WITH PLOTLY
# =============================================================================
# Display sample images from each emotion class to understand
# what the model will be learning to classify.
# =============================================================================

def display_sample_images_plotly(X, y, samples_per_class=4, title="Sample Images"):
    """
    Display sample images from each class using Plotly.

    Args:
        X: Image array (N, H, W, 1) or (N, H, W)
        y: Label array (integer class indices)
        samples_per_class: Number of samples to show per class
        title: Chart title
    """
    fig = make_subplots(
        rows=NUM_CLASSES, cols=samples_per_class,
        subplot_titles=[f"{CLASS_NAMES[i].title()} #{j+1}"
                       for i in range(NUM_CLASSES)
                       for j in range(samples_per_class)],
        vertical_spacing=0.08,
        horizontal_spacing=0.02
    )

    for class_idx, class_name in enumerate(CLASS_NAMES):
        # Get indices for this class
        class_indices = np.where(y == class_idx)[0]

        if len(class_indices) == 0:
            continue

        # Random sample
        sample_indices = np.random.choice(
            class_indices,
            min(samples_per_class, len(class_indices)),
            replace=False
        )

        for j, idx in enumerate(sample_indices):
            img = X[idx]
            # Handle both (H,W,1) and (H,W) shapes
            if len(img.shape) == 3:
                img = img.squeeze()

            fig.add_trace(
                go.Heatmap(
                    z=np.flipud(img),  # Flip for correct orientation
                    colorscale='Gray',
                    showscale=False,
                    hoverinfo='skip'
                ),
                row=class_idx + 1, col=j + 1
            )

    fig.update_layout(
        title=dict(text=f'<b>{title}</b>', x=0.5, font=dict(size=18)),
        height=600,
        width=800,
        template='plotly_white'
    )

    # Hide axes
    fig.update_xaxes(showticklabels=False, showgrid=False)
    fig.update_yaxes(showticklabels=False, showgrid=False)

    fig.show()

    # Print class distribution
    print("\n" + "=" * 50)
    print("CLASS DISTRIBUTION IN DISPLAYED DATA")
    print("=" * 50)
    unique, counts = np.unique(y, return_counts=True)
    for idx, count in zip(unique, counts):
        pct = count / len(y) * 100
        print(f"  {CLASS_NAMES[idx].title():<12}: {count:>6,} ({pct:>5.1f}%)")
    print(f"  {'─'*40}")
    print(f"  {'TOTAL':<12}: {len(y):>6,}")

# Display samples from original training set
print("\n📸 Sample Images from Original Training Set:")
X_train_orig = data_original['X_train']
y_train_orig = data_original['y_train']
display_sample_images_plotly(X_train_orig, y_train_orig, samples_per_class=4,
                             title="Sample Images from Each Emotion Class (Original Dataset)")

# @title
# =============================================================================
# MODEL 0 (BASELINE): SIMPLE CNN ON ORIGINAL DATASET
# =============================================================================
#
# This is our baseline model trained on the original MIT FER dataset.
# It establishes a performance reference point BEFORE any data cleaning.
#
# =============================================================================
# 📊 EXPECTED RESULTS
# =============================================================================
#
# Validation Accuracy: ~65-66%
# Training Accuracy:   ~60-62%
# Overfitting Gap:     Low (~4-5%)
#
# Why relatively LOW accuracy?
# • Original dataset has problematic split ratios (74/25/0.6)
# • Test set is nearly useless (only 128 images = 0.6%)
# • Potential label noise and duplicates
# • Model capacity limited by simple architecture
#
# Why LOW overfitting despite no regularization?
# • Model is underfitting - hasn't learned the data well
# • Large validation set (25%) provides stable estimates
# • Class weights help but can't fix fundamental data issues
#
# =============================================================================

def build_model_0():
    """
    Model 0 (Baseline): Simple CNN for establishing baseline performance.

    Architecture:
    - 3 Conv blocks with increasing filters (32 → 64 → 128)
    - Light dropout (0.25)
    - No augmentation, no regularization

    Purpose: Show performance on problematic original dataset
    """
    model = Sequential([
        # Block 1
        Conv2D(32, (3, 3), padding='same', input_shape=(IMG_SIZE, IMG_SIZE, 1)),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D((2, 2)),
        Dropout(0.25),

        # Block 2
        Conv2D(64, (3, 3), padding='same'),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D((2, 2)),
        Dropout(0.25),

        # Block 3
        Conv2D(128, (3, 3), padding='same'),
        BatchNormalization(),
        Activation('relu'),
        MaxPooling2D((2, 2)),
        Dropout(0.25),

        # Dense layers
        Flatten(),
        Dense(256, activation='relu'),
        Dropout(0.5),
        Dense(NUM_CLASSES, activation='softmax')
    ])

    return model


# Alias for backward compatibility
build_baseline_model = build_model_0

# Build and compile model
model_0 = build_model_0()
model_0.compile(
    optimizer=Adam(learning_rate=0.001),
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

print('✅ Model 0 (Baseline) built and compiled')
print(f'   Parameters: {model_0.count_params():,}')
print()
print('📐 Model Architecture:')
model_0.summary()
print()
print('📊 Expected: ~65-66% validation accuracy (inflated due to data leakage)')

# @title
# =============================================================================
# TRAIN MODEL 0 (BASELINE) ON ORIGINAL DATASET
# =============================================================================
#
# This establishes our performance baseline on the problematic original data.
# Expected: ~65-66% validation accuracy
#
# =============================================================================

TRAIN_MODEL_0 = True  # Set to False to skip

if TRAIN_MODEL_0:
    start_timer('model_0_train')
    print('=' * 70)
    print('🚀 TRAINING MODEL 0 (Baseline) on Original MIT Dataset')
    print('=' * 70)

    # Extract data arrays
    X_train_orig = data_original['X_train']
    y_train_orig = data_original['y_train']
    y_train_orig_cat = data_original['y_train_cat']

    X_val_orig = data_original['X_val']
    y_val_orig_cat = data_original['y_val_cat']

    # Reset random seed for reproducibility
    np.random.seed(SEED)
    tf.random.set_seed(SEED)

    # Compute class weights (function prints them)
    class_weights_orig = compute_class_weights(y_train_orig)

    # Callbacks
    # NOTE: Model 0 learns slowly on messy data - needs higher patience
    model_0_callbacks = [ReduceLROnPlateau(monitor='val_loss', factor=0.5,
                          patience=5, min_lr=1e-6, verbose=1),
        ModelCheckpoint(f'{MODELS_PATH}/model_0_baseline.keras',
                        monitor='val_accuracy', save_best_only=True, verbose=1)
    ]

    # Train
    print('\n🏋️ Training Model 0 on ORIGINAL dataset (with all its problems)...')
    history_0 = model_0.fit(
        X_train_orig, y_train_orig_cat,
        validation_data=(X_val_orig, y_val_orig_cat),
        epochs=75,
        batch_size=BATCH_SIZE,
        class_weight=class_weights_orig,
        callbacks=model_0_callbacks,
        verbose=1
    )

    # Results
    best_val_0 = max(history_0.history['val_accuracy'])
    best_epoch_0 = history_0.history['val_accuracy'].index(best_val_0) + 1
    final_train_0 = history_0.history['accuracy'][best_epoch_0 - 1]
    gap_0 = (final_train_0 - best_val_0) * 100

    print(f'\n✅ MODEL 0 (BASELINE) RESULTS:')
    print(f'   Best validation accuracy: {best_val_0*100:.2f}%')
    print(f'   Training accuracy at best epoch: {final_train_0*100:.2f}%')
    print(f'   Overfitting gap: {gap_0:.1f}%')
    print(f'   Best epoch: {best_epoch_0}')

    print(f'\n⚠️ Note: This accuracy reflects the problematic original dataset:')
    print(f'   • Imbalanced splits (74% train, 25% val, 0.6% test)')
    print(f'   • Test set nearly useless (only 128 images)')
    print(f'   • This is why we need to stratify the dataset!')
    # Record timing
    train_time_0 = stop_timer('model_0_train', 'model_training')
    TIMING_DATA['model_training']['model_0_details'] = {
        'name': 'Model 0 (Baseline)',
        'epochs_configured': 75,
        'epochs_completed': len(history_0.history['accuracy']),
        'parameters': model_0.count_params(),
        'batch_size': BATCH_SIZE,
        'time_seconds': train_time_0,
        'time_per_epoch': train_time_0 / len(history_0.history['accuracy'])
    }
    print(f'\n⏱️ Model 0 training time: {format_time(train_time_0)} ({train_time_0/60:.1f} min)')
else:
    print('⏭️ Skipping Model 0 training (TRAIN_MODEL_0 = False)')
    print('   Expected result: ~65-66% validation accuracy')