print('Kernel probe: cell 5 started')
probe_value = 2 + 2
print('Kernel probe: cell 5 completed, value =', probe_value)

import json
print('Import OK: json')

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
print('Import OK: os and TF log level set')

import random
random.seed(42)
print('Import OK: random (seed=42)')

import shutil
print('Import OK: shutil')

import subprocess
print('Import OK: subprocess')

import sys
print('Import OK: sys', sys.version.split()[0])

import matplotlib.pyplot as plt
print('Import OK: matplotlib')

import numpy as np
print('Import OK: numpy', np.__version__)

import pandas as pd
print('Import OK: pandas', pd.__version__)

import seaborn as sns
sns.set_theme(style='whitegrid')
print('Import OK: seaborn')

import tensorflow as tf
tf.keras.utils.set_random_seed(42)
print('Import OK: tensorflow', tf.__version__)

from pathlib import Path

IS_COLAB = 'google.colab' in sys.modules
GITHUB_REPO_OWNER = 'FrancisBurnet'
GITHUB_REPO_NAME = 'francisburnet'
GITHUB_REPO_BRANCH = 'main'
GITHUB_REPO_URL = f'https://github.com/{GITHUB_REPO_OWNER}/{GITHUB_REPO_NAME}.git'
CAPSTONE_ROOT = Path('Incremental Capstones/Deep Learning Specialization/Capstone Session 10')
PDF_EXPECTED_CLASSES = ['with_mask', 'without_mask', 'mask_worn_incorrect']
PDF_CLASS_ALIASES = {
    'mask_worn_incorrect': {'mask_worn_incorrect', 'mask_weared_incorrect'},
}


def resolve_capstone_dir() -> Path | None:
    current = Path.cwd().resolve()
    capstone_parts = CAPSTONE_ROOT.parts
    for candidate in [current, *current.parents]:
        if len(candidate.parts) >= len(capstone_parts) and candidate.parts[-len(capstone_parts):] == capstone_parts:
            return candidate
        nested_candidate = candidate / CAPSTONE_ROOT
        if nested_candidate.exists():
            return nested_candidate
    return None


def ensure_colab_capstone_dir() -> Path:
    repo_root = Path('/content') / GITHUB_REPO_NAME
    capstone_dir = repo_root / CAPSTONE_ROOT
    data_dir = capstone_dir / 'data'

    if not repo_root.exists():
        subprocess.run(['git', 'clone', '--depth', '1', '--branch', GITHUB_REPO_BRANCH, GITHUB_REPO_URL, str(repo_root)], check=True)

    has_lfs = subprocess.run(['bash', '-lc', 'command -v git-lfs >/dev/null 2>&1'], check=False).returncode == 0
    if not has_lfs:
        subprocess.run(['apt-get', 'update', '-qq'], check=True)
        subprocess.run(['apt-get', 'install', '-y', 'git-lfs'], check=True)

    subprocess.run(['git', 'lfs', 'install', '--local'], cwd=repo_root, check=True)
    subprocess.run(['git', 'lfs', 'pull', '--include', f'{CAPSTONE_ROOT.as_posix()}/data/**'], cwd=repo_root, check=True)

    if not data_dir.exists():
        raise RuntimeError(f'Expected GitHub-backed data directory at {data_dir}, but it was not found after clone + LFS pull.')

    # Verify actual image content was pulled — LFS pointer stubs are <200 bytes
    image_files = [f for f in data_dir.rglob('*') if f.is_file() and f.suffix.lower() in {'.jpg', '.jpeg', '.png'}]
    if not image_files:
        raise RuntimeError(f'No image files found under {data_dir} — LFS pull may not have downloaded content.')
    tiny = [f for f in image_files[:10] if f.stat().st_size < 1000]
    if tiny:
        raise RuntimeError(
            f'LFS pointer file detected — actual image content was not downloaded.\n'
            f'File: {tiny[0]} ({tiny[0].stat().st_size} bytes)\n'
            f'Fix: run `git lfs pull` inside {repo_root}'
        )
    print(f'LFS verified: {len(image_files)} image files found under data/')

    return capstone_dir


