from datasets import load_dataset
from openai import OpenAI  # or your preferred client
from pinecone import Pinecone, ServerlessSpec
import gradio as gr
import json
from tqdm.auto import tqdm

import ast
import os
import pandas as pd


#Setup API Keys
openai_api_key = os.getenv("OPEN_API_KEY")
openai_client = OpenAI(api_key=openai_api_key)
papi_key = os.getenv("PINECONE_KEY")

###Setup Pinecone and create context for the user's query. 

#Initialize Pinecone singleton object with API keys and set an index name.
pc = Pinecone(api_key=papi_key)
INDEX_NAME = 'summarizer'

#Delete any existing indices
if INDEX_NAME in [index.name for index in pc.list_indexes()]:
    pc.delete_index(INDEX_NAME)

#Create a new index for the dataset with shape of the embeddings, cosine similarity, and setup Serverless index with AWS. 
pc.create_index(name=INDEX_NAME, 
    dimension=1536, 
    metric='cosine',
    spec=ServerlessSpec(cloud='aws', region='us-east-1'))
index = pc.Index(INDEX_NAME)

max_articles_num = 500
df = pd.read_csv('./data/wiki.csv', nrows=max_articles_num)

prepped = []

for i, row in tqdm(df.iterrows(), total=df.shape[0]):
    meta = ast.literal_eval(row['metadata'])
    prepped.append({'id':row['id'], 
                    'values':ast.literal_eval(row['values']), 
                    'metadata':meta})
    #Upload data to Pinecone in batches. 
    if len(prepped) >= 250:
        index.upsert(prepped)
        prepped = []

#Use ada embedding model to create vector embeddings for input articles. 
def get_embeddings(articles, model="text-embedding-ada-002"):
   return openai_client.embeddings.create(input = articles, model=model)

def create_prompt(query):
    #Create embedding for input query. 
    embed = get_embeddings([query])
    
    #Search match in Pinecone. 
    res = index.query(vector=embed.data[0].embedding, top_k=3, include_metadata=True)
    
    #Extract Metadata and Text from the matches object returned by Pinecone. 
    contexts = [
        x['metadata']['text'] for x in res['matches']
    ]
    
    prompt_start = (
        "Answer the question based on the context below.\n\n"+
        "Context:\n"
    )
    
    prompt_end = (
        f"\n\nQuestion: {query}\nAnswer:"
    )
    
    #Create a Prompt for OpenAI and provide results of pinecone as part of the prompt. 
    prompt = (
        prompt_start + "\n\n---\n\n".join(contexts) + 
        prompt_end
    )
    return prompt

def summarize(prompt):
    #Use completions API with GPT 3.5 Turbo, and prepared prompt to call OpenAI. 
    res = openai_client.completions.create(
        model="gpt-3.5-turbo-instruct",
        prompt=prompt,
        temperature=0,
        max_tokens=636,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )

    return res.choices[0].text
    

def run_query(query):
    prompt=create_prompt(query)
    result=summarize(prompt)
    return result

# Create the Gradio interface
demo = gr.Interface(fn=run_query, inputs=gr.Textbox(label="User Input", placeholder="Type your question here..."), outputs=gr.Textbox(label="Matching Questions from Vector Database"))

# Launch the app
demo.launch(share=True)