import gradio as gr
import PIL.Image
import shlex
import shutil
import subprocess
from pathlib import Path
import os
import torch
from tqdm import tqdm



def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
    w, h = image.size
    if w == h:
        return image
    elif w > h:
        new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0))
        new_image.paste(image, (0, (w - h) // 2))
        return new_image
    else:
        new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0))
        new_image.paste(image, ((h - w) // 2, 0))
        return new_image



def train_submit(
        prompt, anchor_prompt, concept_type, reg_lambda, iterations, lr, openai_key, save_path, mem_impath=None
    ):
        if not torch.cuda.is_available():
            raise gr.Error('CUDA is not available.')

        torch.cuda.empty_cache()
        original_prompt = prompt
        parameter_group = "cross-attn"
        train_batch_size = 4
        if concept_type == 'style':
            class_data_dir = f'./data/samples_painting/'
            anchor_prompt = f'./assets/painting.txt'
            openai_key = ''
        elif concept_type == 'object':
            os.makedirs('temp', exist_ok=True)
            class_data_dir = f'./temp/{anchor_prompt}'
            name = save_path.split('/')[-1]
            prompt = f'{anchor_prompt}+{prompt}'
            assert openai_key is not None

            if len(openai_key.split('\n')) > 1:
                openai_key = openai_key.split('\n')
                with open(f'./temp/{name}.txt', 'w') as f:
                    for prompt_ in openai_key:
                        f.write(prompt_.strip()+'\n')
                openai_key = ''
                anchor_prompt = f'./temp/{name}.txt'
        elif concept_type == 'memorization':
            os.system("wget https://dl.fbaipublicfiles.com/sscd-copy-detection/sscd_imagenet_mixup.torchscript.pt -P assets/") 
            os.makedirs('temp', exist_ok=True)
            prompt = f'*+{prompt}'
            name = save_path.split('/')[-1]
            train_batch_size = 1
            lr = 5e-7
            parameter_group = "full-weight"

            assert openai_key is not None
            assert mem_impath is not None

            if len(openai_key.split('\n')) > 1:
                openai_key = openai_key.split('\n')
                with open(f'./temp/{name}.txt', 'w') as f:
                    for prompt_ in openai_key:
                        f.write(prompt_.strip()+'\n')
                openai_key = ''
                anchor_prompt = f'./temp/{name}.txt'
            else:
                anchor_prompt = prompt

            print(mem_impath)
            image = PIL.Image.open(mem_impath[0][0].name)
            image = pad_image(image)
            image = image.convert('RGB')
            mem_impath = f"./temp/{original_prompt.lower().replace(' ', '')}.jpg"
            image.save(mem_impath, format='JPEG', quality=100)

            class_data_dir = f"./temp/{original_prompt.lower().replace(' ', '')}"
            

        command = f'''
        accelerate launch concept-ablation-diffusers/train.py \
          --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"  \
          --output_dir={save_path} \
          --class_data_dir={class_data_dir} \
          --class_prompt="{anchor_prompt}"  \
          --caption_target "{prompt}" \
          --concept_type {concept_type} \
          --resolution=512  \
          --train_batch_size={train_batch_size}  \
          --learning_rate={lr}  \
          --max_train_steps={iterations} \
          --scale_lr --hflip \
          --parameter_group {parameter_group} \
          --openai_key "{openai_key}" \
          --enable_xformers_memory_efficient_attention  --num_class_images 500
        '''

        if concept_type == 'style':
            command += f'  --noaug'
        
        if concept_type == 'memorization':
            command += f' --use_8bit_adam  --with_prior_preservation --prior_loss_weight=1.0 --mem_impath {mem_impath}'
            
        with open(f'{save_path}/train.sh', 'w') as f:
            command_s = ' '.join(command.split())
            f.write(command_s)

        res = subprocess.run(shlex.split(command))

        if res.returncode == 0:
            result_message = 'Training Completed!'
        else:
            result_message = 'Training Failed!'
        weight_paths = sorted(Path(save_path).glob('*.bin'))
        print(weight_paths)
        return gr.update(value=result_message), weight_paths[0]


def inference(model_path, prompt, n_steps, generator):
    import sys 
    sys.path.append('concept-ablation/diffusers/.')
    from model_pipeline import CustomDiffusionPipeline
    import torch

    pipe = CustomDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to("cuda")
    image1 = pipe(prompt, num_inference_steps=n_steps, guidance_scale=6., eta=1., generator=generator).images[0]

    pipe.load_model(model_path)
    image2 = pipe(prompt, num_inference_steps=n_steps, guidance_scale=6., eta=1., generator=generator).images[0]

    return image1, image2