import sys
_module = sys.modules[__name__]
del sys
interface = _module
nuScenes = _module
nuScenes_graphs = _module
nuScenes_raster = _module
nuScenes_vector = _module
evaluate = _module
covernet_loss = _module
goal_pred_nll = _module
metric = _module
min_ade = _module
min_fde = _module
miss_rate = _module
mtp_loss = _module
pi_bc = _module
utils = _module
aggregator = _module
concat = _module
global_attention = _module
goal_conditioned = _module
pgp = _module
covernet = _module
decoder = _module
lvm = _module
mtp = _module
multipath = _module
utils = _module
encoder = _module
pgp_encoder = _module
polyline_subgraph = _module
raster_encoder = _module
model = _module
preprocess = _module
train = _module
evaluator = _module
initialization = _module
preprocessor = _module
trainer = _module
utils = _module
visualizer = _module
visualize = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import abc


import torch.utils.data as torch_data


import numpy as np


from typing import Union


from typing import Dict


from typing import Tuple


from typing import List


import torch


import math


import torch.nn as nn


from torch.distributions import Categorical


from sklearn.cluster import KMeans


from scipy.spatial.distance import cdist


from torch.nn.utils.rnn import pack_padded_sequence


from torchvision.models import resnet18


from torchvision.models import resnet34


from torchvision.models import resnet50


from torch.utils.tensorboard import SummaryWriter


import torch.optim


import time


import matplotlib.pyplot as plt


class PredictionAggregator(nn.Module):
    """
    Base class for context aggregators for single agent prediction.
    Aggregates a set of context (map, surrounding agent) encodings and outputs either a single aggregated context vector
    or 'K' selectively aggregated context vectors for multimodal prediction.
    """

    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def forward(self, encodings: Dict) ->Union[Dict, torch.Tensor]:
        """
        Forward pass for prediction aggregator
        :param encodings: Dictionary with target agent and context encodings
        :return agg_encoding: Aggregated context encoding
        """
        raise NotImplementedError()


class Concat(PredictionAggregator):
    """
    Concatenates target agent encoding and all context encodings.
    Set of context encodings needs to be the same size, ideally with a well-defined order.
    """

    def __init__(self):
        super().__init__()

    def forward(self, encodings: Dict) ->torch.Tensor:
        """
        Forward pass for Concat aggregator
        """
        target_agent_enc = encodings['target_agent_encoding']
        context_enc = encodings['context_encoding']
        batch_size = target_agent_enc.shape[0]
        if context_enc['combined'] is not None:
            context_vec = context_enc['combined'].reshape(batch_size, -1)
        else:
            map_vec = context_enc['map'].reshape(batch_size, -1) if context_enc['map'] else torch.empty(batch_size, 0)
            vehicle_vec = context_enc['vehicles'].reshape(batch_size, -1) if context_enc['map'] else torch.empty(batch_size, 0)
            ped_vec = context_enc['pedestrians'].reshape(batch_size, -1) if context_enc['pedestrians'] else torch.empty(batch_size, 0)
            context_vec = torch.cat((map_vec, vehicle_vec, ped_vec), dim=1)
        op = torch.cat((target_agent_enc, context_vec), dim=1)
        return op


class GlobalAttention(PredictionAggregator):
    """
    Aggregate context encoding using scaled dot product attention. Query obtained using target agent encoding,
    Keys and values obtained using map and surrounding agent encodings.
    """

    def __init__(self, args: Dict):
        """
        args to include

        enc_size: int Dimension of encodings generated by encoder
        emb_size: int Size of embeddings used for queries, keys and values
        num_heads: int Number of attention heads

        """
        super().__init__()
        self.query_emb = nn.Linear(args['target_agent_enc_size'], args['emb_size'])
        self.key_emb = nn.Linear(args['context_enc_size'], args['emb_size'])
        self.val_emb = nn.Linear(args['context_enc_size'], args['emb_size'])
        self.mha = nn.MultiheadAttention(args['emb_size'], args['num_heads'])

    def forward(self, encodings: Dict) ->torch.Tensor:
        """
        Forward pass for attention aggregator
        """
        target_agent_enc = encodings['target_agent_encoding']
        context_enc = encodings['context_encoding']
        if context_enc['combined'] is not None:
            combined_enc, combined_masks = context_enc['combined'], context_enc['combined_masks'].bool()
        else:
            combined_enc, combined_masks = self.get_combined_encodings(context_enc)
        query = self.query_emb(target_agent_enc).unsqueeze(0)
        keys = self.key_emb(combined_enc).permute(1, 0, 2)
        vals = self.val_emb(combined_enc).permute(1, 0, 2)
        op, _ = self.mha(query, keys, vals, key_padding_mask=combined_masks)
        op = op.squeeze(0)
        op = torch.cat((target_agent_enc, op), dim=-1)
        return op

    @staticmethod
    def get_combined_encodings(context_enc: Dict) ->Tuple[torch.Tensor, torch.Tensor]:
        """
        Creates a combined set of map and surrounding agent encodings to be aggregated using attention.
        """
        encodings = []
        masks = []
        if 'map' in context_enc:
            encodings.append(context_enc['map'])
            masks.append(context_enc['map_masks'])
        if 'vehicles' in context_enc:
            encodings.append(context_enc['vehicles'])
            masks.append(context_enc['vehicle_masks'])
        if 'pedestrians' in context_enc:
            encodings.append(context_enc['pedestrians'])
            masks.append(context_enc['pedestrian_masks'])
        combined_enc = torch.cat(encodings, dim=1)
        combined_masks = torch.cat(masks, dim=1).bool()
        return combined_enc, combined_masks


