import glob
from argparse import Namespace
from configs import paths_config
from models.StyleCLIP.mapper.scripts.inference import run
from scripts.run_pti import run_PTI

meta_data = {
    'afro': ['afro', False, False, True],
    'angry': ['angry', False, False, True],
    'Beyonce': ['beyonce', False, False, False],
    'bobcut': ['bobcut', False, False, True],
    'bowlcut': ['bowlcut', False, False, True],
    'curly hair': ['curly_hair', False, False, True],
    'Hilary Clinton': ['hilary_clinton', False, False, False],
    'Jhonny Depp': ['depp', False, False, False],
    'mohawk': ['mohawk', False, False, True],
    'purple hair': ['purple_hair', False, False, False],
    'surprised': ['surprised', False, False, True],
    'Taylor Swift': ['taylor_swift', False, False, False],
    'trump': ['trump', False, False, False],
    'Mark Zuckerberg': ['zuckerberg', False, False, False]
}


def styleclip_edit(use_multi_id_G, run_id, use_wandb, edit_types):
    images_dir = paths_config.input_data_path
    pretrained_mappers = paths_config.style_clip_pretrained_mappers
    data_dir_name = paths_config.input_data_id
    if run_id == '':
        run_id = run_PTI(run_name='', use_wandb=use_wandb, use_multi_id_training=False)
    images = glob.glob(f"{images_dir}/*.jpeg")
    w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
    for image_name in images:
        image_name = image_name.split(".")[0].split("/")[-1]
        embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
        latent_path = f'{embedding_dir}/0.pt'
        for edit_type in set(meta_data.keys()).intersection(edit_types):
            edit_id = meta_data[edit_type][0]
            args = {
                "exp_dir": f'{paths_config.styleclip_output_dir}',
                "checkpoint_path": f"{pretrained_mappers}/{edit_id}.pt",
                "couple_outputs": False,
                "mapper_type": "LevelsMapper",
                "no_coarse_mapper": meta_data[edit_type][1],
                "no_medium_mapper": meta_data[edit_type][2],
                "no_fine_mapper": meta_data[edit_type][3],
                "stylegan_size": 1024,
                "test_batch_size": 1,
                "latents_test_path": latent_path,
                "test_workers": 1,
                "run_id": run_id,
                "image_name": image_name,
                'edit_name': edit_type,
                "data_dir_name": data_dir_name
            }

            run(Namespace(**args), run_id, image_name, use_multi_id_G)