import numpy as np
import json, codecs
from sklearn.metrics import classification_report
from itertools import groupby
import torch
from torch.nn import functional as F


# Evaluation
def _f1(_p, _r):
    if _p == 0 or _r == 0:
        return 0
    return 2 * _p * _r / (_p + _r)


def _prf(gold, pred):
    if len(pred) == 0:
        p = 0.0
    else:
        p = len(gold & pred) / len(pred)
    r = len(gold & pred) / len(gold)
    return p, r, _f1(p, r)


def metrics(data, ids=None, log=True):
    if ids != None:
        data = [sample for sample in data if sample['annotation_id'] in ids]
    gold = [sample['label_target'] for sample in data]
    pred = [sample['predicted_target'] for sample in data]
    if log:
        print(classification_report(gold, pred, digits= 3))
    report = classification_report(gold, pred, output_dict=True)
    return_dict = dict()
    for k in report:
        if isinstance (report[k], dict):
            for k2 in report[k]:
                return_dict[k + '-' + k2] = report[k][k2]
        else:
            return_dict[k] = report[k]
    return return_dict


def split_list(lst, split_token='[CLS]'):
    return [[split_token] + list(group) for k, group in groupby(lst, lambda x: x == split_token) if not k]


def load_jsonl_sentselecting(src, rationale_tau=None):
    with codecs.open(src, 'r', encoding='utf-8') as f_in:
        data = []
        for line in f_in.readlines():
            sample = json.loads(line.strip())
            if 'all_logits' in sample:
                sample['all_logits'] = np.asarray(sample['all_logits'])
            else: sample['all_logits'] = None
            sample['logits'] = np.asarray(sample['logits'])
            sample['text'] = split_list(sample['text'])
            sample['text'] = [' '.join(sent).replace(' ##', '') for sent in sample['text']]
            if 'meta' in sample:
                sample['rationales'] = [sent for rationale in sample['meta']['rationale_idx'] for sent in rationale]
            else:
                sample['rationales'] = None
            
            if 'rationale_probabilities' in sample:
                if rationale_tau is not None:
                    # Recalculate confidences using the same tau as during training
                    tensor = torch.FloatTensor(sample['all_logits'].max(axis=1))
                    sample['rationale_probabilities'] = F.softmax(tensor / rationale_tau, dim=0).numpy()
                else:
                    sample['rationale_probabilities'] = np.asarray(sample['rationale_probabilities'])
                sample['selected_sentences'] = [sorted(sents) for sents in sample['selected_sentences']]
            else:
                sample['rationale_probabilities'] = None
                
            
            data.append(sample)
        return data


