from typing import Any, Callable, Iterator, List, Mapping, Optional

import numpy as np
import sacrebleu
from tqdm import tqdm

from datalabs.operations.aggregate.aggregating import Aggregating, aggregating
from datalabs.operations.featurize import get_gender_bias
from datalabs.operations.operation import dataset_operation, DatasetOperation


class TextMatchingAggregating(Aggregating, DatasetOperation):
    def __init__(
        self,
        name: str = None,
        func: Callable[..., Any] = None,
        resources: Optional[Mapping[str, Any]] = None,
        contributor: str = None,
        processed_fields: List = ["text1", "text2"],
        generated_field: str = None,
        task="text-matching",
        description=None,
    ):
        super().__init__(
            name=name,
            func=func,
            resources=resources,
            contributor=contributor,
            task=task,
            description=description,
        )
        self._type = "TextMatchingAggregating"
        self.processed_fields = ["text1", "text2"]
        if isinstance(processed_fields, str):
            self.processed_fields[0] = processed_fields
        else:
            self.processed_fields = processed_fields
        self.generated_field = generated_field
        self._data_type = "Dataset"


class text_matching_aggregating(aggregating, dataset_operation):
    def __init__(
        self,
        name: Optional[str] = None,
        resources: Optional[Mapping[str, Any]] = None,
        contributor: str = None,
        processed_fields: List = ["text1", "text2"],
        generated_field: str = None,
        task="text-matching",
        description=None,
    ):
        super().__init__(
            name=name,
            resources=resources,
            contributor=contributor,
            description=description,
        )
        self.processed_fields = processed_fields
        self.generated_field = generated_field
        self.task = task

    def __call__(self, *param_arg):
        if callable(self.name):
            tf_class = TextMatchingAggregating(name=self.name.__name__, func=self.name)
            return tf_class(*param_arg)
        else:
            f = param_arg[0]
            name = self.name or f.__name__
            tf_cls = TextMatchingAggregating(
                name=name,
                func=f,
                resources=self.resources,
                contributor=self.contributor,
                processed_fields=self.processed_fields,
                generated_field=self.generated_field,
                task=self.task,
                description=self.description,
            )
            return tf_cls


def get_similarity_by_sacrebleu(text1, text2):
    # pip install sacrebleu
    references = [text1]
    hypothesis = text2
    score = sacrebleu.sentence_bleu(hypothesis, references).score

    return score


