from typing import Any

num_runs: int = 104
samples_per_run: int = 50_000


def get_config(index: int) -> dict[str, Any]:
    return {
        'num_steps': 40,
        'time_steps': [0, 1, 40],
        'model_name': 'edm-imagenet-64x64-cond-adm',
        'model_load_path': 'PATH/results/lines-gan-imagenet-64x64-4-bs-16-br-32-conditional-disc-resume-ema-ef0f402c-115e-4f8a-be90-45c286e75f63.0/checkpoints/last.tmp.1/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667.pth',
        'noises_folder': 'PATH/imagenet-64x64/second-phase/noises',
        'labels_folder': 'PATH/imagenet-64x64/second-phase/labels',
        'output_folder': 'PATH/imagenet-64x64/second-phase/ef0f402c-115e-4f8a-be90-45c286e75f63_[0,1,40]',
        'num_samples': samples_per_run,
        'batch_size': 50,
        'num_classes': 1000,
        '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)]