class GoalConditioned(GlobalAttention):
    """
    Goal conditioned aggregator with following functionality.
    1) Predicts goal probabilities over lane nodes
    2) Samples goals
    3) Outputs goal conditioned encodings for N samples to pass on to the trajectory decoder
    """

    def __init__(self, args):
        """
        args to include

        for aggregating map and agent context
        enc_size: int Dimension of encodings generated by encoder
        emb_size: int Size of embeddings used for queries, keys and values
        num_heads: int Number of attention heads

        for goal prediction
        'pre_train': bool, whether the model is being pre-trained using ground truth goals.
        'context_enc_size': int, size of node encoding
        'target_agent_enc_size': int, size of target agent encoding
        'goal_h1_size': int, size of first layer of goal prediction header
        'goal_h2_size': int, size of second layer of goal prediction header
        'num_samples': int, number of goals to sample
        """
        super(GoalConditioned, self).__init__(args)
        self.goal_h1 = nn.Linear(args['context_enc_size'] + args['target_agent_enc_size'], args['goal_h1_size'])
        self.goal_h2 = nn.Linear(args['goal_h1_size'], args['goal_h2_size'])
        self.goal_op = nn.Linear(args['goal_h2_size'], 1)
        self.num_samples = args['num_samples']
        self.leaky_relu = nn.LeakyReLU()
        self.log_softmax = nn.LogSoftmax(dim=1)
        self.pre_train = args['pre_train']

    def forward(self, encodings: Dict) ->Dict:
        """
        Forward pass for goal conditioned aggregator
        :param encodings: dictionary with encoder outputs
        :return: outputs, dictionary with
            'agg_encoding': aggregated encodings
            'goal_log_probs':  log probabilities over nodes corresponding to predicted goals
        """
        target_agent_encoding = encodings['target_agent_encoding']
        node_encodings = encodings['context_encoding']['combined']
        node_masks = encodings['context_encoding']['combined_masks']
        goal_log_probs = self.compute_goal_probs(target_agent_encoding, node_encodings, node_masks)
        if self.pre_train and self.training:
            max_nodes = node_masks.shape[1]
            goals = encodings['node_seq_gt'][:, -1].unsqueeze(1).repeat(1, self.num_samples).long() - max_nodes
        else:
            goals = Categorical(torch.exp(goal_log_probs).unsqueeze(1).repeat(1, self.num_samples, 1)).sample()
        agg_enc = super(GoalConditioned, self).forward(encodings)
        agg_enc = agg_enc.unsqueeze(1).repeat(1, self.num_samples, 1)
        batch_indices = torch.arange(agg_enc.shape[0]).unsqueeze(1).repeat(1, self.num_samples)
        goal_encodings = node_encodings[batch_indices, goals]
        agg_enc = torch.cat((agg_enc, goal_encodings), dim=2)
        outputs = {'agg_encoding': agg_enc, 'goal_log_probs': goal_log_probs}
        return outputs

    def compute_goal_probs(self, target_agent_encoding, node_encodings, node_masks):
        """
        Forward pass for goal prediction header
        :param target_agent_encoding: tensor encoding the target agent's past motion
        :param node_encodings: tensor of node encodings provided by the encoder
        :param node_masks: masks indicating whether a node exists for a given index in the tensor
        :return:
        """
        max_nodes = node_encodings.shape[1]
        target_agent_enc_size = target_agent_encoding.shape[-1]
        node_enc_size = node_encodings.shape[-1]
        target_agent_encoding = target_agent_encoding.unsqueeze(1).repeat(1, max_nodes, 1)
        enc = torch.cat((target_agent_encoding, node_encodings), dim=2)
        masks_goal = ~node_masks.unsqueeze(-1).bool()
        enc_batched = torch.masked_select(enc, masks_goal).reshape(-1, target_agent_enc_size + node_enc_size)
        goal_ops_ = self.goal_op(self.leaky_relu(self.goal_h2(self.leaky_relu(self.goal_h1(enc_batched)))))
        goal_ops = torch.zeros_like(masks_goal).float()
        goal_ops = goal_ops.masked_scatter_(masks_goal, goal_ops_).squeeze(-1)
        goal_log_probs = self.log_softmax(goal_ops + torch.log(1 - node_masks))
        return goal_log_probs


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class PGP(PredictionAggregator):
    """
    Policy header + selective aggregator from "Multimodal trajectory prediction conditioned on lane graph traversals"
    1) Outputs edge probabilities corresponding to pi_route
    2) Samples pi_route to output traversed paths
    3) Selectively aggregates context along traversed paths
    """

    def __init__(self, args: Dict):
        """
        args to include
        'pre_train': bool, whether the model is being pre-trained using ground truth node sequence.
        'node_enc_size': int, size of node encoding
        'target_agent_enc_size': int, size of target agent encoding
        'pi_h1_size': int, size of first layer of policy header
        'pi_h2_size': int, size of second layer of policy header
        'emb_size': int, embedding size for attention layer for aggregating node encodings
        'num_heads: int, number of attention heads
        'num_samples': int, number of sampled traversals (and encodings) to output
        """
        super().__init__()
        self.pre_train = args['pre_train']
        self.pi_h1 = nn.Linear(2 * args['node_enc_size'] + args['target_agent_enc_size'] + 2, args['pi_h1_size'])
        self.pi_h2 = nn.Linear(args['pi_h1_size'], args['pi_h2_size'])
        self.pi_op = nn.Linear(args['pi_h2_size'], 1)
        self.pi_h1_goal = nn.Linear(args['node_enc_size'] + args['target_agent_enc_size'], args['pi_h1_size'])
        self.pi_h2_goal = nn.Linear(args['pi_h1_size'], args['pi_h2_size'])
        self.pi_op_goal = nn.Linear(args['pi_h2_size'], 1)
        self.leaky_relu = nn.LeakyReLU()
        self.log_softmax = nn.LogSoftmax(dim=2)
        self.horizon = args['horizon']
        self.num_samples = args['num_samples']
        self.pos_enc = PositionalEncoding1D(args['node_enc_size'])
        self.query_emb = nn.Linear(args['target_agent_enc_size'], args['emb_size'])
        self.key_emb = nn.Linear(args['node_enc_size'], args['emb_size'])
        self.val_emb = nn.Linear(args['node_enc_size'], args['emb_size'])
        self.mha = nn.MultiheadAttention(args['emb_size'], args['num_heads'])

    def forward(self, encodings: Dict) ->Dict:
        """
        Forward pass for PGP aggregator
        :param encodings: dictionary with encoder outputs
        :return: outputs: dictionary with
            'agg_encoding': aggregated encodings along sampled traversals
            'pi': discrete policy (probabilities over outgoing edges) for graph traversal
        """
        target_agent_encoding = encodings['target_agent_encoding']
        node_encodings = encodings['context_encoding']['combined']
        node_masks = encodings['context_encoding']['combined_masks']
        s_next = encodings['s_next']
        edge_type = encodings['edge_type']
        pi = self.compute_policy(target_agent_encoding, node_encodings, node_masks, s_next, edge_type)
        if self.pre_train and self.training:
            sampled_traversals = encodings['node_seq_gt'].unsqueeze(1).repeat(1, self.num_samples, 1).long()
        else:
            init_node = encodings['init_node']
            sampled_traversals = self.sample_policy(torch.exp(pi), s_next, init_node)
        agg_enc = self.aggregate(sampled_traversals, node_encodings, target_agent_encoding)
        outputs = {'agg_encoding': agg_enc, 'pi': pi}
        return outputs

    def aggregate(self, sampled_traversals, node_encodings, target_agent_encoding) ->torch.Tensor:
        batch_size = node_encodings.shape[0]
        max_nodes = node_encodings.shape[1]
        unique_traversals = [torch.unique(i, dim=0, return_counts=True) for i in sampled_traversals]
        traversals_batched = torch.cat([i[0] for i in unique_traversals], dim=0)
        counts_batched = torch.cat([i[1] for i in unique_traversals], dim=0)
        batch_idcs = torch.cat([(n * torch.ones(len(i[1])).long()) for n, i in enumerate(unique_traversals)])
        batch_idcs = batch_idcs.unsqueeze(1).repeat(1, self.horizon)
        dummy_enc = torch.zeros_like(node_encodings)
        node_encodings = torch.cat((node_encodings, dummy_enc), dim=1)
        node_enc_selected = node_encodings[batch_idcs, traversals_batched]
        pos_enc = self.pos_enc(torch.zeros_like(node_enc_selected))
        node_enc_selected += pos_enc
        target_agent_enc_batched = target_agent_encoding[batch_idcs[:, 0]]
        query = self.query_emb(target_agent_enc_batched).unsqueeze(0)
        keys = self.key_emb(node_enc_selected).permute(1, 0, 2)
        vals = self.val_emb(node_enc_selected).permute(1, 0, 2)
        key_padding_mask = torch.as_tensor(traversals_batched >= max_nodes)
        att_op, _ = self.mha(query, keys, vals, key_padding_mask)
        att_op = att_op.squeeze(0).repeat_interleave(counts_batched, dim=0).view(batch_size, self.num_samples, -1)
        agg_enc = torch.cat((target_agent_encoding.unsqueeze(1).repeat(1, self.num_samples, 1), att_op), dim=-1)
        return agg_enc

    def sample_policy(self, pi, s_next, init_node) ->torch.Tensor:
        """
        Sample graph traversals using discrete policy.
        :param pi: tensor with probabilities corresponding to the policy
        :param s_next: look-up table for next node for a given source node and edge
        :param init_node: initial node to start the policy at
        :return:
        """
        with torch.no_grad():
            batch_size = pi.shape[0]
            max_nodes = pi.shape[1]
            batch_idcs = torch.arange(batch_size, device=device).unsqueeze(1).repeat(1, self.num_samples).view(-1)
            sampled_traversals = torch.zeros(batch_size, self.num_samples, self.horizon, device=device).long()
            pi_dummy = torch.zeros_like(pi)
            pi_dummy[:, :, -1] = 1
            s_next_dummy = torch.zeros_like(s_next)
            s_next_dummy[:, :, -1] = max_nodes + torch.arange(max_nodes).unsqueeze(0).repeat(batch_size, 1)
            pi = torch.cat((pi, pi_dummy), dim=1)
            s_next = torch.cat((s_next, s_next_dummy), dim=1)
            pi_s = init_node.unsqueeze(1).repeat(1, self.num_samples, 1).view(-1, max_nodes)
            s = Categorical(pi_s).sample()
            sampled_traversals[:, :, 0] = s.reshape(batch_size, self.num_samples)
            for n in range(1, self.horizon):
                pi_s = pi[batch_idcs, s]
                a = Categorical(pi_s).sample()
                s = s_next[batch_idcs, s, a].long()
                sampled_traversals[:, :, n] = s.reshape(batch_size, self.num_samples)
        return sampled_traversals

    def compute_policy(self, target_agent_encoding, node_encodings, node_masks, s_next, edge_type) ->torch.Tensor:
        """
        Forward pass for policy header
        :param target_agent_encoding: tensor encoding the target agent's past motion
        :param node_encodings: tensor of node encodings provided by the encoder
        :param node_masks: masks indicating whether a node exists for a given index in the tensor
        :param s_next: look-up table for next node for a given source node and edge
        :param edge_type: look-up table with edge types
        :return pi: tensor with probabilities corresponding to the policy
        """
        batch_size = node_encodings.shape[0]
        max_nodes = node_encodings.shape[1]
        max_nbrs = s_next.shape[2] - 1
        node_enc_size = node_encodings.shape[2]
        target_agent_enc_size = target_agent_encoding.shape[1]
        src_node_enc = node_encodings.unsqueeze(2).repeat(1, 1, max_nbrs, 1)
        dst_idcs = s_next[:, :, :-1].reshape(batch_size, -1).long()
        batch_idcs = torch.arange(batch_size).unsqueeze(1).repeat(1, max_nodes * max_nbrs)
        dst_node_enc = node_encodings[batch_idcs, dst_idcs].reshape(batch_size, max_nodes, max_nbrs, node_enc_size)
        target_agent_enc = target_agent_encoding.unsqueeze(1).unsqueeze(2).repeat(1, max_nodes, max_nbrs, 1)
        edge_enc = torch.cat((torch.as_tensor(edge_type[:, :, :-1] == 1, device=device).unsqueeze(3).float(), torch.as_tensor(edge_type[:, :, :-1] == 2, device=device).unsqueeze(3).float()), dim=3)
        enc = torch.cat((target_agent_enc, src_node_enc, dst_node_enc, edge_enc), dim=3)
        enc_goal = torch.cat((target_agent_enc[:, :, 0, :], src_node_enc[:, :, 0, :]), dim=2)
        masks = torch.sum(edge_enc, dim=3, keepdim=True).bool()
        masks_goal = ~node_masks.unsqueeze(-1).bool()
        enc_batched = torch.masked_select(enc, masks).reshape(-1, target_agent_enc_size + 2 * node_enc_size + 2)
        enc_goal_batched = torch.masked_select(enc_goal, masks_goal).reshape(-1, target_agent_enc_size + node_enc_size)
        pi_ = self.pi_op(self.leaky_relu(self.pi_h2(self.leaky_relu(self.pi_h1(enc_batched)))))
        pi = torch.zeros_like(masks).float()
        pi = pi.masked_scatter_(masks, pi_).squeeze(-1)
        pi_goal_ = self.pi_op_goal(self.leaky_relu(self.pi_h2_goal(self.leaky_relu(self.pi_h1_goal(enc_goal_batched)))))
        pi_goal = torch.zeros_like(masks_goal).float()
        pi_goal = pi_goal.masked_scatter_(masks_goal, pi_goal_)
        pi = torch.cat((pi, pi_goal), dim=-1)
        op_masks = torch.log(torch.as_tensor(edge_type != 0).float())
        pi = self.log_softmax(pi + op_masks)
        return pi


