import gradio as gr
import random
from time import time, ctime

def promptgen2(choice, num, artist):
  t = time()
  print(ctime(t))

  if choice == "Prompt Generator v0.1(Better quality)":
    prompt = open('pr1.txt').read().splitlines()
  elif choice == "Prompt Generator v0.2(More tags)":
    prompt = open('pr2.txt').read().splitlines()

  if int(num) < 1 or int(num) > 20:
    num = 10
  
  if int(artist) < 0 or int(artist) > 40:
    artist = 2

  vocab = len(prompt)
  generated = []
  artists_num = 0
  while len(sorted(set(generated), key=lambda d: generated.index(d))) < num:
    rand = random.randint(0, vocab)
    if prompt[rand-1].startswith('art by') and artists_num < artist:
        artists_num +=1 
        generated.append(prompt[rand-1])
    elif not prompt[rand-1].startswith('art by'):
        generated.append(prompt[rand-1])
  print(', '.join(set(generated)) + '\n')
  return ', '.join(set(generated))

def promptgen(choice, num, artist):
  try:
      gen_prompt = open('gen_prompt.txt').read().splitlines()
  except:
      pass
  t = time()
  print(ctime(t))

  if choice == "Prompt Generator v0.1(Better quality)":
    prompt = open('pr1.txt').read().splitlines()
  elif choice == "Prompt Generator v0.2(More tags)":
    prompt = open('pr2.txt').read().splitlines()

  if int(num) < 1 or int(num) > 20:
    num = 10
  
  if int(artist) < 0 or int(artist) > 40:
    artist = 2

  vocab = len(prompt)
  generated = []
  artists_num = 0
  with open('gen_prompt.txt', 'w') as f:
      while len(sorted(set(generated), key=lambda d: generated.index(d))) < num:
        rand = random.randint(0, vocab)
        if prompt[rand-1].startswith('art by') and artists_num < artist:
            artists_num +=1 
            generated.append(prompt[rand-1])
        elif not prompt[rand-1].startswith('art by'):
            generated.append(prompt[rand-1])
      print(', '.join(set(generated)) + '\n')
      f.write(f"{str(ctime(t))} - {', '.join(set(generated))}")
      return ', '.join(set(generated))

demo = gr.Blocks()

with demo:
  gr.HTML(
        """
            <div style="text-align: center; margin: 0 auto;">
              <div style="display: inline-flex;align-items: center;gap: 0.8rem;font-size: 1.75rem;">
                <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
                  Simple Prompt Generator v0.4 (Gradio Demo)
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                Simple prompt generation script for Midjourney, DALLe, Stable and Disco diffusion and etc neural networks. <br> <p>More examples in <a class='link-info' href="https://github.com/WiNE-iNEFF/Simple_Prompt_Generator" target="_blank">Github</a> and <a class='link-info' href="https://wine-ineff.github.io/Simple_Prompt_Generator/" target="_blank">Project site</a></p>
              </p>
              <center>
                <img style="display: inline-block, margin-right: 1%;" src='https://visitor-badge.glitch.me/badge?page_id=WiNE-iNEFF_Simple_Prompt_Generator&left_text=GithubVisitors' alt='visitor badge'><img style="display: inline-block;" src='https://visitor-badge.glitch.me/badge?page_id=WiNE-iNEFF_HF_Simple_Prompt_Generator&left_text=HuggingFaceVisitors' alt='visitor badge'>
              </center>
            </div>
        """
    )
  with gr.Column():
    model_size = gr.Radio(["Prompt Generator v0.1(Better quality)", "Prompt Generator v0.2(More tags)"], label="Model Variant", value="Prompt Generator v0.1(Better quality)")
    number = gr.Number(value="10", label="Num of tag (MAX 20)", show_label=True)
    artist = gr.Number(value="2", label="Num of artist (Standart 2)", show_label=True)
    out = gr.Textbox(lines=4, label="Generated Prompts")
  greet_btn = gr.Button("Generate")
  greet_btn.click(fn=promptgen, inputs=[model_size, number, artist], outputs=out)
  gr.HTML(
            """
                <div class="footer">
                    <div style='text-align: center;'>Simple Prompt Generator by <a href='https://twitter.com/wine_ineff' target='_blank'>Artsem Holub (WiNE-iNEFF)</a><br>More information about this demo and script your can find in <a class='link-info' href="https://github.com/WiNE-iNEFF/Simple_Prompt_Generator" target="_blank">Github</a> and <a class='link-info' href="https://wine-ineff.github.io/Simple_Prompt_Generator/" target="_blank">Project site</a></div>
               </div>
           """
        )

demo.queue(concurrency_count=4)
demo.launch()