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

__all__ = ['evaluater']

# Cell
import openai, pandas as pd, numpy as np, datetime, json, time
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 = {}
    prompts = {}
    kwargs = {
        "engine":"davinci",
        "temperature":0,
        "max_tokens":10,
        "stop":"\n",
    }

    def addPrompt(self, label:str, definition:str):
        """
        set definition of label to definition
        """
        try:
            self.definitions[label]["definition"] = definition
        except:
            self.definitions[label]= {"definition":definition}

    def removePrompt(self, label:str):
        """
        set definition of label to definition)
        """
        del self.definitions[label]

    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=["data", "label"]

    def generateResponses(self, queries, prompt, maxQuerySize=1500, verbose=False):
        """
        run a prompt against queries.
        prompt should be formated as "str {} str" to be able to use format
        to stick each query in
        """
        results = []
        for query in queries:
            #r = self.query(prompt=prompt.format(query[:maxQuerySize]), verbose=verbose, **self.kwargs)
            if verbose:
                print(r)
            results.append(r)
        return results

    def train_test_split(self, trainSize=0, testSize=0):
        """
        generates train and test set
        train sets with openai api can be small; max of 10 fills up the prompt window generally
        """
        if testSize==0:
            print("give it a train and test")
            return -1

        working = self.examples.copy()

        train = pd.DataFrame()
        test = pd.DataFrame()

        for label in working.label.unique():
            newEntries = working[working.label==label]
            newEntries = newEntries.sample(min(len(newEntries), trainSize))
            train = pd.concat([train, newEntries])
            working.drop(newEntries.index, inplace=True)
        self.train = train
        self.test = test

    def optimizeSearchCutoff(self, df, term, trainSize=200, testSize=200):
        """
        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
        """



    def useTermPrompt(self, n=20, overwrite=True, inplace=True, fewShot = 0, labels=[], verbose=False):
        """
        for each label in labels, add to the
        """
        if len(labels) == 0:
            labels = [x for x in self.definitions]

        if n > 0:
            targets = self.test.sample(n)
        else:
            targets = self.test.copy()


        all_responses = {}
        for label in labels:
                prompt = self.definitions[label]["prompt"]
                print(prompt)
                responses = self.generateResponses(targets["data"].to_list(), prompt, verbose=verbose)

                with open("{}/{}".format(self.outdir, datetime.datetime.now().strftime("%Y%m%d%H%m%S") + ".json"), "w") as fh:
                    json.dump({"prompt":prompt, "data":{"queries":targets["data"].to_list(), "preds":responses}}, fh, indent=4)

                all_responses[label]=responses

                targets[label] = responses
                self.labeled = targets
        self.labeled = targets


