// Small, separate synthetic training/validation sets demonstrate the workflow.
// Real image preprocessing and label mappings belong to the experiment.
import aner.nn
import aner.tensor
import aner.parallel

let train_rows: Tensor = [
    [-1.0, -0.7, -0.9, -0.8],
    [1.0, 0.7, 0.9, 0.8],
    [-0.8, -0.3, -0.6, -0.4],
    [0.8, 0.3, 0.6, 0.4],
    [-0.2, -0.1, -0.4, -0.3]]
let train_images = tensor.reshape(train_rows, List(5, 1, 2, 2))
let train_labels: Tensor = [[0.0], [1.0], [0.0], [1.0], [0.0]]
let valid_rows: Tensor = [
    [-0.6, -0.4, -0.7, -0.5],
    [0.6, 0.4, 0.7, 0.5],
    [0.3, 0.2, 0.5, 0.4]]
let valid_images = tensor.reshape(valid_rows, List(3, 1, 2, 2))
let valid_labels: Tensor = [[0.0], [1.0], [1.0]]

let model = nn.cnn(1, 2, 2, seed: 31)
    .conv2d(outputs: 8, kernel: 1).relu().head(classes: 2)
    .on("cpu", precision: "Float64")
let head_fit = model.freeze_backbone().fit(train_images, train_labels,
    epochs: 12, batch: 3, rate: 0.03, clip_norm: 1.0,
    weight_decay: 0.0001, shuffle: false,
    execution: parallel.local(workers: 2))
let fitted = head_fit.freeze_backbone(false).fit(train_images, train_labels,
    epochs: 20, batch: 3, rate: 0.01, clip_norm: 1.0,
    weight_decay: 0.0001, shuffle: true)

let validation = fitted.evaluate(valid_images, valid_labels, batch: 2)
print(validation.summary())
print(validation.loss())
print(validation.accuracy())
print(validation.samples())
print(fitted.probabilities(valid_images).cols())
print(fitted.predict(valid_images).rows())
// An independently loaded next batch retains Adam moments and cumulative steps.
let next_batch = fitted.train_step(train_images, train_labels,
    rate: 0.001, clip_norm: 1.0, weight_decay: 0.0001)
print(next_batch.updates())
// Explicit new-data phase preserves learned weights while resetting moments:
// let adapted = fitted.reset_optimizer().fit(new_images, new_labels, epochs: 2)
// Save to a fresh path for exact parameter/optimizer continuation:
// fitted.save("cnn-workflow.anercnn")