import os
os.environ["OPENAI_API_KEY"] = os.getenv('api_key')
import openai
import gradio as gr
from PIL import Image
import base64 
import requests
import httpx

def get_completion(prompt, model="gpt-3.5-turbo"):
    messages = [{"role": "user", "content": prompt}]
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0, # this i the degree of randomness of the model's output
    )
    return response.choices[0].message["content"]
def topic_diagram(topic):
  prompt_sum = f"""
  Imagine you are Charles Darwin, an expert in generating diagrams that explain complex concepts. Your goal is to create a diagram that visually represents a specific concept related to {topic}. The diagram should effectively communicate the relationships, processes, or hierarchies involved in the concept.
  Generate code for the diagram in mermaid syntax surrounded by ```.
  Output:
  Mermaid code: ```code```
"""

  response_sum = get_completion(prompt_sum)
  return response_sum
import re

def extract_text_between_backticks(text):
    pattern = r"```(.+?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    return matches

def generate_diagram(mermaid_code):
    graphbytes = mermaid_code.encode("ascii")
    base64_bytes = base64.b64encode(graphbytes)
    base64_string = base64_bytes.decode("ascii")
    return "https://mermaid.ink/img/" + base64_string


def visualize_diagram(input_text):
    llm_output = topic_diagram(input_text)
    mermaid_code = extract_text_between_backticks(llm_output)
    print (mermaid_code)
    url = generate_diagram(mermaid_code[0])
    print ("Url",url)
    img = Image.open(requests.get(url, headers={'content-type': 'image/png'},stream=True).raw)
    print ("Image",img)
    return img

input_text = gr.inputs.Textbox(label="Enter your question")
output_image = gr.outputs.Image(type="pil", label="Diagram")

gr.Interface(
    fn=visualize_diagram,
    inputs=input_text,
    outputs=output_image,
    title="Darwin's Diagram Visualizer",
    description="Enter your question and see the generated diagram.",
    allow_flagging=False
).launch()