""" Run rejection framework."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import os
import ast
import math
import numpy as np
import pandas as pd

import ndcgprime
import conformal
import tensorflow as tf

from multiprocessing import  Pool
from functools import partial,reduce

import warnings
warnings.filterwarnings('ignore')

flags = tf.flags
FLAGS = flags.FLAGS

flags.DEFINE_string(
    "data_dir", None,
    "The input data dir. Should contain the .txt files"
)

flags.DEFINE_float(
    "annotation_score_threshold", 2,
    "A threshold to decide whether a review is relevant or not."
)

flags.DEFINE_integer(
    "k", 10,
    "Top k reviews"
)


def parallelize(data, func, num_of_processes=8):
    data_split = np.array_split(data, num_of_processes)
    pool = Pool(num_of_processes)
    data = pd.concat(pool.map(func, data_split))
    pool.close()
    pool.join()
    return data

def run_on_subset(func, data_subset):
    return data_subset.apply(func, axis=1)

def parallelize_on_rows(data, func, num_of_processes=8):
    return parallelize(data, partial(run_on_subset, func), num_of_processes)


def baseline(train,test):
    '''The baseline is to tune a threshold directly on predicted scores (or probabilities).'''
    
    models = ["FLTR","BertQA"]
    train_copy = train.copy()
    for model in models:
        best_threshold = 0
        highest_NDCG = 0
        
        conf = [item[1] for sublist in train[model+'_score'] for item in sublist]
        epsilons = [np.percentile(conf,i) for i in [5+i*5 for i in range(0,19)]]
        
        for e in epsilons:
            cp = train_copy.copy()
            cp[model+'_pre'] = cp[model+'_score'].apply(lambda x: [item[1] if item[1]>=e else 0 for item in x])

            cp_a = cp[cp['answerable']==1]
            cp_na = cp[cp['answerable']==0]

            ndcg_a = ndcgprime.NDCGPrime_PQA_beta(cp_a,by=model+'_pre',k=FLAGS.k)
            ndcg_na = ndcgprime.NDCGPrime_PQA_beta(cp_na,by=model+'_pre',k=FLAGS.k)
            
            # Using Geometric Mean to avoid overfitting either on 
            # answerable questions or unanswerble questions.
            ndcgp = math.sqrt(ndcg_a*ndcg_na)  
                
            if ndcgp > highest_NDCG:
                best_threshold = e
                highest_NDCG = ndcgp
      
        train[model+'_pre']= train[model+'_score'].apply(lambda x: [item[1] if item[1]>=best_threshold else 0 for item in x])
        test[model+'_pre']= test[model+'_score'].apply(lambda x: [item[1] if item[1]>=best_threshold else 0 for item in x])
    return train,test

def tunning_threshold_on_confidence(train,test,by,k):
    ''' Tunnig theshold on confidence scores generated by IMCP.'''
    train = train.reset_index(drop=True)
    epsilons = [round(0.05*i,2) for i in range(0,20)]
      
    best_theshold = 0
    highest_NDCGP = 0
               
    for e in epsilons:
        
        def leave_one_out_tunning(row):
            cp = train.copy()
            idx = row["index"]
            test_idx = cp['index'].isin([idx])
            train_cp = cp[~test_idx]
            test_cp = cp[test_idx]
            test_predictions = conformal.MICP(train_cp,test_cp,by,epsilon=e)
            return test_predictions
    
        train[by+'_conformal'] = train.apply(leave_one_out_tunning,axis=1)
        ndcg_a = ndcgprime.NDCGPrime_PQA_beta(train[train['answerable']==1],by = by+'_conformal',k=k)
        ndcg_na = ndcgprime.NDCGPrime_PQA_beta(train[train['answerable']==0],by = by+'_conformal',k=k)
        ndcgp = math.sqrt(ndcg_a*ndcg_na)
            
        if ndcgp > highest_NDCGP:
                highest_NDCGP = ndcgp
                best_theshold = e
      
    print('Best e: %.3f' % best_theshold)
    test = test.reset_index(drop=True)
    test_predictions = conformal.MICP(train,test,by,best_theshold)        
    return [list(test_predictions)]

def filter_by_FLTR(row):
    r = row['FLTR_score']
    r = [item for item  in r]
    indices = list(range(len(r)))
    indices.sort(key=lambda x: r[x],reverse=True)
    indices = indices[:10]
      
    BertQA_scores = row['BertQA_score']
    r = [value if index in indices else 0 for index,value in enumerate(BertQA_scores)]
    row['BertQA_score'] = r
    return row


def leave_one_out(row):
    k = FLAGS.k
    cp = data.copy()
    idx = row["index"]
    test_idx = cp['index'].isin([idx])
    train = cp[~test_idx]
    test = cp[test_idx]
    train,test = baseline(train,test)
    
    models = ['FLTR','BertQA']
    NDCGP = {}
    
    for model in models:     
        col_name = model+"_CP"+"_tunning"
        train_temp = train.copy()
        test_temp = test.copy()    
        test[col_name+'_predictions'] = tunning_threshold_on_confidence(train_temp,test_temp, by=model,k=k)
        
        
    predictions = [model+"_CP_tunning"+'_predictions' for model in models]+[
        model+'_original_pre' for model in models]+[
        model+'_pre' for model in models]
                
    for pre in predictions:
        if pre not in NDCGP:
            cp = test.copy()
            temp = [ndcgprime.NDCGPrime_PQA_beta(cp,by=pre,k=k)]
            NDCGP[pre] = temp
            
    for pre in predictions: 
        row[pre] = test[pre]
        
    row['NDCGP'] = NDCGP
    
    return row


def main(_):
    
    models = ['FLTR','BertQA']
    global data
    data_path = os.path.join(FLAGS.data_dir,'test_predictions.txt')
    data = pd.read_csv(data_path,sep='\t',encoding='utf-8',
                       converters={'annotation_score':ast.literal_eval,
                                   'reviews':ast.literal_eval,
                                   'FLTR_score':ast.literal_eval,
                                   'BertQA_score':ast.literal_eval,})
    data['ave_annot_score'] = data['annotation_score'].apply(lambda x: [np.mean(item) for item in x])
    data['label'] = data['ave_annot_score'].apply(lambda x: [1 if item>=FLAGS.annotation_score_threshold 
                                                             else 0 for item in x ])
    data['ave_annot_score'] = data['ave_annot_score'].apply(lambda x: [item if item>=FLAGS.annotation_score_threshold 
                                                                       else 0 for item in x ])
    data['answerable'] = data['ave_annot_score'].apply(lambda x: 1 if sum(x)>0 else 0)
    data = data.apply(filter_by_FLTR,axis=1)
    data['BertQA_score'] = data['BertQA_score'].apply(lambda x: [[1-item,item] for item in x])
    data['FLTR_score'] = data['FLTR_score'].apply(lambda x: [[1-item,item] for item in x])
    data['FLTR_original_pre'] = data['FLTR_score'].apply(lambda x: [item[1] for item in x])
    data['BertQA_original_pre'] = data['BertQA_score'].apply(lambda x: [item[1] for item in x])
    
    ''' multi-processing data'''
    df = parallelize_on_rows(data,leave_one_out)      
    
    
    A = df[df['answerable'] == 1]
    NA = df[df['answerable'] == 0]

    NDCGP_A = A['NDCGP'].tolist()
    NDCGP_NA = NA['NDCGP'].tolist()

    NDCGP_a = {}
    NDCGP_na = {}

    for item in NDCGP_A:
        for key in item.keys():
            if key in NDCGP_a:
                temp = NDCGP_a[key]
                NDCGP_a[key] = temp+item[key]
            else:
                NDCGP_a[key] = item[key]

    for item in NDCGP_NA:
        for key in item.keys():
            if key in NDCGP_na:
                temp = NDCGP_na[key]
                NDCGP_na[key] = temp+item[key]
            else:
                NDCGP_na[key] = item[key]
    
    k = FLAGS.k
    for key in NDCGP_a.keys():
        print(key, "Geometric Mean    NDCG`@",str(k),": %.3f" % math.sqrt(np.mean(NDCGP_a[key])*np.mean(NDCGP_na[key])))

    for key in NDCGP_a.keys():
        print(key, "Geometric Mean    A NDCG`@",str(k),": %.3f" % np.mean(NDCGP_a[key]))

    for key in NDCGP_na.keys():
        print(key, "Geometric Mean   NA NDCG`@",str(k),": %.3f" % np.mean(NDCGP_na[key]))
          
if __name__ == "__main__":
    flags.mark_flag_as_required("data_dir")
    tf.app.run()

