import boto3
from botocore.exceptions import ClientError
import pprint
import time
import streamlit as st
from sentence_transformers import CrossEncoder

#model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", max_length=512)
####### Add this Kendra Rescore ranking
#kendra_ranking = boto3.client("kendra-ranking",region_name = 'us-east-1')
#print("Create a rescore execution plan.")

# Provide a name for the rescore execution plan
#name = "MyRescoreExecutionPlan"
# Set your required additional capacity units
# Don't set capacity units if you don't require more than 1 unit given by default
#capacity_units = 2

# try:
#     rescore_execution_plan_response = kendra_ranking.create_rescore_execution_plan(
#         Name = name,
#         CapacityUnits = {"RescoreCapacityUnits":capacity_units}
#     )

#     pprint.pprint(rescore_execution_plan_response)

#     rescore_execution_plan_id = rescore_execution_plan_response["Id"]

#     print("Wait for Amazon Kendra to create the rescore execution plan.")

#     while True:
#         # Get the details of the rescore execution plan, such as the status
#         rescore_execution_plan_description = kendra_ranking.describe_rescore_execution_plan(
#             Id = rescore_execution_plan_id
#         )
#         # When status is not CREATING quit.
#         status = rescore_execution_plan_description["Status"]
#         print(" Creating rescore execution plan. Status: "+status)
#         time.sleep(60)
#         if status != "CREATING":
#             break

# except ClientError as e:
#         print("%s" % e)

# print("Program ends.")
#########################

@st.cache_resource
def re_rank(self_, rerank_type, search_type, question, answers):
   
    ans = []
    ids = []
    ques_ans = []
    query = question[0]['question']
    for i in answers[0]['answer']:
        if(self_ == "search"):
            
            ans.append({
                    "Id": i['id'],
                    "Body": i["desc"],
                    "OriginalScore": i['score'],
                    "Title":i["desc"]
                    })
            ids.append(i['id'])
            ques_ans.append((query,i["desc"]))
        
        else:
            ans.append({'text':i})
            
            ques_ans.append((query,i))
        
            

    re_ranked = [{}]
    ####### Add this Kendra Rescore ranking
    # if(rerank_type == 'Kendra Rescore'):
    #     rescore_response = kendra_ranking.rescore(
    #         RescoreExecutionPlanId = 'b2a4d4f3-98ff-4e17-8b69-4c61ed7d91eb',
    #         SearchQuery = query,
    #         Documents = ans
    #     )
    #     re_ranked[0]['answer']=[]
    #     for result in rescore_response["ResultItems"]:

    #         pos_ = ids.index(result['DocumentId'])

    #         re_ranked[0]['answer'].append(answers[0]['answer'][pos_])
    #     re_ranked[0]['search_type']=search_type,
    #     re_ranked[0]['id'] = len(question)
    #     return re_ranked
        
    # if(rerank_type == 'Cross Encoder'):

    #     scores = model.predict(
    #                 ques_ans
    #                     )
        
    #     index__ = 0
    #     for i in ans:
    #         i['new_score'] = scores[index__]
    #         index__ = index__+1

    #     ans_sorted = sorted(ans, key=lambda d: d['new_score'],reverse=True) 
        
        
    #     def retreive_only_text(item):
    #         return item['text']
            
    #     if(self_ == 'rag'):
    #         return list(map(retreive_only_text, ans_sorted)) 

       
    #     re_ranked[0]['answer']=[]
    #     for j in ans_sorted:
    #         pos_ = ids.index(j['Id'])
    #         re_ranked[0]['answer'].append(answers[0]['answer'][pos_])
    #     re_ranked[0]['search_type']= search_type,
    #     re_ranked[0]['id'] = len(question)
    #     return re_ranked