def eval_sentselecting(data, log=False):
    distinct_sentence_precisions = []
    distinct_sentence_recalls = []
    joint_sentence_precisions = []
    joint_sentence_recalls = []
    gold_targets = []
    pred_targets = []
    
    true_true = 0
    true_partially_true = 0
    correct_targets = 0
    
    # this is not used in single sentence
    is_single_sent = sum([len(sample['selected_sentences']) for sample in data]) == 0
    
    for sample in data:
        gold_rationales = [set(rat) for rat in sample['meta']['rationale_idx']]
        if is_single_sent:
            predicted_rationales = set([sample['rationale_probabilities'].argmax()])
        else:
            idx = sample['rationale_probabilities'].argmax()
            predicted_rationales = set(sample['selected_sentences'][idx])
            
        predicted_target = sample['predicted_target']
        true_target = sample['label_target']
        
        gold_targets.append(true_target)
        pred_targets.append(predicted_target)
        
        # joint
        joint_gold_rationales = set([s for rat in gold_rationales for s in rat])
        p, r, _ = _prf(joint_gold_rationales, predicted_rationales)
        joint_sentence_precisions.append(p)
        joint_sentence_recalls.append(r)

        # distinct
        best_p = -1
        best_r = -1
        best_f1 = -1
        for rat in gold_rationales:
            p, r, f1 = _prf(rat, predicted_rationales)
            if f1 > best_f1:
                best_f1 = f1
                best_r = r
                best_p = p
        distinct_sentence_precisions.append(best_p)
        distinct_sentence_recalls.append(best_r)
        
        # Right for right reasons
        is_correct_target = predicted_target == true_target
        if is_correct_target:
            correct_targets += 1
            tt = False
            tpt = False
            for rat in gold_rationales:
                if predicted_rationales == rat:
                    tt = True
                    tpt = True
                    break
                elif len(predicted_rationales & rat) > 0:
                    tpt = True
            if tt:
                true_true += 1
            if tpt:
                true_partially_true += 1
            
    macro_sent_joint_r = sum(joint_sentence_recalls) / len(joint_sentence_recalls)
    macro_sent_joint_p = sum(joint_sentence_precisions) / len(joint_sentence_precisions)
    macro_sent_joint_f1 = _f1(macro_sent_joint_p, macro_sent_joint_r)
    if log:
        print('Sentence (Joint)')
        print('Macro Recall', macro_sent_joint_r)
        print('Macro Precision', macro_sent_joint_p)
        print('Macro F1', macro_sent_joint_f1)
        print()
        
    macro_sent_distinct_r = sum(distinct_sentence_recalls) / len(distinct_sentence_recalls)
    macro_sent_distinct_p = sum(distinct_sentence_precisions) / len(distinct_sentence_precisions)
    macro_sent_distinct_f1 = _f1(macro_sent_distinct_p, macro_sent_distinct_r)
    
    if log:
        print('Sentence (Distinct)')
        print('Macro Recall', macro_sent_distinct_r)
        print('Macro Precision', macro_sent_distinct_p)
        print('Macro F1', macro_sent_distinct_f1)
        print()
    
        # Right for right
        print('Right for Right Reasons')
        print('Plain acc', correct_targets / len(data))
        print('Exactly right', true_true / len(data))
        print('Partially right', true_partially_true / len(data))
        print()
    
        # Target:
        print('Target')
        print(classification_report(gold_targets, pred_targets, digits=3))
    classification_dict = classification_report(gold_targets, pred_targets, output_dict=True)
    
    results = {
        'rationale_p': macro_sent_distinct_p,
        'rationale_r': macro_sent_distinct_r,
        'rationale_f1': macro_sent_distinct_f1,
        'acc': correct_targets / len(data),
        'acc_em': true_true / len(data),
        'acc_pm': true_partially_true / len(data),
        'f1a': classification_dict['macro avg']['f1-score']
    }
    return results


