import torch


def get_plan_feats(plans, scene_tensor, agent_tensor):
    """
    Returns location coordinates, map and agent features for a given batch of plans

    Inputs
    plans: Sequences of row and column values on grid. shape: (Batchsize, horizon, 2)
    scene_tensor: Tensor of scene features: (Batchsize, C_s, H, W)
    agent_tensor: Tensor of agent features: (Batchsize, C_a, H, W)

    Output
    scene_feats: Scene features along plan (Batchsize, horizon, C_s)
    agent_feats: Agent features along plan (Batchsize, horizon, C_a)
    """
    h = scene_tensor.shape[2]
    scene_tensor = scene_tensor.reshape(scene_tensor.shape[0], scene_tensor.shape[1], -1)
    agent_tensor = agent_tensor.reshape(agent_tensor.shape[0], agent_tensor.shape[1], -1)
    plans = plans[:, :, 0] * h + plans[:, :, 1]
    plans_s = plans[:, None, :].repeat(1, scene_tensor.shape[1], 1).long()
    plans_a = plans[:, None, :].repeat(1, agent_tensor.shape[1], 1).long()
    scene_feats = torch.gather(scene_tensor, 2, plans_s)
    agent_feats = torch.gather(agent_tensor, 2, plans_a)
    scene_feats = scene_feats.permute(0, 2, 1)
    agent_feats = agent_feats.permute(0, 2, 1)
    return scene_feats, agent_feats