# AUTOGENERATED! DO NOT EDIT! File to edit: 02_evaluateFromSearch.ipynb (unless otherwise specified).

__all__ = ['evaluater']

# Cell
import openai, pandas as pd, numpy as np, datetime, json, time, os
from OpenAISurveyWrapper import wrapper

# Cell
class evaluater(wrapper.wrapper):
    """
    wrapper on the API to load in survey definitions
    a format for displaying

    self.myFormat: a function to format a prompt
    """

    labeled = pd.DataFrame()
    raw = pd.DataFrame()
    examples = pd.DataFrame()
    searchTerms = {}

    def addSearchTerm(self, label:str, searchTerm:str):
        """
        k, this is badly named, but I want to be able to evaluate how search does
        """
        try:
            self.searchTerms[label]["searchTerm"] = searchTerm
        except:
            self.searchTerms[label]= {"searchTerm":searchTerm}

    def removeSearchTerm(self, term:str):
        """
        delete search term
        """
        del self.searchTerms[term]

    def loadExamplesFromDataFrame(self, examples:dict):
        """
        takes a 2 column dataframe; first column should be "text", second "label"
        """
        if len(examples.columns)!=2:
            raise Exception("need 2 columns, first should be data second label")
        self.examples = examples
        self.examples.columns=["text", "label"]


    def printSearchTerms(self):
        """
        print off what I'm going to be searching
        """
        for term in self.searchTerms:
            print("{}: {}".format(term, self.searchTerms[term]["searchTerm"]))

    def optimizeSingleSearchCutoff(self, df:pd.DataFrame, term:str, outputDir:str, trainSize=200, testSize=200, sleepTime:float=2):
        """
        run a search against the df texts to get most similar
        the API returns the same similarity scale when the same term is queried multiple times against different docs

        saves the output by the name of the search query to the outputDir
        """

        if df.columns[0] != "text" or df.columns[1] !="label":
            print("pass a df with columns label and text")
            print(df.columns)
            return

        responses = []
        trgts = df[:trainSize + testSize].copy()
        startIdx = 0
        endIdx = 200
        while startIdx < len(trgts):
            #responses = responses + trgts[startIdx:endIdx].text.to_list()

            responses = responses + openai.Engine("davinci").search(documents=\
                        trgts[startIdx:endIdx].text.to_list(), query=term)["data"]
            startIdx += 200
            endIdx += 200
            time.sleep(sleepTime)

        trgts["seach_sim"] = [float(x["score"]) for x in responses]

        cutoff_f1 = pd.DataFrame()
        for i in range(0, 700):
            trgts["search_pos"] = trgts["seach_sim"].apply(lambda x: 1 if x > i else 0)
            cutoff_f1.at["f1", i] = f1_score(trgts.label[:trainSize], trgts.search_pos[:trainSize])

        cutoff = cutoff_f1.T.idxmax()[0]

        cutoff_f1_test = pd.DataFrame()
        for i in range(0, 700):
            trgts["search_pos"] = trgts["seach_sim"].apply(lambda x: 1 if x > i else 0)
            cutoff_f1_test.at["f1", i] = f1_score(trgts.label[trainSize:], trgts.search_pos[trainSize:])

        if outputDir != "":
            os.makedirs(outputDir, exist_ok=True)
            trgts.to_pickle("{}/{}.pkl".format(outputDir, term))

        return cutoff_f1_test[cutoff][0]

    def generateSearchScoresForLabel(self, outFileName="", trainSize=200, testSize=200):
        """
        autogeneratres the search scores
        """

        results = pd.DataFrame()
        for term in self.searchTerms:
            search_term = self.searchTerms[term]["searchTerm"]
            print(search_term)
            myResult = self.optimizeSingleSearchCutoff(self.examples, search_term, outFileName, trainSize=trainSize, testSize=testSize)
            print(myResult)
            results.at[search_term, "testF1"] = myResult

        self.results = results
        return results