"""Fast checks for the synthetic depth-scaling task plumbing."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from experiments.run import build, get_args, load_task


def check_task(dataset, expected_in, expected_out, extra):
    argv = [
        "run.py", "--dataset", dataset, "--device", "cpu",
        "--task_train_examples", "64", "--task_test_examples", "32",
        "--batch_size", "16", "--depth", "2", "--width", "12",
    ] + extra
    old = sys.argv
    try:
        sys.argv = argv
        args = get_args()
    finally:
        sys.argv = old
    train, test, n_in, n_out = load_task(args, "cpu")
    assert (n_in, n_out) == (expected_in, expected_out)
    args.n_in, args.n_out = n_in, n_out
    x, y = next(iter(train))
    assert x.shape == (16, n_in)
    assert y.min() >= 0 and y.max() < n_out
    for mode in ("bp", "dfa", "sdil"):
        args.mode = mode
        net, _ = build(args, "cpu")
        assert net.logits(x).shape == (16, n_out)
    assert sum(batch_y.numel() for _, batch_y in test) == 32
    print(f"{dataset}: input={n_in}, classes={n_out}, fixed task seed={args.task_seed}")


if __name__ == "__main__":
    torch.manual_seed(0)
    check_task("teacher", 20, 4,
               ["--task_n_in", "20", "--task_classes", "4",
                "--teacher_depth", "3", "--teacher_width", "10"])
    check_task("hierarchical", 16, 4,
               ["--task_levels", "4", "--task_classes", "4"])
    check_task("tentmap", 5, 2,
               ["--task_levels", "4", "--task_n_in", "5"])
    print("ALL SYNTHETIC TASK CHECKS PASSED")