class PredictionDecoder(nn.Module):
    """
    Base class for decoders for single agent prediction.
    Outputs K trajectories and/or their probabilities
    """

    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def forward(self, agg_encoding: Union[torch.Tensor, Dict]) ->Union[torch.Tensor, Dict]:
        """
        Forward pass for prediction decoder
        :param agg_encoding: Aggregated context encoding
        :return outputs: K Predicted trajectories and/or their probabilities/scores
        """
        raise NotImplementedError()


def cluster_traj(k: int, traj: torch.Tensor):
    """
    clusters sampled trajectories to output K modes.
    :param k: number of clusters
    :param traj: set of sampled trajectories, shape [batch_size, num_samples, traj_len, 2]
    :return: traj_clustered:  set of clustered trajectories, shape [batch_size, k, traj_len, 2]
             scores: scores for clustered trajectories (basically 1/rank), shape [batch_size, k]
    """
    batch_size = traj.shape[0]
    num_samples = traj.shape[1]
    traj_len = traj.shape[2]
    data = traj[:, :, 0::3, :]
    data = data.reshape(batch_size, num_samples, -1).detach().cpu().numpy()
    cluster_ops = ray.get([cluster_and_rank.remote(k, data_slice) for data_slice in data])
    cluster_lbls = [cluster_op['lbls'] for cluster_op in cluster_ops]
    cluster_counts = [cluster_op['counts'] for cluster_op in cluster_ops]
    cluster_ranks = [cluster_op['ranks'] for cluster_op in cluster_ops]
    lbls = torch.as_tensor(cluster_lbls, device=device).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, traj_len, 2).long()
    traj_summed = torch.zeros(batch_size, k, traj_len, 2, device=device).scatter_add(1, lbls, traj)
    cnt_tensor = torch.as_tensor(cluster_counts, device=device).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, traj_len, 2)
    traj_clustered = traj_summed / cnt_tensor
    scores = 1 / torch.as_tensor(cluster_ranks, device=device)
    scores = scores / torch.sum(scores, dim=1)[0]
    return traj_clustered, scores