@text_matching_aggregating(
    name="get_statistics",
    contributor="datalab",
    task="text-matching, natural-language-inference",
    description="Calculate the overall statistics (e.g., average length) of a given "
    "text pair classification datasets. e,g. natural language inference",
)
def get_statistics(samples: Iterator):
    """
        Input:
        samples: [{
         "text1":
         "text2":
        }]
        Output:
            dict:

        usage:
        you can test it with following code:

    from datalabs import load_dataset
    from aggregate.text_matching import *
    dataset = load_dataset('sick')
    res = dataset['test'].apply(get_statistics)
    print(next(res))

    """
    # for hate speech
    # from hatesonar import Sonar
    # sonar = Sonar()

    sample_infos = []

    text1_lengths = []
    text2_lengths = []
    labels_to_number = {}
    vocab = {}
    number_of_tokens = 0
    gender_results = []
    # hatespeech = {
    #                  "hate_speech":{"ratio":0,"texts":[]},
    #                     "offensive_language":{"ratio":0,"texts":[]},
    #                     "neither":{"ratio":0,"texts":[]}}
    text1_divided_text2 = []
    similarities = []

    for sample in tqdm(samples):

        text1, text2, label = sample["text1"], sample["text2"], sample["label"]

        similarity_of_text_pair = get_similarity_by_sacrebleu(text1, text2)
        similarities.append(similarity_of_text_pair)

        # average length of text1
        text1_length = len(text1.split(" "))
        text1_lengths.append(text1_length)

        # average length of text2
        text2_length = len(text2.split(" "))
        text2_lengths.append(text2_length)

        # text1/text2
        text1_divided_text2.append(len(text1.split(" ")) / len(text2.split(" ")))

        # label info
        if label in labels_to_number.keys():
            labels_to_number[label] += 1
        else:
            labels_to_number[label] = 1

        # update the number of tokens
        number_of_tokens += len(text1.split())
        number_of_tokens += len(text2.split())

        # Vocabulary info
        for w in (text1 + text2).split(" "):

            if w in vocab.keys():
                vocab[w] += 1
            else:
                vocab[w] = 1

        # Gender info
        gender_result1 = get_gender_bias.func(text1)
        gender_result2 = get_gender_bias.func(text2)
        gender_results.append(gender_result1["gender_bias_info"])
        gender_results.append(gender_result2["gender_bias_info"])

        # hataspeech
        # results = sonar.ping(text=text1)
        # class_1 = results['top_class']
        # confidence = 0
        # for value in results['classes']:
        #     if value['class_name'] == class_1:
        #         confidence = value['confidence']
        #         break
        #
        # hatespeech[class_1]["ratio"] += 1
        # if class_1 != "neither":
        #     hatespeech[class_1]["texts"].append(text1)

        # results = sonar.ping(text=text2)
        # class_2 = results['top_class']
        # confidence = 0
        # for value in results['classes']:
        #     if value['class_name'] == class_2:
        #         confidence = value['confidence']
        #         break
        #
        # hatespeech[class_2]["ratio"] += 1
        # if class_2 != "neither":
        #     hatespeech[class_2]["texts"].append(text2)

        sample_info = {
            "text1": text1,
            "text2": text2,
            "label": label,
            "text1_length": text1_length,
            "text2_length": text2_length,
            "text1_gender": gender_result1,
            "text2_gender": gender_result2,
            # "text1_hate_speech_class":class_1,
            # "text2_hate_speech_class":class_2,
            "text1_divided_text2": len(text1.split(" ")) / len(text2.split(" ")),
            "similarity_of_text_pair": similarity_of_text_pair,
        }
        if len(sample_infos) < 10000:
            sample_infos.append(sample_info)
    # ------------------ Dataset-level ----------------
    # get vocabulary
    vocab_sorted = dict(sorted(vocab.items(), key=lambda item: item[1], reverse=True))

    # compute dataset-level gender_ratio
    gender_ratio = {
        "word": {"male": 0, "female": 0},
        "single_name": {"male": 0, "female": 0},
    }
    for result in gender_results:
        res_word = result["word"]  # noqa
        gender_ratio["word"]["male"] += result["word"]["male"]
        gender_ratio["word"]["female"] += result["word"]["female"]
        gender_ratio["single_name"]["male"] += result["single_name"]["male"]
        gender_ratio["single_name"]["female"] += result["single_name"]["female"]

    n_gender = gender_ratio["word"]["male"] + gender_ratio["word"]["female"]
    if n_gender != 0:
        gender_ratio["word"]["male"] /= n_gender
        gender_ratio["word"]["female"] /= n_gender
    else:
        gender_ratio["word"]["male"] = 0
        gender_ratio["word"]["female"] = 0

    n_gender = (
        gender_ratio["single_name"]["male"] + gender_ratio["single_name"]["female"]
    )
    if n_gender != 0:
        gender_ratio["single_name"]["male"] /= n_gender
        gender_ratio["single_name"]["female"] /= n_gender
    else:
        gender_ratio["single_name"]["male"] = 0
        gender_ratio["single_name"]["female"] = 0

    # get ratio of hate_speech:offensive_language:neither
    # for k,v in hatespeech.items():
    #     hatespeech[k]["ratio"] /= 2* len(samples)

    res = {
        "dataset-level": {
            "length_info": {
                "max_text1_length": np.max(text1_lengths),
                "min_text1_length": np.min(text1_lengths),
                "average_text1_length": np.average(text1_lengths),
                "max_text2_length": np.max(text2_lengths),
                "min_text2_length": np.min(text2_lengths),
                "average_text2_length": np.average(text2_lengths),
                "text1_divided_text2": np.average(text1_divided_text2),
            },
            "label_info": {
                "ratio": min(labels_to_number.values())
                * 1.0
                / max(labels_to_number.values()),
                "distribution": labels_to_number,
            },
            "vocabulary_info": vocab_sorted,
            "number_of_samples": len(samples),
            "number_of_tokens": number_of_tokens,
            "gender_info": gender_ratio,
            "average_similarity": np.average(similarities),
            # "hatespeech_info": hatespeech,
        },
        "sample-level": sample_infos,
    }

    return res
