"""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, split = 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", "fa", "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
    assert split["task_seed"] == args.task_seed
    print(f"{dataset}: input={n_in}, classes={n_out}, fixed task seed={args.task_seed}")


def check_synthetic_validation():
    argv = [
        "run.py", "--dataset", "teacher", "--device", "cpu",
        "--task_train_examples", "80", "--task_test_examples", "32",
        "--val_examples", "20", "--eval_split", "validation",
        "--batch_size", "16", "--depth", "2", "--width", "12",
        "--task_n_in", "20", "--task_classes", "4",
        "--teacher_depth", "3", "--teacher_width", "10",
    ]
    old = sys.argv
    try:
        sys.argv = argv
        args = get_args()
    finally:
        sys.argv = old
    train, validation, _, _, split = load_task(args, "cpu")
    assert sum(y.numel() for _, y in train) == 60
    assert sum(y.numel() for _, y in validation) == 20
    assert split["evaluation_split"] == "validation"
    assert split["validation_index_sha256"]
    assert split["split_from_training_only"] is True
    print("synthetic validation: 60/20; test generator untouched")


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"])
    check_synthetic_validation()
    print("ALL SYNTHETIC TASK CHECKS PASSED")