if IS_COLAB:
    CAPSTONE_DIR = ensure_colab_capstone_dir()
else:
    CAPSTONE_DIR = resolve_capstone_dir()

if CAPSTONE_DIR is None:
    raise RuntimeError(f'Could not resolve capstone directory from {CAPSTONE_ROOT}')

OUTPUT_ROOT = CAPSTONE_DIR
OUTPUT_MODE = 'permanent capstone outputs'
OUTPUT_DISPLAY = (CAPSTONE_ROOT / 'outputs').as_posix()

SOURCE_DATA_DIR = CAPSTONE_DIR / 'data'
if not SOURCE_DATA_DIR.exists():
    raise RuntimeError(f'Expected staged dataset at {SOURCE_DATA_DIR}, but it does not exist.')

AVAILABLE_CLASS_FOLDERS = sorted([path.name for path in SOURCE_DATA_DIR.iterdir() if path.is_dir()])
if len(AVAILABLE_CLASS_FOLDERS) < 2:
    raise RuntimeError(f'Expected at least two class folders under {SOURCE_DATA_DIR}, found {AVAILABLE_CLASS_FOLDERS}.')

CLASS_NAME_LOOKUP = {}
for expected_name in PDF_EXPECTED_CLASSES:
    alias_set = PDF_CLASS_ALIASES.get(expected_name, {expected_name})
    for candidate_name in AVAILABLE_CLASS_FOLDERS:
        if candidate_name in alias_set:
            CLASS_NAME_LOOKUP[expected_name] = candidate_name
            break

ordered_classes = [CLASS_NAME_LOOKUP[name] for name in PDF_EXPECTED_CLASSES if name in CLASS_NAME_LOOKUP]
remaining_classes = [name for name in AVAILABLE_CLASS_FOLDERS if name not in ordered_classes]
CLASS_NAMES = ordered_classes + remaining_classes
NUM_CLASSES = len(CLASS_NAMES)
MISSING_PDF_CLASSES = [name for name in PDF_EXPECTED_CLASSES if name not in CLASS_NAME_LOOKUP]

OUTPUTS_DIR = (OUTPUT_ROOT / 'outputs').resolve()
PLOTS_DIR = OUTPUTS_DIR / 'plots'
GENERATED_SPLIT_DIR = OUTPUTS_DIR / 'generated_split'
OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)
PLOTS_DIR.mkdir(parents=True, exist_ok=True)
IMAGE_SIZE = (128, 128)
BATCH_SIZE = 32
# None = use all available images (balanced to the smallest class at split time)
TRAIN_IMAGES_PER_CLASS = None
TEST_IMAGES_PER_CLASS = None

print('Runtime:', 'Google Colab' if IS_COLAB else 'Notebook runtime')
print('GitHub source repo:', f'https://github.com/{GITHUB_REPO_OWNER}/{GITHUB_REPO_NAME}/tree/{GITHUB_REPO_BRANCH}')
print('Capstone artifact path:', CAPSTONE_ROOT.as_posix())
print('Source data (GitHub-backed):', f'{CAPSTONE_ROOT.as_posix()}/data')
print('Output mode:', OUTPUT_MODE)
print('Output target:', OUTPUT_DISPLAY)
print('PDF expected classes:', PDF_EXPECTED_CLASSES)
print('Available source classes:', AVAILABLE_CLASS_FOLDERS)
print('Missing PDF classes from source:', MISSING_PDF_CLASSES if MISSING_PDF_CLASSES else 'None')
print('Classes used for this run:', CLASS_NAMES)

# Clear any previous generated split so re-runs start clean
if GENERATED_SPLIT_DIR.exists():
    shutil.rmtree(GENERATED_SPLIT_DIR)

