# Base Dependencies

# -----------------

import numpy as np

from typing import Dict, List

# PyTorch Dependencies

# ---------------------

import torch

from torch import Tensor

# Auxiliar Functions

# -------------------

def sort_batch(

    batch: Dict[str, List[List[float]]], lengths: List[List[float]]

) -> Dict[str, List[List[float]]]:

    """

    Sort a minibatch by the length of the sequences with the longest sequences first

    return the sorted batch targes and sequence lengths. This way the output can be used by pack_padded_sequences(...)

    Args:

        batch (Dict[str,  List[List[float]]]): batch of data

    Return:

        Dict[str,  List[List[float]]]: batch of data ordered in descending order of sequence length.

    """

    perm_idx = np.argsort(-lengths)

    for key in batch.keys():

        batch[key] = batch[key][perm_idx]

    return batch

def pad_seqs(

    seqs: List[List[float]], lengths: List[int], padding_idx: int

) -> List[List[float]]:

    """Pads sequences

    Args:

        seqs (List[List[float]]): sequences of different lengths

        lengths (List[int]): length of each sequence

        padding_idx (int): value used for padding

    Returns:

        List[List[float]]: padded sequences

    """

    batch_size = len(lengths)

    max_length = max(lengths)

    padded_seqs = np.full(

        shape=(batch_size, max_length), fill_value=padding_idx, dtype=np.int32

    )

    for i, l in enumerate(lengths):

        padded_seqs[i, 0:l] = seqs[i]

    return padded_seqs

def pad_and_sort_batch(batch: Dict, padding_idx: int, rd_max: int) -> Dict[str, Tensor]:

    """

    DataLoaderBatch should be a list of (sequence, target, length) tuples...

    Returns a padded tensor of sequences sorted from longest to shortest,

    """

    for key in ["char_length", "seq_length", "label"]:

        batch[key] = np.array(batch[key])

    for key in ["e1", "e2"]:

        seqs = batch[key]

        # pad entities apart to avoid unnecessary padding

        lengths = list(map(lambda x: len(x), seqs))

        batch[key] = pad_seqs(seqs, lengths, padding_idx)

    for key in ["rd1", "rd2"]:

        seqs = batch[key]

        # pad relative distance with maximum value

        batch[key] = pad_seqs(seqs, batch["seq_length"], rd_max)

    for key in ["sent", "iob", "pos", "dep"]:

        seqs = batch[key]

        # pad other features with the common padding index

        batch[key] = pad_seqs(seqs, batch["seq_length"], padding_idx)

    return sort_batch(batch, batch["seq_length"])

def custom_collate(data: Dict[str, List[List[float]]]):

    """Separates the inputs and the targets

    Args:

        data (Dict[str, List[List[float]]]): batch of data

    Returns:

        Tuple[Dict[str, Tensor], Tensor]: inputs and targets.

    """

    inputs = {}

    targets = torch.from_numpy(data[0]["label"]).long()

    for key, value in data[0].items():

        if key != "label":

            inputs[key] = torch.from_numpy(value).long()

    return inputs, targets