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
import plotly.express as px

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')

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
    )
    
    print(resp)
    
    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)

    fig = px.treemap(df, path=['name', 'level1', 'level2'], 
                 values='count', color='level1')
    fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))
    fig.show()
    fig.write_image('temp1.png')
    img1 = Image.open("temp1.png")


    return img,img1

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)')
        with gr.Column():
            textbox2 = gr.Textbox(placeholder="Enter OpenAI API Key...", lines=1,label='Your API Key (Mandatory)')
    with gr.Row() as row:
        with gr.Column():
            btn = gr.Button("Generate")    
        with gr.Column():
            examples = gr.Examples(examples=['English Premier League','Heavy metal music','Face recognition','Arsenal Football Club',],
                           inputs=[textbox1])
    with gr.Row() as row:    
        with gr.Column():
            output_image1 = gr.components.Image(label="Your Mind Map as Graph")
        with gr.Column():
            output_image2 = gr.components.Image(label="Your Mind Map as Tree Map")


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

demo.launch()