# Count available images per class and use ALL of them (80/20 split per class)
source_files_by_class = {}
for class_name in CLASS_NAMES:
    files = sorted([p for p in (SOURCE_DATA_DIR / class_name).iterdir() if p.is_file()])
    random.shuffle(files)
    source_files_by_class[class_name] = files

split_manifest = {'source_counts': {}, 'generated_counts': {}}
for split_name in ['train', 'test']:
    for class_name in CLASS_NAMES:
        (GENERATED_SPLIT_DIR / split_name / class_name).mkdir(parents=True, exist_ok=True)

for class_name in CLASS_NAMES:
    files = source_files_by_class[class_name]
    n_total = len(files)
    n_test = max(1, int(n_total * 0.20))  # 20% to test
    n_train = n_total - n_test            # 80% to train
    
    split_manifest['source_counts'][class_name] = n_total
    split_manifest['generated_counts'][class_name] = {'train': n_train, 'test': n_test}
    
    train_files = files[:n_train]
    test_files = files[n_train:]
    
    for file_path in train_files:
        shutil.copy2(file_path, GENERATED_SPLIT_DIR / 'train' / class_name / file_path.name)
    for file_path in test_files:
        shutil.copy2(file_path, GENERATED_SPLIT_DIR / 'test' / class_name / file_path.name)

print(f'Using ALL available images (80/20 train/test split per class):')
for class_name in CLASS_NAMES:
    train_ct = split_manifest['generated_counts'][class_name]['train']
    test_ct = split_manifest['generated_counts'][class_name]['test']
    print(f'  {class_name}: {train_ct} train + {test_ct} test')

with open(OUTPUTS_DIR / 'session_10_split_manifest.json', 'w', encoding='utf-8') as handle:
    json.dump(split_manifest, handle, indent=2)
split_manifest

# No rescale here — the Lambda(preprocess) layer inside each model handles
# resnet50/efficientnet preprocessing on raw 0-255 pixel values.
train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    validation_split=0.2,
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.08,
    zoom_range=0.15,
    horizontal_flip=True,
    brightness_range=(0.85, 1.15),
    fill_mode='nearest',
)
test_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

train_generator = train_datagen.flow_from_directory(
    GENERATED_SPLIT_DIR / 'train',
    target_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    class_mode='categorical',
    subset='training',
    shuffle=True,
)
validation_generator = train_datagen.flow_from_directory(
    GENERATED_SPLIT_DIR / 'train',
    target_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    class_mode='categorical',
    subset='validation',
    shuffle=False,
)
test_generator = test_datagen.flow_from_directory(
    GENERATED_SPLIT_DIR / 'test',
    target_size=IMAGE_SIZE,
    batch_size=BATCH_SIZE,
    class_mode='categorical',
    shuffle=False,
)
class_indices = train_generator.class_indices
class_labels = {index: label for label, index in class_indices.items()}
class_indices

callbacks = [
    tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=2, verbose=0),
    tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=4, restore_best_weights=True, verbose=0),
]

HAS_GPU = bool(tf.config.list_physical_devices('GPU'))
if IS_COLAB and not HAS_GPU:
    # Colab CPU sessions are often slower than local runs.
    PRETRAIN_EPOCHS = 6
    FINE_TUNE_EPOCHS = 8
    MAX_TRAIN_STEPS = 20
    MAX_VAL_STEPS = 8
elif IS_COLAB and HAS_GPU:
    PRETRAIN_EPOCHS = 8
    FINE_TUNE_EPOCHS = 12
    MAX_TRAIN_STEPS = None
    MAX_VAL_STEPS = None
else:
    PRETRAIN_EPOCHS = 15
    FINE_TUNE_EPOCHS = 25
    MAX_TRAIN_STEPS = None
    MAX_VAL_STEPS = None

