import numpy as np
import torch


def flatten_dataset(x, y, indices, visits_length, case="los"):
    x_flat = []
    y_flat = []
    len_list = []
    for i in indices:
        len_list.append(visits_length[i])
        for v in range(visits_length[i]):
            x_flat.append(x[i][v])
            if case == "los":
                y_flat.append(y[i][v][1])
            elif case == "outcome":
                y_flat.append(y[i][v].tolist())
    return np.array(x_flat), np.array(y_flat), np.array(len_list)


def numpy_dataset(x, y, x_lab_length):
    x = x.numpy()
    y = y.numpy()
    x_lab_length = x_lab_length.numpy()
    y_los = y[:, :, 1]
    y_outcome = y[:, 0, 0]
    return x, y_outcome, y_los, x_lab_length