from collections import namedtuple
import random
import unicodedata

import torch
from torch.nn.utils.rnn import pad_sequence

# prefix used by the wordpiece tokenizer to indicate that the token continues the previous word
WORDPIECE_PREFIX = '##'
# probability of masking a token during BERT training
P_MASK = 0.15
# probability of replacing a masked token with a random token
P_TOKEN_REPLACE = 0.1
# probability of using the original token instead of masking it
P_TOKEN_ORIGINAL = 0.1
# a range of reasonable token ids to use for replacement during model training
TOKEN_REPLACE_RANGE = (1000, 29999)
# attention mask value that tells the model to not ignore the token
NOT_MASKED = 1
# token type for a single input sequence
TOKEN_TYPE_A = 0
# padding value used for attention_mask
MASK_PAD = 0
# padding value used for token_type_ids
TOKEN_TYPE_PAD = 1
# label to use for tokens to ignore when computing masked language modeling training loss
LABEL_IGNORE = -100
# maximum number of input tokens BERT supports
BERT_MAX_TOKENS = 512

# used to represent inputs to the BERT model
BertInput = namedtuple(
    typename='BertInput',
    field_names=['input_ids', 'attention_mask', 'token_type_ids', 'labels', 'masked_idxs'],
)

# all the configuration options
Config = namedtuple(
    'Config',
    [
        'doc_key',
        'summary_key',
        'summaries_key',
        'model_name',
        'measure',
        'gap',
        'min_token_length_normal',
        'min_token_length_lead',
        'min_token_length_followup',
        'device',
        'random_seed',
        'inference_batch_size',
        'inference_mask_evenly',
        'filler_token',
        'help_sep',
        'finetune_batch_size',
        'finetune_epochs',
        'finetune_mask_evenly',
        'finetune_chunk_size',
        'finetune_chunk_stride',
        'learning_rate',
        'warmup_steps',
    ]
)

# the default configuration options that can reproduce the paper's results and don't require a GPU
Defaults = Config(
    doc_key='doc',
    summary_key='summary',
    summaries_key='summaries',
    model_name='bert-base-uncased',
    measure='relative',
    gap=6,
    min_token_length_normal=4,
    min_token_length_lead=2,
    min_token_length_followup=100,
    device='cpu',
    random_seed=1,
    inference_batch_size=1,
    inference_mask_evenly=True,
    filler_token='.',
    help_sep='',
    finetune_batch_size=1,
    finetune_epochs=10,
    finetune_chunk_size=64,
    finetune_chunk_stride=32,
    learning_rate=5e-5,
    finetune_mask_evenly=False,
    warmup_steps=0,
)


def batch_data(data, batch_size):
    """Given a list, batch that list into chunks of size batch_size

    Args:
        data (List): list to be batched
        batch_size (int): size of each batch

    Returns:
        batches (List[List]): a list of lists, each inner list of size batch_size except possibly
            the last one.
    """
    batches = [
        data[i:i + batch_size]
        for i in range(0, len(data), batch_size)
    ]
    return batches


def is_token_large_enough(token, next_token, min_token_lengths):
    """Determine if a token is large enough according to min_token_lengths

    Args:
        token (str): a wordpiece token
        next_token (str): the next wordpiece token in the sequence
        min_token_lengths (Tuple[int, int, int]): minimum token lengths for normal tokens, lead
            tokens, and followup tokens

    Returns:
        large_enough (bool): whether or not the token is large enough
    """
    min_normal, min_lead, min_followup = min_token_lengths
    token_size = len(token)

    if token.startswith(WORDPIECE_PREFIX):
        token_size -= len(WORDPIECE_PREFIX)
        return token_size >= min_followup
    elif next_token.startswith(WORDPIECE_PREFIX):
        return token_size >= min_lead
    else:
        return token_size >= min_normal


def mask_tokens_evenly(tokens, gap, min_token_lengths, mask_token):
    """Produce several maskings for the given tokens where each masking is created by masking every
    "gap" tokens, as long as the token is large enough according to min_token_lengths.

    Args:
        tokens (List[str]): a sequence of wordpiece tokens
        gap (int): the spacing in-between masked tokens
        min_token_lengths (Tuple[int, int, int]): minimum token lengths for normal tokens, lead
            tokens, and followup tokens
        mask_token (str): wordpiece token to use for masking

    Returns:
        masked_inputs (List[List[str]]): a list of token sequences, where each token sequence
            contains masked tokens separated by "gap" tokens.
        all_answers (List[Dict[int, str]]): a list of "answer" dicts, where each answer dict maps
            token indices corresponding to masked tokens back to their original token.
    """
    gap = min(gap, len(tokens))
    masked_inputs = []
    all_answers = []
    for modulus in range(gap):
        masked_input = []
        answers = {}
        for idx, token in enumerate(tokens):
            next_token = '' if idx + 1 == len(tokens) else tokens[idx + 1]
            large_enough = is_token_large_enough(token, next_token, min_token_lengths)

            if idx % gap == modulus and large_enough:
                masked_input.append(mask_token)
                answers[idx] = token
            else:
                masked_input.append(token)

        if len(answers) > 0:
            masked_inputs.append(masked_input)
            all_answers.append(answers)

    return masked_inputs, all_answers


