"""Fast checks for frozen train/validation/test protocol plumbing."""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.data import get_dataset_splits
from sdil.core import SDILConfig, SDILNet
from experiments.run import (calibration_work_per_event, fixed_training_probe,
                             residual_lesion_report)


def loader_labels(loader):
    return torch.cat([labels.cpu() for _, labels in loader])


def main():
    first = get_dataset_splits(
        "mnist", batch_size=256, device="cpu", val_examples=1000, split_seed=2027)
    second = get_dataset_splits(
        "mnist", batch_size=256, device="cpu", val_examples=1000, split_seed=2027)
    train, validation, test, n_in, n_out, metadata = first
    _, validation_2, _, _, _, metadata_2 = second

    assert (n_in, n_out) == (784, 10)
    assert metadata["train_examples"] == 59000
    assert metadata["validation_examples"] == 1000
    assert metadata["test_examples"] == 10000
    assert metadata["validation_index_sha256"] == metadata_2["validation_index_sha256"]
    assert metadata["validation_class_counts"] == {str(i): 100 for i in range(10)}
    assert torch.equal(loader_labels(validation), loader_labels(validation_2))
    assert sum(y.numel() for _, y in train) == 59000
    assert sum(y.numel() for _, y in validation) == 1000
    assert sum(y.numel() for _, y in test) == 10000

    fmnist = get_dataset_splits(
        "fmnist", batch_size=256, device="cpu", val_examples=5000, split_seed=2027)
    f_train, f_validation, f_test, f_n_in, f_n_out, f_metadata = fmnist
    assert (f_n_in, f_n_out) == (784, 10)
    assert f_metadata["train_examples"] == 55000
    assert f_metadata["validation_examples"] == 5000
    assert f_metadata["validation_class_counts"] == {str(i): 500 for i in range(10)}
    assert sum(y.numel() for _, y in f_train) == 55000
    assert sum(y.numel() for _, y in f_validation) == 5000
    assert sum(y.numel() for _, y in f_test) == 10000

    probe_loader = get_dataset_splits(
        "mnist", batch_size=256, device="cpu", val_examples=1000,
        split_seed=2027)[0]
    generator_state = probe_loader.g.get_state().clone()
    probe_x, probe_y = fixed_training_probe(probe_loader, 137, "cpu")
    assert torch.equal(generator_state, probe_loader.g.get_state())
    assert torch.equal(probe_x, probe_loader.x[:137])
    assert torch.equal(probe_y, probe_loader.y[:137])

    net = SDILNet([10, 8, 8, 3], device="cpu")
    simultaneous = calibration_work_per_event(
        net, SDILConfig(pert_ndirs=2, pert_mode="simultaneous"))
    layerwise = calibration_work_per_event(
        net, SDILConfig(pert_ndirs=2, pert_mode="layerwise"))
    overridden = calibration_work_per_event(
        net, SDILConfig(pert_ndirs=1, pert_mode="simultaneous"),
        pert_mode="layerwise", pert_ndirs=2)
    assert simultaneous == {
        "batch_loss_evaluations": 4,
        "forward_equivalent_batches": 5.0,
        "perturbation_batch_expansion": 4,
    }
    assert layerwise["batch_loss_evaluations"] == 9
    assert abs(layerwise["forward_equivalent_batches"] - 11.0 / 3.0) < 1e-12
    assert overridden == layerwise

    lesion_net = SDILNet([1, 4, 4, 4, 4, 2], act="relu", residual=True, device="cpu")
    lesion_x = torch.linspace(0, 1, 16).view(-1, 1)
    lesion_y = torch.zeros(16, dtype=torch.long)
    lesion = residual_lesion_report(
        lesion_net, [(lesion_x, lesion_y)], lesion_x[:8], fraction=1.0 / 3.0)
    assert lesion["interior_layers"] == [1, 2, 3]
    assert lesion["lesioned_layers"] == [3]
    assert len(lesion["branch_to_skip_rms"]) == 3
    assert 0.0 <= lesion["lesion_eval_acc"] <= 1.0
    print("validation split hash:", metadata["validation_index_sha256"])
    print("train/validation/test: 59000/1000/10000; stratification exact")
    print("FashionMNIST recovery split: 55000/5000/10000; stratification exact")
    print("training-prefix diagnostic probe: deterministic; shuffle state unchanged")
    print("simultaneous/layerwise calibration cost accounting: exact")
    print("final-third residual-block lesion: exact block selection and finite report")
    print("ALL PROTOCOL SMOKE CHECKS PASSED")


if __name__ == "__main__":
    main()