import argparse
from typing import Any

from scripts.plots.scripts.grid.main import run as run_main
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_main(
        name=name,
        model_name='edm-afhqv2-64x64-uncond-vp',
        edm_model_path='PATH/state_dicts/edm-afhqv2-64x64-uncond-vp.pth',
        vgg_model_path='PATH/results/lines-afhqv2-64x64-1-loss-lpips-vgg-64-bs-512-br-1-0fd31e38-e3b7-40d2-9862-d5eb3cb3e0e0.0/checkpoints/50000/ema/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667.pth',
        gan_model_path='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_0.66667.pth',
        save_folder='samples-plot-grid',
        image_height=64,
        image_width=64,
        image_channels=3,
        num_samples=500,
        num_steps=40,
        two_time_steps=[0, 1],
        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()
