import torch


def extract(training_data, valid_indices, t, off_policy, act_dim):
    """Extract from the training data all the elements for the given timestep as a tuple of
    stacked tensors.

    Args:
        training_data (list): List of TraceTrainingData objects from which to extract.
        valid_indices (list): List of indices corresponding to train_data object for which t is
            a valid timestep.
        t (int): Time-step for which the data has to be extracted.
        off_policy (bool): Whether the training_data objects come from old policies.
        act_dim (int): Dimension of actions.

    Returns:
        A tuple (actions, policies, rewards, q_values, values, rhos, avg_p) where each element is
        a tensor of stacked elements from training_data[valid_indices] at time-step t.
        rhos contain the importance samples computed, or ones if not off_policy.
    """
    # Extract elements at position -t from training data
    actions = torch.cat(tuple(training_data[i].actions[-t] for i in valid_indices), dim=0)
    policies = torch.cat(tuple(training_data[i].policies[-t] for i in valid_indices), dim=0)
    q_values = torch.cat(tuple(training_data[i].q_values[-t] for i in valid_indices), dim=0)
    rewards = torch.cat(tuple(training_data[i].rewards[-t] for i in valid_indices), dim=0)
    values = torch.cat(tuple(training_data[i].values[-t] for i in valid_indices), dim=0)
    avg_p = torch.cat(tuple(training_data[i].avg_policies[-t] for i in valid_indices), dim=0)
    if off_policy:
        old_policies = torch.cat(
            tuple(training_data[i].old_policies[-t] for i in valid_indices), dim=0)
        rhos = policies.detach() / old_policies
    else:
        rhos = torch.ones(actions.shape[0], act_dim)
    return actions, policies, rewards, q_values, values, rhos, avg_p