class LVM(PredictionDecoder):

    def __init__(self, args):
        """
        Latent variable conditioned decoder.

        args to include:
        agg_type: 'combined' or 'sample_specific'. Whether we have a single aggregated context vector or sample-specific
        num_samples: int Number of trajectories to sample
        op_len: int Length of predicted trajectories
        lv_dim: int Dimension of latent variable
        encoding_size: int Dimension of encoded scene + agent context
        hidden_size: int Size of output mlp hidden layer
        num_clusters: int Number of final clustered trajectories to output

        """
        super().__init__()
        self.agg_type = args['agg_type']
        self.num_samples = args['num_samples']
        self.op_len = args['op_len']
        self.lv_dim = args['lv_dim']
        self.hidden = nn.Linear(args['encoding_size'] + args['lv_dim'], args['hidden_size'])
        self.op_traj = nn.Linear(args['hidden_size'], args['op_len'] * 2)
        self.leaky_relu = nn.LeakyReLU()
        self.num_clusters = args['num_clusters']

    def forward(self, inputs: Union[Dict, torch.Tensor]) ->Dict:
        """
        Forward pass for latent variable model.

        :param inputs: aggregated context encoding,
         shape for combined encoding: [batch_size, encoding_size]
         shape if sample specific encoding: [batch_size, num_samples, encoding_size]
        :return: predictions
        """
        if type(inputs) is torch.Tensor:
            agg_encoding = inputs
        else:
            agg_encoding = inputs['agg_encoding']
        if self.agg_type == 'combined':
            agg_encoding = agg_encoding.unsqueeze(1).repeat(1, self.num_samples, 1)
        elif len(agg_encoding.shape) != 3 or agg_encoding.shape[1] != self.num_samples:
            raise Exception('Expected ' + str(self.num_samples) + 'encodings for each train/val data')
        batch_size = agg_encoding.shape[0]
        z = torch.randn(batch_size, self.num_samples, self.lv_dim, device=device)
        agg_encoding = torch.cat((agg_encoding, z), dim=2)
        h = self.leaky_relu(self.hidden(agg_encoding))
        traj = self.op_traj(h)
        traj = traj.reshape(batch_size, self.num_samples, self.op_len, 2)
        traj_clustered, probs = cluster_traj(self.num_clusters, traj)
        predictions = {'traj': traj_clustered, 'probs': probs}
        if type(inputs) is dict:
            for key, val in inputs.items():
                if key != 'agg_encoding':
                    predictions[key] = val
        return predictions


def bivariate_gaussian_activation(ip: torch.Tensor) ->torch.Tensor:
    """
    Activation function to output parameters of bivariate Gaussian distribution
    """
    mu_x = ip[..., 0:1]
    mu_y = ip[..., 1:2]
    sig_x = ip[..., 2:3]
    sig_y = ip[..., 3:4]
    rho = ip[..., 4:5]
    sig_x = torch.exp(sig_x)
    sig_y = torch.exp(sig_y)
    rho = torch.tanh(rho)
    out = torch.cat([mu_x, mu_y, sig_x, sig_y, rho], dim=-1)
    return out


class MTP(PredictionDecoder):

    def __init__(self, args):
        """
        Prediction decoder for MTP

        args to include:
        num_modes: int number of modes K
        op_len: int prediction horizon
        hidden_size: int hidden layer size
        encoding_size: int size of context encoding
        use_variance: Whether to output variance params along with mean predicted locations
        """
        super().__init__()
        self.agg_type = args['agg_type']
        self.num_modes = args['num_modes']
        self.op_len = args['op_len']
        self.use_variance = args['use_variance']
        self.op_dim = 5 if self.use_variance else 2
        self.hidden = nn.Linear(args['encoding_size'], args['hidden_size'])
        self.traj_op = nn.Linear(args['hidden_size'], args['op_len'] * self.op_dim * self.num_modes)
        self.prob_op = nn.Linear(args['hidden_size'], self.num_modes)
        self.leaky_relu = nn.LeakyReLU(0.01)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, agg_encoding: torch.Tensor) ->Dict:
        """
        Forward pass for MTP
        :param agg_encoding: aggregated context encoding
        :return predictions: dictionary with 'traj': K predicted trajectories and
            'probs': K corresponding probabilities
        """
        h = self.leaky_relu(self.hidden(agg_encoding))
        batch_size = h.shape[0]
        traj = self.traj_op(h)
        probs = self.log_softmax(self.prob_op(h))
        traj = traj.reshape(batch_size, self.num_modes, self.op_len, self.op_dim)
        probs = probs.squeeze(dim=-1)
        traj = bivariate_gaussian_activation(traj) if self.use_variance else traj
        predictions = {'traj': traj, 'probs': probs}
        return predictions


