from typing import Any


def get_config(index: int, sample_steps: list[int], ema_power: float) -> dict[str, Any]:
    return {
        'num_steps': 40,
        'time_steps': sample_steps,
        'image_height': 64,
        'image_width': 64,
        'image_channels': 3,
        'calculate_fid': True,
        'calculate_isc': True,
        'fid_reference_path': 'PATH/fid-refs/imagenet-64x64.npz',
        'model_name': 'edm-imagenet-64x64-cond-adm',
        'model_path': f'PATH/results/lines-imagenet-64x64-9-loss-lpips-vgg-64-bs-32-br-64-ae3da90b-fee8-423a-935c-4feb477edde9.0/checkpoints/last/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_{ema_power}.pth',
        'num_classes': 1000,
        'fid_save_path': f'results_fidelity/fid/lines-imagenet-64x64-9-loss-lpips-vgg-64-bs-32-br-64-ae3da90b-fee8-423a-935c-4feb477edde9.0/checkpoints/last/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_{ema_power}/{"[" + str(",".join(map(str, sample_steps))) + "]"}/{index}.pickle',
        'isc_save_path': f'results_fidelity/isc/lines-imagenet-64x64-9-loss-lpips-vgg-64-bs-32-br-64-ae3da90b-fee8-423a-935c-4feb477edde9.0/checkpoints/last/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_{ema_power}/{"[" + str(",".join(map(str, sample_steps))) + "]"}/{index}.pickle',
        'batch_size': 10
    }


def get_configs() -> list[dict[str, Any]]:
    configs: list[dict[str, Any]] = []
    for sample_steps in [[0, 40], [0, 1, 40], [0, 1, 2, 40], [0, 3, 5, 40]]:
        for index in range(10):
            configs.append(get_config(index, sample_steps, round(2 / 3, 5)))
    return configs
