import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px

import chromadb
chroma_client = chromadb.Client()

collection = chroma_client.create_collection(name="emails")

df.loc[4, 'text']

for i in df.index:
    collection.add(
     
        documents = df.loc[i, 'text'],

      
        metadatas = [{"sender": df.loc[i, 'sender'],
                     "recipient1": df.loc[i, 'recipient1'],
                     "recipient2": df.loc[i, 'recipient2'],
                     "recipient3": df.loc[i, 'recipient3'],
                     "subject": df.loc[i, 'Subject'],
                     "folder": df.loc[i, 'folder'],
                     "date": str(df.loc[i, 'date'])
                     }],

     
        ids = str(i)
    )

collection.get(
    ids=["140"]
)

results = collection.query(
    query_texts = ["this is a document"],
    n_results = 2,
    include = ['distances', 'metadatas', 'documents']
)
results


from chromadb.utils import embedding_functions



sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="paraphrase-MiniLM-L3-v2")

collection_minilm = chroma_client.create_collection(name="emails_minilm", embedding_function=sentence_transformer_ef)


for i in df.index:
    print(i)
    collection_minilm.add(
       
        documents = df.loc[i, 'text'],

        metadatas = [{"sender": df.loc[i, 'sender'],
                     "recipient1": df.loc[i, 'recipient1'],
                     "recipient2": df.loc[i, 'recipient2'],
                     "recipient3": df.loc[i, 'recipient3'],
                     "subject": df.loc[i, 'Subject'],
                     "folder": df.loc[i, 'folder'],
                     "date": str(df.loc[i, 'date'])
                     }],

    
        ids = str(i)
    )

results = collection_minilm.query(
    query_texts = ["this is a document"],
    n_results = 2,
    include = ['distances', 'metadatas', 'documents']
)
results





import gradio as gr


def query_chromadb(question,numberOfResults):
    results = collection_minilm.query(
          n_results = numberOfResults,
      )

    return results['documents'][0]

iface = gr.Interface(
    fn=query_chromadb,
    inputs=["text","number"],
    outputs="text",
    title="Email Dataset Interface",
    description="Insert the question or the key word to find the topic correlated in the dataset"
)

iface.launch(share=True)



import ast

def create_output(dictionary, number):

    dictionary_ids = str(dictionary['ids'])

  
    dictionary_ids_clean = dictionary_ids.strip("[]")

    dictionary_ids_clean = dictionary_ids_clean.replace("'", "")

   
    dictionary_ids_list = dictionary_ids_clean.split(", ")

    string_results = "";


    for n in range(number):
      t = collection_minilm.get(
              ids=[dictionary_ids_list[n]]
          )


      id = str(t["ids"])
      doc = str(t["documents"])
      metadata = str(t["metadatas"])

      dictionary_metadata = ast.literal_eval(metadata.strip("[]"))

      string_results_old = string_results

      string_temp = """---------------
      SUBJECT: """ + dictionary_metadata['subject'] + """"
      MESSAGE: """ + "\n" + doc + """
      ---------------"""

      string_results = string_results_old + string_temp

    return string_results

def query_chromadb_advanced(question,numberOfResults):
    results = collection_minilm.query(
          query_texts = question,
          n_results = numberOfResults,
      )

    return create_output(results, numberOfResults)


result_advance = query_chromadb_advanced("bank", 4)

print(result_advance)

iface = gr.Interface(
    fn=query_chromadb_advanced,
    inputs=["text","number"],
    outputs="text",
    title="Email Dataset Interface",
    description="Insert the question or the key word to find the topic correlated in the dataset"
)

iface.launch(share=True)