import pandas as pd
import time
from sagemaker.pytorch.model import PyTorchPredictor
from sagemaker.deserializers import JSONDeserializer
from sagemaker.serializers import JSONSerializer
from sagemaker.predictor_async import AsyncPredictor
from sagemaker.predictor import Predictor

def get_probs(endpoint_key, sentences, questions, append_data_back = None, threshold = 0.6):
    
    if type(questions) != list:
        raise TypeError("questions parameter must be of type list, with string values inside of that list!")
    
    list_of_probs_for_all_cats = []
    for question in questions:
    
        data = {"sentences" : sentences,
                     "questions": [question]}
        
        predictor = Predictor(endpoint_name=endpoint_key, serializer=JSONSerializer(), deserializer=JSONDeserializer())
        async_predictor = AsyncPredictor(predictor)
        
        
        result = async_predictor.predict_async(data, "s3://ask-t5-stage/async_inference/ip")
        
        list_of_probs_for_specific_cat = []
        
        while True:
            try:
                results = result.get_result()
                for result in results:
                    list_of_probs_for_specific_cat.append(result[1])
                
                list_of_probs_for_all_cats.append(list_of_probs_for_specific_cat)
                break
            except Exception as e:
                if 'Inference could still be running' in str(e):
                    time.sleep(5)
                else:
                    raise(str(e))
    
    transpose_list_of_probs_for_all_cats = [list(i) for i in zip(*list_of_probs_for_all_cats)]
    probs_df = pd.DataFrame( transpose_list_of_probs_for_all_cats)
    probs_df.columns = questions
    
    """nelow code up until the if isinstance line is responsible for creating a thrid column that has the 
    respective categories assigned to the respective sentences """
    label_results = []
    for row in transpose_list_of_probs_for_all_cats:
        row_result = []
        for i, element in enumerate(row):
            if element > threshold:
                row_result.append(questions[i])
        label_results.append(row_result)
    probs_df['labels'] = label_results        

    
    if isinstance(append_data_back, pd.DataFrame):
    

        append_data_back.reset_index(drop=True, inplace=True)
        probs_df.reset_index(drop=True, inplace=True)
        original_df_with_probs = pd.concat([append_data_back, probs_df], axis=1)
        
        return original_df_with_probs
    
    elif isinstance(append_data_back, list):
        df_with_sentences = pd.DataFrame(append_data_back, columns = ['sentences'])
        df_with_sentences = pd.concat([df_with_sentences, probs_df], axis=1)
        return df_with_sentences
    
    elif append_data_back == None:
        
        return probs_df
    
    
    else:
        raise TypeError("append_data_back parameter must be of type pandas dataframe or list!")
        