Callbacks

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.callbacks import (
    EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard, CSVLogger,
    LearningRateScheduler, TerminateOnNaN, BackupAndRestore,
)
import datetime, math, os

# 1) Build a model
model = keras.Sequential([
    keras.layers.Input(shape=(28, 28)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10),
])
model.compile(
    optimizer='adam',
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['sparse_categorical_accuracy'],
)

# 2) EarlyStopping — stop when val loss stops improving
es = EarlyStopping(
    monitor='val_loss',
    patience=5,                  # wait 5 epochs without improvement
    min_delta=1e-4,              # 'improvement' threshold
    restore_best_weights=True,    # roll back to the best checkpoint
    verbose=1,
)

# 3) ModelCheckpoint — save best (and/or periodically)
ckpt_dir = './ckpt'
os.makedirs(ckpt_dir, exist_ok=True)
mc_best = ModelCheckpoint(
    filepath=f'{ckpt_dir}/best.keras',
    monitor='val_loss',
    save_best_only=True,
    verbose=1,
)
mc_every = ModelCheckpoint(
    filepath=f'{ckpt_dir}/epoch-{{epoch:03d}}.keras',
    save_freq='epoch',
)

# 4) ReduceLROnPlateau — back off LR when stuck
rlr = ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.5,                  # halve the LR
    patience=3,                  # ... after 3 epochs of no improvement
    min_lr=1e-6,
    verbose=1,
)

# 5) Learning-rate schedules — explicit cosine
def cosine_lr(epoch, lr, *, base_lr=1e-3, warmup=5, total=50):
    if epoch < warmup:
        return base_lr * (epoch + 1) / max(warmup, 1)
    progress = (epoch - warmup) / max(total - warmup, 1)
    return base_lr * 0.5 * (1 + math.cos(math.pi * progress))

ls = LearningRateScheduler(cosine_lr, verbose=1)

# 6) TensorBoard — metrics + graphs + per-batch histograms
logdir = './tb/' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
tb = TensorBoard(
    log_dir=logdir,
    histogram_freq=1,            # log weight histograms each epoch
    profile_batch=(2, 4),         # profiler — slow, use sparingly
    update_freq='epoch',
    write_graph=True,
    write_images=False,
)

# Launch:  tensorboard --logdir ./tb

# 7) CSV log + NaN guard + crash recovery
csv = CSVLogger('./train.csv', append=True)
nan = TerminateOnNaN()
backup = BackupAndRestore(backup_dir='./backup')   # resume from last completed epoch

# 8) Fit with the lot
history = model.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    epochs=50,
    batch_size=64,
    callbacks=[es, mc_best, mc_every, rlr, ls, tb, csv, nan, backup],
    verbose=2,
)

# 9) Custom callback — Slack/Discord ping on val loss improvement
import urllib.request, json

class NotifyOnBest(keras.callbacks.Callback):
    def __init__(self, webhook_url, monitor='val_loss', mode='min'):
        super().__init__()
        self.webhook = webhook_url
        self.monitor = monitor
        self.mode = mode
        self.best = math.inf if mode == 'min' else -math.inf

    def on_epoch_end(self, epoch, logs=None):
        v = logs.get(self.monitor)
        if v is None: return
        improved = v < self.best if self.mode == 'min' else v > self.best
        if improved:
            self.best = v
            self._notify(f'epoch {epoch+1}: {self.monitor}={v:.4f} ✅')

    def on_train_end(self, logs=None):
        self._notify(f'Training complete. Best {self.monitor}: {self.best:.4f}')

    def _notify(self, text):
        body = json.dumps({'text': text}).encode()
        req = urllib.request.Request(self.webhook, data=body, headers={'Content-Type': 'application/json'})
        try: urllib.request.urlopen(req, timeout=5).read()
        except Exception as e: print('notify failed:', e)

notify = NotifyOnBest(os.environ['SLACK_WEBHOOK'])

# 10) Custom callback — gradient norm logger (debugging)
class GradNormLogger(keras.callbacks.Callback):
    def __init__(self, sample_x, sample_y): super().__init__(); self.x, self.y = sample_x, sample_y
    def on_epoch_end(self, epoch, logs=None):
        with tf.GradientTape() as tape:
            preds = self.model(self.x, training=True)
            loss  = self.model.compiled_loss(self.y, preds)
        grads = tape.gradient(loss, self.model.trainable_weights)
        gnorm = tf.linalg.global_norm(grads)
        print(f'epoch {epoch+1} grad_norm={float(gnorm):.4f}')

# 11) Callback ordering matters
# Keras processes callbacks in the order you pass them.
# Logging-style callbacks last (after model state is finalized).
# EarlyStopping before BackupAndRestore so resume picks the latest non-stopped state.

# 12) Class-based custom callback — full lifecycle
class Demo(keras.callbacks.Callback):
    def on_train_begin(self, logs=None): ...
    def on_train_end  (self, logs=None): ...
    def on_epoch_begin(self, epoch, logs=None): ...
    def on_epoch_end  (self, epoch, logs=None): ...
    def on_train_batch_begin(self, batch, logs=None): ...
    def on_train_batch_end  (self, batch, logs=None): ...
    def on_test_begin (self, logs=None): ...
    def on_test_end   (self, logs=None): ...
    def on_predict_begin(self, logs=None): ...
    def on_predict_end  (self, logs=None): ...

# 13) Production checklist
#   ✓ ModelCheckpoint(save_best_only=True) somewhere durable
#   ✓ EarlyStopping with restore_best_weights=True
#   ✓ ReduceLROnPlateau OR an explicit schedule
#   ✓ BackupAndRestore for long runs (resume after preemption)
#   ✓ CSVLogger or TensorBoard for post-hoc analysis
#   ✓ TerminateOnNaN — catch divergence early
#   ✓ Notification callback so you don't refresh logs at midnight

# 14) Common bugs
#   • EarlyStopping fires before model has converged → patience too small
#   • restore_best_weights=False (default) → final weights aren't the best
#   • ModelCheckpoint saving on training loss, not val → overfits to train
#   • ReduceLROnPlateau + LearningRateScheduler conflict → pick one
#   • TensorBoard profile_batch on every batch → massive logs, slow training
#   • Custom callback accessing self.model.optimizer.lr — set via assign for new TF versions

from tensorflow.keras import callbacks
cbs = [
    callbacks.EarlyStopping(patience=3, restore_best_weights=True),
    callbacks.ModelCheckpoint('best.keras', save_best_only=True),
    callbacks.TensorBoard('logs/'),
]