#!/usr/bin/env python
# coding: utf-8

import os
import openai
import gradio as gr

import torch
from diffusers import StableDiffusionPipeline
from torch import autocast

from contextlib import nullcontext
#from PIL import Image
#from torchvision import transforms



openai.api_key = os.getenv('openaikey')
authtoken = os.getenv('authtoken')

device = "cuda" if torch.cuda.is_available() else "cpu"
context = autocast if device == "cuda" else nullcontext
dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained("stale2000/sd-dnditem", torch_dtype=dtype, use_auth_token=authtoken)  
pipe = pipe.to(device)

disable_safety = True

if disable_safety:
  def null_safety(images, **kwargs):
      return images, False
  pipe.safety_checker = null_safety

def create_files():
    directory = 'C:\\Users\\brcwa\\OneDrive\\Desktop\\destinyCaptures\\dnd\\fullcaptionsimple\\full'
    for filename2 in os.listdir(directory):
        if not filename2.endswith('txt'):
            continue
        f = os.path.join(directory, filename2)
        # checking if it is a file
        if os.path.isfile(f):
            text_file = open(f, "r")
            lines = text_file.read()
            print(lines.split(',')[1] + "," + lines.split(',')[1])

#create_files()

def createGPTPrompt(item_type, description):
    return item_type.split(",")[0].split(" ")[-1] + " of " + description

def convert_lines(lines):
    key_arr = []
    key_hash = {}
    for line in lines:
        key = line.split(",")[0]
        val = line.split(",")[1]
        key_arr.append(key)
        key_hash[key] = val

    return key_arr, key_hash

def predict(dropdown, style_dropdown, manual_gpt_replacement, manual_sd_prompt, n_samples, history=[]):

    # gpt3
    sd_input = ""
    gpt_input = ""

    description = style_dropdown
    if manual_sd_prompt != '':
        gpt_input = manual_gpt_replacement
    else:
        gpt_input = "Describe the mechanics of a 5th Edition DnD item called '" + createGPTPrompt(dropdown, description) + "' :"
    
    if manual_sd_prompt != '':
        sd_input = manual_sd_prompt
    else:
        sd_input = "dnditem, " + dropdown + ", " + style_hashmap[style_dropdown]  + ", circle inner background and white outerbackground"

    
    response = openai.Completion.create(
    model="text-davinci-003",
    prompt=gpt_input,
    temperature=0.9,
    max_tokens=150,
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0.6)
    
    # tokenize the new input sentence
    responseText = response["choices"][0]["text"]
    history.append((sd_input, responseText))


    #img generation
    scale = 5.5

    #with autocast("cuda"):
    #  images = pipe(n_samples*[prompt], guidance_scale=scale).images

    with context("cuda"):
        images = pipe(n_samples*[sd_input], guidance_scale=scale, num_inference_steps=40).images

    
    return history, history, images



#inputText = gr.Textbox(placeholder="input query")
manual_gpt_query = gr.Textbox(placeholder="Input any query here, to replace the gpt query builder entirely.")
manual_sd_prompt = gr.Textbox(placeholder="Input any query here, to replace the gpt query builder entirely.")



choiceArr = ["a pair of boots", "a cloak", "a pair of gloves", "a helmet", "a necklace", "a ring", "a robe", "a rod", "a shield", "a staff", "a sword", "a wand"]
dropdown = gr.Dropdown(label= "Item Type", choices=choiceArr, value="a pair of boots")



text_file = open("styles.txt", "r")
lines = text_file.read().split('\n')
dropdown_arr, style_hashmap = convert_lines(lines)


style_dropdown = gr.Dropdown(label= "Item Ability and Style", choices=dropdown_arr, value="elvenkind")

output_img = gr.Gallery(label="Generated image")
output_img.style(grid=2)

step_slide = gr.Slider(1, 4, value=2, step=1),
slide = gr.Slider(label="Number of Images Generated", minimum=1, maximum=4, value=2, step=1)
gr.Interface(fn=predict,
             inputs=[dropdown, style_dropdown, manual_gpt_query,manual_sd_prompt,slide,'state'],
            
             outputs=["chatbot",'state', output_img]).launch()