FINE_TUNE_LR = 1e-5
print('Training profile:', {
    'is_colab': IS_COLAB,
    'has_gpu': HAS_GPU,
    'pretrain_epochs': PRETRAIN_EPOCHS,
    'fine_tune_epochs': FINE_TUNE_EPOCHS,
    'max_train_steps': MAX_TRAIN_STEPS,
    'max_val_steps': MAX_VAL_STEPS,
})

def build_transfer_model(base_model_name):
    if base_model_name == 'EfficientNetB0':
        preprocess = tf.keras.applications.efficientnet.preprocess_input
        base_model = tf.keras.applications.EfficientNetB0(include_top=False, weights='imagenet', input_shape=(128, 128, 3))
        dropout_rate = 0.2
        fine_tune_at = max(0, len(base_model.layers) - 40)
    elif base_model_name == 'ResNet50':
        preprocess = tf.keras.applications.resnet50.preprocess_input
        base_model = tf.keras.applications.ResNet50(include_top=False, weights='imagenet', input_shape=(128, 128, 3))
        dropout_rate = 0.5
        fine_tune_at = max(0, len(base_model.layers) - 30)
    else:
        raise ValueError(base_model_name)

    base_model.trainable = False
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(128, 128, 3)),
        tf.keras.layers.Lambda(preprocess),
        base_model,
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dropout(dropout_rate),
        tf.keras.layers.Dense(NUM_CLASSES, activation='softmax'),
    ])
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    model._capstone_fine_tune_at = fine_tune_at
    return model

def fine_tune_model(model, base_model_name):
    base_model = next(layer for layer in model.layers if isinstance(layer, tf.keras.Model))
    fine_tune_at = getattr(model, '_capstone_fine_tune_at', max(0, len(base_model.layers) - 30))
    base_model.trainable = True
    for layer in base_model.layers[:fine_tune_at]:
        layer.trainable = False
    for layer in base_model.layers[fine_tune_at:]:
        if isinstance(layer, tf.keras.layers.BatchNormalization):
            layer.trainable = False
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=FINE_TUNE_LR),
        loss='categorical_crossentropy',
        metrics=['accuracy'],
    )
    return model

histories = {}
evaluations = []
prediction_examples = {}
trained_models = {}

steps_per_epoch = None if MAX_TRAIN_STEPS is None else min(len(train_generator), MAX_TRAIN_STEPS)
validation_steps = None if MAX_VAL_STEPS is None else min(len(validation_generator), MAX_VAL_STEPS)

for model_name in ['EfficientNetB0', 'ResNet50']:
    tf.keras.backend.clear_session()
    train_generator.reset()
    validation_generator.reset()
    test_generator.reset()
    model = build_transfer_model(model_name)
    frozen_history = model.fit(
        train_generator,
        validation_data=validation_generator,
        epochs=PRETRAIN_EPOCHS,
        steps_per_epoch=steps_per_epoch,
        validation_steps=validation_steps,
        callbacks=callbacks,
        verbose=0,
    )
    model = fine_tune_model(model, model_name)
    fine_tune_history = model.fit(
        train_generator,
        validation_data=validation_generator,
        epochs=FINE_TUNE_EPOCHS,
        steps_per_epoch=steps_per_epoch,
        validation_steps=validation_steps,
        callbacks=callbacks,
        verbose=0,
    )
    histories[model_name] = pd.concat([
        pd.DataFrame(frozen_history.history),
        pd.DataFrame(fine_tune_history.history),
    ], ignore_index=True)
    test_loss, test_accuracy = model.evaluate(test_generator, verbose=0)
    test_generator.reset()
    probabilities = model.predict(test_generator, verbose=0)
    predicted_labels = np.argmax(probabilities, axis=1)
    true_labels = test_generator.classes.copy()
    evaluations.append({
        'model': model_name,
        'epochs_ran': int(len(histories[model_name])),
        'test_loss': float(test_loss),
        'test_accuracy': float(test_accuracy),
    })
    prediction_examples[model_name] = {
        'filenames': test_generator.filenames[:10],
        'true_labels': [class_labels[index] for index in true_labels[:10]],
        'predicted_labels': [class_labels[index] for index in predicted_labels[:10]],
    }
    trained_models[model_name] = model

