// Run in a fresh directory; the continuation example creates dropout.anermodel.
// Partitioning, initialization, shuffling and dropout have separate seeds.
import aner.dataset
import aner.nn
import aner.ml

let parts = dataset.iris().partition(validation: 0.2, test: 0.2, seed: 81)
let network = nn.sequential(seed: 42)
    .dense(inputs: 4, outputs: 16, initialization: "he").relu()
    .dropout(probability: 0.2)
    .dense(inputs: 16, outputs: 3, initialization: "xavier")
let plan = ml.classifier(network, standardize: true, optimizer: "adam")
    .dropout_seed(seed: 2026)
let config = ml.training(epochs: 30, batch: 16, rate: 0.01, seed: 91)
    .validation(parts.validation()).select_best(true)
let fitted = plan.fit(parts.train(), config: config)
print(fitted.summary())
// Validation and prediction use the complete network, with dropout disabled.
print(fitted.evaluate(parts.test()).summary())

// Omit dropout_seed for a fresh automatic seed. Preserve its summary/checkpoint
// to recover the exact integer seed, stream protocol and continuation position.
let automatic = ml.classifier(network, standardize: true, optimizer: "adam")
    .fit(parts.train(), config: ml.training(epochs: 2, batch: 16, rate: 0.01, seed: 91))
print(automatic.summary())

// Checkpoint continuation restores the seed, sample schedule and optimizer state.
// The same pattern works with optimizer: "sgd".
let partial = plan.fit(parts.train(),
    config: ml.training(epochs: 2, batch: 16, rate: 0.01, seed: 91))
partial.save("dropout.anermodel")
let restored = ml.load_classifier("dropout.anermodel")
let continued = plan.warm_start(restored).fit(parts.train(),
    config: ml.training(epochs: 2, batch: 16, rate: 0.01, seed: 91))
let uninterrupted = plan.fit(parts.train(),
    config: ml.training(epochs: 4, batch: 16, rate: 0.01, seed: 91))
print(continued.training_loss() == uninterrupted.training_loss())