""" Holds scoring functions for coref and quotes """

import itertools
import pdb

from sklearn.metrics import cohen_kappa_score

from annotated_span import group_annotations, match_spans, match_annotated_spans, normalize_annotations_to_id, match_links


def score_attribution_labels(matched_labels, mismatched_labels, mismatched_extraction_labels, metric='cohen'):
    """ Scores annotation labels for lists of tuples:
            matched_labels: (span1, span2) for spans that match and have the same annotation
            mismatched_labels: (span1, span2) for spans that match but do not have the same annotation
            mismatched_extraction_labels: (span1, span2) for spans that match but one has a NULL label (wasn't extracted)

        Returns score_on_matched_extractions, score_on_all_extractions
    """

    if len(matched_labels) == 0:
        if len(mismatched_labels) == 0 and len(mismatched_extraction_labels) == 0: # nothing predicted
            return (1.0,1.0)
        return (0.0,0.0)

    # Build corresponding lists of attributions
    matched_labels1 = list(list(zip(*matched_labels))[0])
    matched_labels2 = list(list(zip(*matched_labels))[1])
    if len(mismatched_labels) == 0:
        mismatched_labels1 = []
        mismatched_labels2 = []
    else:
        mismatched_labels1 = list(list(zip(*mismatched_labels))[0])
        mismatched_labels2 = list(list(zip(*mismatched_labels))[1])
    if len(mismatched_extraction_labels) == 0:
        mismatched_ext_labels1 = []
        mismatched_ext_labels2 = []
    else:
        mismatched_ext_labels1 = list(list(zip(*mismatched_extraction_labels))[0])
        mismatched_ext_labels2 = list(list(zip(*mismatched_extraction_labels))[1])

    # Calculate attribution agreement on matched extractions
    ext_labels1 = matched_labels1 + mismatched_labels1
    ext_labels2 = matched_labels2 + mismatched_labels2
    
    # Normalize character names to IDs
    annotations1 = matched_labels1 + mismatched_labels1 + mismatched_ext_labels1
    annotations2 = matched_labels2 + mismatched_labels2 + mismatched_ext_labels2
    char2id = normalize_annotations_to_id(annotations1, annotations2)

    vals1 = [char2id[span.annotation] for span in ext_labels1]
    vals2 = [char2id[span.annotation] for span in ext_labels2]
    score_on_matched_extractions = cohen_kappa_score(vals1, vals2)

    # Calculate attribution agreement on all extractions
    labels1 = matched_labels1 + mismatched_labels1 + mismatched_ext_labels1
    labels2 = matched_labels2 + mismatched_labels2 + mismatched_ext_labels2
    vals1 = [char2id[span.annotation] for span in labels1]
    vals2 = [char2id[span.annotation] for span in labels2]
    score_on_all_extractions = cohen_kappa_score(vals1, vals2)

    return score_on_matched_extractions, score_on_all_extractions


#def span_attribution_confusion_matrix(baseline_spans, experimental_spans, gold_spans):
#        a = 0 # Both correct
#        b = 0 # Baseline correct, experiment incorrect
#        c = 0 # Baseline incorrect, experiment correct
#        d = 0 # Both incorrect
#        for baseline_quote, experimental_quote, gold_quote in zip(self.ordered_predictions['quotes']['baseline'], self.ordered_predictions['quotes']['experimental'], self.ordered_predictions['quotes']['gold']):
#            if utils.characters_match(baseline_quote.speaker, gold_quote.speaker) and utils.characters_match(experimental_quote.speaker, gold_quote.speaker):
#                a += 1
#            elif utils.characters_match(baseline_quote.speaker, gold_quote.speaker) and not utils.characters_match(experimental_quote.speaker, gold_quote.speaker):
#                b += 1
#            elif not utils.characters_match(baseline_quote.speaker, gold_quote.speaker) and utils.characters_match(experimental_quote.speaker, gold_quote.speaker):
#                c += 1
#            else:
#                d += 1
#                
#        table = [[a, b],
#                 [c, d]]