class TrajectoryDataset(torch_data.Dataset):
    """
    Base class for trajectory datasets.
    """

    def __init__(self, mode: str, data_dir: str):
        """
        Initialize trajectory dataset.
        :param mode: Mode of operation of dataset
        :param data_dir: Directory to store extracted pre-processed data
        """
        if mode not in ['compute_stats', 'extract_data', 'load_data']:
            raise Exception('Dataset mode needs to be one of {compute_stats, extract_data or load_data}')
        self.mode = mode
        self.data_dir = data_dir
        if mode != 'load_data' and not os.path.isdir(self.data_dir):
            os.mkdir(self.data_dir)

    @abc.abstractmethod
    def __len__(self) ->int:
        """
        Returns size of dataset
        """
        raise NotImplementedError()

    def __getitem__(self, idx: int) ->Union[Dict, int]:
        """
        Get data point, based on mode of operation of dataset.
        :param idx: data index
        """
        if self.mode == 'compute_stats':
            return self.compute_stats(idx)
        elif self.mode == 'extract_data':
            self.extract_data(idx)
            return 0
        else:
            return self.load_data(idx)

    @abc.abstractmethod
    def compute_stats(self, idx: int) ->Dict:
        """
        Function to compute dataset statistics like max surrounding agents, max nodes, max edges etc.
        :param idx: data index
        """
        raise NotImplementedError()

    def extract_data(self, idx: int):
        """
        Function to extract data. Bulk of the dataset functionality will be implemented by this method.
        :param idx: data index
        """
        inputs = self.get_inputs(idx)
        ground_truth = self.get_ground_truth(idx)
        data = {'inputs': inputs, 'ground_truth': ground_truth}
        self.save_data(idx, data)

    @abc.abstractmethod
    def get_inputs(self, idx: int) ->Dict:
        """
        Extracts model inputs.
        :param idx: data index
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_ground_truth(self, idx: int) ->Dict:
        """
        Extracts ground truth 'labels' for training.
        :param idx: data index
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def load_data(self, idx: int) ->Dict:
        """
        Function to load extracted data.
        :param idx: data index
        :return data: Dictionary with pre-processed data
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def save_data(self, idx: int, data: Dict):
        """
        Function to save extracted pre-processed data.
        :param idx: data index
        :param data: Dictionary with pre-processed data
        """
        raise NotImplementedError()


class SingleAgentDataset(TrajectoryDataset):
    """
    Base class for single agent dataset. While we implicitly model all surrounding agents in the scene, predictions
    are made for a single target agent at a time.
    """

    @abc.abstractmethod
    def get_map_representation(self, idx: int) ->Union[np.ndarray, Dict]:
        """
        Extracts map representation
        :param idx: data index
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_surrounding_agent_representation(self, idx: int) ->Union[np.ndarray, Dict]:
        """
        Extracts surrounding agent representation
        :param idx: data index
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_target_agent_representation(self, idx: int) ->Union[np.ndarray, Dict]:
        """
        Extracts target agent representation
        :param idx: data index
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_target_agent_future(self, idx: int) ->Union[np.ndarray, Dict]:
        """
        Extracts future trajectory for target agent
        :param idx: data index
        """
        raise NotImplementedError()


def k_means_anchors(k, ds: SingleAgentDataset):
    """
    Extracts anchors for multipath/covernet using k-means on train set trajectories
    """
    prototype_traj = ds[0]['ground_truth']['traj']
    traj_len = prototype_traj.shape[0]
    traj_dim = prototype_traj.shape[1]
    ds_size = len(ds)
    trajectories = np.zeros((ds_size, traj_len, traj_dim))
    for i, data in enumerate(ds):
        trajectories[i] = data['ground_truth']['traj']
    clustering = KMeans(n_clusters=k).fit(trajectories.reshape((ds_size, -1)))
    anchors = np.zeros((k, traj_len, traj_dim))
    for i in range(k):
        anchors[i] = np.mean(trajectories[clustering.labels_ == i], axis=0)
    anchors = torch.from_numpy(anchors).float()
    return anchors


class Multipath(MTP):

    def __init__(self, args):
        """
        Prediction decoder for Multipath. Almost identical to MTP, but predicts residuals with respect to anchors,
        include the same arguments
        """
        super().__init__(args)
        self.anchors = nn.Parameter(torch.zeros(self.num_modes, self.op_len, 2), requires_grad=False)

    def generate_anchors(self, ds: SingleAgentDataset):
        """
        Function to initialize anchors
        :param ds: train dataset for single agent trajectory prediction
        """
        self.anchors = nn.Parameter(k_means_anchors(self.num_modes, ds))

    def forward(self, agg_encoding: torch.Tensor) ->Dict:
        """
        Forward pass for Multipath
        :param agg_encoding: aggregated context encoding
        :return predictions: dictionary with 'traj': K predicted trajectories and
            'probs': K corresponding probabilities
        """
        predictions = super().forward(agg_encoding)
        predictions['traj'][..., :2] += self.anchors.unsqueeze(0)
        return predictions


class PredictionEncoder(nn.Module):
    """
    Base class for encoders for single agent prediction.
    """

    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def forward(self, inputs: Dict) ->Dict:
        """
        Abstract method for forward pass. Returns dictionary of encodings. Should typically include
        1) target agent encoding, 2) context encoding: encodes map and surrounding agents.

        Context encodings will typically be a set of features (agents or parts of the map),
        with shape: [batch_size, set_size, feature_dim],
        sometimes along with masks for some set elements to account for varying set sizes

        :param inputs: Dictionary with
            'target_agent_representation': target agent history
            'surrounding_agent_representation': surrounding agent history
            'map_representation': HD map representation
        :return encodings: Dictionary with input encodings
        """
        raise NotImplementedError()


class GAT(nn.Module):
    """
    GAT layer for aggregating local context at each lane node. Uses scaled dot product attention using pytorch's
    multihead attention module.
    """

    def __init__(self, in_channels, out_channels):
        """
        Initialize GAT layer.
        :param in_channels: size of node encodings
        :param out_channels: size of aggregated node encodings
        """
        super().__init__()
        self.query_emb = nn.Linear(in_channels, out_channels)
        self.key_emb = nn.Linear(in_channels, out_channels)
        self.val_emb = nn.Linear(in_channels, out_channels)
        self.att = nn.MultiheadAttention(out_channels, 1)

    def forward(self, node_encodings, adj_mat):
        """
        Forward pass for GAT layer
        :param node_encodings: Tensor of node encodings, shape [batch_size, max_nodes, node_enc_size]
        :param adj_mat: Bool tensor, adjacency matrix for edges, shape [batch_size, max_nodes, max_nodes]
        :return:
        """
        queries = self.query_emb(node_encodings.permute(1, 0, 2))
        keys = self.key_emb(node_encodings.permute(1, 0, 2))
        vals = self.val_emb(node_encodings.permute(1, 0, 2))
        att_op, _ = self.att(queries, keys, vals, attn_mask=~adj_mat)
        return att_op.permute(1, 0, 2)


