import numpy as np
import torch


def flat_x_y_collate(batch, *args, **kwargs) -> tuple:
    """
    Collate function for a classical model with flat x and y feature vectors
    """
    # Flatten each item in the batch
    flattened_x = list()
    flattened_y = list()
    for i in range(len(batch)):
        item = batch[i]
        flattened_x_item = list()
        flattened_y_item = list()
        for j, (category, windows) in enumerate(item.items()):
            for feature in windows.T:
                if category == "target":
                    flattened_y_item.append(feature)
                else:
                    flattened_x_item.append(feature)
        # create numpy arrays with t_1_f1, t_1_f2, t_2_f1, t_2_f2 and so on
        reordered_x = np.empty((len(flattened_x_item) * len(flattened_x_item[0])))
        reordered_y = np.empty((len(flattened_y_item) * len(flattened_y_item[0])))
        for j in range(len(flattened_x_item)):
            for k in range(len(flattened_x_item[j])):
                reordered_x[j + k * len(flattened_x_item)] = flattened_x_item[j][k]

        for j in range(len(flattened_y_item)):
            for k in range(len(flattened_y_item[j])):
                reordered_y[j + k * len(flattened_y_item)] = flattened_y_item[j][k]

        # append to the list
        flattened_x.append(reordered_x)
        flattened_y.append(reordered_y)

    return np.array(flattened_x), np.array(flattened_y)


def simple_x_y_collate(batch):
    """
    Collate function for LSTM model
    """
    collated_x = list()
    collated_y = list()
    for item in batch:
        current_x = list()
        current_y = list()
        for category, windows in item.items():
            if category == "target_features":
                current_y.append(windows)
            else:
                current_x.append(windows)

        collated_x.append(np.concatenate(current_x, axis=1))
        collated_y.append(np.concatenate(current_y, axis=1))

    # convert to numpy arrays
    collated_x = np.array(collated_x)
    collated_y = np.array(collated_y)

    # convert to torch tensors
    return torch.tensor(collated_x, dtype=torch.float32), torch.tensor(collated_y, dtype=torch.float32)


def collate_with_padding(batch,
                         padding_value: float = 0.0, ):
    """
    Collate the batch with padding
    """

    # create simple lists for x and y
    collated_x = list()
    collated_y = list()
    for item in batch:
        current_x = list()
        current_y = list()
        for category, windows in item.items():
            if category == "target_features":
                current_y.append(windows)
            else:
                current_x.append(windows)

        collated_x.append(np.concatenate(current_x, axis=1))
        collated_y.append(np.concatenate(current_y, axis=1))

    # get the max length of the x and y
    max_x_length = max([x.shape[0] for x in collated_x])
    max_y_length = max([y.shape[0] for y in collated_y])
    # pad the x and y
    padded_x = list()
    padded_y = list()
    for x, y in zip(collated_x, collated_y):
        padded_x.append(
            np.pad(x, ((max_x_length - x.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))
        padded_y.append(
            np.pad(y, ((max_y_length - y.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))

    return torch.tensor(padded_x, dtype=torch.float32), torch.tensor(padded_y, dtype=torch.float32)