def mask_tokens_randomly(tokens, min_token_lengths, mask_token):
    """Produce several maskings for the given tokens by randomly choosing tokens to mask

    Args:
        tokens (List[str]): a sequence of wordpiece tokens
        min_token_lengths (Tuple[int, int, int]): minimum token lengths for normal tokens, lead
            tokens, and followup tokens
        mask_token (str): wordpiece token to use for masking

    Returns:
        masked_inputs (List[List[str]]): a list of token sequences, where each token sequence
            contains masked tokens chosen randomly.
        all_answers (List[Dict[int, str]]): a list of "answer" dicts, where each answer dict maps
            token indices corresponding to masked tokens back to their original token.
    """
    n_mask = max(int(len(tokens) * P_MASK), 1)

    token_positions = []
    for idx, token in enumerate(tokens):
        next_token = '' if idx + 1 == len(tokens) else tokens[idx + 1]
        if is_token_large_enough(token, next_token, min_token_lengths):
            token_positions.append(idx)
    random.shuffle(token_positions)

    all_inputs, all_answers = [], []
    while len(token_positions) > 0:
        positions_to_mask = token_positions[:n_mask]
        token_positions = token_positions[n_mask:]

        inputs, answers = [], {}
        for idx, token in enumerate(tokens):
            if idx in positions_to_mask:
                inputs.append(mask_token)
                answers[idx] = token
            else:
                inputs.append(token)

        all_inputs.append(inputs)
        all_answers.append(answers)

    return all_inputs, all_answers


def stack_tensor(input_list, pad_value, device):
    """Given a batch of inputs, stack them into a single tensor on the given device, padding them
    at the back with pad_value to make sure they are all the same length.

    Args:
        input_list (List[List[int]]): a list of input sequences
        pad_value (int): the value to use for padding input sequences to make them the same length
        device (str): torch device (usually "cpu" or "cuda")

    Returns:
        stacked_tensor (torch.LongTensor): a tensor of dimensions (batch size) x (seq length)
    """
    tensor_list = [
        torch.LongTensor(inputs) for inputs in input_list
    ]
    stacked_tensor = pad_sequence(
        sequences=tensor_list,
        batch_first=True,
        padding_value=pad_value
    ).to(device)

    return stacked_tensor


def get_input_tensors(input_batch, device, tokenizer):
    """Given a list of BertInputs, return the relevant tensors that are fed into BERT.

    Args:
        input_batch (List[BertInput]): a batch of model inputs
        device (str): torch device (usually "cpu" or "cuda")
        tokenizer (BertTokenizer): the wordpiece tokenizer used for BERT

    Returns:
        input_ids (torch.LongTensor): ids corresponding to input tokens
        attention_mask (torch.LongTensor): tells BERT about parts of the input to ignore
        token_type_ids (torch.LongTensor): used to differentiate input segments
        labels (torch.LongTensor): contains the original token ids for tokens that were masked
    """
    input_ids_list = [inputs.input_ids for inputs in input_batch]
    attention_mask_list = [inputs.attention_mask for inputs in input_batch]
    token_type_ids_list = [inputs.token_type_ids for inputs in input_batch]
    labels_list = [inputs.labels for inputs in input_batch]

    id_pad, = tokenizer.convert_tokens_to_ids([tokenizer.pad_token])
    input_ids = stack_tensor(input_ids_list, pad_value=id_pad, device=device)
    attention_mask = stack_tensor(attention_mask_list, pad_value=MASK_PAD, device=device)
    token_type_ids = stack_tensor(token_type_ids_list, pad_value=TOKEN_TYPE_PAD, device=device)

    if labels_list[0] is not None:
        labels = stack_tensor(labels_list, pad_value=LABEL_IGNORE, device=device)
    else:
        labels = None

    return input_ids, attention_mask, token_type_ids, labels


def determine_correctness(outputs, answers):
    """Given dicts corresponding to predicted tokens and actual tokens at different indices, return
    a list of bools for whether or not those predictions were correct.

    Args:
        outputs (List[Dict[int, str]]): each list represents a different input masking, and each
            dict maps indices to model predictions
        answers (List[Dict[int, str]]): each list represents a different input masking, and each
            dict maps indices to original tokens

    Returns:
        correctness (List[bool]): a list of values that are True if the model made a correct
            prediction and False otherwise
    """
    correctness = []
    for output, answer in zip(outputs, answers):
        for idx, actual_token in answer.items():
            predicted_token = output[idx]
            correctness.append(predicted_token == actual_token)

    return correctness


def measure_relative(S):
    """Calculate the "measure-relative" score as defined in the paper

    Args:
        S (List[List[int]]): accuracy counts as defined in the paper

    Returns:
        score (float): measure-relative score
    """
    return (S[0][1] - S[1][0]) / (S[0][0] + S[1][1] + S[0][1] + S[1][0])


def measure_improve(S):
    """Calculate the "measure-improve" score as defined in the paper

    Args:
        S (List[List[int]]): accuracy counts as defined in the paper

    Returns:
        score (float): measure-improve score
    """
    return S[0][1] / (S[0][0] + S[1][1] + S[0][1])


def clean_text(text):
    """Return a cleaned version of the input text

    Args:
        text (str): dirty text

    Returns:
        text (str): cleaned text
    """
    text = unicodedata.normalize('NFKD', text)
    return text
