"""Run from the project root with --config experiment.json --mode overfit/train."""
import numpy as np
from common import arguments, settings, load_data, print_scores
from network_numpy import initialize, forward, loss_signal, backward, step, scores

# region one_step

def train_step(params, X_batch, y_batch, cfg):
    cache = forward(params, X_batch)
    loss, dZ2 = loss_signal(cache, y_batch, cfg['loss'])
    grads = backward(params, cache, dZ2)
    step(params, grads, cfg['learning_rate'])
    return loss

# endregion one_step


def main():
    args = arguments()
    cfg = settings(args.config)
    X, y = load_data(args.config, cfg, 'train')
    B = cfg['batch_size']
    if B > len(X):
        raise ValueError('Batch is larger than the training set')
    params = initialize(cfg)
    if args.mode == 'overfit':
        xb, yb = X[:B], y[:B]
        print_scores('Fixed batch BEFORE', *scores(params, xb, yb, cfg['loss']), len(xb))
        for _ in range(cfg['overfit_steps']):
            train_step(params, xb, yb, cfg)
        print_scores('Fixed batch AFTER', *scores(params, xb, yb, cfg['loss']), len(xb))
        print('This checks memorization of one batch, not generalization.')
        return
    print_scores('Training BEFORE', *scores(params, X, y, cfg['loss']), len(X))
    rng = np.random.default_rng(cfg['shuffle_seed'])
    updates = 0
    # region batches
    for epoch in range(cfg['epochs']):
        order = rng.permutation(len(X)) if cfg['shuffle'] else np.arange(len(X))
        for start in range(0, len(X), B):
            indices = order[start:start + B]
            train_step(params, X[indices], y[indices], cfg)
            updates += 1
    # endregion batches
    print(f'Completed {cfg["epochs"]} epochs and {updates} updates.')
    print_scores('Training AFTER', *scores(params, X, y, cfg['loss']), len(X))
    X_test, y_test = load_data(args.config, cfg, 'test')
    print_scores('Held-out, no updates', *scores(params, X_test, y_test, cfg['loss']), len(X_test))
    print('Small teaching subset; this is not a full-MNIST benchmark. Weights are not saved.')


if __name__ == '__main__':
    main()