#!/usr/bin/env python3
"""
🏆 CAPSTONE: End-to-End Production ML Pipeline
Combines ALL techniques from Pages 1-9:
  - Page 1: TensorFlow basics, tensors
  - Page 2: Keras compile/fit/callbacks
  - Page 3: CNN, augmentation, transfer learning
  - Page 4: tf.data pipeline, mixed precision, cache
  - Page 5: (NLP concepts referenced)
  - Page 6: (Transformer concepts referenced)
  - Page 7: Custom training concepts, gradient clipping
  - Page 8: (GAN concepts referenced)
  - Page 9: SavedModel, TFLite, Docker deployment
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import time
import os

# ═══════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════
DATA_DIR = "data/train"
IMG_SIZE = (224, 224)
BATCH_SIZE = 32
EPOCHS_PHASE1 = 15
EPOCHS_PHASE2 = 10
LR_PHASE1 = 1e-3
LR_PHASE2 = 1e-5
MODEL_DIR = "saved_model/capstone"
TFLITE_PATH = "capstone.tflite"

# ═══════════════════════════════════════════════════
# STEP 1: DATA PIPELINE (Page 4)
# ═══════════════════════════════════════════════════
print("📊 Step 1: Loading data...")

train_ds = keras.utils.image_dataset_from_directory(
    DATA_DIR, image_size=IMG_SIZE, batch_size=BATCH_SIZE,
    validation_split=0.2, subset="training", seed=42,
    label_mode="int")

val_ds = keras.utils.image_dataset_from_directory(
    DATA_DIR, image_size=IMG_SIZE, batch_size=BATCH_SIZE,
    validation_split=0.2, subset="validation", seed=42,
    label_mode="int")

NUM_CLASSES = len(train_ds.class_names)
print(f"  Classes: {train_ds.class_names} ({NUM_CLASSES})")

# Optimize pipeline (Page 4: cache + prefetch)
AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

# ═══════════════════════════════════════════════════
# STEP 2: MODEL (Page 3 + 4)
# ═══════════════════════════════════════════════════
print("🧠 Step 2: Building model...")

# Mixed precision (Page 4)
tf.keras.mixed_precision.set_global_policy('mixed_float16')

# Data augmentation (Page 3)
augmentation = keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.15),
    layers.RandomZoom(0.1),
    layers.RandomContrast(0.1),
    layers.RandomBrightness(0.1),
], name="augmentation")

# Transfer learning backbone (Page 3)
base_model = keras.applications.EfficientNetB0(
    input_shape=(*IMG_SIZE, 3),
    include_top=False,
    weights="imagenet"
)
base_model.trainable = False  # Phase 1: freeze backbone

# Full model
model = keras.Sequential([
    augmentation,
    layers.Rescaling(1./255),
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.BatchNormalization(),
    layers.Dropout(0.3),
    layers.Dense(128, activation="relu"),
    layers.Dropout(0.2),
    layers.Dense(NUM_CLASSES, activation="softmax", dtype="float32")
], name="capstone_classifier")

model.summary()
print(f"  Total params: {model.count_params():,}")

# ═══════════════════════════════════════════════════
# STEP 3: TRAINING — Phase 1 (Page 2 + 7)
# ═══════════════════════════════════════════════════
print(f"\n🏋️ Step 3a: Phase 1 — Train head (backbone frozen)...")

model.compile(
    optimizer=keras.optimizers.Adam(LR_PHASE1, clipnorm=1.0),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
    jit_compile=True  # XLA for speed (Page 4)
)

callbacks_p1 = [
    keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=5, restore_best_weights=True),
    keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=3),
    keras.callbacks.TensorBoard(log_dir="logs/phase1"),
]

start = time.time()
history_p1 = model.fit(
    train_ds, validation_data=val_ds,
    epochs=EPOCHS_PHASE1, callbacks=callbacks_p1)
p1_time = time.time() - start
print(f"  Phase 1 done in {p1_time:.0f}s")

# ═══════════════════════════════════════════════════
# STEP 3b: Phase 2 — Fine-tune (Page 3)
# ═══════════════════════════════════════════════════
print(f"\n🔧 Step 3b: Phase 2 — Fine-tune top backbone layers...")

base_model.trainable = True
for layer in base_model.layers[:-20]:
    layer.trainable = False

trainable = sum(1 for l in model.layers if hasattr(l, 'trainable') and l.trainable)
print(f"  Unfrozen top layers for fine-tuning")

model.compile(
    optimizer=keras.optimizers.Adam(LR_PHASE2, clipnorm=1.0),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
    jit_compile=True
)

callbacks_p2 = [
    keras.callbacks.EarlyStopping(
        monitor="val_accuracy", patience=3, restore_best_weights=True),
    keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=2, min_lr=1e-7),
    keras.callbacks.ModelCheckpoint(
        "best_model.keras", save_best_only=True, monitor="val_accuracy"),
    keras.callbacks.TensorBoard(log_dir="logs/phase2"),
]

start = time.time()
history_p2 = model.fit(
    train_ds, validation_data=val_ds,
    epochs=EPOCHS_PHASE2, callbacks=callbacks_p2)
p2_time = time.time() - start
print(f"  Phase 2 done in {p2_time:.0f}s")

# ═══════════════════════════════════════════════════
# STEP 4: EVALUATION
# ═══════════════════════════════════════════════════
print("\n📊 Step 4: Evaluation...")

val_loss, val_acc = model.evaluate(val_ds, verbose=0)
print(f"  Val Loss: {val_loss:.4f}")
print(f"  Val Accuracy: {val_acc:.1%}")

# Per-class accuracy
y_true, y_pred = [], []
for images, labels in val_ds:
    preds = model.predict(images, verbose=0)
    y_true.extend(labels.numpy())
    y_pred.extend(np.argmax(preds, axis=1))

y_true, y_pred = np.array(y_true), np.array(y_pred)
print("\n  Per-class accuracy:")
for i, name in enumerate(train_ds.class_names):
    mask = y_true == i
    if mask.sum() > 0:
        acc = (y_pred[mask] == i).mean()
        print(f"    {name:15s}: {acc:.1%} ({mask.sum()} samples)")

# ═══════════════════════════════════════════════════
# STEP 5: EXPORT & DEPLOY (Page 9)
# ═══════════════════════════════════════════════════
print("\n🚀 Step 5: Export & Deploy...")

# 5a. SavedModel (for TF Serving)
os.makedirs(MODEL_DIR, exist_ok=True)
model.save(f"{MODEL_DIR}/1")
sm_size = sum(os.path.getsize(os.path.join(dp, f))
    for dp, dn, filenames in os.walk(f"{MODEL_DIR}/1")
    for f in filenames) / (1024*1024)
print(f"  SavedModel: {sm_size:.1f} MB → {MODEL_DIR}/1/")

# 5b. TFLite (for mobile — dynamic range quantization)
converter = tf.lite.TFLiteConverter.from_saved_model(f"{MODEL_DIR}/1")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()
with open(TFLITE_PATH, "wb") as f:
    f.write(tflite_model)
tflite_size = len(tflite_model) / (1024*1024)
print(f"  TFLite:     {tflite_size:.1f} MB → {TFLITE_PATH}")
print(f"  Compression: {sm_size/tflite_size:.1f}×")

# 5c. Verify TFLite accuracy
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_d = interpreter.get_input_details()
output_d = interpreter.get_output_details()
print(f"  TFLite input:  {input_d[0]['shape']} {input_d[0]['dtype']}")
print(f"  TFLite output: {output_d[0]['shape']} {output_d[0]['dtype']}")

# ═══════════════════════════════════════════════════
# FINAL REPORT
# ═══════════════════════════════════════════════════
print(f"""
{'='*60}
🏆 CAPSTONE PROJECT COMPLETE!
{'='*60}
📊 Data:       {NUM_CLASSES} classes
🧠 Model:      EfficientNetB0 + custom head
⚡ Training:   Phase 1 ({p1_time:.0f}s) + Phase 2 ({p2_time:.0f}s)
🎯 Accuracy:   {val_acc:.1%}
💾 SavedModel: {sm_size:.1f} MB → TF Serving ready
📱 TFLite:     {tflite_size:.1f} MB → Mobile ready
🐳 Docker:     docker run -p 8501:8501 ...
{'='*60}
""")