from typing import Any

num_runs: int = 21
samples_per_run: int = 50_000


def get_config(index: int) -> dict[str, Any]:
    return {
        'num_steps': 18,
        'time_steps': [0, 1, 18],
        'model_name': 'edm-cifar10-32x32-cond-vp',
        'model_load_path': 'PATH/results/lines-gan-cifar10-32x32-5-bs-32-br-16-e50fbed6-d2db-45be-94ee-0de3b5c4ab60.0/checkpoints/best_beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667.pth.tmp.1',
        'noises_folder': 'PATH/cifar10-32x32/second-phase/noises',
        'labels_folder': 'PATH/cifar10-32x32/second-phase/labels',
        'output_folder': 'PATH/cifar10-32x32/second-phase/e50fbed6-d2db-45be-94ee-0de3b5c4ab60_[0,1,18]',
        'num_samples': samples_per_run,
        'batch_size': 200,
        'num_classes': 10,
        'start_noise_index': index * samples_per_run,
        'start_label_index': index * samples_per_run,
        'start_output_index': index * samples_per_run
    }


def get_configs() -> list[dict[str, Any]]:
    return [get_config(i) for i in range(num_runs)]
