import re
import string
import numpy as np
from typing import List, Tuple, Dict, Union
from tqdm import tqdm

from nltk.translate.bleu_score import corpus_bleu
from nltk.tokenize import sent_tokenize, word_tokenize


def decode_all(eval_preds, tokenizer, use_additional_metrics=False):
    # In this group of tasks we suppose that golden labels are reference sentences
    # and predictions are summaries / translations generated by model, in human-readable format
    predictions, labels, *inputs = eval_preds
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [label.strip() for label in decoded_labels]

    # For old versions of transformers, `inputs` are not passed
    decoded_texts = None
    if use_additional_metrics and len(inputs) > 0:
        input_ids = inputs[0]
        input_ids = np.where(input_ids != -100, input_ids, tokenizer.pad_token_id)
        decoded_texts = tokenizer.batch_decode(input_ids, skip_special_tokens=True)
        decoded_texts = [text.strip() for text in decoded_texts]

    return decoded_preds, decoded_labels, decoded_texts


def calculate_bleu(
    predictions: List[str], labels: List[str], aggregate: Union[bool, str] = True
) -> Dict[str, np.ndarray]:
    aggregation = _get_aggregation(aggregate)
    return {
        "bleu": aggregation(
            [pair_bleu(pred, label) for pred, label in tqdm(zip(predictions, labels))]
        )
    }


def split_into_words(text):
    return re.sub(
        "\s+",
        " ",
        text.translate(str.maketrans("", "", string.punctuation))
        .replace("\n", " ")
        .lower(),
    ).split()


def pair_bleu(references, prediction):
    """
    Compute the bleu score between two given texts.
    A smoothing function is used to avoid zero scores when
    there are no common higher order n-grams between the
    texts.
    """
    tok_ref = [word_tokenize(sent) for sent in sent_tokenize(references)]
    tok_pred = [word_tokenize(sent) for sent in sent_tokenize(prediction)]
    score = 0
    for c_cent in tok_pred:
        try:
            score += corpus_bleu(
                [tok_ref], [c_cent], smoothing_function=smoothing_function
            )
        except KeyError:
            score = 0.0
    try:
        score /= len(tok_pred)
    except ZeroDivisionError:
        score = 0.0

    return score


def smoothing_function(p_n, references, hypothesis, hyp_len):
    """
    Smooth-BLEU (BLEUS) as proposed in the paper:
    Chin-Yew Lin, Franz Josef Och. ORANGE: a method for evaluating automatic
    evaluation metrics for machine translation. COLING 2004.
    """
    smoothed_p_n = []
    for i, p_i in enumerate(p_n, start=1):
        # Smoothing is not applied for unigrams
        if i > 1:
            # If hypothesis length is lower than the current order, its value equals (0 + 1) / (0 + 1) = 0
            if hyp_len < i:
                assert p_i.denominator == 1
                smoothed_p_n.append(1)
            # Otherwise apply smoothing
            else:
                smoothed_p_i = (p_i.numerator + 1) / (p_i.denominator + 1)
                smoothed_p_n.append(smoothed_p_i)
        else:
            smoothed_p_n.append(p_i)
    return smoothed_p_n


def _get_aggregation(aggregate):
    if aggregate == True:
        return np.mean
    elif isinstance(aggregate, str):
        return getattr(np, aggregate)
    elif aggregate == False:
        return None
    return aggregate

def calculate_additional_hf_metrics(additional_metrics, decoded_labels, decoded_preds):
    result = {}
    for add_metric in additional_metrics:
        if add_metric.name != "sacrebleu":
            add_metrics_result = add_metric.compute(
                predictions=decoded_preds, references=decoded_labels
            )
        else:
            add_metrics_result = add_metric.compute(
                predictions=decoded_preds,
                references=[[lab] for lab in decoded_labels],
            )
            add_metrics_result["sacrebleu"] = add_metrics_result.pop("score")
        result.update(add_metrics_result)
    return result