def links(mention_cluster):
    """ Returns a set of all the links in an entity between lists of AnnotatedSpans 
        Links are tuples (chap_id, start_token_id, end_token_id)
    """
    
    if len(mention_cluster) == 1: # self-link
        links = {(mention_cluster[0].get_location(), 
                mention_cluster[0].get_location())}

    else:
        links = {(mention1.get_location(), mention2.get_location()) for mention1, 
                    mention2 in itertools.combinations(mention_cluster, 2)}

    n_mentions = len(mention_cluster)
    if len(mention_cluster) != 1:
        if len(links) != (n_mentions * (n_mentions-1))/2:
            pdb.set_trace()
        
    return links


def lea(clusters, ref_clusters):
    """ Generic LEA calculcation between 2 sets of clusters..
        Input clusters are of form {'annotation': [AnnotatedSpan, ...]}
        Returns a numerator and denominator to be used for recall and precision.
    """

    cluster_resolutions = {}
    cluster_sizes = {}
    
    for ref_cluster_name, ref_mentions in ref_clusters.items():
        ref_links = links(ref_mentions)
        common_links = 0
        for cluster_name, mentions in clusters.items():
            cluster_links = links(mentions)
            overlap = match_links(cluster_links, ref_links)
            common_links += len(overlap)
            
        cluster_resolution = common_links/len(ref_links)
        cluster_resolutions[ref_cluster_name] = cluster_resolution
        cluster_sizes[ref_cluster_name] = len(ref_mentions)
        
    num = sum([cluster_sizes[c] * cluster_resolutions[c] for c in ref_clusters])
    den = sum(cluster_sizes.values())

    return num, den


def lea_recall(predicted_clusters, gold_clusters):
    """ Calculates LEA recall between predicted mention clusters and gold mention 
        clusters for a fic.
        Input clusters are of form {'annotation': [AnnotatedSpan, ...]}
    """
    if len(predicted_clusters) == 0:
        return 0
    num, den = lea(predicted_clusters, gold_clusters)
    return num/den


def lea_precision(predicted_clusters, gold_clusters):
    """ Calculates LEA precision between predicted mention clusters and gold mention 
        clusters for a fic.
        Input clusters are of form {'annotation': [AnnotatedSpan, ...]}
    """
    if len(predicted_clusters) == 0:
        return 0
    num, den = lea(gold_clusters, predicted_clusters)
    return num/den


#def lea_recall(predicted_clusters, gold_clusters):
#    """ Calculates LEA recall between predicted mention clusters and gold mention clusters.
#        Input clusters are of form {'annotation': [AnnotatedSpan, ...]}
#    """
#
#    cluster_resolutions = {}
#    cluster_sizes = {}
#    
#    for gold_cluster_name, gold_mentions in gold_clusters.items():
#        gold_links = links(gold_mentions)
#        cluster_resolution = 0
#        for predicted_cluster, predicted_mentions in predicted_clusters.items():
#            predicted_links = links(predicted_mentions)
#            cluster_resolution += len(predicted_links.intersection(gold_links))
#            
#        cluster_resolution = cluster_resolution/len(gold_links)
#        cluster_resolutions[gold_cluster_name] = cluster_resolution
#        cluster_sizes[gold_cluster_name] = len(gold_mentions)
#        
#    # take importance (size) of clusters into account
#    if sum(cluster_sizes.values()) == 0: # no predicted clusters
#        fic_recall = 0
#    else:
#        fic_recall = sum([cluster_sizes[c] * cluster_resolutions[c] for c in gold_clusters])/sum(cluster_sizes.values())
#
#    return fic_recall
#
#
#def lea_precision(predicted_clusters, gold_clusters):
#    
#    cluster_resolutions = {}
#    cluster_sizes = {}
#    
#    for predicted_cluster_name, predicted_mentions in predicted_clusters.items():
#        predicted_links = links(predicted_mentions)
#        cluster_resolution = 0
#        for gold_cluster, gold_mentions in gold_clusters.items():
#            gold_links = links(gold_mentions)
#            cluster_resolution += len(predicted_links.intersection(gold_links))
#        
#        cluster_resolution = cluster_resolution/len(predicted_links)
#        cluster_resolutions[predicted_cluster_name] = cluster_resolution
#        cluster_sizes[predicted_cluster_name] = len(predicted_mentions)
#        
#    # take importance (size) of clusters into account
#    if sum(cluster_sizes.values()) == 0: # no predicted clusters
#        fic_precision = 0
#    else:
#        fic_precision = sum([cluster_sizes[c] * cluster_resolutions[c] for c in predicted_clusters])/sum(cluster_sizes.values())
#        
#    return fic_precision


def lea_moosavi_recall(predicted_clusters, gold_clusters):
    """ Matching original LEA implementation at
        https://github.com/ns-moosavi/coval/blob/master/coval/eval/evaluator.py
    """
    #clusters = {m.get_location(): k for k,m in ref_clusters.items()}
    gold2pred = get_markable_assignments(gold_clusters, predicted_clusters)

    if len(predicted_clusters) == 0:
        return 0
    num, den = lea_moosavi(gold_clusters, predicted_clusters, gold2pred)
    return num/den


def lea_moosavi_precision(predicted_clusters, gold_clusters):
    """ Calculates LEA precision between predicted mention clusters and gold mention 
        clusters for a fic.
        Input clusters are of form {'annotation': [AnnotatedSpan, ...]}
    """
    pred2gold = get_markable_assignments(predicted_clusters, gold_clusters)
    if len(predicted_clusters) == 0:
        return 0
    num, den = lea_moosavi(predicted_clusters, gold_clusters, pred2gold)
    return num/den


def lea_moosavi(clusters, ref_clusters, mapping):
    """ Matching original LEA implementation at
        https://github.com/ns-moosavi/coval/blob/master/coval/eval/evaluator.py
    """
    num, den = 0, 0

    for cluster_name, c in clusters.items():
        if len(c) == 1: # singleton
            all_links = 1
            if c[0].get_location() in mapping and len(
                    ref_clusters[mapping[c[0].get_location()]]) == 1:
                common_links = 1
            else:
                common_links = 0
        else:
            common_links = 0
            links_common = set()
            all_links = len(c) * (len(c) - 1) / 2.0
            for i, m in enumerate(c):
                if m.get_location() in mapping:
                    for m2 in c[i + 1:]:
                        if m2.get_location() in mapping and mapping[
                            m.get_location()] == mapping[m2.get_location()]:
                            common_links += 1
                            links_common.add((m,m2))
        num += len(c) * common_links / float(all_links)
        den += len(c)
    
    return num, den


def get_markable_assignments(inp_clusters, out_clusters):
    """ Returns a mapping of mentions in inp_clusters to their cluster number
        in out_clusters
        From https://github.com/ns-moosavi/coval/blob/87071a6293dc2e786dcfe2ed78e9369c17e41b3b/coval/arrau/reader.py#L178 """
    markable_cluster_ids = {}
    out_dic = {}
    for cluster_name, cluster in out_clusters.items():
        for m in cluster:
            out_dic[m.get_location()] = cluster_name

    for cluster_name, cluster in inp_clusters.items():
        for im in cluster:
            for om in out_dic:
                if im.get_location() == om:
                    markable_cluster_ids[im.get_location()] = out_dic[om]
                    break

    return markable_cluster_ids


def f_score(precision, recall):
    if precision + recall == 0:
        return 0
    return 2 * (precision * recall)/(precision + recall)


def calculate_lea(predicted_mentions, gold_mentions):
    """ Calculate LEA coreference evaluation from gold and predicted AnnotatedSpans """

    predicted_clusters = group_annotations(predicted_mentions)
    gold_clusters = group_annotations(gold_mentions)

    #recall = lea_recall(predicted_clusters, gold_clusters)
    #precision = lea_precision(predicted_clusters, gold_clusters)
    recall = lea_moosavi_recall(predicted_clusters, gold_clusters)
    precision = lea_moosavi_precision(predicted_clusters, gold_clusters)
    f1 = f_score(precision, recall)

    return f1, precision, recall


def coref_scores(predicted_mentions, gold_mentions, exact_match=True):
    """ Returns coref scores """

    scores = {}
    scores['lea_f1'], scores['lea_precision'], scores['lea_recall'] = calculate_lea(
        predicted_mentions, gold_mentions)
    return scores


def quote_extraction_scores(predicted_quotes, gold_quotes, matched_pred_quotes, matched_gold_quotes, false_positives, false_negatives):
    extraction_precision = min(len(matched_pred_quotes)/len(predicted_quotes), 1)
    extraction_recall = len(matched_gold_quotes)/len(gold_quotes)
    extraction_f1 = f_score(extraction_precision, extraction_recall)

    return extraction_f1, extraction_precision, extraction_recall 


def quote_attribution_scores(predicted_quotes, gold_quotes, correct_attributions, incorrect_attributions):
    # Quote attribution accuracy on all predicted quotes.
    # If the predicted quote is not a real quote span, is not a match
    attribution_precision = len(correct_attributions)/len(predicted_quotes)
    attribution_recall = len(correct_attributions)/len(gold_quotes)
    attribution_f1 = f_score(attribution_precision, attribution_recall)

    return attribution_f1, attribution_precision, attribution_recall


def quote_scores(predicted_quotes, gold_quotes, exact_match=True):
    """ Returns quote extraction and attribution scores """
    scores = {}

    # Edge cases
    if len(predicted_quotes) == 0 or len(gold_quotes) == 0:
        matched_pred_quotes, matched_gold_quotes, false_positives, false_negatives, correct_attributions, incorrect_attributions = [], [], [], [], [], []
        if len(predicted_quotes) == 0:
            if len(gold_quotes) == 0:
                scores['extraction_precision'] = 1
                scores['attribution_precision'] = 1
            else:
                scores['extraction_precision'] = 0
                scores['attribution_precision'] = 0
        if len(gold_quotes) == 0:
            scores['extraction_recall'] = 1 # everyone gets perfect recall if there are no quotes
            scores['attribution_recall'] = 1 # everyone gets perfect recall if there are no quotes
        else:
            scores['extraction_recall'] = 0
            scores['attribution_recall'] = 0
        scores['extraction_f1'] = f_score(scores['extraction_precision'], scores['extraction_recall'])
        scores['attribution_f1'] = f_score(scores['attribution_precision'], scores['attribution_recall'])

    else:
        # Precision, recall of the quote extraction (the markables)
        matched_pred_quotes, matched_gold_quotes, false_positives, false_negatives = \
            match_spans(predicted_quotes, gold_quotes, exact=exact_match)
        scores['extraction_f1'], scores['extraction_precision'], scores['extraction_recall'] = quote_extraction_scores(predicted_quotes, gold_quotes, matched_pred_quotes, matched_gold_quotes, false_positives, false_negatives)

        # Quote attribution precision, recall, f1
        correct_attributions, incorrect_attributions = match_annotated_spans(
            matched_pred_quotes, matched_gold_quotes, matched=True, 
            incorrect_extractions=false_positives)
        (scores['attribution_f1'], 
            scores['attribution_precision'], 
            scores['attribution_recall']) = quote_attribution_scores(
            predicted_quotes, gold_quotes, correct_attributions,
            incorrect_attributions)

    groups = {}
    groups['predicted_quotes'] = predicted_quotes
    groups['gold_quotes'] = gold_quotes
    groups['matched_pred_quotes'] = matched_pred_quotes
    groups['matched_gold_quotes'] = matched_gold_quotes
    groups['false_positives'] = false_positives
    groups['false_negatives'] = false_negatives
    groups['correct_attributions'] = correct_attributions
    groups['incorrect_attributions'] = incorrect_attributions

    return scores, groups
