import argparse
from overrides import overrides
from typing import List

from sacrerouge.commands import MetricSetupSubcommand
from sacrerouge.data import MetricsDict
from sacrerouge.data.types import ReferenceType, SummaryType
from sacrerouge.metrics import Metric, ReferenceBasedMetric

try:
    from bert_score import score, create_idf_dict
except ImportError:
    BERTSCORE_INSTALLED = False

    @Metric.register('bertscore')
    class BertScore(ReferenceBasedMetric):
        def __init__(self, *args, **kwargs):
            pass

        def score_multi_all(self, *args, **kwargs):
            raise NotImplementedError('Please install the "bert_score" python library library to use BertScore')
else:
    BERTSCORE_INSTALLED = True

    @Metric.register('bertscore')
    class BertScore(ReferenceBasedMetric):
        def __init__(self,
                     model_type: str = None,
                     num_layers: int = None,
                     idf: bool = False,
                     nthreads: int = 4,
                     batch_size: int = 64,
                     lang: str = 'en',
                     verbose: bool = False) -> None:
            super().__init__()
            self.model_type = model_type
            self.num_layers = num_layers
            self.idf = idf
            self.nthreads = nthreads
            self.batch_size = batch_size
            self.lang = lang
            self.verbose = verbose

        def _flatten_summary(self, summary: SummaryType) -> str:
            if isinstance(summary, list):
                return ' '.join(summary)
            return summary

        def _get_unique_references(self, references_list: List[List[str]]) -> List[str]:
            unique_references = set()
            for references in references_list:
                for reference in references:
                    unique_references.add(reference)
            return list(unique_references)

        def _run(self,
                 summaries_list: List[List[SummaryType]],
                 references_list: List[List[SummaryType]]) -> List[List[MetricsDict]]:
            summaries_list = [[self._flatten_summary(summary) for summary in summaries] for summaries in summaries_list]
            references_list = [[self._flatten_summary(reference) for reference in references] for references in
                               references_list]

            # Create the candidate and reference lists for passing to the scoring function
            input_candidates = []
            input_references = []
            empty_inputs = set()
            for i, (summaries, references) in enumerate(zip(summaries_list, references_list)):
                for j, summary in enumerate(summaries):
                    if len(summary) == 0:
                        empty_inputs.add((i, j))
                    else:
                        input_candidates.append(summary)
                        input_references.append(references)

            # If we are using IDF weighting, then the default script will create
            # the IDF score based on all of the references, but we could potentially
            # have duplicate references, which would make "incorrect" counts. Instead,
            # create the idf dict based on unique references
            if self.idf:
                unique_references = self._get_unique_references(references_list)
                idf = create_idf_dict(unique_references, model_type=self.model_type, nthreads=self.nthreads,
                                      lang=self.lang)
            else:
                idf = False

            # Score the summaries
            precisions, recalls, f1s = score(input_candidates, input_references,
                                             model_type=self.model_type,
                                             num_layers=self.num_layers,
                                             idf=idf,
                                             nthreads=self.nthreads,
                                             batch_size=self.batch_size,
                                             lang=self.lang,
                                             verbose=self.verbose)

            # Remap the scores to the summaries
            index = 0
            metrics_lists = []
            for i, summaries in enumerate(summaries_list):
                metrics_lists.append([])
                for j, summary in enumerate(summaries):
                    if (i, j) in empty_inputs:
                        precision, recall, f1 = 0.0, 0.0, 0.0
                    else:
                        precision = precisions[index].item()
                        recall = recalls[index].item()
                        f1 = f1s[index].item()
                        index += 1

                    metrics_lists[-1].append(MetricsDict({
                        'bertscore': {
                            'precision': precision,
                            'recall': recall,
                            'f1': f1,
                        }
                    }))

            return metrics_lists

        def score_multi_all(self,
                            summaries_list: List[List[SummaryType]],
                            references_list: List[List[ReferenceType]]) -> List[List[MetricsDict]]:
            return self._run(summaries_list, references_list)


@MetricSetupSubcommand.register('bertscore')
class BertScoreSetupSubcommand(MetricSetupSubcommand):
    @overrides
    def add_subparser(self, parser: argparse._SubParsersAction):
        description = 'Setup the BERTScore metric'
        self.parser = parser.add_parser('bertscore', description=description, help=description)
        self.parser.set_defaults(subfunc=self.run)

    @overrides
    def run(self, args):
        try:
            import bert_score
            print('BertScore setup success')
        except ImportError:
            print('Please pip install "bert_score" to complete setup')