import argparse
import os
#os.environ["PL_GLOBAL_SEED"]="1"
import string
import json
#import cv2
from datetime import datetime
from tqdm import tqdm
from ldm.generate import Generate
from random import choice
import PIL
from PIL.Image import Resampling
import math
import pdb
import random
from scipy.special import expit

import itertools
#os.environ["PL_GLOBAL_SEED"]="1"

#def permutation(pop_size, chrom_length):
#    return itertools.permutation(pop_size, chrom_length)

_t2i = Generate()

def get_folder_name(prompt = ""):
    now = datetime.now()

    prompt_parts = prompt.split()
    first_word = prompt_parts[0] if len(prompt_parts) > 0 else ""

    h, m, s = (str(t).ljust(2, "0") for t in [now.hour, now.minute, now.second])

    return f"{first_word}-{now.year}-{now.month}-{now.day}-{h}{m}{s}"

def get_vid_path(prompt = ""):
    return os.path.join(".", "outputs", "vid-samples", get_folder_name(prompt))

def prompt2vid(**config):
    prompt = config["prompt"]
    n_frames = config["n_frames"]
    #initial_image = config["init_img"]
    fps = config["fps"] if "fps" in config else 30.0
    cfg_scale = config["cfg_scale"] if "cfg_scale" in config else 7.5
#    strength = config["strength"] if "strength" in config else 0.8
    zoom_speed = config["zoom_speed"] if "zoom_speed" in config else 2.0
    vid_path = get_vid_path(prompt)
    frames_path = os.path.join(vid_path, "frames")
    os.makedirs(frames_path, exist_ok=True)

    pixels_sizes = (
#        64,
#        128,

        # dont go under here becuase it just produces junk
#        256,
        512,
    )


    #stepsize = 85
    #stepsize = 5
    #stepsize = 35
    seed = 111


    fixed_seed = seed
    # many_iterations64
    many_iterations=1
    #many_iterations=2
    for stepsize in (25,):
        for pixel in pixels_sizes:
            #seed = seed + 1
#            video_writer = cv2.VideoWriter(os.path.join(vid_path, f"videox{pixel}.mp4"), 1, fps, (pixel, pixel))
            #logfile = open(os.path.join(vid_path, f"videox{pixel}.log"),"w")
            stepno = 0
            order = list(range(0,1132))
            #order = (1,)
            #random.shuffle(order)
            #print(order)
            for parameter in order:
                #state = random.getstate()
                #random.seed(parameter)
                #rslice = random.sample(order,n_frames) # we want the slice to have a different sequence
                #random.setstate(state)

                for i in tqdm(range(n_frames), desc=f"Creating Video {stepsize} "):
                    #print(stepsize)
                    if stepno > 100:
                        stepno = 1

                    vars = [ ]


                    #     wei =  abs(math.sin(weight * stepno ))
                    #     val = x *math.exp(stepno)
                    #     if val > 4294967295:
                    #         val = val % 4294967295
                    #     vars.append(( val , wei))

                    stepno =         stepno + 1
                    #stepsize =stepsize + 1 # how much do we increase
                    mutations = []


                    #@if i > 0:
                    mutations.append(parameter)


                    co = 0.2
                    step = (co * 2 )/n_frames
                    mutation = (-1 * co) + (i * step)

                    data = dict(
                        #                strength=strength,
                        seed=fixed_seed,
                        variation_amount=0,
                        width=pixel,
                        height=pixel,
                        cfg_scale=7,
                        #latent_channels=8,
                        #downsampling_factor=10,
                        #strength=0.85,
                        steps=stepsize,
                        iterations=many_iterations,
                        #with_variations=vars,
                        mutations=mutations,
                        mutation=mutation,
                    )
                    images = _t2i.prompt2image(prompt,**data)

                    for (j,image) in enumerate(images):
                        next_frame_filename = os.path.join(frames_path, f"p{parameter}s{i}c{j}s{stepsize}x{pixel}.png")
                        #print(j,image)
                        image[0].save(next_frame_filename)
                        #for x in range(4):# write the same frame again
                        #logfile.write(json.dumps() + "\n")
#                        video_writer.write(cv2.imread(next_frame_filename))
                        logfile2 = open(os.path.join(frames_path,  f"p{parameter}s{i}c{j}s{stepsize}x{pixel}.json"),"w")

                        logfile2.write(json.dumps(
                            dict(
                                seed=image[1],
                                #mutatations=image[2],
                                #data=image,
                                fn=next_frame_filename,
                                prompt=prompt,
                                data=data)                  )+ "\n")
                        logfile2.close()

                #logfile.close()

                #cv2.destroyAllWindows()
                #video_writer.release()

#pdb.set_trace()
prompt2vid(   prompt = "cube. block. wooden. spinning. rotating. bright red. low poly. 3d. letter a. alphabet. simple. focal length. ",
              n_frames=4          )