import json
import os
import shutil
from argparse import ArgumentParser
import time
from time import gmtime, strftime
from PIL import Image
import torch
from torch import autocast
from tqdm.auto import tqdm
import numpy as np
import piexif
import piexif.helper

from diffusers import (
    DDIMScheduler,
    #PNDMScheduler,
)

from diffusers import StableDiffusionImg2ImgPipeline 

t = torch.cuda.get_device_properties(0).total_memory
if t <=8500000000:
    print("Not enough GPU memory to generate pictures")
else:
    file1 = open("prompt.txt","r+")
    text = file1.read()
    print(text)

    ##Generating with stable diffusion
    device = "cuda"
    model_path = "CompVis/stable-diffusion-v1-4"

    # Using DDIMScheduler as anexample,this also works with PNDMScheduler
    # uncomment this line if you want to use it.

    #scheduler = PNDMScheduler.from_config(model_path, subfolder="scheduler", use_auth_token=True)
    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model_path,
        scheduler=scheduler,
        revision="fp16", 
        torch_dtype=torch.float16,
        use_auth_token=True
    ).to(device)
    if t <= 10500000000:
        pipe.enable_attention_slicing()
        
    def dummy_checker(images, **kwargs):
        return images, False
    pipe.safety_checker = dummy_checker

    def preprocess(image):
        w, h = image.size
        w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
        image = image.resize((w, h), Image.Resampling.LANCZOS)
        image = np.array(image).astype(np.float32) / 255.0
        image = image[None].transpose(0, 3, 1, 2)
        image = torch.from_numpy(image)
        return 2.*image - 1.

    path_img = []

    startStrength = 0.86
    deltaStrength = 0.02
    endStrength = 0.941

    startScale = 7.5
    deltaScale = 2.5
    endScale = 15.0

    startSeed = 1022
    endSeed = 1024

    directory_in = "./data/input"
    if not os.path.exists(directory_in):
        os.makedirs(directory_in)
    for root, subdirectories, files in os.walk(directory_in):
        for filename in files:
            if filename.endswith(".png"):
                path_img.append(os.path.join(root, filename))

    print("Found " +str(len(path_img)) +" pictures")
    start_time = time.time()

    counterr=0
    allwork=0
    directory="./data/out/" + strftime("%Y-%m-%d_%H-%M-%S", gmtime()) + "/"
    if not os.path.exists(directory):
        os.makedirs(directory)
    with open(directory+"prompt.txt", 'w') as f:
        f.write(text)  
    shutil.copytree(directory_in, directory+"input")
    for i in path_img:
        print(i)
        if not os.path.exists(directory+str(counterr)):
            os.makedirs(directory+str(counterr))

        init_img = Image.open(i)
        init_img = init_img.resize((768, 512))
        init_image = preprocess(init_img)
        prompt = text

        for seed in range(endSeed-startSeed+1):
            generator = torch.Generator(device=device).manual_seed(startSeed+seed)
            guidance_scale = startScale
            while guidance_scale <= endScale:
                strenght = startStrength
                while strenght <= endStrength:
                    with autocast("cuda"):
                        image = pipe(prompt=prompt, init_image=init_image, strength=strenght, guidance_scale=guidance_scale, generator=generator)["sample"][0]
                    image.save(directory+str(counterr)+"/" + str(allwork) +".jpg")

                    exif_dict = piexif.load(directory+str(counterr)+"/" + str(allwork) +".jpg")
                    userCommentAsDict = {}
                    userCommentAsDict['Seed']=str(seed+startSeed)
                    userCommentAsDict['Strength']=str(strenght)
                    userCommentAsDict['Guidance_scale']=str(guidance_scale)
                    user_comment = piexif.helper.UserComment.dump(json.dumps(userCommentAsDict))
                    exif_dict["Exif"][piexif.ExifIFD.UserComment] = user_comment
                    exif_bytes = piexif.dump(exif_dict)
                    piexif.insert(exif_bytes, directory+str(counterr)+"/" + str(allwork) +".jpg")

                    allwork+=1

                    strenght+=deltaStrength
                guidance_scale+=deltaScale
                
        
        counterr+=1

    print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds")