# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TODO: Add a description here."""

from collections import defaultdict
from typing import List, Dict, Tuple, NamedTuple

import datasets
import evaluate
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizer

from .prediction import Prediction


_CITATION = """\
@inproceedings{Hu:et-al:2020,
  author = {Hu, Jennifer and Gauthier, Jon and Qian, Peng and Wilcox, Ethan and Levy, Roger},
  title = {A systematic assessment of syntactic generalization in neural language models},
  booktitle = {Proceedings of the Association of Computational Linguistics},
  year = {2020}
}
"""

# TODO: Add description of the module here
_DESCRIPTION = """
"""


# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Runs SyntaxGym evaluations on the given model and test suite.
Args:
    suite (Dataset): SyntaxGym test suite loaded as a Dataset.
    model_id (str): model used for calculating surprisals
            NOTE: The SyntaxGym evaluations are only well-defined for causal language models.
                    This includes models such as gpt2, causal variations of bert,
                    causal versions of t5, and more (the full list can be found
                    in the AutoModelForCausalLM documentation here:
                    https://huggingface.co/docs/transformers/master/en/model_doc/auto#transformers.AutoModelForCausalLM )
Returns:
    prediction_results: A list of prediction results per item. A list of lists,
            one per item, containing the boolean prediction result for each
            prediction in the test suite,
    region_totals: A list of total surprisals for each region (nested within
            condition and item). A list of dictionaries (one per item), each
            mapping tuples (condition_name, region_number) to a float
            total surprisal value (i.e. negative log-2 probability).
Examples:
    TODO

    >>> my_new_module = evaluate.load("cpllab/syntaxgym")
    >>> ...
"""


SUITE_DATASET_CONDITION_SPEC = {
    "condition_name": datasets.Value("string"),
    "content": datasets.Value("string"),
    "regions": datasets.Sequence({
        "region_number": datasets.Value("int32"),
        "content": datasets.Value("string")
    })
}


SUITE_DATASET_SPEC = {
    "suite_name": datasets.Value("string"),
    "item_number": datasets.Value("int32"),
    "conditions": datasets.Sequence(SUITE_DATASET_CONDITION_SPEC),
    "predictions": datasets.Sequence(datasets.Value("string")),
}


class SyntaxGymMetricSuiteResult(NamedTuple):
    """
    Evaluation results for a single suite.
    """
    suite_name: str
    prediction_results: List[List[bool]]
    region_totals: List[Dict[Tuple[str, int], float]]

    @property
    def accuracy(self) -> float:
        return np.array(self.prediction_results).all(axis=1).mean(axis=0)


SyntaxGymMetricResult = Dict[str, SyntaxGymMetricSuiteResult]


def prepare_tokenizer(model, batch_size, add_start_token=True) -> Tuple[PreTrainedTokenizer, Dict]:
    """
    Load and prepare a tokenizer for SyntaxGym evaluation.

    Returns:
        tokenizer:
        tokenizer_kwargs: suggested kwargs for any tokenizer calls
    """
    tokenizer = AutoTokenizer.from_pretrained(model.name_or_path)

    # if batch_size > 1 (which generally leads to padding being required), and
    # if there is not an already assigned pad_token, assign an existing
    # special token to also be the padding token
    if tokenizer.pad_token is None and batch_size > 1:
        existing_special_tokens = list(tokenizer.special_tokens_map_extended.values())
        # check that the model already has at least one special token defined
        assert (
            len(existing_special_tokens) > 0
        ), "If batch_size > 1, model must have at least one special token to use for padding. Please use a different model or set batch_size=1."
        # assign one of the special tokens to also be the pad token
        tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})

    if add_start_token:
        # leave room for <BOS> token to be added:
        assert (
            tokenizer.bos_token is not None
        ), "Input model must already have a BOS token if using add_start_token=True. Please use a different model, or set add_start_token=False"
        max_tokenized_len = model.config.max_length - 1
    else:
        max_tokenized_len = model.config.max_length

    tokenizer_kwargs = {
        "add_special_tokens": False,
        "padding": True,
        "max_length": max_tokenized_len
    }
    return tokenizer, tokenizer_kwargs


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class SyntaxGym(evaluate.EvaluationModule):
    """
    Defines SyntaxGym evaluation logic for causal language models.
    """

    def _info(self):
        seq = datasets.Sequence
        features = datasets.Features({
            "dataset": SUITE_DATASET_SPEC
        })
        return evaluate.EvaluationModuleInfo(
            module_type="metric",
            description="TODO",
            citation=_CITATION,
            inputs_description="TODO",
            features=features,
            homepage="https://syntaxgym.org",
            codebase_urls=["https://github.com/cpllab/syntaxgym-core"],
        )

    def _compute(self, dataset, model_id, batch_size=8, add_start_token=False, device=None) -> SyntaxGymMetricResult:
        if device is not None:
            assert device in ["gpu", "cpu", "cuda"]
            if device == "gpu":
                device = "cuda"
        else:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        model = AutoModelForCausalLM.from_pretrained(model_id)
        model = model.to(device)
        model.eval()

        tokenizer, tokenizer_kwargs = prepare_tokenizer(model, batch_size, add_start_token)

        results = {}
        result_keys = ["prediction_results", "region_totals"]
        # TODO batch all items together
        for item in datasets.logging.tqdm(dataset):
            result_single = self._compute_single(item, tokenizer, tokenizer_kwargs,
                                                 model, device)

            suite_name = item["suite_name"]
            if suite_name not in results:
                results[suite_name] = SyntaxGymMetricSuiteResult(suite_name, [], [])
            for k in result_keys:
                getattr(results[suite_name], k).append(result_single[k])

        return results

    def _compute_single(self, item, tokenizer, tokenizer_kwargs, model, device):
        tokenized = tokenizer(item["conditions"]["content"],
                              return_tensors="pt",
                              return_offsets_mapping=True,
                              **tokenizer_kwargs).to(device)

        # input_ids: B * T
        input_ids = tokenized["input_ids"]
        assert input_ids.ndim == 2

        # Compute sentence level surprisals.
        with torch.no_grad():
            # Pre-softmax predictive distribution B * T * V
            logits = model(input_ids).logits
            surprisals = -logits.log_softmax(dim=2) / np.log(2)

        # surprisals: B * T * V
        assert surprisals.ndim == 3

        # Get surprisals of expected words.
        surps_shifted = surprisals[:, :-1, :]
        expected_ids = input_ids[:, 1:]

        # reindexed surprisals: B * (T - 1)
        surprisals = torch.gather(surps_shifted, 2, expected_ids.unsqueeze(2)) \
            .squeeze(2)

        #### aggregate
        condition_names = item["conditions"]["condition_name"]
        region_totals = {condition_name: defaultdict(float)
                         for condition_name in condition_names}
        region2tokens = self.compute_region_token_mapping(
            item, input_ids, tokenized["offset_mapping"])

        for i, (i_cond, i_inputs) in enumerate(zip(condition_names, input_ids)):
            for region_number, region_tokens in region2tokens[i_cond].items():
                for token in region_tokens:
                    if token == 0:
                        # surprisal not defined. pass.
                        continue
                    elif token <= surprisals.shape[1]:
                        region_totals[i_cond][region_number] += surprisals[i, token - 1]
                    else:
                        # TODO don't think this is an issue, just should clean
                        # up the aggregation output
                        assert token == surprisals.shape[1], \
                            "%s %s" % (token, surprisals.shape[1])

        region_totals = {(condition_name, region_number): float(total)
                         for condition_name, totals in region_totals.items()
                         for region_number, total in totals.items()}

        results = {
            "prediction_results": [
                Prediction(i, formula, "sum").formula(region_totals)
                for i, formula in enumerate(item["predictions"])
            ],

            "region_totals": region_totals
        }
        return results

    def get_region_edges(self, item, condition_idx):
        """
        Get left edge of each region as a character index.
        """
        # NB this is coupled with `condition_to_string` logic of course

        regions = item["conditions"]["regions"][condition_idx]

        idx = 0
        ret = []
        for r_idx, region_content in enumerate(regions["content"]):
            ret.append(idx)

            region_size = len(region_content)
            if region_content.strip() != "" and r_idx != 0 and not region_content.startswith(","):
                # Add joining space
                region_size += 1

            idx += region_size

        return ret

    def compute_region_token_mapping(self, item, input_ids: torch.LongTensor,
                                     offset_mapping: List[Tuple[int, int]]
                                     ) -> Dict[str, Dict[int, List[int]]]:
        # input_ids: B * T
        # offset_mapping: B * T * 2
        # assumes batch is sorted according to item's condition_name order

        condition_names = item["conditions"]["condition_name"]
        region2tokens = {cond: defaultdict(list) for cond in condition_names}

        max_long = torch.iinfo(torch.int64).max

        for i_cond, (i_tokens, i_offsets) in enumerate(zip(input_ids, offset_mapping)):
            region_edges = self.get_region_edges(item, i_cond)

            t_cursor, r_cursor = 0, 0
            while t_cursor < i_tokens.shape[0]:
                # token = i_tokens[t_cursor]
                token_char_start, token_char_end = i_offsets[t_cursor]

                if token_char_start == token_char_end == 0:
                    # This is a padding token. Skip.
                    # TODO what about BOS/EOS? some models incorporate them
                    t_cursor += 1
                    continue

                region_start = region_edges[r_cursor]
                region_end = region_edges[r_cursor + 1] \
                    if r_cursor + 1 < len(region_edges) else max_long

                # NB region boundaries are left edges, hence the >= here.
                if token_char_start >= region_end:
                    r_cursor += 1
                    continue

                region2tokens[condition_names[i_cond]][r_cursor + 1].append(t_cursor)
                t_cursor += 1

        return region2tokens