import torch


def compute_edge_feats_dict(edge_ixs, det_df, fps, use_cuda):
    """
    Computes a dictionary of edge features among pairs of detections
    Args:
        edge_ixs: Edges tensor with shape (2, num_edges)
        det_df: processed detections datafrmae
        fps: fps for the given sequence
        use_cuda: bool, determines whether operations must be performed in GPU
    Returns:
        Dict where edge key is a string referring to the attr name, and each val is a tensor of shape (num_edges)
        with vals of that attribute for each edge.

    """
    device = torch.device("cuda" if torch.cuda.is_available() and use_cuda else "cpu")
    row, col = edge_ixs

    secs_time_dists = torch.from_numpy(det_df['frame'].values).float().to(device) / fps

    bb_height = torch.from_numpy(det_df['bb_height'].values).float().to(device)
    bb_width = torch.from_numpy(det_df['bb_width'].values).float().to(device)

    feet_x = torch.from_numpy(det_df['feet_x'].values).float().to(device)
    feet_y = torch.from_numpy(det_df['feet_y'].values).float().to(device)

    mean_bb_heights = (bb_height[row] + bb_height[col]) / 2

    edge_feats_dict = {'secs_time_dists': secs_time_dists[col] - secs_time_dists[row],

                       'norm_feet_x_dists': (feet_x[col] - feet_x[row]) / mean_bb_heights,
                       'norm_feet_y_dists': (feet_y[col] - feet_y[row]) / mean_bb_heights,

                       'bb_height_dists': torch.log(bb_height[col] / bb_height[row]),
                       'bb_width_dists': torch.log(bb_width[col] / bb_width[row])}

    return edge_feats_dict