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 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 %i pictures", len(path_img))
    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)
            strenght = startStrength
            while strenght <= endStrength:
                guidance_scale = startScale
                while guidance_scale <= endScale:
                    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
                    
                    guidance_scale+=deltaScale
                strenght+=deltaStrength
        
        counterr+=1

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