import pandas as pd
from glob import glob
from scipy import spatial 
from collections import defaultdict

import tiktoken
import openai
import gradio as gr
from tenacity import retry, stop_after_attempt, wait_random_exponential



#df = pd.read_json('https://www.dropbox.com/scl/fi/uh964d1k6woc9wo3l2slc/dyf_w_embeddings.json?rlkey=j23j5338n4e88kvvsmj7s7aff&dl=1')
df = pd.read_json('rw7.json')

GPT_MODEL = 'gpt-3.5-turbo'
EMBEDDING_MODEL = "text-embedding-ada-002"
@retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(6))
def ask_naive(query):
    messages = [
        {"role": "system", "content": "You are a collegee sociology professor. Provide a very brief answer to this student question."},
        {"role": "user", "content": query},
    ]
    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo',
        messages=messages,
    )
    
    response_message = response["choices"][0]["message"]["content"]
    return response_message

# search function

# search function
def strings_ranked_by_relatedness(
    query: str,
    df: pd.DataFrame,
    relatedness_fn=lambda x, y: 1 - spatial.distance.cosine(x, y),
    top_n: int = 100
) -> tuple[list[str], list[float]]:
    """Returns a list of strings and relatednesses, sorted from most related to least."""
    query_embedding_response = openai.Embedding.create(
        model=EMBEDDING_MODEL,
        input=query,
    )
    query_embedding = query_embedding_response["data"][0]["embedding"]
    strings_and_relatednesses = [
        (row["text"], relatedness_fn(query_embedding, row["embedding"]))
        for i, row in df.iterrows()
    ]
    strings_and_relatednesses.sort(key=lambda x: x[1], reverse=True)
    strings, relatednesses = zip(*strings_and_relatednesses)
    return strings[:top_n]

def num_tokens(text: str) -> int:
    """Return the number of tokens in a string."""
    encoding = tiktoken.encoding_for_model('gpt-3.5-turbo')
    return len(encoding.encode(text))

def build_resources(psuedo_answer):
    related_book_selections = strings_ranked_by_relatedness(psuedo_answer, df, top_n=15)
    message  = 'Real World Sociology selections:\n'
    for selection in related_book_selections:
        if (
            num_tokens(message + selection)
            > 3000
        ):
            break
        else:
            message += '\n' + selection
    print(num_tokens(message))
    return message
    
@retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(6))
def respond(question, textbook_samples):
    messages = [
        {"role": "system", "content": "You are a college profesor who excels at explaining topics to students. Start with a direct answer to the question. Then, definition/overview of the concept's essence; break it down into understandable pieces; use clear language and structure. Where approriate, provide connections  and comparisions to  related terms. "},
        {"role": "user", "content": f"""Use markdown and emphasize important phrases in bold. Respond to the following question: {question}.
        
        When contructing the question, use the following information from the textbook.
        {textbook_samples}
        """ }
        ]
    
    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo',
        n=1,
        messages=messages)
    return response["choices"][0]["message"]["content"]
        
def ask(query):
    psuedo_answer = ask_naive(query)
    resources = build_resources(psuedo_answer)
    response = respond(query, resources)
    
    return response



intro_text = '''
# Ask the textbook

This app responds to your questions by looking up the most relevant selections from the textbook, and asking ChatGPT to respond based on the selections. It can take up to 30 seconds to respond.
'''

outro_text = '''
**Caveats:** Like all apps that employ large language models, this one has the possiblitiy for bias and confabulation. 

**Behind the Scenes**

This app uses sentence embeddings and a large language model to craft the response. Behind the scenes, it involves the following steps:

1. Each page from the textbook (or segment of the article if it's long) is converted into a fixed-length vector representation using OpenAI's text-embedding-ada-002 model. These representations are stored in a dataframe.
2. Your question is embedded using the same text-embedding-ada-002 model to convert it into a fixed-length vector representation.
3. To find the most relevant articles to your question, cosine similarity is calculated between the query vector and all the page vectors. The pages with the highest cosine similarity are retrieved as the top matches.
5. All of the relevant texts (from Step 3), along with the original search query, are passed to OpenAI's ChatGPT 3.5 model with specific instructions to answer the question using  the supplied texts.

'''



block = gr.Blocks(theme = 'bethecloud/storj_theme')

with block:
    gr.Markdown(intro_text)
        
    # Define the input and output blocks
    input_block = gr.Textbox(label='Question')
    research_btn = gr.Button(value="Ask the textbook")
    output_block = gr.Markdown(label="Response")
    research_btn.click(ask, inputs=input_block, outputs=output_block)
    gr.Examples(["What is the difference beween organic and mechnical solidarity?", 
                 ], inputs=[input_block])
    gr.Markdown(outro_text)



# Launch the interface
block.launch()