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,
        'reference_path': 'PATH/afhqv2-64x64-edm.npz',
        'model_name': 'edm-afhqv2-64x64-uncond-vp',
        'model_path': f'PATH/results/lines-gan-afhqv2-64x64-1-bs-32-br-16-a54bdffd-4d4f-4b01-82b6-124625059812.0/checkpoints/30000/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_{ema_power}.pth',
        'output_save_path': f'results_fid/lines-gan-afhqv2-64x64-1-bs-32-br-16-a54bdffd-4d4f-4b01-82b6-124625059812.0/30000/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': 50
    }


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