def eval_sentselecting(data, log=False):
    distinct_sentence_precisions = []
    distinct_sentence_recalls = []
    joint_sentence_precisions = []
    joint_sentence_recalls = []
    gold_targets = []
    pred_targets = []
    
    true_true = 0
    true_partially_true = 0
    correct_targets = 0
    
    # this is not used in single sentence
    is_single_sent = sum([len(sample['selected_sentences']) for sample in data]) == 0
    
    for sample in data:
        gold_rationales = [set(rat) for rat in sample['meta']['rationale_idx']]
        if is_single_sent:
            predicted_rationales = set([sample['rationale_probabilities'].argmax()])
        else:
            idx = sample['rationale_probabilities'].argmax()
            predicted_rationales = set(sample['selected_sentences'][idx])
            
        predicted_target = sample['predicted_target']
        true_target = sample['label_target']
        
        gold_targets.append(true_target)
        pred_targets.append(predicted_target)
        
        # joint
        joint_gold_rationales = set([s for rat in gold_rationales for s in rat])
        p, r, _ = _prf(joint_gold_rationales, predicted_rationales)
        joint_sentence_precisions.append(p)
        joint_sentence_recalls.append(r)
        
        
        # distinct
        best_p = -1
        best_r = -1
        best_f1 = -1
        for rat in gold_rationales:
            p, r, f1 = _prf(rat, predicted_rationales)
            if f1 > best_f1:
                best_f1 = f1
                best_r = r
                best_p = p
        distinct_sentence_precisions.append(best_p)
        distinct_sentence_recalls.append(best_r)
        
        # Right for right reasons
        is_correct_target = predicted_target == true_target
        if is_correct_target:
            correct_targets += 1
            tt = False
            tpt = False
            for rat in gold_rationales:
                if predicted_rationales == rat:
                    tt = True
                    tpt = True
                    break
                elif len(predicted_rationales & rat) > 0:
                    tpt = True
            if tt:
                true_true += 1
            if tpt:
                true_partially_true += 1
            
    macro_sent_joint_r = sum(joint_sentence_recalls) / len(joint_sentence_recalls)
    macro_sent_joint_p = sum(joint_sentence_precisions) / len(joint_sentence_precisions)
    macro_sent_joint_f1 = _f1(macro_sent_joint_p, macro_sent_joint_r)
    if log:
        print('Sentence (Joint)')
        print('Macro Recall', macro_sent_joint_r)
        print('Macro Precision', macro_sent_joint_p)
        print('Macro F1', macro_sent_joint_f1)
        print()
        
    macro_sent_distinct_r = sum(distinct_sentence_recalls) / len(distinct_sentence_recalls)
    macro_sent_distinct_p = sum(distinct_sentence_precisions) / len(distinct_sentence_precisions)
    macro_sent_distinct_f1 = _f1(macro_sent_distinct_p, macro_sent_distinct_r)
    
    if log:
        print('Sentence (Distinct)')
        print('Macro Recall', macro_sent_distinct_r)
        print('Macro Precision', macro_sent_distinct_p)
        print('Macro F1', macro_sent_distinct_f1)
        print()
    
        # Right for right
        print('Right for Right Reasons')
        print('Plain acc', correct_targets / len(data))
        print('Exactly right', true_true / len(data))
        print('Partially right', true_partially_true / len(data))
        print()
    
        # Target:
        print('Target')
        print(classification_report(gold_targets, pred_targets, digits=3))
    classification_dict = classification_report(gold_targets, pred_targets, output_dict=True)
    
    results = {
        'rationale_p': macro_sent_distinct_p,
        'rationale_r': macro_sent_distinct_r,
        'rationale_f1': macro_sent_distinct_f1,
        'acc': correct_targets / len(data),
        'acc_em': true_true / len(data),
        'acc_pm': true_partially_true / len(data),
        'f1a': classification_dict['macro avg']['f1-score']
    }
    return results

import numpy as np
import json, codecs
from sklearn.metrics import classification_report
from itertools import groupby
import torch
from torch.nn import functional as F


# Evaluation
def _f1(_p, _r):
    if _p == 0 or _r == 0:
        return 0
    return 2 * _p * _r / (_p + _r)


def _prf(gold, pred):
    if len(pred) == 0:
        p = 0.0
    else:
        p = len(gold & pred) / len(pred)
    r = len(gold & pred) / len(gold)
    return p, r, _f1(p, r)


def metrics(data, ids=None, log=True):
    if ids != None:
        data = [sample for sample in data if sample['annotation_id'] in ids]
    gold = [sample['label_target'] for sample in data]
    pred = [sample['predicted_target'] for sample in data]
    if log:
        print(classification_report(gold, pred, digits= 3))
    report = classification_report(gold, pred, output_dict=True)
    return_dict = dict()
    for k in report:
        if isinstance (report[k], dict):
            for k2 in report[k]:
                return_dict[k + '-' + k2] = report[k][k2]
        else:
            return_dict[k] = report[k]
    return return_dict


def split_list(lst, split_token='[CLS]'):
    return [[split_token] + list(group) for k, group in groupby(lst, lambda x: x == split_token) if not k]


def load_jsonl_sentselecting(src, rationale_tau=None):
    with codecs.open(src, 'r', encoding='utf-8') as f_in:
        data = []
        for line in f_in.readlines():
            sample = json.loads(line.strip())
            if 'all_logits' in sample:
                sample['all_logits'] = np.asarray(sample['all_logits'])
            else: sample['all_logits'] = None
            sample['logits'] = np.asarray(sample['logits'])
            sample['text'] = split_list(sample['text'])
            sample['text'] = [' '.join(sent).replace(' ##', '') for sent in sample['text']]
            if 'meta' in sample:
                sample['rationales'] = [sent for rationale in sample['meta']['rationale_idx'] for sent in rationale]
            else:
                sample['rationales'] = None
            
            if 'rationale_probabilities' in sample:
                if rationale_tau is not None:
                    # Recalculate confidences using the same tau as during training
                    tensor = torch.FloatTensor(sample['all_logits'].max(axis=1))
                    sample['rationale_probabilities'] = F.softmax(tensor / rationale_tau, dim=0).numpy()
                else:
                    sample['rationale_probabilities'] = np.asarray(sample['rationale_probabilities'])
                sample['selected_sentences'] = [sorted(sents) for sents in sample['selected_sentences']]
            else:
                sample['rationale_probabilities'] = None
                
            
            data.append(sample)
        return data