evaluations_df = pd.DataFrame(evaluations).sort_values('test_accuracy', ascending=False).reset_index(drop=True)
evaluations_df

for model_name, history_df in histories.items():
    history_df.to_csv(OUTPUTS_DIR / f'{model_name.lower()}_history.csv', index=False)
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    axes[0].plot(history_df['accuracy'], label='train')
    axes[0].plot(history_df['val_accuracy'], label='validation')
    axes[0].set_title(f'{model_name} Accuracy by Epoch')
    axes[0].legend()
    axes[1].plot(history_df['loss'], label='train')
    axes[1].plot(history_df['val_loss'], label='validation')
    axes[1].set_title(f'{model_name} Loss by Epoch')
    axes[1].legend()
    fig.tight_layout()
    fig.savefig(PLOTS_DIR / f'{model_name.lower()}_training_history.png', dpi=150)
    plt.show()
    plt.close(fig)

fig, ax = plt.subplots(figsize=(10, 5))
evaluations_df.plot(x='model', y='test_accuracy', kind='bar', ax=ax, color=['#1f77b4', '#ff7f0e'])
ax.set_title('Session 10 Model Comparison')
ax.set_ylim(0, 1.0)
fig.tight_layout()
fig.savefig(PLOTS_DIR / 'model_comparison.png', dpi=150)
plt.show()
plt.close(fig)

best_model_name = evaluations_df.iloc[0]['model']
example_info = prediction_examples[best_model_name]
figure, axes = plt.subplots(2, 5, figsize=(16, 7))
for axis, filename, true_label, predicted_label in zip(axes.flatten(), example_info['filenames'], example_info['true_labels'], example_info['predicted_labels']):
    image_path = GENERATED_SPLIT_DIR / 'test' / filename
    image = tf.keras.utils.load_img(image_path, target_size=IMAGE_SIZE)
    axis.imshow(image)
    axis.set_title(f'T:{true_label}\nP:{predicted_label}', fontsize=10)
    axis.axis('off')
figure.tight_layout()
figure.savefig(PLOTS_DIR / 'best_model_prediction_examples.png', dpi=150)
plt.show()
plt.close(figure)

evaluations_df.to_csv(OUTPUTS_DIR / 'session_10_model_comparison.csv', index=False)
summary = {
    'source_counts': split_manifest['source_counts'],
    'generated_counts': split_manifest['generated_counts'],
    'class_indices': class_indices,
    'pdf_expected_classes': PDF_EXPECTED_CLASSES,
    'available_source_classes': AVAILABLE_CLASS_FOLDERS,
    'missing_pdf_classes_from_source': MISSING_PDF_CLASSES,
    'source_governance_notes': [
        'PDF directions are the source of truth for task sequence and deliverables (Task A, Task B, Task C).',
        'GitHub-backed dataset assets are the source of truth for notebook inputs per Project_DEV_Rules_PROMPT.md.',
        'When PDF-expected classes are missing from source files, execution uses available classes and records the gap explicitly.',
        'The generated train/test split is a fixed stratified sample to keep transfer-learning runs executable on the current CPU-only environment.',
    ],
    'model_results': evaluations,
    'best_model': evaluations_df.iloc[0].to_dict(),
    'best_model_example_predictions': prediction_examples[best_model_name],
}
with open(OUTPUTS_DIR / 'session_10_summary.json', 'w', encoding='utf-8') as handle:
    json.dump(summary, handle, indent=2)
summary

import json
import shutil
import sys
import types

# tensorflowjs imports optional modules even for plain Keras layers export.
# On this Windows + py3.12 setup they may be unavailable, so provide no-op stubs.
for optional_module in ('tensorflow_decision_forests', 'tensorflow_hub'):
    if optional_module not in sys.modules:
        sys.modules[optional_module] = types.ModuleType(optional_module)

