import json
import os
import shutil
import time
from time import strftime
from PIL import Image
import torch
from torch import autocast
import numpy as np
import piexif
import piexif.helper

from diffusers import (
    #DDIMScheduler,
    PNDMScheduler,
    StableDiffusionImg2ImgPipeline
)

t = torch.cuda.get_device_properties(0).total_memory
if t <= 6400000000:
    print("Running with less than 6gb memory. Working is not garanted")

file1 = open("prompt.txt", "r+")
text = file1.read()
print(text)

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

scheduler = PNDMScheduler.from_config(model_path, subfolder="scheduler")
#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
).to(device)
pipe.enable_attention_slicing()


def dummy_checker(images, **kwargs):
    return images, False


pipe.safety_checker = dummy_checker


def preprocess(image):
    w, h = image.size
    # resize to integer multiple of 32
    w, h = map(lambda x: x - x % 32, (w, h))
    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") + "/"
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")