def eval_sentselecting(data, log=False):
    distinct_sentence_precisions = []
    distinct_sentence_recalls = []
    joint_sentence_precisions = []
    joint_sentence_recalls = []
    gold_targets = []
    pred_targets = []
    
    true_true = 0
    true_partially_true = 0
    correct_targets = 0
    
    # this is not used in single sentence
    is_single_sent = sum([len(sample['selected_sentences']) for sample in data]) == 0
    
    for sample in data:
        gold_rationales = [set(rat) for rat in sample['meta']['rationale_idx']]
        if is_single_sent:
            predicted_rationales = set([sample['rationale_probabilities'].argmax()])
        else:
            idx = sample['rationale_probabilities'].argmax()
            predicted_rationales = set(sample['selected_sentences'][idx])
            
        predicted_target = sample['predicted_target']
        true_target = sample['label_target']
        
        gold_targets.append(true_target)
        pred_targets.append(predicted_target)
        
        # joint
        joint_gold_rationales = set([s for rat in gold_rationales for s in rat])
        p, r, _ = _prf(joint_gold_rationales, predicted_rationales)
        joint_sentence_precisions.append(p)
        joint_sentence_recalls.append(r)

        # distinct
        best_p = -1
        best_r = -1
        best_f1 = -1
        for rat in gold_rationales:
            p, r, f1 = _prf(rat, predicted_rationales)
            if f1 > best_f1:
                best_f1 = f1
                best_r = r
                best_p = p
        distinct_sentence_precisions.append(best_p)
        distinct_sentence_recalls.append(best_r)
        
        # Right for right reasons
        is_correct_target = predicted_target == true_target
        if is_correct_target:
            correct_targets += 1
            tt = False
            tpt = False
            for rat in gold_rationales:
                if predicted_rationales == rat:
                    tt = True
                    tpt = True
                    break
                elif len(predicted_rationales & rat) > 0:
                    tpt = True
            if tt:
                true_true += 1
            if tpt:
                true_partially_true += 1
            
    macro_sent_joint_r = sum(joint_sentence_recalls) / len(joint_sentence_recalls)
    macro_sent_joint_p = sum(joint_sentence_precisions) / len(joint_sentence_precisions)
    macro_sent_joint_f1 = _f1(macro_sent_joint_p, macro_sent_joint_r)
    if log:
        print('Sentence (Joint)')
        print('Macro Recall', macro_sent_joint_r)
        print('Macro Precision', macro_sent_joint_p)
        print('Macro F1', macro_sent_joint_f1)
        print()
        
    macro_sent_distinct_r = sum(distinct_sentence_recalls) / len(distinct_sentence_recalls)
    macro_sent_distinct_p = sum(distinct_sentence_precisions) / len(distinct_sentence_precisions)
    macro_sent_distinct_f1 = _f1(macro_sent_distinct_p, macro_sent_distinct_r)
    
    if log:
        print('Sentence (Distinct)')
        print('Macro Recall', macro_sent_distinct_r)
        print('Macro Precision', macro_sent_distinct_p)
        print('Macro F1', macro_sent_distinct_f1)
        print()
    
        # Right for right
        print('Right for Right Reasons')
        print('Plain acc', correct_targets / len(data))
        print('Exactly right', true_true / len(data))
        print('Partially right', true_partially_true / len(data))
        print()
    
        # Target:
        print('Target')
        print(classification_report(gold_targets, pred_targets, digits=3))
    classification_dict = classification_report(gold_targets, pred_targets, output_dict=True)
    
    results = {
        'rationale_p': macro_sent_distinct_p,
        'rationale_r': macro_sent_distinct_r,
        'rationale_f1': macro_sent_distinct_f1,
        'acc': correct_targets / len(data),
        'acc_em': true_true / len(data),
        'acc_pm': true_partially_true / len(data),
        'f1a': classification_dict['macro avg']['f1-score']
    }
    return results