# Stub JAX import path required by tensorflowjs.converters.jax_conversion
if 'jax' not in sys.modules:
    jax_mod = types.ModuleType('jax')
    jax_experimental = types.ModuleType('jax.experimental')
    jax2tf_mod = types.ModuleType('jax.experimental.jax2tf')
    jax_experimental.jax2tf = jax2tf_mod
    jax_mod.experimental = jax_experimental
    sys.modules['jax'] = jax_mod
    sys.modules['jax.experimental'] = jax_experimental
    sys.modules['jax.experimental.jax2tf'] = jax2tf_mod

import tensorflowjs as tfjs

# ── Export ResNet50 to TF.js ─────────────────────────────────────────────────
trained_resnet50 = trained_models['ResNet50']
resnet_base   = next(l for l in trained_resnet50.layers if isinstance(l, tf.keras.Model))
gap_layer     = next(l for l in trained_resnet50.layers if isinstance(l, tf.keras.layers.GlobalAveragePooling2D))
dropout_layer = next(l for l in trained_resnet50.layers if isinstance(l, tf.keras.layers.Dropout))
dense_layer   = next(l for l in trained_resnet50.layers if isinstance(l, tf.keras.layers.Dense))

export_model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(128, 128, 3)),
    resnet_base,
    gap_layer,
    dropout_layer,
    dense_layer,
])

# Temp export dir
TFJS_EXPORT_DIR = OUTPUTS_DIR / 'tfjs_resnet50'
if TFJS_EXPORT_DIR.exists():
    shutil.rmtree(TFJS_EXPORT_DIR)
TFJS_EXPORT_DIR.mkdir(parents=True)

tfjs.converters.save_keras_model(export_model, str(TFJS_EXPORT_DIR))

# Write companion metadata.json with correct 3-class labels and preprocessing note
resnet50_eval = next(e for e in evaluations if e['model'] == 'ResNet50')
tfjs_metadata = {
    'modelType': 'capstone-resnet50',
    'labels': CLASS_NAMES,
    'imageSize': list(IMAGE_SIZE),
    'source': 'Session 10 capstone TensorFlow workflow',
    'preprocessing': 'resnet50_caffe',
    'preprocessingNote': 'Pass raw 0-255 RGB pixels; JS must do: toFloat() -> reverse(-1) [RGB->BGR] -> subtract [103.939, 116.779, 123.68]',
    'testLoss': resnet50_eval['test_loss'],
    'testAccuracy': resnet50_eval['test_accuracy'],
    'epochsRan': resnet50_eval['epochs_ran'],
}
with open(TFJS_EXPORT_DIR / 'metadata.json', 'w', encoding='utf-8') as fh:
    json.dump(tfjs_metadata, fh, indent=2)

# ── Deploy to web public assets ──────────────────────────────────────────────
if not IS_COLAB:
    capstone_root = CAPSTONE_DIR
    web_model_dir = None
    for parent in [capstone_root, *capstone_root.parents]:
        candidate = parent / 'web' / 'public' / 'assets' / 'models' / 'session-10-capstone-resnet50'
        if candidate.parent.exists():
            web_model_dir = candidate
            break
    if web_model_dir is not None:
        web_model_dir.mkdir(parents=True, exist_ok=True)
        # Non-destructive overwrite avoids Windows lock issues when directory is open/in use.
        shutil.copytree(str(TFJS_EXPORT_DIR), str(web_model_dir), dirs_exist_ok=True)
        print(f'Deployed to: {web_model_dir}')
    else:
        print('Web model dir not found — skipping deploy (run manually from Colab or local).')

print('TF.js export complete.')
print('Labels:', tfjs_metadata['labels'])
print('Test accuracy:', round(tfjs_metadata['testAccuracy'] * 100, 1), '%')