class PGPEncoder(PredictionEncoder):

    def __init__(self, args: Dict):
        """
        GRU based encoder from PGP. Lane node features and agent histories encoded using GRUs.
        Additionally, agent-node attention layers infuse each node encoding with nearby agent context.
        Finally GAT layers aggregate local context at each node.

        args to include:

        target_agent_feat_size: int Size of target agent features
        target_agent_emb_size: int Size of target agent embedding
        taret_agent_enc_size: int Size of hidden state of target agent GRU encoder

        node_feat_size: int Size of lane node features
        node_emb_size: int Size of lane node embedding
        node_enc_size: int Size of hidden state of lane node GRU encoder

        nbr_feat_size: int Size of neighboring agent features
        nbr_enb_size: int Size of neighboring agent embeddings
        nbr_enc_size: int Size of hidden state of neighboring agent GRU encoders

        num_gat_layers: int Number of GAT layers to use.
        """
        super().__init__()
        self.target_agent_emb = nn.Linear(args['target_agent_feat_size'], args['target_agent_emb_size'])
        self.target_agent_enc = nn.GRU(args['target_agent_emb_size'], args['target_agent_enc_size'], batch_first=True)
        self.node_emb = nn.Linear(args['node_feat_size'], args['node_emb_size'])
        self.node_encoder = nn.GRU(args['node_emb_size'], args['node_enc_size'], batch_first=True)
        self.nbr_emb = nn.Linear(args['nbr_feat_size'] + 1, args['nbr_emb_size'])
        self.nbr_enc = nn.GRU(args['nbr_emb_size'], args['nbr_enc_size'], batch_first=True)
        self.query_emb = nn.Linear(args['node_enc_size'], args['node_enc_size'])
        self.key_emb = nn.Linear(args['nbr_enc_size'], args['node_enc_size'])
        self.val_emb = nn.Linear(args['nbr_enc_size'], args['node_enc_size'])
        self.a_n_att = nn.MultiheadAttention(args['node_enc_size'], num_heads=1)
        self.mix = nn.Linear(args['node_enc_size'] * 2, args['node_enc_size'])
        self.leaky_relu = nn.LeakyReLU()
        self.gat = nn.ModuleList([GAT(args['node_enc_size'], args['node_enc_size']) for _ in range(args['num_gat_layers'])])

    def forward(self, inputs: Dict) ->Dict:
        """
        Forward pass for PGP encoder
        :param inputs: Dictionary with
            target_agent_representation: torch.Tensor, shape [batch_size, t_h, target_agent_feat_size]
            map_representation: Dict with
                'lane_node_feats': torch.Tensor, shape [batch_size, max_nodes, max_poses, node_feat_size]
                'lane_node_masks': torch.Tensor, shape [batch_size, max_nodes, max_poses, node_feat_size]

                (Optional)
                's_next': Edge look-up table pointing to destination node from source node
                'edge_type': Look-up table with edge type

            surrounding_agent_representation: Dict with
                'vehicles': torch.Tensor, shape [batch_size, max_vehicles, t_h, nbr_feat_size]
                'vehicle_masks': torch.Tensor, shape [batch_size, max_vehicles, t_h, nbr_feat_size]
                'pedestrians': torch.Tensor, shape [batch_size, max_peds, t_h, nbr_feat_size]
                'pedestrian_masks': torch.Tensor, shape [batch_size, max_peds, t_h, nbr_feat_size]
            agent_node_masks:  Dict with
                'vehicles': torch.Tensor, shape [batch_size, max_nodes, max_vehicles]
                'pedestrians': torch.Tensor, shape [batch_size, max_nodes, max_pedestrians]

            Optionally may also include the following if edges are defined for graph traversal
            'init_node': Initial node in the lane graph based on track history.
            'node_seq_gt': Ground truth node sequence for pre-training

        :return:
        """
        target_agent_feats = inputs['target_agent_representation']
        target_agent_embedding = self.leaky_relu(self.target_agent_emb(target_agent_feats))
        _, target_agent_enc = self.target_agent_enc(target_agent_embedding)
        target_agent_enc = target_agent_enc.squeeze(0)
        lane_node_feats = inputs['map_representation']['lane_node_feats']
        lane_node_masks = inputs['map_representation']['lane_node_masks']
        lane_node_embedding = self.leaky_relu(self.node_emb(lane_node_feats))
        lane_node_enc = self.variable_size_gru_encode(lane_node_embedding, lane_node_masks, self.node_encoder)
        nbr_vehicle_feats = inputs['surrounding_agent_representation']['vehicles']
        nbr_vehicle_feats = torch.cat((nbr_vehicle_feats, torch.zeros_like(nbr_vehicle_feats[:, :, :, 0:1])), dim=-1)
        nbr_vehicle_masks = inputs['surrounding_agent_representation']['vehicle_masks']
        nbr_vehicle_embedding = self.leaky_relu(self.nbr_emb(nbr_vehicle_feats))
        nbr_vehicle_enc = self.variable_size_gru_encode(nbr_vehicle_embedding, nbr_vehicle_masks, self.nbr_enc)
        nbr_ped_feats = inputs['surrounding_agent_representation']['pedestrians']
        nbr_ped_feats = torch.cat((nbr_ped_feats, torch.ones_like(nbr_ped_feats[:, :, :, 0:1])), dim=-1)
        nbr_ped_masks = inputs['surrounding_agent_representation']['pedestrian_masks']
        nbr_ped_embedding = self.leaky_relu(self.nbr_emb(nbr_ped_feats))
        nbr_ped_enc = self.variable_size_gru_encode(nbr_ped_embedding, nbr_ped_masks, self.nbr_enc)
        nbr_encodings = torch.cat((nbr_vehicle_enc, nbr_ped_enc), dim=1)
        queries = self.query_emb(lane_node_enc).permute(1, 0, 2)
        keys = self.key_emb(nbr_encodings).permute(1, 0, 2)
        vals = self.val_emb(nbr_encodings).permute(1, 0, 2)
        attn_masks = torch.cat((inputs['agent_node_masks']['vehicles'], inputs['agent_node_masks']['pedestrians']), dim=2)
        att_op, _ = self.a_n_att(queries, keys, vals, attn_mask=attn_masks)
        att_op = att_op.permute(1, 0, 2)
        lane_node_enc = self.leaky_relu(self.mix(torch.cat((lane_node_enc, att_op), dim=2)))
        adj_mat = self.build_adj_mat(inputs['map_representation']['s_next'], inputs['map_representation']['edge_type'])
        for gat_layer in self.gat:
            lane_node_enc += gat_layer(lane_node_enc, adj_mat)
        lane_node_masks = ~lane_node_masks[:, :, :, 0].bool()
        lane_node_masks = lane_node_masks.any(dim=2)
        lane_node_masks = ~lane_node_masks
        lane_node_masks = lane_node_masks.float()
        encodings = {'target_agent_encoding': target_agent_enc, 'context_encoding': {'combined': lane_node_enc, 'combined_masks': lane_node_masks, 'map': None, 'vehicles': None, 'pedestrians': None, 'map_masks': None, 'vehicle_masks': None, 'pedestrian_masks': None}}
        if 'init_node' in inputs:
            encodings['init_node'] = inputs['init_node']
            encodings['node_seq_gt'] = inputs['node_seq_gt']
            encodings['s_next'] = inputs['map_representation']['s_next']
            encodings['edge_type'] = inputs['map_representation']['edge_type']
        return encodings

    @staticmethod
    def variable_size_gru_encode(feat_embedding: torch.Tensor, masks: torch.Tensor, gru: nn.GRU) ->torch.Tensor:
        """
        Returns GRU encoding for a batch of inputs where each sample in the batch is a set of a variable number
        of sequences, of variable lengths.
        """
        masks_for_batching = ~masks[:, :, :, 0].bool()
        masks_for_batching = masks_for_batching.any(dim=-1).unsqueeze(2).unsqueeze(3)
        feat_embedding_batched = torch.masked_select(feat_embedding, masks_for_batching)
        feat_embedding_batched = feat_embedding_batched.view(-1, feat_embedding.shape[2], feat_embedding.shape[3])
        seq_lens = torch.sum(1 - masks[:, :, :, 0], dim=-1)
        seq_lens_batched = seq_lens[seq_lens != 0].cpu()
        if len(seq_lens_batched) != 0:
            feat_embedding_packed = pack_padded_sequence(feat_embedding_batched, seq_lens_batched, batch_first=True, enforce_sorted=False)
            _, encoding_batched = gru(feat_embedding_packed)
            encoding_batched = encoding_batched.squeeze(0)
            masks_for_scattering = masks_for_batching.squeeze(3).repeat(1, 1, encoding_batched.shape[-1])
            encoding = torch.zeros(masks_for_scattering.shape, device=device)
            encoding = encoding.masked_scatter(masks_for_scattering, encoding_batched)
        else:
            batch_size = feat_embedding.shape[0]
            max_num = feat_embedding.shape[1]
            hidden_state_size = gru.hidden_size
            encoding = torch.zeros((batch_size, max_num, hidden_state_size), device=device)
        return encoding

    @staticmethod
    def build_adj_mat(s_next, edge_type):
        """
        Builds adjacency matrix for GAT layers.
        """
        batch_size = s_next.shape[0]
        max_nodes = s_next.shape[1]
        max_edges = s_next.shape[2]
        adj_mat = torch.diag(torch.ones(max_nodes, device=device)).unsqueeze(0).repeat(batch_size, 1, 1).bool()
        dummy_vals = torch.arange(max_nodes, device=device).unsqueeze(0).unsqueeze(2).repeat(batch_size, 1, max_edges)
        dummy_vals = dummy_vals.float()
        s_next[edge_type == 0] = dummy_vals[edge_type == 0]
        batch_indices = torch.arange(batch_size).unsqueeze(1).unsqueeze(2).repeat(1, max_nodes, max_edges)
        src_indices = torch.arange(max_nodes).unsqueeze(0).unsqueeze(2).repeat(batch_size, 1, max_edges)
        adj_mat[batch_indices[:, :, :-1], src_indices[:, :, :-1], s_next[:, :, :-1].long()] = True
        adj_mat = adj_mat | torch.transpose(adj_mat, 1, 2)
        return adj_mat


class PolylineSubgraphs(PredictionEncoder):

    def __init__(self, args: Dict):
        """
        Polyline subgraph encoder from VectorNet (Gao et al., CVPR 2020).
        Has N encoder layers. Each layer encodes every feature in a polyline using an MLP with shared
        weights, followed by a permutation invariant aggregation operator (element-wise max used in the paper).
        Aggregated vector is concatenated with each independent feature encoding.
        Layer is repeated N times. Final encodings are passed through the permutation invariant
        aggregation operator to give polyline encodings.

        args to include
            'num_layers': int Number of repeated encoder layers
            'mlp_size':  int Width of MLP hidden layer
            'lane_feat_size': int Lane feature dimension
            'agent_feat_size': int Agent feature dimension

        """
        super().__init__()
        self.num_layers = args['num_layers']
        self.mlp_size = args['mlp_size']
        self.lane_feat_size = args['lane_feat_size']
        self.agent_feat_size = args['agent_feat_size']
        """
        Note: I'm not completely sure if VectorNet uses different MLPs for agents, map polylines and map polygons.
        The paper doesn't seem to mention this clearly. However, agents and map polylines will typically have different 
        attribute features. At least the first linear layer has to be different. 
        Shouldn't affect the global attention aggregator. All final feats will have the same dimensions.
        """
        lane_encoders = [nn.Linear(self.lane_feat_size + 2, self.mlp_size)]
        for n in range(1, self.num_layers):
            lane_encoders.append(nn.Linear(self.mlp_size * 2, self.mlp_size))
        self.lane_encoders = nn.ModuleList(lane_encoders)
        target_agent_encoders = [nn.Linear(self.agent_feat_size + 2, self.mlp_size)]
        for n in range(1, self.num_layers):
            target_agent_encoders.append(nn.Linear(self.mlp_size * 2, self.mlp_size))
        self.target_agent_encoders = nn.ModuleList(target_agent_encoders)
        surrounding_vehicle_encoders = [nn.Linear(self.agent_feat_size + 2, self.mlp_size)]
        for n in range(1, self.num_layers):
            surrounding_vehicle_encoders.append(nn.Linear(self.mlp_size * 2, self.mlp_size))
        self.surrounding_vehicle_encoders = nn.ModuleList(surrounding_vehicle_encoders)
        surrounding_ped_encoders = [nn.Linear(self.agent_feat_size + 2, self.mlp_size)]
        for n in range(1, self.num_layers):
            surrounding_ped_encoders.append(nn.Linear(self.mlp_size * 2, self.mlp_size))
        self.surrounding_ped_encoders = nn.ModuleList(surrounding_ped_encoders)
        self.layer_norm = nn.LayerNorm(self.mlp_size)
        self.relu = nn.ReLU()

    def forward(self, inputs: Dict) ->Dict:
        target_agent_feats = inputs['target_agent_representation']
        lane_feats = inputs['map_representation']['lane_node_feats']
        lane_masks = inputs['map_representation']['lane_node_masks']
        vehicle_feats = inputs['surrounding_agent_representation']['vehicles']
        vehicle_masks = inputs['surrounding_agent_representation']['vehicle_masks']
        ped_feats = inputs['surrounding_agent_representation']['pedestrians']
        ped_masks = inputs['surrounding_agent_representation']['pedestrian_masks']
        target_agent_feats = self.convert2vectornet_feat_format(target_agent_feats.unsqueeze(1))
        target_agent_enc, _ = self.encode(self.target_agent_encoders, target_agent_feats, torch.zeros_like(target_agent_feats))
        target_agent_enc = target_agent_enc.squeeze(1)
        lane_feats = self.convert2vectornet_feat_format(lane_feats)
        lane_masks = lane_masks[:, :, :-1, :]
        lane_enc, lane_masks = self.encode(self.lane_encoders, lane_feats, lane_masks)
        vehicle_feats = self.convert2vectornet_feat_format(vehicle_feats)
        vehicle_masks = vehicle_masks[:, :, :-1, :]
        vehicle_enc, vehicle_masks = self.encode(self.surrounding_vehicle_encoders, vehicle_feats, vehicle_masks)
        ped_feats = self.convert2vectornet_feat_format(ped_feats)
        ped_masks = ped_masks[:, :, :-1, :]
        ped_enc, ped_masks = self.encode(self.surrounding_ped_encoders, ped_feats, ped_masks)
        encodings = {'target_agent_encoding': target_agent_enc, 'context_encoding': {'combined': None, 'combined_masks': None, 'map': lane_enc, 'vehicles': vehicle_enc, 'pedestrians': ped_enc, 'map_masks': lane_masks, 'vehicle_masks': vehicle_masks, 'pedestrian_masks': ped_masks}}
        return encodings

    def encode(self, encoder_layers: nn.ModuleList, input_feats: torch.Tensor, masks: torch.Tensor) ->Tuple[torch.Tensor, torch.Tensor]:
        """
        Applies encoding layers to a given set of input feats
        """
        masks = masks[..., 0]
        masks[masks == 1] = -math.inf
        encodings = input_feats
        for n in range(len(encoder_layers)):
            encodings = self.relu(self.layer_norm(encoder_layers[n](encodings)))
            encodings = encodings + masks.unsqueeze(-1)
            agg_enc, _ = torch.max(encodings, dim=2)
            encodings = torch.cat((encodings, agg_enc.unsqueeze(2).repeat(1, 1, encodings.shape[2], 1)), dim=3)
            encodings[encodings == -math.inf] = 0
        agg_encoding, _ = torch.max(encodings, dim=2)
        masks[masks == -math.inf] = 1
        return agg_encoding, masks[..., 0]

    @staticmethod
    def convert2vectornet_feat_format(feats: torch.Tensor) ->torch.Tensor:
        """
        Helper function to convert a tensor of node features to the vectornet format.
        By default the datasets return node features of the format [x, y, attribute feats...].
        Vectornet uses the following format [x, y, x_next, y_next, attribute_feats]
        :param feats: Tensor of feats, shape [batch_size, max_polylines, max_len, feat_dim]
        :return: Tensor of updated feats, shape [batch_size, max_polylines, max_len, feat_dim + 2]
        """
        xy = feats[:, :, :-1, :2]
        xy_next = feats[:, :, 1:, :2]
        attr = feats[:, :, :-1, 2:]
        feats = torch.cat((xy, xy_next, attr), dim=3)
        return feats


class RasterEncoder(PredictionEncoder):

    def __init__(self, args: Dict):
        """
        CNN encoder for raster representation of HD maps and surrounding agent trajectories.

        args to include
            'backbone': str CNN backbone to use (resnet18, resnet34 or resnet50)
            'input_channels': int Size of scene features at each grid cell
            'use_positional_encoding: bool Whether or not to add positional encodings to final set of features
            'target_agent_feat_size': int Size of target agent state
        """
        super().__init__()
        resnet_backbones = {'resnet18': resnet18, 'resnet34': resnet34, 'resnet50': resnet50}
        resnet_model = resnet_backbones[args['backbone']](pretrained=False)
        conv1_new = nn.Conv2d(args['input_channels'], 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        modules = list(resnet_model.children())[:-2]
        modules[0] = conv1_new
        self.backbone = nn.Sequential(*modules)
        num_channels = 2048 if self.backbone == 'resnet50' else 512
        self.use_pos_enc = args['use_positional_encoding']
        if self.use_pos_enc:
            self.pos_enc = PositionalEncodingPermute2D(num_channels)
        self.target_agent_encoder = nn.Linear(args['target_agent_feat_size'], args['target_agent_enc_size'])
        self.relu = nn.ReLU()

    def forward(self, inputs: Dict) ->Dict:
        """
        Forward pass for raster encoder
        :param inputs: Dictionary with
            target_agent_representation: torch.Tensor with target agent state, shape[batch_size, target_agent_feat_size]
            surrounding_agent_representation: Rasterized BEV representation, shape [batch_size, 3, H, W]
            map_representation: Rasterized BEV representation, shape [batch_size, 3, H, W]
        :return encodings: Dictionary with
            'target_agent_encoding': torch.Tensor of shape [batch_size, 3],
            'context_encoding': torch.Tensor of shape [batch_size, N, backbone_feat_dim]
        """
        target_agent_representation = inputs['target_agent_representation']
        surrounding_agent_representation = inputs['surrounding_agent_representation']
        map_representation = inputs['map_representation']
        rasterized_input = torch.cat((map_representation, surrounding_agent_representation), dim=1)
        context_encoding = self.backbone(rasterized_input)
        if self.use_pos_enc:
            context_encoding = context_encoding + self.pos_enc(context_encoding)
        context_encoding = context_encoding.view(context_encoding.shape[0], context_encoding.shape[1], -1)
        context_encoding = context_encoding.permute(0, 2, 1)
        target_agent_enc = self.relu(self.target_agent_encoder(target_agent_representation))
        encodings = {'target_agent_encoding': target_agent_enc, 'context_encoding': {'combined': context_encoding, 'combined_masks': torch.zeros_like(context_encoding[..., 0]), 'map': None, 'vehicles': None, 'pedestrians': None, 'map_masks': None, 'vehicle_masks': None, 'pedestrian_masks': None}}
        return encodings


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (MTP,
     lambda: ([], {'args': _mock_config(agg_type=4, num_modes=4, op_len=4, use_variance=4, encoding_size=4, hidden_size=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Multipath,
     lambda: ([], {'args': _mock_config(agg_type=4, num_modes=4, op_len=4, use_variance=4, encoding_size=4, hidden_size=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
]

class Test_nachiket92_PGP(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

