import os
from PIL import Image
import torchvision.transforms.functional as f
from utils import load_face_generator
from omegaconf import OmegaConf
import random
import sys

def generate_face_image(
    anything_facemaker,
    class_concept, 
    face_img_pil=None, 
    controlnet_conditioning_scale=1.0,
    strength=0.95,
):
    face_img_pil = f.center_crop(
        f.resize(face_img_pil, 512), 512).convert('RGB')
    prompt = anything_facemaker.prompt_template.format(class_concept)
    # # There are four ways to generate a image by now.
    # pure_generate = anything_facemaker.generate(prompt=prompt, image=face_img_pil, do_inversion=False)
    # inversion = anything_facemaker.generate(prompt=prompt, image=face_img_pil, strength=strength, do_inversion=True)

    if controlnet_conditioning_scale == None:
        init_face_pil = anything_facemaker.generate(prompt=prompt)
        return init_face_pil
    
    if strength is None:
        pure_control = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=False,
                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale)
        init_face_pil = pure_control
    else:
        control_inversion = anything_facemaker.face_control_generate(prompt=prompt, face_img_pil=face_img_pil, do_inversion=True, 
                                                                 strength=strength,
                                                                 controlnet_conditioning_scale=controlnet_conditioning_scale)
        init_face_pil = control_inversion
    return init_face_pil


def experiment(anything_facemaker, concepts_path, face_img_path, output_dir,
               controlnet_conditioning_scale=1., strength=0.95):
    os.makedirs(output_dir, exist_ok=True)
    face_img_pil = Image.open(face_img_path)
    face_img_pil = f.center_crop(
        f.resize(face_img_pil, 512), 512).convert('RGB')
    with open(concepts_path) as fr:
        concepts = fr.read().split('\n')
    concepts = [concept for concept in concepts if len(concept)!=0]
    random.shuffle(concepts)
    for concept in concepts[:4]:
        save_path = os.path.join(output_dir, f'{concept}.png')
        if os.path.exists(save_path):
            continue
        init_face_pil = generate_face_image(
            anything_facemaker,
            class_concept=concept,
            face_img_pil=face_img_pil,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            strength=strength,
        )

        save_path = os.path.join(output_dir, f'{concept}.png')
        init_face_pil.save(save_path)



if __name__=='__main__':
    # run this in repo path:
    # PYTHONPATH=.:$PYTHONPATH python experiments/experiment.py 
    model_config_path = 'resources/models.yaml'
    # model_config_path = 'resources/models_personality.yaml'
    model_config = OmegaConf.load(model_config_path)['models']
    gameicon_config = model_config['GameIconInstitute_mode']

    # face_img_path = 'resources/images/faces/0.jpg'
    face_img_dir='resources/images/faces'
    faces = os.listdir(face_img_dir)
    controlnet_conditioning_scale=1.
    strength=0.95

    for model, model_info in model_config.items():
        
        anything_facemaker = load_face_generator(
            model_dir=model_info['model_dir'],
            lora_path=model_info['lora_path'],
            prompt_template=model_info['prompt_template'],
            negative_prompt=model_info['negative_prompt']
        )
        output_dir = os.path.join(sys.argv[1], model)
        os.makedirs(output_dir, exist_ok=True)
        # concept test, with control and inversion
        input_dir = 'resources/prompts'
        for dir, folders, files in os.walk(input_dir):
            for file in files:
                input_file = os.path.join(dir, file)
                file_output_dir = os.path.join(output_dir, file)
                print(f'input_file: {input_file}')
                print(f'file_output_dir: {file_output_dir}')
                face_img_path = os.path.join(face_img_dir, random.choice(faces))
                experiment(anything_facemaker, input_file, face_img_path, output_dir=file_output_dir,
                        controlnet_conditioning_scale=controlnet_conditioning_scale,
                        strength=strength)

    # # concept, with control and inversion
    # experiment(anything_facemaker, 'resources/concepts.txt', face_img_path, output_dir='results/concepts/control_inversion',
    #         controlnet_conditioning_scale=controlnet_conditioning_scale,
    #         strength=strength)

    # # concept test, no control no inversion
    # experiment(anything_facemaker, 'resources/concepts_test.txt', face_img_path, output_dir='results/concepts_test/generate',
    #         controlnet_conditioning_scale=None,
    #         strength=strength)

    # # concept, no control no inversion
    # experiment(anything_facemaker, 'resources/concepts.txt', face_img_path, output_dir='results/concepts/generate',
    #         controlnet_conditioning_scale=None,
    #         strength=strength)