import argparse
from typing import Any

from scripts.plots.scripts.samples.main_single_step import run as run_single_step
from utils.logger.logger import Logger
from utils.utils import get_class_name


def run(
        name: str,
        seed_entropy_noises: int = None
) -> None:
    Logger.debug(
        f'{get_class_name(run)} - '
        f'name: {name}, seed_entropy_noises: {seed_entropy_noises}'
    )
    run_single_step(
        name=name,
        model_name='edm-cifar10-32x32-uncond-vp',
        gan_model_path='PATH/results/lines-gan-cifar10-32x32-uncond-13-bs-32-br-16-83b461b5-b32e-41b3-b31c-ceb9c739fe43.0/checkpoints/10000/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.75.pth',
        save_folder='samples-plot',
        image_height=32,
        image_width=32,
        image_channels=3,
        num_samples=100,
        rows=10,
        cols=10,
        num_steps=18,
        batch_size=50,
        seed_entropy_noises=seed_entropy_noises,
        num_processes=8
    )
    Logger.debug(f'{get_class_name(run)} - done')


def run_from_config(config: dict[str, Any]) -> None:
    run(**config)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', type=str, required=True)
    parser.add_argument('--seed_entropy_noises', type=int, default=None)
    return parser.parse_args()


def get_config_from_args(args: argparse.Namespace) -> dict[str, Any]:
    return vars(args)


def main() -> None:
    args = parse_args()
    config = get_config_from_args(args)
    run_from_config(config)


if __name__ == '__main__':
    main()
