import json
import requests
import gradio as gr
import random
import time
import os
import datetime
from datetime import datetime
from PIL import Image
from PIL import ImageOps
from PIL import Image, ImageDraw, ImageFont
import json
import io
from PIL import Image
import openai
import pandas as pd
from graphviz import Digraph

HRA_TOKEN=os.getenv("HRA_TOKEN")

"""
headers = {'Content-type': 'application/json', 'Accept': 'text/plain'}
url_hraprompts='https://us-central1-createinsightsproject.cloudfunctions.net/gethrahfprompts'

data={"prompt_type":'mindmap_prompt',"hra_token":HRA_TOKEN}
try:
    r = requests.post(url_hraprompts, data=json.dumps(data), headers=headers)
except requests.exceptions.ReadTimeout as e:
    print(e)
print(r.content)


prompt_text=str(r.content, 'UTF-8')
"""
prompt_text="Given a topic make a mindmap. The mindmap should cover various aspects not limited to concept, features, companies, types, applications, challenges etc. Output the mindmap in json format. Topic is :"

print(prompt_text)

def getmindmap(topic,openapikey):
    print('*******************')
    dateforfilesave=datetime.today().strftime("%d-%m-%Y %I:%M%p")
    print(topic)
    print(dateforfilesave)

    os.environ['OPENAI_API_KEY'] = str(openapikey)
    openai.api_key=str(openapikey)
    
    prompt=prompt_text+topic
    resp=openai.Completion.create(
      model="text-davinci-003",
      prompt=prompt,
      max_tokens=4000,
      temperature=0
    )
    
    
    df=pd.DataFrame(json.loads(resp['choices'][0]['text']))
    
    df['level1']=df['children'].apply(lambda x: x['name'])
    df['level1_tmp']=df['children'].apply(lambda x: x['children'])
    s = df.pop('level1_tmp').explode().to_frame()
    df = pd.merge(df.reset_index(),
        s.reset_index(),on='index'
            
    )
    df['level2']=df['level1_tmp'].apply(lambda x: x['name'])
    df['count']=[1]*len(df)

    dot = Digraph()

    dot.graph_attr['rankdir'] = 'LR' 
    for item in list(set(df['level1'].tolist())):
      dot.edge(str(list(set(df["name"].tolist()))[0]), str(item), label='')
    for item in list(set(df['level1'].tolist())):
      tempdf=df[df['level1']==item]
      for stuff in tempdf['level2'].tolist():
        dot.edge(str(item), str(stuff), label='',)

    r=requests.get('https://quickchart.io/graphviz?format=png&graph='+dot.source)
    dataBytesIO = io.BytesIO(r.content)
    img=Image.open(dataBytesIO)
    img.seek(0)
    name='temp.png'
    img.save(name)

    return img

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Mind Map Generator</center></h1>")
    gr.Markdown(
        """Enter a topic  and get a quick Mind Map. Use examples as a guide. \n\nNote: Error typically occurs with wrong OpenAI API key & sometimes with ChatGPT's inability to give a structured mindmap"""
        )
    with gr.Row() as row:
        with gr.Column():
            textbox1 = gr.Textbox(placeholder="Enter topic for Mind Map...", lines=1,label='Your topic (Mandatory)')
            textbox2 = gr.Textbox(placeholder="Enter OpenAI API Key...", lines=1,label='Your API Key (Mandatory)')
            btn = gr.Button("Generate")    
            examples = gr.Examples(examples=['English Premier League','Heavy metal music','Face recognition','Arsenal Football Club',],
                           inputs=[textbox1])
        with gr.Column():
            output_image1 = gr.components.Image(label="Your Mind Map")


    btn.click(getmindmap,inputs=[textbox1,textbox2], outputs=[output_image1])
    

demo.launch()