import torch


def apply_plot_args_to_trajectories(args, obs_traj, gt, predictions):
    """
    Apply command line arguments to the trajectories to be plotted
    :param args: command line arguments containing the several changes/restrictions to the trajectories
    :param obs_traj: Tensor of shape [obs_traj_len, num_peds, 2]. Past trajectories of all pedestrians
    :param gt: Tensor of shape [pred_traj_len, num_peds, 2]. Real Future (ground truth) trajectories of all pedestrians
    :param predictions: Tensor of shape [num_methods, pred_traj_len, num_peds, 2]. Predicted trajectories by
    'num_methods' different models, for all pedestrians
    :return: the three aforementioned tensors, with the arguments applied.
    """
    squeeze_unsqueeze = obs_traj.ndim < 3  # for the case of just one pedestrian - shape (traj_len, 2)
    if squeeze_unsqueeze:
        obs_traj, gt, predictions = obs_traj.unsqueeze(1), gt.unsqueeze(1), predictions.unsqueeze(2)
    num_peds = obs_traj.shape[1]
    seq_mask = torch.ones(num_peds, device=obs_traj.device).to(torch.bool)
    if args.ignore_neighbours_past:
        seq_mask *= ~torch.all(torch.isnan(gt[:, :, 0]), dim=0)
    obs_traj, gt, predictions = obs_traj[:, seq_mask], gt[:, seq_mask], predictions[:, :, seq_mask]
    if args.switch_x_y:
        obs_traj[:, :, [0, 1]] = obs_traj[:, :, [1, 0]]
        gt[:, :, [0, 1]] = gt[:, :, [1, 0]]
        predictions[:, :, :, [0, 1]] = predictions[:, :, :, [1, 0]]
    if args.invert_x:
        obs_traj[:, :, 0] = - obs_traj[:, :, 0]
        gt[:, :, 0] = - gt[:, :, 0]
        predictions[:, :, :, 0] = - predictions[:, :, :, 0]
    if args.invert_y:
        obs_traj[:, :, 1] = - obs_traj[:, :, 1]
        gt[:, :, 1] = - gt[:, :, 1]
        predictions[:, :, :, 1] = - predictions[:, :, :, 1]
    if squeeze_unsqueeze:
        return obs_traj.squeeze(1), gt.squeeze(1), predictions.squeeze(2)
    return obs_traj, gt, predictions