import numpy as np

import train_omniglot


def fake_omniglot_batches(shots, batch_classes):
    def generator():
        while True:
            yield np.array([[0.0, 1.0], [1.0, 0.0]], dtype=float), np.array([0, 1], dtype=int)

    return generator()


def passthrough_tqdm(iterable, **kwargs):
    return iterable


def test_run_few_shot_does_not_depend_on_global_args(monkeypatch) -> None:
    monkeypatch.setattr(train_omniglot, "omniglot_batches", fake_omniglot_batches)
    monkeypatch.setattr(train_omniglot, "tqdm", passthrough_tqdm)
    monkeypatch.delattr(train_omniglot, "args", raising=False)

    summary_data = train_omniglot.run_few_shot(
        shots=1,
        num_batches=2,
        num_excit=2,
        num_inhib=0,
        num_mod=0,
        p_feedforward=0.0,
        p_feedback=0.0,
        log_every=1,
    )

    assert summary_data["shots"] == 1
    assert summary_data["num_batches"] == 2
    assert summary_data["total_slots"] >= 1