def eval_sentselecting(data, log=False):
    distinct_sentence_precisions = []
    distinct_sentence_recalls = []
    joint_sentence_precisions = []
    joint_sentence_recalls = []
    gold_targets = []
    pred_targets = []
    
    true_true = 0
    true_partially_true = 0
    correct_targets = 0
    
    # this is not used in single sentence
    is_single_sent = sum([len(sample['selected_sentences']) for sample in data]) == 0
    
    for sample in data:
        gold_rationales = [set(rat) for rat in sample['meta']['rationale_idx']]
        if is_single_sent:
            predicted_rationales = set([sample['rationale_probabilities'].argmax()])
        else:
            idx = sample['rationale_probabilities'].argmax()
            predicted_rationales = set(sample['selected_sentences'][idx])
            
        predicted_target = sample['predicted_target']
        true_target = sample['label_target']
        
        gold_targets.append(true_target)
        pred_targets.append(predicted_target)
        
        # joint
        joint_gold_rationales = set([s for rat in gold_rationales for s in rat])
        p, r, _ = _prf(joint_gold_rationales, predicted_rationales)
        joint_sentence_precisions.append(p)
        joint_sentence_recalls.append(r)
        
        
        # distinct
        best_p = -1
        best_r = -1
        best_f1 = -1
        for rat in gold_rationales:
            p, r, f1 = _prf(rat, predicted_rationales)
            if f1 > best_f1:
                best_f1 = f1
                best_r = r
                best_p = p
        distinct_sentence_precisions.append(best_p)
        distinct_sentence_recalls.append(best_r)
        
        # Right for right reasons
        is_correct_target = predicted_target == true_target
        if is_correct_target:
            correct_targets += 1
            tt = False
            tpt = False
            for rat in gold_rationales:
                if predicted_rationales == rat:
                    tt = True
                    tpt = True
                    break
                elif len(predicted_rationales & rat) > 0:
                    tpt = True
            if tt:
                true_true += 1
            if tpt:
                true_partially_true += 1
            
    macro_sent_joint_r = sum(joint_sentence_recalls) / len(joint_sentence_recalls)
    macro_sent_joint_p = sum(joint_sentence_precisions) / len(joint_sentence_precisions)
    macro_sent_joint_f1 = _f1(macro_sent_joint_p, macro_sent_joint_r)
    if log:
        print('Sentence (Joint)')
        print('Macro Recall', macro_sent_joint_r)
        print('Macro Precision', macro_sent_joint_p)
        print('Macro F1', macro_sent_joint_f1)
        print()
        
    macro_sent_distinct_r = sum(distinct_sentence_recalls) / len(distinct_sentence_recalls)
    macro_sent_distinct_p = sum(distinct_sentence_precisions) / len(distinct_sentence_precisions)
    macro_sent_distinct_f1 = _f1(macro_sent_distinct_p, macro_sent_distinct_r)
    
    if log:
        print('Sentence (Distinct)')
        print('Macro Recall', macro_sent_distinct_r)
        print('Macro Precision', macro_sent_distinct_p)
        print('Macro F1', macro_sent_distinct_f1)
        print()
    
        # Right for right
        print('Right for Right Reasons')
        print('Plain acc', correct_targets / len(data))
        print('Exactly right', true_true / len(data))
        print('Partially right', true_partially_true / len(data))
        print()
    
        # Target:
        print('Target')
        print(classification_report(gold_targets, pred_targets, digits=3))
    classification_dict = classification_report(gold_targets, pred_targets, output_dict=True)
    
    results = {
        'rationale_p': macro_sent_distinct_p,
        'rationale_r': macro_sent_distinct_r,
        'rationale_f1': macro_sent_distinct_f1,
        'acc': correct_targets / len(data),
        'acc_em': true_true / len(data),
        'acc_pm': true_partially_true / len(data),
        'f1a': classification_dict['macro avg']['f1-score']
    }
    return results

# Evaluation ##', '') for sent in sample['text']] # Recalculate confidences using the same tau as during training # this is not used in single sentence # joint # distinct # Right for right reasons # Right for right # Target: # this is not used in single sentence # joint # distinct # Right for right reasons # Right for right # Target: