

from __future__ import division
from __future__ import absolute_import
#typing
from collections import Counter

#overrides
from sklearn.utils.linear_assignment_ import linear_assignment
import numpy as np
import torch

from allennlp.training.metrics.metric import Metric

class ConllCorefScores(Metric):
    def __init__(self)        :
        self.scorers = [Scorer(m) for m in (Scorer.muc, Scorer.b_cubed, Scorer.ceafe)]

    #overrides
    def __call__(self,  # type: ignore
                 top_spans              ,
                 antecedent_indices              ,
                 predicted_antecedents              ,
                 metadata_list                      ):
        u"""
        Parameters
        ----------
        top_spans : ``torch.Tensor``
            (start, end) indices for all spans kept after span pruning in the model.
            Expected shape: (batch_size, num_spans, 2)
        antecedent_indices : ``torch.Tensor``
            For each span, the indices of all allowed antecedents for that span.  This is
            independent of the batch dimension, as it's just based on order in the document.
            Expected shape: (num_spans, num_antecedents)
        predicted_antecedents: ``torch.Tensor``
            For each span, this contains the index (into antecedent_indices) of the most likely
            antecedent for that span.
            Expected shape: (batch_size, num_spans)
        metadata_list : ``List[Dict[str, Any]]``
            A metadata dictionary for each instance in the batch.  We use the "clusters" key from
            this dictionary, which has the annotated gold coreference clusters for that instance.
        """
        top_spans, antecedent_indices, predicted_antecedents = self.unwrap_to_tensors(top_spans,
                                                                                      antecedent_indices,
                                                                                      predicted_antecedents)
        for i, metadata in enumerate(metadata_list):
            gold_clusters, mention_to_gold = self.get_gold_clusters(metadata[u"clusters"])
            predicted_clusters, mention_to_predicted = self.get_predicted_clusters(top_spans[i],
                                                                                   antecedent_indices,
                                                                                   predicted_antecedents[i])
            for scorer in self.scorers:
                scorer.update(predicted_clusters, gold_clusters, mention_to_predicted, mention_to_gold)

    #overrides
    def get_metric(self, reset       = False)                              :
        metrics = (lambda e: e.get_precision(), lambda e: e.get_recall(), lambda e: e.get_f1())
        precision, recall, f1_score = tuple(sum(metric(e) for e in self.scorers) / len(self.scorers)
                                            for metric in metrics)
        if reset:
            self.reset()
        return precision, recall, f1_score

    #overrides
    def reset(self):
        self.scorers = [Scorer(metric) for metric in (Scorer.muc, Scorer.b_cubed, Scorer.ceafe)]

    @staticmethod
    def get_gold_clusters(gold_clusters):
        gold_clusters = [tuple(tuple(m) for m in gc) for gc in gold_clusters]
        mention_to_gold = {}
        for gold_cluster in gold_clusters:
            for mention in gold_cluster:
                mention_to_gold[mention] = gold_cluster
        return gold_clusters, mention_to_gold

    @staticmethod
    def get_predicted_clusters(top_spans              ,
                               antecedent_indices              ,
                               predicted_antecedents              ):                                            
                                                                                                  
                                                                                                               
        # Pytorch 0.4 introduced scalar tensors, so our calls to tuple() and such below don't
        # actually give ints unless we convert to numpy first.  So we do that here.
        top_spans = top_spans.numpy()  # (num_spans, 2)
        antecedent_indices = antecedent_indices.numpy()  # (num_spans, num_antecedents)
        predicted_antecedents = predicted_antecedents.numpy()  # (num_spans,)

        predicted_clusters_to_ids                             = {}
        clusters                              = []
        for i, predicted_antecedent in enumerate(predicted_antecedents):
            if predicted_antecedent < 0:
                continue

            # Find predicted index in the antecedent spans.
            predicted_index = antecedent_indices[i, predicted_antecedent]
            # Must be a previous span.
            assert i > predicted_index
            antecedent_span                  = tuple(top_spans[predicted_index])  # type: ignore

            # Check if we've seen the span before.
            if antecedent_span in list(predicted_clusters_to_ids.keys()):
                predicted_cluster_id      = predicted_clusters_to_ids[antecedent_span]
            else:
                # We start a new cluster.
                predicted_cluster_id = len(clusters)
                clusters.append([antecedent_span])
                predicted_clusters_to_ids[antecedent_span] = predicted_cluster_id

            mention                  = tuple(top_spans[i])  # type: ignore
            clusters[predicted_cluster_id].append(mention)
            predicted_clusters_to_ids[mention] = predicted_cluster_id

        # finalise the spans and clusters.
        final_clusters = [tuple(cluster) for cluster in clusters]
        # Return a mapping of each mention to the cluster containing it.
        mention_to_cluster                                                     = dict((
                mention, final_clusters[cluster_id])
                for mention, cluster_id in list(predicted_clusters_to_ids.items()))

        return final_clusters, mention_to_cluster


ConllCorefScores = Metric.register(u"conll_coref_scores")(ConllCorefScores)

class Scorer(object):
    u"""
    Mostly borrowed from <https://github.com/clarkkev/deep-coref/blob/master/evaluation.py>
    """
    def __init__(self, metric):
        self.precision_numerator = 0
        self.precision_denominator = 0
        self.recall_numerator = 0
        self.recall_denominator = 0
        self.metric = metric

    def update(self, predicted, gold, mention_to_predicted, mention_to_gold):
        if self.metric == self.ceafe:
            p_num, p_den, r_num, r_den = self.metric(predicted, gold)
        else:
            p_num, p_den = self.metric(predicted, mention_to_gold)
            r_num, r_den = self.metric(gold, mention_to_predicted)
        self.precision_numerator += p_num
        self.precision_denominator += p_den
        self.recall_numerator += r_num
        self.recall_denominator += r_den

    def get_f1(self):
        precision = 0 if self.precision_denominator == 0 else\
            self.precision_numerator / float(self.precision_denominator)
        recall = 0 if self.recall_denominator == 0 else\
            self.recall_numerator / float(self.recall_denominator)
        return 0 if precision + recall == 0 else 2 * precision * recall / (precision + recall)

    def get_recall(self):
        if self.recall_numerator == 0:
            return 0
        else:
            return self.recall_numerator / float(self.recall_denominator)

    def get_precision(self):
        if self.precision_numerator == 0:
            return 0
        else:
            return self.precision_numerator / float(self.precision_denominator)

    def get_prf(self):
        return self.get_precision(), self.get_recall(), self.get_f1()

    @staticmethod
    def b_cubed(clusters, mention_to_gold):
        u"""
        Averaged per-mention precision and recall.
        <https://pdfs.semanticscholar.org/cfe3/c24695f1c14b78a5b8e95bcbd1c666140fd1.pdf>
        """
        numerator, denominator = 0, 0
        for cluster in clusters:
            if len(cluster) == 1:
                continue
            gold_counts = Counter()
            correct = 0
            for mention in cluster:
                if mention in mention_to_gold:
                    gold_counts[tuple(mention_to_gold[mention])] += 1
            for cluster2, count in list(gold_counts.items()):
                if len(cluster2) != 1:
                    correct += count * count
            numerator += correct / float(len(cluster))
            denominator += len(cluster)
        return numerator, denominator

    @staticmethod
    def muc(clusters, mention_to_gold):
        u"""
        Counts the mentions in each predicted cluster which need to be re-allocated in
        order for each predicted cluster to be contained by the respective gold cluster.
        <http://aclweb.org/anthology/M/M95/M95-1005.pdf>
        """
        true_p, all_p = 0, 0
        for cluster in clusters:
            all_p += len(cluster) - 1
            true_p += len(cluster)
            linked = set()
            for mention in cluster:
                if mention in mention_to_gold:
                    linked.add(mention_to_gold[mention])
                else:
                    true_p -= 1
            true_p -= len(linked)
        return true_p, all_p

    @staticmethod
    def phi4(gold_clustering, predicted_clustering):
        u"""
        Subroutine for ceafe. Computes the mention F measure between gold and
        predicted mentions in a cluster.
        """
        return 2 * len([mention for mention in gold_clustering if mention in predicted_clustering])\
               / float(len(gold_clustering) + len(predicted_clustering))

    @staticmethod
    def ceafe(clusters, gold_clusters):
        u"""
        Computes the  Constrained EntityAlignment F-Measure (CEAF) for evaluating coreference.
        Gold and predicted mentions are aligned into clusterings which maximise a metric - in
        this case, the F measure between gold and predicted clusters.

        <https://www.semanticscholar.org/paper/On-Coreference-Resolution-Performance-Metrics-Luo/de133c1f22d0dfe12539e25dda70f28672459b99>
        """
        clusters = [cluster for cluster in clusters if len(cluster) != 1]
        scores = np.zeros((len(gold_clusters), len(clusters)))
        for i, gold_cluster in enumerate(gold_clusters):
            for j, cluster in enumerate(clusters):
                scores[i, j] = Scorer.phi4(gold_cluster, cluster)
        matching = linear_assignment(-scores)
        similarity = sum(scores[matching[:, 0], matching[:, 1]])
        return similarity, len(clusters), similarity, len(gold_clusters)
