backups/dev/utils/visualization.py

import math
import os
import torch
import pickle
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import numpy.typing as npt
import fnmatch
import io
import seaborn as sns
import matplotlib.axes as Axes
import matplotlib.transforms as mtransforms
from PIL import Image
from functools import wraps
from typing import Sequence, Union, Optional
from tqdm import tqdm
from typing import List, Literal
from argparse import ArgumentParser
from scipy.ndimage.filters import gaussian_filter
from matplotlib.patches import FancyBboxPatch, Polygon, Rectangle, Circle
from matplotlib.collections import LineCollection
from torch_geometric.data import HeteroData, Dataset
from waymo_open_dataset.protos import scenario_pb2

from dev.utils.func import CONSOLE
from dev.modules.attr_tokenizer import Attr_Tokenizer
from dev.datasets.preprocess import TokenProcessor, cal_polygon_contour, AGENT_TYPE
from dev.datasets.scalable_dataset import WaymoTargetBuilder


__all__ = ['plot_occ_grid', 'plot_interact_edge', 'plot_map_edge', 'plot_insert_grid', 'plot_binary_map',
           'plot_map_token', 'plot_prob_seed', 'plot_scenario', 'get_heatmap', 'draw_heatmap', 'plot_val', 'plot_tokenize']


def safe_run(func):

    @wraps(func)
    def wrapper1(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            print(e)
            return

    @wraps(func)
    def wrapper2(*args, **kwargs):
        return func(*args, **kwargs)

    if int(os.getenv('DEBUG', 0)):
        return wrapper2
    else:
        return wrapper1


@safe_run
def plot_occ_grid(scenario_id, occ, gt_occ=None, save_path='', mode='agent', prefix=''):

    def generate_box_edges(matrix, find_value=1):
        y, x = np.where(matrix == find_value)
        edges = []

        for xi, yi in zip(x, y):
            edges.append([(xi - 0.5, yi - 0.5), (xi + 0.5, yi - 0.5)])
            edges.append([(xi + 0.5, yi - 0.5), (xi + 0.5, yi + 0.5)])
            edges.append([(xi + 0.5, yi + 0.5), (xi - 0.5, yi + 0.5)])
            edges.append([(xi - 0.5, yi + 0.5), (xi - 0.5, yi - 0.5)])
    
        return edges

    os.makedirs(save_path, exist_ok=True)
    n = int(math.sqrt(occ.shape[-1]))

    plot_n = 3
    plot_t = 5

    occ_list = []
    for i in range(plot_n):
        for j in range(plot_t):
            occ_list.append(occ[i, j].reshape(n, n))

    occ_gt_list = []
    if gt_occ is not None:
        for i in range(plot_n):
            for j in range(plot_t):
                occ_gt_list.append(gt_occ[i, j].reshape(n, n))

    row_labels = [f'n={n}' for n in range(plot_n)]
    col_labels = [f't={t}' for t in range(plot_t)]

    fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
    plt.subplots_adjust(wspace=0.1, hspace=0.1)

    for i, ax in enumerate(axes.flat):
        # NOTE: do not set vmin and vamx!
        ax.imshow(occ_list[i], cmap='viridis', interpolation='nearest')
        ax.axis('off')

        if occ_gt_list:
            gt_edges = generate_box_edges(occ_gt_list[i])
            gts = LineCollection(gt_edges, colors='blue', linewidths=0.5)
            ax.add_collection(gts)
            insert_edges = generate_box_edges(occ_gt_list[i], find_value=-1)
            inserts = LineCollection(insert_edges, colors='red', linewidths=0.5)
            ax.add_collection(inserts)

        ax.add_patch(plt.Rectangle((-0.5, -0.5), occ_list[i].shape[1], occ_list[i].shape[0],
                                linewidth=2, edgecolor='black', facecolor='none'))

    for i, ax in enumerate(axes[:, 0]):
        ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
                    fontsize=12, ha="right", va="center", rotation=0)

    for j, ax in enumerate(axes[0, :]):
        ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
                    fontsize=12, ha="center", va="bottom")

    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_occ_{mode}.png'), dpi=500, bbox_inches='tight')
    plt.close()


@safe_run
def plot_interact_edge(edge_index, scenario_ids, batch_sizes, num_seed, num_step, save_path='interact_edge_map',
                        **kwargs):

    num_batch = len(scenario_ids)
    batches = torch.cat([
        torch.arange(num_batch).repeat_interleave(repeats=batch_sizes, dim=0),
        torch.arange(num_batch).repeat_interleave(repeats=num_seed, dim=0),
    ], dim=0).repeat(num_step).numpy()

    num_agent = batch_sizes.sum() + num_seed * num_batch
    batch_sizes = torch.nn.functional.pad(batch_sizes, (1, 0), mode='constant', value=0)
    ptr = torch.cumsum(batch_sizes, dim=0)
    # assume difference scenarios and different timestep have the same number of seed agents
    ptr_seed = torch.tensor(np.array([0] + [num_seed] * num_batch), device=ptr.device)

    all_av_index = None
    if 'av_index' in kwargs:
        all_av_index = kwargs.pop('av_index').cpu() - ptr[:-1]

    is_bos = np.zeros((batch_sizes.sum(), num_step)).astype(np.bool_)
    if 'is_bos' in kwargs:
        is_bos = kwargs.pop('is_bos').cpu().numpy()

    src_index = torch.unique(edge_index[1])
    for idx, src in enumerate(tqdm(src_index)):

        src_batch = batches[src]

        src_row = src % num_agent
        if src_row // batch_sizes.sum() > 0:
            seed_row = src_row % batch_sizes.sum() - ptr_seed[src_batch]
            src_row = batch_sizes[src_batch + 1] + seed_row
        else:
            src_row = src_row - ptr[src_batch]

        src_col = src // (num_agent)
        src_mask = np.zeros((batch_sizes[src_batch + 1] + num_seed, num_step))
        src_mask[src_row, src_col] = 1

        tgt_mask = np.zeros((src_mask.shape[0], num_step))
        tgt_index = edge_index[0, edge_index[1] == src]
        for tgt in tgt_index:

            tgt_batch = batches[tgt]

            tgt_row = tgt % num_agent
            if tgt_row // batch_sizes.sum() > 0:
                seed_row = tgt_row % batch_sizes.sum() - ptr_seed[tgt_batch]
                tgt_row = batch_sizes[tgt_batch + 1] + seed_row
            else:
                tgt_row = tgt_row - ptr[tgt_batch]

            tgt_col = tgt // num_agent
            tgt_mask[tgt_row, tgt_col] = 1
            assert tgt_batch == src_batch

        selected_step = tgt_mask.sum(axis=0) > 0
        if selected_step.sum() > 1:
            print(f"\nidx={idx}", src.item(), src_row.item(), src_col.item())
            print(selected_step)
            print(edge_index[:, edge_index[1] == src].tolist())

        if all_av_index is not None:
            kwargs['av_index'] = int(all_av_index[src_batch])

        t = kwargs.get('t', src_col)
        n = kwargs.get('n', 0)
        is_bos_batch = is_bos[ptr[src_batch] : ptr[src_batch + 1]]
        plot_binary_map(src_mask, tgt_mask, save_path, suffix=f'_{scenario_ids[src_batch]}_{t:02d}_{n:02d}_{idx:04d}',
                        is_bos=is_bos_batch, **kwargs)


@safe_run
def plot_map_edge(edge_index, pos_a, data, save_path='map_edge_map'):

    map_points = data['map_point']['position'][:, :2].cpu().numpy()
    token_pos = data['pt_token']['position'][:, :2].cpu().numpy()
    token_heading = data['pt_token']['orientation'].cpu().numpy()
    num_pt = token_pos.shape[0]
    
    agent_index = torch.unique(edge_index[1])
    for i in tqdm(agent_index):
        xy = pos_a[i].cpu().numpy()
        pt_index = edge_index[0, edge_index[1] == i].cpu().numpy()
        pt_index = pt_index % num_pt

        plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
        _, ax = plt.subplots()
        ax.set_axis_off()

        plot_map_token(ax, map_points, token_pos[pt_index], token_heading[pt_index], colors='blue')

        ax.scatter(xy[0], xy[1], s=0.5, c='red', edgecolors='none')

        os.makedirs(save_path, exist_ok=True)
        plt.savefig(os.path.join(save_path, f'map_{i}.png'), dpi=600, bbox_inches='tight')
        plt.close()


def get_heatmap(x, y, prob, s=3, bins=1000):
    heatmap, xedges, yedges = np.histogram2d(x, y, bins=bins, weights=prob, density=True)

    heatmap = gaussian_filter(heatmap, sigma=s)

    extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
    return heatmap.T, extent


@safe_run
def draw_heatmap(vector, vector_prob, gt_idx):
    fig, ax = plt.subplots(figsize=(10, 10))
    vector_prob = vector_prob.cpu().numpy()

    for j in range(vector.shape[0]):
        if j in gt_idx:
            color = (0, 0, 1)
        else:
            grey_scale = max(0, 0.9 - vector_prob[j])
            color = (0.9, grey_scale, grey_scale)

        # if lane[j, k, -1] == 0: continue
        x0, y0, x1, y1, = vector[j, :4]
        ax.plot((x0, x1), (y0, y1), color=color, linewidth=2)

    return plt


@safe_run
def plot_insert_grid(scenario_id, prob, grid, ego_pos, map, save_path='', prefix='', inference=False, indices=None,
                     all_t_in_one=False):

    """
        prob: float array of shape (num_step, num_grid)
        grid: float array of shape (num_grid, 2)
    """

    os.makedirs(save_path, exist_ok=True)

    n = int(math.sqrt(prob.shape[1]))

    # grid = grid[:, np.newaxis] + ego_pos[np.newaxis, ...]
    for t in range(ego_pos.shape[0]):

        plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
        _, ax = plt.subplots()

        # plot probability
        prob_t = prob[t].reshape(n, n)
        plt.imshow(prob_t, cmap='viridis', interpolation='nearest')

        if indices is not None:
            indice = indices[t]

            if isinstance(indice, (int, float, np.int_)):
                indice = [indice]

            for _indice in indice:
                if _indice == -1: continue

                row = _indice // n
                col = _indice % n

                rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
                ax.add_patch(rect)

        ax.grid(False)
        ax.set_aspect('equal', adjustable='box')

        plt.title('Prob of Rel Position Grid')
        plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_heat_map_{t}.png'), dpi=300, bbox_inches='tight')
        plt.close()

        if all_t_in_one:
            break


@safe_run
def plot_insert_grid(scenario_id, prob, indices=None, save_path='', prefix='', inference=False):

    """
        prob: float array of shape (num_seed, num_step, num_grid)
        grid: float array of shape (num_grid, 2)
    """

    os.makedirs(save_path, exist_ok=True)

    n = int(math.sqrt(prob.shape[-1]))

    plot_n = 3
    plot_t = 5

    prob_list = []
    for i in range(plot_n):
        for j in range(plot_t):
            prob_list.append(prob[i, j].reshape(n, n))

    indice_list = []
    if indices is not None:
        for i in range(plot_n):
            for j in range(plot_t):
                indice_list.append(indices[i, j])

    row_labels = [f'n={n}' for n in range(plot_n)]
    col_labels = [f't={t}' for t in range(plot_t)]

    fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
    fig.suptitle('Prob of Insert Position Grid')
    plt.subplots_adjust(wspace=0.1, hspace=0.1)

    for i, ax in enumerate(axes.flat):
        ax.imshow(prob_list[i], cmap='viridis', interpolation='nearest')
        ax.axis('off')

        if indice_list:
            row = indice_list[i] // n
            col = indice_list[i] % n
            rect = Rectangle((col - .5, row - .5), 1, 1, edgecolor='red', facecolor='none', lw=2)
            ax.add_patch(rect)

        ax.add_patch(plt.Rectangle((-0.5, -0.5), prob_list[i].shape[1], prob_list[i].shape[0],
                                linewidth=2, edgecolor='black', facecolor='none'))

    for i, ax in enumerate(axes[:, 0]):
        ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
                    fontsize=12, ha="right", va="center", rotation=0)

    for j, ax in enumerate(axes[0, :]):
        ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
                    fontsize=12, ha="center", va="bottom")

    ax.grid(False)
    ax.set_aspect('equal', adjustable='box')

    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_insert_map.png'), dpi=500, bbox_inches='tight')
    plt.close()


@safe_run
def plot_binary_map(src_mask, tgt_mask, save_path='', suffix='', av_index=None, is_bos=None, **kwargs):

    from matplotlib.colors import ListedColormap
    os.makedirs(save_path, exist_ok=True)

    fig, axes = plt.subplots(1, 2, figsize=(10, 8))

    title = []
    if kwargs.get('t', None) is not None:
        t = kwargs['t']
        title.append(f't={t}')

    if kwargs.get('n', None) is not None:
        n = kwargs['n']
        title.append(f'n={n}')

    plt.title(' '.join(title))

    cmap = ListedColormap(['white', 'green'])
    axes[0].imshow(src_mask, cmap=cmap, interpolation='nearest')

    cmap = ListedColormap(['white', 'orange'])
    axes[1].imshow(tgt_mask, cmap=cmap, interpolation='nearest')

    if av_index is not None:
        rect = Rectangle((-0.5, av_index - 0.5), src_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
        axes[0].add_patch(rect)
        rect = Rectangle((-0.5, av_index - 0.5), tgt_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
        axes[1].add_patch(rect)

    if is_bos is not None:
        rows, cols = np.where(is_bos)
        for row, col in zip(rows, cols):
            rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
            axes[0].add_patch(rect)
            rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
            axes[1].add_patch(rect)

    for ax in axes:
        ax.set_xticks(range(src_mask.shape[1] + 1), minor=False)
        ax.set_yticks(range(src_mask.shape[0] + 1), minor=False)
        ax.grid(which='major', color='gray', linestyle='--', linewidth=0.5)

    plt.savefig(os.path.join(save_path, f'map{suffix}.png'), dpi=300, bbox_inches='tight')
    plt.close()


@safe_run
def plot_prob_seed(scenario_id, prob, save_path, prefix='', indices=None):

    os.makedirs(save_path, exist_ok=True)

    plt.figure(figsize=(8, 5))
    plt.imshow(prob, cmap='viridis', aspect='auto')
    plt.colorbar()

    plt.title('Seed Probability')

    if indices is not None:

        for col in range(indices.shape[1]):
            for row in indices[:, col]:

                if row == -1: continue

                rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
                plt.gca().add_patch(rect)

    plt.tight_layout()
    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_prob_seed.png'), dpi=300, bbox_inches='tight')
    plt.close()


@safe_run
def plot_raw():
    plt.figure(figsize=(30, 30))
    plt.rcParams['axes.facecolor']='white'

    data_path = '/u/xiuyu/work/dev4/data/waymo/scenario/training'
    os.makedirs("data/vis/raw/0/", exist_ok=True)
    file_list = os.listdir(data_path)

    for cnt_file, file in enumerate(file_list):
        file_path = os.path.join(data_path, file)
        dataset = tf.data.TFRecordDataset(file_path, compression_type='')
        for scenario_idx, data in enumerate(dataset):
            scenario = scenario_pb2.Scenario()
            scenario.ParseFromString(bytearray(data.numpy()))
            tqdm.write(f"scenario id: {scenario.scenario_id}")

            # draw maps
            for i in range(len(scenario.map_features)):

                # draw lanes
                if str(scenario.map_features[i].lane) != '':
                    line_x = [z.x for z in scenario.map_features[i].lane.polyline]
                    line_y = [z.y for z in scenario.map_features[i].lane.polyline]
                    plt.scatter(line_x, line_y, c='g', s=5)
                    plt.text(line_x[0], line_y[0], str(scenario.map_features[i].id), fontdict={'family': 'serif', 'size': 20, 'color': 'green'})

                # draw road_edge
                if str(scenario.map_features[i].road_edge) != '':
                    road_edge_x = [polyline.x for polyline in scenario.map_features[i].road_edge.polyline]
                    road_edge_y = [polyline.y for polyline in scenario.map_features[i].road_edge.polyline]
                    plt.scatter(road_edge_x, road_edge_y)
                    plt.text(road_edge_x[0], road_edge_y[0], scenario.map_features[i].road_edge.type, fontdict={'family': 'serif', 'size': 20, 'color': 'black'})
                    if scenario.map_features[i].road_edge.type == 2:
                        plt.scatter(road_edge_x, road_edge_y, c='k')
                    elif scenario.map_features[i].road_edge.type == 3:
                        plt.scatter(road_edge_x, road_edge_y, c='purple')
                        print(scenario.map_features[i].road_edge)
                    else:
                        plt.scatter(road_edge_x, road_edge_y, c='k')

                # draw road_line
                if str(scenario.map_features[i].road_line) != '':
                    road_line_x = [j.x for j in scenario.map_features[i].road_line.polyline]
                    road_line_y = [j.y for j in scenario.map_features[i].road_line.polyline]
                    if scenario.map_features[i].road_line.type == 7:
                        plt.plot(road_line_x, road_line_y, c='y')
                    elif scenario.map_features[i].road_line.type == 8:
                        plt.plot(road_line_x, road_line_y, c='y') 
                    elif scenario.map_features[i].road_line.type == 6:
                        plt.plot(road_line_x, road_line_y, c='y')
                    elif scenario.map_features[i].road_line.type == 1:
                        for i in range(int(len(road_line_x) / 7)):
                            plt.plot(road_line_x[i * 7 : 5 + i * 7], road_line_y[i * 7 : 5 + i * 7], color='w')
                    elif scenario.map_features[i].road_line.type == 2:
                        plt.plot(road_line_x, road_line_y, c='w')
                    else:
                        plt.plot(road_line_x, road_line_y, c='w')
            
            # draw tracks
            scenario_has_invalid_tracks = False
            for i in range(len(scenario.tracks)):
                traj_x = [center.center_x for center in scenario.tracks[i].states]
                traj_y = [center.center_y for center in scenario.tracks[i].states]
                head = [center.heading for center in scenario.tracks[i].states]
                valid = [center.valid for center in scenario.tracks[i].states]
                print(valid)
                if i == scenario.sdc_track_index:
                    plt.scatter(traj_x[0], traj_y[0], s=140, c='r', marker='s')
                    plt.scatter([x for x, v in zip(traj_x, valid) if v],
                                [y for y, v in zip(traj_y, valid) if v], s=14, c='r')
                    plt.scatter([x for x, v in zip(traj_x, valid) if not v],
                                [y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
                else:
                    plt.scatter(traj_x[0], traj_y[0], s=140, c='k', marker='s')
                    plt.scatter([x for x, v in zip(traj_x, valid) if v],
                                [y for y, v in zip(traj_y, valid) if v], s=14, c='b')
                    plt.scatter([x for x, v in zip(traj_x, valid) if not v],
                                [y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
                if valid.count(False) > 0:
                    scenario_has_invalid_tracks = True
            if scenario_has_invalid_tracks:
                plt.savefig(f"scenario_{scenario_idx}_{scenario.scenario_id}.png")
                plt.clf()
                breakpoint()
        break


colors = [
    ('#1f77b4', '#1a5a8a'),  # blue
    ('#2ca02c', '#217721'),  # green
    ('#ff7f0e', '#cc660b'),  # orange
    ('#9467bd', '#6f4a91'),  # purple
    ('#d62728', '#a31d1d'),  # red
    ('#000000', '#000000'),  # black
]

@safe_run
def plot_gif():
    data_path = "/u/xiuyu/work/dev4/data/waymo_processed/training"
    os.makedirs("data/vis/processed/0/gif", exist_ok=True)
    file_list = os.listdir(data_path)

    for scenario_idx, file in tqdm(enumerate(file_list), leave=False, desc="Scenario"):

        fig, ax = plt.subplots()
        ax.set_axis_off()

        file_path = os.path.join(data_path, file)
        data = pickle.load(open(file_path, "rb"))
        scenario_id = data['scenario_id']

        save_path = os.path.join("data/vis/processed/0/gif",
                                 f"scenario_{scenario_idx}_{scenario_id}.gif")
        if os.path.exists(save_path):
            tqdm.write(f"Skipped {save_path}.")
            continue

        # draw maps
        ax.scatter(data['map_point']['position'][:, 0],
                   data['map_point']['position'][:, 1], s=0.2, c='black', edgecolors='none')

        # draw agents
        agent_data = data['agent']
        av_index = agent_data['av_index']
        position = agent_data['position'] # (num_agent, 91, 3)
        heading = agent_data['heading'] # (num_agent, 91)
        shape = agent_data['shape'] # (num_agent, 91, 3)
        category = agent_data['category'] # (num_agent,)
        valid_mask = (position[..., 0] != 0) & (position[..., 1] != 0) # (num_agent, 91)

        num_agent = valid_mask.shape[0]
        num_timestep = position.shape[1]
        is_av = np.arange(num_agent) == int(av_index)

        is_blue = valid_mask.sum(axis=1) == num_timestep
        is_green = ~valid_mask[:, 0] & valid_mask[:, -1]
        is_orange = valid_mask[:, 0] & ~valid_mask[:, -1]
        is_purple = (valid_mask.sum(axis=1) != num_timestep
                    ) & (~is_green) & (~is_orange)
        agent_colors = np.zeros((num_agent,))
        agent_colors[is_blue] = 1
        agent_colors[is_green] = 2
        agent_colors[is_orange] = 3
        agent_colors[is_purple] = 4
        agent_colors[is_av] = 5

        veh_mask = category == 1
        ped_mask = category == 2
        cyc_mask = category == 3
        shape[veh_mask, :, 1] = 1.8
        shape[veh_mask, :, 0] = 1.8
        shape[ped_mask, :, 1] = 0.5
        shape[ped_mask, :, 0] = 0.5
        shape[cyc_mask, :, 1] = 1.0
        shape[cyc_mask, :, 0] = 1.0

        fig_paths = []
        for tid in tqdm(range(num_timestep), leave=False, desc="Timestep"):
            current_valid_mask = valid_mask[:, tid]
            xs = position[current_valid_mask, tid, 0]
            ys = position[current_valid_mask, tid, 1]
            widths = shape[current_valid_mask, tid, 1]
            lengths = shape[current_valid_mask, tid, 0]
            angles = heading[current_valid_mask, tid]
            current_agent_colors = agent_colors[current_valid_mask]

            drawn_agents = []
            contours = cal_polygon_contour(xs, ys, angles, widths, lengths) # (num_agent, 4, 2)
            contours = np.concatenate([contours, contours[:, 0:1]], axis=1) # (num_agent, 5, 2)
            for x, y, width, length, angle, color_type in zip(
                xs, ys, widths, lengths, angles, current_agent_colors):
                agent = plt.Rectangle((x, y), width, length, angle=((angle + np.pi / 2) / np.pi * 360) % 360,
                                      linewidth=0.2,
                                      facecolor=colors[int(color_type) - 1][0],
                                      edgecolor=colors[int(color_type) - 1][1])
                ax.add_patch(agent)
                drawn_agents.append(agent)
            plt.gca().set_aspect('equal', adjustable='box')
            # for contour, color_type in zip(contours, agent_colors):
            #     drawn_agent = ax.plot(contour[:, 0], contour[:, 1])
            #     drawn_agents.append(drawn_agent)

            fig_path = os.path.join("data/vis/processed/0/",
                                    f"scenario_{scenario_idx}_{scenario_id}_{tid}.png")
            plt.savefig(fig_path, dpi=600)
            fig_paths.append(fig_path)

            for drawn_agent in drawn_agents:
                drawn_agent.remove()

        plt.close()

        # generate gif
        import imageio.v2 as imageio
        images = []
        for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
            images.append(imageio.imread(fig_path))
        imageio.mimsave(save_path, images, duration=0.1)


@safe_run
def plot_map_token(ax: Axes, map_points: npt.NDArray, token_pos: npt.NDArray, token_heading: npt.NDArray, colors: Union[str, npt.NDArray]=None):

    plot_map(ax, map_points)

    x, y = token_pos[:, 0], token_pos[:, 1]
    u = np.cos(token_heading)
    v = np.sin(token_heading)

    if colors is None:
        colors = np.random.rand(x.shape[0], 3)
    ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=0.2, color=colors, width=0.005,
               headwidth=0.2, headlength=2)
    ax.scatter(x, y, color='blue', s=0.2, edgecolors='none')
    ax.axis("equal")


@safe_run
def plot_map(ax: Axes, map_points: npt.NDArray, color='black'):
    ax.scatter(map_points[:, 0], map_points[:, 1], s=0.5, c=color, edgecolors='none')

    xmin = np.min(map_points[:, 0])
    xmax = np.max(map_points[:, 0])
    ymin = np.min(map_points[:, 1])
    ymax = np.max(map_points[:, 1])
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)


@safe_run
def plot_agent(ax: Axes, xy: Sequence[float], heading: float, type: str, state, is_av: bool=False,
               pl2seed_radius: float=25., attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], **kwargs):

    if type == 'veh':
        length = 4.3
        width = 1.8
        size = 1.0
    elif type == 'ped':
        length = 0.5
        width = 0.5
        size = 0.1
    elif type == 'cyc':
        length = 1.9
        width = 0.5
        size = 0.3
    else:
        raise ValueError(f"Unsupported agent type {type}")

    if kwargs.get('label', None) is not None:
        ax.text(
            xy[0] + 1.5, xy[1] + 1.5,
            kwargs['label'], fontsize=2, color="darkred", ha="center", va="center"
        )

    patch = FancyBboxPatch([-length / 2, -width / 2], length, width, linewidth=.2, **kwargs)
    transform = (
        mtransforms.Affine2D().rotate(heading).translate(xy[0], xy[1])
        + ax.transData
    )
    patch.set_transform(transform)

    kwargs['label'] = None
    angles = [0, 2 * np.pi / 3, np.pi, 4 * np.pi / 3]
    pts = np.stack([size * np.cos(angles), size * np.sin(angles)], axis=-1)
    center_patch = Polygon(pts, zorder=10., linewidth=.2, **kwargs)
    center_patch.set_transform(transform)

    ax.add_patch(patch)
    ax.add_patch(center_patch)

    if is_av:

        if attr_tokenizer is not None:

            circle_patch = Circle(
                (xy[0], xy[1]), pl2seed_radius, linewidth=0.5, edgecolor='gray', linestyle='--', facecolor='none'
            )
            ax.add_patch(circle_patch)

            grid = attr_tokenizer.get_grid(torch.tensor(np.array(xy)).float(),
                                           torch.tensor(np.array([heading])).float()).numpy()[0] # (num_grid, 2)
            ax.scatter(grid[:, 0], grid[:, 1], s=0.3, c='blue', edgecolors='none')
            ax.text(grid[0, 0], grid[0, 1], 'Front', fontsize=2, color='darkred', ha='center', va='center')
            ax.text(grid[-1, 0], grid[-1, 1], 'Back', fontsize=2, color='darkred', ha='center', va='center')

        if enter_index:
            for i in enter_index:
                ax.plot(grid[int(i), 0], grid[int(i), 1], marker='x', color='red', markersize=1)

    return patch, center_patch


@safe_run
def plot_all(map, xs, ys, angles, types, colors, is_avs, pl2seed_radius: float=25.,
             attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], labels: list=[], **kwargs):

    plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
    _, ax = plt.subplots()
    ax.set_axis_off()

    plot_map(ax, map)

    if not labels:
        labels = [None] * xs.shape[0]

    for x, y, angle, type, color, label, is_av in zip(xs, ys, angles, types, colors, labels, is_avs):
        assert type in ('veh', 'ped', 'cyc'), f"Unsupported type {type}."
        plot_agent(ax, [x, y], angle.item(), type, None, is_av, facecolor=color, edgecolor='k', label=label,
                   pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, enter_index=enter_index)

    ax.grid(False)
    ax.set_aspect('equal', adjustable='box')

    # ! set plot limit if need
    if kwargs.get('limit_size', None):
        cx = float(xs[is_avs])
        cy = float(ys[is_avs])

        lx, ly = kwargs['limit_size']
        xmin, xmax = cx - lx, cx + lx
        ymin, ymax = cy - ly, cy + ly

        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)

    # ax.legend(loc='best', frameon=True)

    pil_image = None
    if kwargs.get('save_path', None):
        plt.savefig(kwargs['save_path'], dpi=600, bbox_inches="tight")

    else:
        # ！convert to PIL image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=600, bbox_inches='tight')
        buf.seek(0)
        pil_image = Image.open(buf).convert('RGB')

    plt.close()

    return pil_image


@safe_run
def plot_file(gt_folder: str,
              folder: Optional[str] = None,
              files: Optional[str] = None,
              save_gif: bool = True,
              batch_idx: Optional[int] = None,
              time_idx: Optional[List[int]] = None,
              limit_size: Optional[List[int]] = None,
              **kwargs,
    ) -> List[Image.Image]:

    from dev.metrics.compute_metrics import _unbatch

    shift = 5

    if files is None:
        assert os.path.exists(folder), f'Path {folder} does not exist.'
        files = list(fnmatch.filter(os.listdir(folder), 'idx_*_rollouts.pkl'))
        CONSOLE.log(f'Found {len(files)} rollouts files from {folder}.')

    if folder is None:
        assert os.path.exists(files), f'Path {files} does not exist.'
        folder = os.path.dirname(files)
        files = [files]

    parent, folder_name = os.path.split(folder.rstrip(os.sep))
    if save_gif:
        save_path = os.path.join(parent, f'{folder_name}_plots')
    else:
        save_path = None

    file_outs = []
    for file in (pbar := tqdm(files, leave=False, desc='Plotting files ...')):
        pbar.set_postfix(file=file)

        with open(os.path.join(folder, file), 'rb') as f:
            preds = pickle.load(f)

        scenario_ids = preds['_scenario_id']
        agent_batch = preds['agent_batch']
        agent_id = _unbatch(preds['agent_id'], agent_batch)
        preds_traj = _unbatch(preds['pred_traj'], agent_batch)
        preds_head = _unbatch(preds['pred_head'], agent_batch)
        preds_type = _unbatch(preds['pred_type'], agent_batch)
        if 'pred_state' in preds:
            preds_state = _unbatch(preds['pred_state'], agent_batch)
        else:
            preds_state = tuple([torch.ones((*traj.shape[:2], traj.shape[2] // shift)) for traj in preds_traj])  # [n_agent, n_rollout, n_step2Hz]
        preds_valid = _unbatch(preds['pred_valid'], agent_batch)

        # ! fetch certain scenario
        if batch_idx is not None:
            scenario_ids = scenario_ids[batch_idx : batch_idx + 1]
            agent_id = (agent_id[batch_idx],)
            preds_traj = (preds_traj[batch_idx],)
            preds_head = (preds_head[batch_idx],)
            preds_type = (preds_type[batch_idx],)
            preds_state = (preds_state[batch_idx],)
            preds_valid = (preds_valid[batch_idx],)

        scenario_outs = []
        for i, scenario_id in enumerate(scenario_ids):
            n_agent, n_rollouts = preds_traj[0].shape[:2]

            rollout_outs = []
            for j in range(n_rollouts):  # 1
                pred = dict(scenario_id=[scenario_id],
                            pred_traj=preds_traj[i][:, j],
                            pred_head=preds_head[i][:, j],
                            pred_state=(
                                torch.cat([torch.zeros(n_agent, 1), preds_state[i][:, j].repeat_interleave(repeats=shift, dim=-1)],
                                          dim=1)
                            ),
                            pred_type=preds_type[i][:, j],
                    )

                # NOTE: hard code!!!
                if 'av_id' in preds:
                    av_index = agent_id[i][:, 0].tolist().index(preds['av_id'])
                else:
                    av_index = n_agent - 1

                # ! load logged data
                data_path = os.path.join(gt_folder, 'validation', f'{scenario_id}.pkl')
                with open(data_path, 'rb') as f:
                    data = pickle.load(f)

                rollout_outs.append(
                    plot_val(data, pred,
                             av_index=av_index,
                             save_path=save_path,
                             save_gif=save_gif,
                             time_idx=time_idx,
                             limit_size=limit_size,
                             **kwargs
                    )
                )

            scenario_outs.append(rollout_outs)
        file_outs.append(scenario_outs)

    return file_outs


@safe_run
def plot_val(data: Union[dict, str], pred: dict, av_index: int, save_path: str, suffix: str='',
             pl2seed_radius: float=75., attr_tokenizer=None, **kwargs):

    if isinstance(data, str):
        assert data.endswith('.pkl'), f'Got invalid data path {data}.'
        assert os.path.exists(data), f'Path {data} does not exist.'
        with open(data, 'rb') as f:
            data = pickle.load(f)

    map_point = data['map_point']['position'].cpu().numpy()

    scenario_id = pred['scenario_id'][0]
    pred_traj = pred['pred_traj'].cpu().numpy() # (num_agent, num_future_step, 2)
    pred_type = list(map(lambda i: AGENT_TYPE[i], pred['pred_type'].tolist()))
    pred_state = pred['pred_state'].cpu().numpy()
    pred_head = pred['pred_head'].cpu().numpy()
    ids = np.arange(pred_traj.shape[0])

    if 'agent_labels' in pred:
        kwargs.update(agent_labels=pred['agent_labels'])

    return plot_scenario(scenario_id, map_point, pred_traj, pred_head, pred_state, pred_type,
                         av_index=av_index, ids=ids, save_path=save_path, suffix=suffix,
                         pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, **kwargs)


@safe_run
def plot_scenario(scenario_id: str,
                  map_data: npt.NDArray,
                  traj: npt.NDArray,
                  heading: npt.NDArray,
                  state: npt.NDArray,
                  types: List[str],
                  av_index: int,
                  color_type: Literal['state', 'type', 'seed', 'insert']='seed',
                  state_type: List[str]=['invalid', 'valid', 'enter', 'exit'],
                  plot_enter: bool=False,
                  suffix: str='',
                  pl2seed_radius: float=25.,
                  attr_tokenizer: Attr_Tokenizer=None,
                  enter_index: List[list] = [],
                  save_gif: bool=True,
                  tokenized: bool=False,
                  agent_labels: List[List[Optional[str]]] = [],
                  **kwargs):

    num_historical_steps = 11
    shift = 5
    num_agent, num_timestep = traj.shape[:2]

    if tokenized:
        num_historical_steps = 2
        shift = 1

    if (
        'save_path' in kwargs
        and kwargs['save_path'] != ''
        and kwargs['save_path'] != None
    ):
        os.makedirs(kwargs['save_path'], exist_ok=True)
        save_id = int(max([0] + list(map(lambda fname: int(fname.split("_")[-1]),
                                        filter(lambda fname: fname.startswith(scenario_id)
                                               and os.path.isdir(os.path.join(kwargs['save_path'], fname)),
                                        os.listdir(kwargs['save_path'])))))) + 1
        os.makedirs(f"{kwargs['save_path']}/{scenario_id}_{str(save_id).zfill(3)}", exist_ok=True)

        if save_id > 1:
            try:
                import shutil
                shutil.rmtree(f"{kwargs['save_path']}/{scenario_id}_{str(save_id - 1).zfill(3)}")
            except:
                pass

    visible_mask = state != state_type.index('invalid')
    if not plot_enter:
        visible_mask &= (state != state_type.index('enter'))

    last_valid_step = visible_mask.shape[1] - 1 - torch.argmax(torch.Tensor(visible_mask).flip(dims=[1]).long(), dim=1)
    ids = None
    if 'ids' in kwargs:
        ids = kwargs['ids']
        last_valid_step = {int(ids[i]): int(last_valid_step[i]) for i in range(len(ids))}

    # agent colors
    agent_colors = np.zeros((num_agent, num_timestep, 3))

    agent_palette = sns.color_palette('husl', n_colors=7)
    state_colors = {state: np.array(agent_palette[i]) for i, state in enumerate(state_type)}
    seed_colors = {seed: np.array(agent_palette[i]) for i, seed in enumerate(['existing', 'entered', 'exited'])}

    if color_type == 'state':
        for t in range(state.shape[1]):
            agent_colors[state[:, t] == state_type.index('invalid'), t * shift : (t + 1) * shift] = state_colors['invalid']
            agent_colors[state[:, t] == state_type.index('valid'), t * shift : (t + 1) * shift] = state_colors['valid']
            agent_colors[state[:, t] == state_type.index('enter'), t * shift : (t + 1) * shift] = state_colors['enter']
            agent_colors[state[:, t] == state_type.index('exit'), t * shift : (t + 1) * shift] = state_colors['exit']

    if color_type == 'seed':
        agent_colors[:, :] = seed_colors['existing']
        is_exited = np.any(state[:, num_historical_steps - 1:] == state_type.index('exit'), axis=-1)
        is_entered = np.any(state[:, num_historical_steps - 1:] == state_type.index('enter'), axis=-1)
        is_entered[av_index + 1:] = True  # NOTE: hard code, need improvment
        agent_colors[is_exited, :] = seed_colors['exited']
        agent_colors[is_entered, :] = seed_colors['entered']

    if color_type == 'insert':
        agent_colors[:, :] = seed_colors['exited']
        agent_colors[av_index + 1:] = seed_colors['existing']

    agent_colors[av_index, :] = np.array(agent_palette[-1])
    is_av = np.zeros_like(state[:, 0]).astype(np.bool_)
    is_av[av_index] = True

    # ! get timesteps to plot
    timesteps = list(range(num_timestep))
    if kwargs.get('time_idx', None) is not None:
        time_idx = kwargs['time_idx']
        assert set(time_idx).issubset(set(timesteps)), f'Got invalid time_idx: {time_idx=} v.s. {timesteps=}'
        timesteps = sorted(time_idx)

    # ! get plot limits
    limit_size = kwargs.get('limit_size', None)
    if limit_size is not None:
        assert len(limit_size) == 2, f'Got invalid `limit_size`: {limit_size=}'

    # ! plot all
    pil_images = []
    fig_paths = []
    for tid in tqdm(timesteps, leave=False, desc="Plot ..."):
        mask_t = visible_mask[:, tid]
        xs = traj[mask_t, tid, 0]
        ys = traj[mask_t, tid, 1]
        angles = heading[mask_t, tid]
        colors = agent_colors[mask_t, tid]
        types_t = [types[i] for i, mask in enumerate(mask_t) if mask]
        if ids is not None:
            ids_t = ids[mask_t]
        is_av_t = is_av[mask_t]
        enter_index_t = enter_index[tid] if enter_index else None
        labels = []
        if agent_labels:
            labels = [agent_labels[i][tid // shift] for i in range(len(agent_labels)) if mask_t[i]]

        fig_path = None
        if kwargs.get('save_path', None) is not None:
            save_path = kwargs['save_path']
            fig_path = os.path.join(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}", f"{tid}.png")
            fig_paths.append(fig_path)

        pil_images.append(
            plot_all(map_data, xs, ys, angles, types_t,
                     colors=colors,
                     save_path=fig_path,
                     is_avs=is_av_t,
                     pl2seed_radius=pl2seed_radius,
                     attr_tokenizer=attr_tokenizer,
                     enter_index=enter_index_t,
                     labels=labels,
                     limit_size=limit_size,
            )
        )

    # generate gif
    if fig_paths and save_gif:
        os.makedirs(os.path.join(save_path, 'gifs'), exist_ok=True)
        images = []
        gif_path = f"{save_path}/gifs/{scenario_id}_{str(save_id).zfill(3)}.gif"
        for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
            images.append(Image.open(fig_path))
        try:
            images[0].save(gif_path, save_all=True, append_images=images[1:], duration=100, loop=0)
            tqdm.write(f"Saved gif at {gif_path}")
            try:
                import shutil
                shutil.rmtree(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}")
                os.remove(f"{save_path}/gifs/{scenario_id}_{str(save_id - 1).zfill(3)}.gif")
            except:
                pass
        except Exception as e:
            tqdm.write(f"{e}! Failed to save gif at {gif_path}")

    return pil_images


def match_token_map(data):

    # init map token
    argmin_sample_len = 3
    map_token_traj_path = '/u/xiuyu/work/dev4/dev/tokens/map_traj_token5.pkl'

    map_token_traj = pickle.load(open(map_token_traj_path, 'rb'))
    map_token = {'traj_src': map_token_traj['traj_src'], }
    traj_end_theta = np.arctan2(map_token['traj_src'][:, -1, 1] - map_token['traj_src'][:, -2, 1],
                                map_token['traj_src'][:, -1, 0] - map_token['traj_src'][:, -2, 0])
    indices = torch.linspace(0, map_token['traj_src'].shape[1]-1, steps=argmin_sample_len).long()
    map_token['sample_pt'] = torch.from_numpy(map_token['traj_src'][:, indices]).to(torch.float)
    map_token['traj_end_theta'] = torch.from_numpy(traj_end_theta).to(torch.float)
    map_token['traj_src'] = torch.from_numpy(map_token['traj_src']).to(torch.float)

    traj_pos = data['map_save']['traj_pos'].to(torch.float)
    traj_theta = data['map_save']['traj_theta'].to(torch.float)
    pl_idx_list = data['map_save']['pl_idx_list']
    token_sample_pt = map_token['sample_pt'].to(traj_pos.device)
    token_src = map_token['traj_src'].to(traj_pos.device)
    max_traj_len = map_token['traj_src'].shape[1]
    pl_num = traj_pos.shape[0]

    pt_token_pos = traj_pos[:, 0, :].clone()
    pt_token_orientation = traj_theta.clone()
    cos, sin = traj_theta.cos(), traj_theta.sin()
    rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
    rot_mat[..., 0, 0] = cos
    rot_mat[..., 0, 1] = -sin
    rot_mat[..., 1, 0] = sin
    rot_mat[..., 1, 1] = cos
    traj_pos_local = torch.bmm((traj_pos - traj_pos[:, 0:1]), rot_mat.view(-1, 2, 2))
    distance = torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1))
    pt_token_id = torch.argmin(distance, dim=1)

    noise = False
    if noise:
        topk_indices = torch.argsort(torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1)), dim=1)[:, :8]
        sample_topk = torch.randint(0, topk_indices.shape[-1], size=(topk_indices.shape[0], 1), device=topk_indices.device)
        pt_token_id = torch.gather(topk_indices, 1, sample_topk).squeeze(-1)

    # cos, sin = traj_theta.cos(), traj_theta.sin()
    # rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
    # rot_mat[..., 0, 0] = cos
    # rot_mat[..., 0, 1] = sin
    # rot_mat[..., 1, 0] = -sin
    # rot_mat[..., 1, 1] = cos
    # token_src_world = torch.bmm(token_src[None, ...].repeat(pl_num, 1, 1, 1).reshape(pl_num, -1, 2),
    #                             rot_mat.view(-1, 2, 2)).reshape(pl_num, token_src.shape[0], max_traj_len, 2) + traj_pos[:, None, [0], :]
    # token_src_world_select = token_src_world.view(-1, 1024, 11, 2)[torch.arange(pt_token_id.view(-1).shape[0]), pt_token_id.view(-1)].view(pl_num, max_traj_len, 2)

    pl_idx_full = pl_idx_list.clone()
    token2pl = torch.stack([torch.arange(len(pl_idx_list), device=traj_pos.device), pl_idx_full.long()])
    count_nums = []
    for pl in pl_idx_full.unique():
        pt = token2pl[0, token2pl[1, :] == pl]
        left_side = (data['pt_token']['side'][pt] == 0).sum()
        right_side = (data['pt_token']['side'][pt] == 1).sum()
        center_side = (data['pt_token']['side'][pt] == 2).sum()
        count_nums.append(torch.Tensor([left_side, right_side, center_side]))
    count_nums = torch.stack(count_nums, dim=0)
    num_polyline = int(count_nums.max().item())
    traj_mask = torch.zeros((int(len(pl_idx_full.unique())), 3, num_polyline), dtype=bool)
    idx_matrix = torch.arange(traj_mask.size(2)).unsqueeze(0).unsqueeze(0)
    idx_matrix = idx_matrix.expand(traj_mask.size(0), traj_mask.size(1), -1)
    counts_num_expanded = count_nums.unsqueeze(-1)
    mask_update = idx_matrix < counts_num_expanded
    traj_mask[mask_update] = True

    data['pt_token']['traj_mask'] = traj_mask
    data['pt_token']['position'] = torch.cat([pt_token_pos, torch.zeros((data['pt_token']['num_nodes'], 1),
                                                                        device=traj_pos.device, dtype=torch.float)], dim=-1)
    data['pt_token']['orientation'] = pt_token_orientation
    data['pt_token']['height'] = data['pt_token']['position'][:, -1]
    data[('pt_token', 'to', 'map_polygon')] = {}
    data[('pt_token', 'to', 'map_polygon')]['edge_index'] = token2pl  # (2, num_points)
    data['pt_token']['token_idx'] = pt_token_id
    return data


@safe_run
def plot_tokenize(data, save_path: str):

    shift = 5
    token_size = 2048
    pl2seed_radius = 75

    # transformation
    transform = WaymoTargetBuilder(num_historical_steps=11,
                                   num_future_steps=80,
                                   max_num=32,
                                   training=False)

    grid_range = 150.
    grid_interval = 3.
    angle_interval = 3.
    attr_tokenizer = Attr_Tokenizer(grid_range=grid_range,
                                    grid_interval=grid_interval,
                                    radius=pl2seed_radius,
                                    angle_interval=angle_interval)

    # tokenization
    token_processor = TokenProcessor(token_size,
                                     training=False,
                                     predict_motion=True,
                                     predict_state=True,
                                     predict_map=True,
                                     state_token={'invalid': 0, 'valid': 1, 'enter': 2, 'exit': 3},
                                     pl2seed_radius=pl2seed_radius)
    CONSOLE.log(f"Loaded token processor with token_size: {token_size}")

    # preprocess
    data: HeteroData = transform(data)
    tokenized_data = token_processor(data)
    CONSOLE.log(f"Keys in tokenized data:\n{tokenized_data.keys()}")

    # plot
    agent_data = tokenized_data['agent']
    map_data = tokenized_data['map_point']
    # CONSOLE.log(f"Keys in agent data:\n{agent_data.keys()}")

    av_index = agent_data['av_index']
    raw_traj = agent_data['position'][..., :2].contiguous()  # [n_agent, n_step, 2]
    raw_heading = agent_data['heading']  # [n_agent, n_step]

    traj = agent_data['traj_pos'][..., :2].contiguous()  # [n_agent, n_step, 6, 2]
    traj = traj[:, :, 1:, :].flatten(1, 2)
    traj = torch.cat([raw_traj[:, :1], traj], dim=1)
    heading = agent_data['traj_heading']  # [n_agent, n_step, 6]
    heading = heading[:, :, 1:].flatten(1, 2)
    heading = torch.cat([raw_heading[:, :1], heading], dim=1)

    agent_state = agent_data['state_idx'].repeat_interleave(repeats=shift, dim=-1)
    agent_state = torch.cat([torch.zeros_like(agent_state[:, :1]), agent_state], dim=1)
    agent_type = agent_data['type']
    ids = np.arange(raw_traj.shape[0])

    return plot_scenario(
                scenario_id=tokenized_data['scenario_id'],
                map_data=tokenized_data['map_point']['position'].numpy(),
                traj=raw_traj.numpy(),
                heading=raw_heading.numpy(),
                state=agent_state.numpy(),
                types=list(map(lambda i: AGENT_TYPE[i], agent_type.tolist())),
                av_index=av_index,
                ids=ids,
                save_path=save_path,
                pl2seed_radius=pl2seed_radius,
                attr_tokenizer=attr_tokenizer,
                color_type='state',
        )


def get_metainfos(folder: str):

    import pandas as pd

    assert os.path.exists(folder), f'Path {folder} does not exist.'
    files = list(fnmatch.filter(os.listdir(folder), 'idx_*_rollouts.pkl'))
    CONSOLE.log(f'Found {len(files)} rollouts files from {folder}.')

    metainfos_path = f'{os.path.normpath(folder)}_metainfos.parquet'
    csv_path = f'{os.path.normpath(folder)}_metainfos.csv'

    if not os.path.exists(metainfos_path):

        data = []
        for file in tqdm(files):
            pkl_data = pickle.load(open(os.path.join(folder, file), 'rb'))
            data.extend((file, scenario_id, index) for index, scenario_id in enumerate(pkl_data['_scenario_id']))

        df = pd.DataFrame(data, columns=('rollout_file', 'scenario_id', 'index'))
        df.to_parquet(metainfos_path)
        df.to_csv(csv_path)
        CONSOLE.log(f'Successfully saved to {metainfos_path}.')

    else:
        CONSOLE.log(f'File {metainfos_path} already exists!')
        return


def plot_comparison(methods: List[str], rollouts_paths: List[str], gt_folders: List[str],
                    save_path: str, scenario_ids: Optional[List[str]] = None):
    import pandas as pd
    from collections import defaultdict

    # ! hyperparameter
    fps = 10

    plot_time = [1, 6, 12, 18, 24, 30]
    # plot_time = [1, 5, 10, 15, 20, 25]
    time_idx = [int(time * fps) for time in plot_time]

    limit_size = [75, 60]  # [width, height]

    # ! load metainfos
    metainfos = defaultdict(dict)
    for method, rollout_path in zip(methods, rollouts_paths):
        meta_info_path = f'{os.path.normpath(rollout_path)}_metainfos.parquet'
        metainfos[method]['df'] = pd.read_parquet(meta_info_path)
        CONSOLE.log(f'Loaded {method=} with {len(metainfos[method]["df"]["scenario_id"])=}.')
    common_scenarios = set(metainfos['ours']['df']['scenario_id'])
    for method, meta_info in metainfos.items():
        if method == 'ours':
            continue
        common_scenarios &= set(meta_info['df']['scenario_id'])
    for method, meta_info in metainfos.items():
        df = metainfos[method]['df']
        metainfos[method]['df'] = df[df['scenario_id'].isin(common_scenarios)]
    CONSOLE.log(f'Filter and get {len(common_scenarios)=}.')

    # ! load data and plot
    if scenario_ids is None:
        scenario_ids = metainfos['ours']['df']['scenario_id'].tolist()
    CONSOLE.log(f'Plotting {len(scenario_ids)=} ...')

    for scenario_id in (pbar := tqdm(scenario_ids)):
        pbar.set_postfix(scenario_id=scenario_id)

        figures = dict()
        for method, rollout_path, gt_folder in zip(methods, rollouts_paths, gt_folders):
            df = metainfos[method]['df']
            _df = df.loc[df['scenario_id'] == scenario_id]
            batch_idx = int(_df['index'].tolist()[0])
            rollout_file = _df['rollout_file'].tolist()[0]
            figures[method] = plot_file(
                                gt_folder=gt_folder,
                                files=os.path.join(rollout_path, rollout_file),
                                save_gif=False,
                                batch_idx=batch_idx,
                                time_idx=time_idx,
                                limit_size=limit_size,
                                color_type='insert',
                            )[0][0][0]

        # ! plot figures
        border = 5
        padding_x = 20
        padding_y = 50

        img_width, img_height = figures['ours'][0].size
        img_width = img_width + 2 * border
        img_height = img_height + 2 * border
        n_col = len(time_idx)
        n_row = len(methods)

        W = n_col * img_width + (n_col - 1) * padding_x
        H = n_row * img_height + (n_row - 1) * padding_y

        canvas = Image.new('RGB', (W, H), 'white')
        for i_row, (method, method_figures) in enumerate(figures.items()):
            for i_col, method_figure in enumerate(method_figures):
                x = i_col * (img_width + padding_x)
                y = i_row * (img_height + padding_y)

                padded_figure = Image.new('RGB', (img_width, img_height), 'black')
                padded_figure.paste(method_figure, (border, border))

                canvas.paste(padded_figure, (x, y))

        canvas.save(
            os.path.join(save_path, f'{scenario_id}.png')
        )


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--data_path', type=str, default='/u/xiuyu/work/dev4/data/waymo_processed')
    parser.add_argument('--tfrecord_dir', type=str, default='validation_tfrecords_splitted')
    # plot tokenized data
    parser.add_argument('--save_folder', type=str, default='plot_gt')
    parser.add_argument('--split', type=str, default='validation')
    parser.add_argument('--scenario_id', type=str, default=None)
    parser.add_argument('--plot_tokenize', action='store_true')
    # plot generated rollouts
    parser.add_argument('--plot_file', action='store_true')
    parser.add_argument('--folder_path', type=str, default=None)
    parser.add_argument('--file_path', type=str, default=None)
    # metainfos
    parser.add_argument('--get_metainfos', action='store_true')
    # plot comparison
    parser.add_argument('--plot_comparison', action='store_true')
    parser.add_argument('--comparison_folder', type=str, default='comparisons')
    args = parser.parse_args()

    if args.plot_tokenize:

        scenario_id = "74ad7b76d5906d39"
        # scenario_id = "1d60300bc06f4801"
        data_path = os.path.join(args.data_path, args.split, f"{scenario_id}.pkl")
        data = pickle.load(open(data_path, "rb"))
        data['tfrecord_path'] = os.path.join(args.tfrecord_dir, f'{scenario_id}.tfrecords')
        CONSOLE.log(f"Loaded scenario {scenario_id}")

        save_path = os.path.join(args.data_path, args.save_folder, args.split)
        os.makedirs(save_path, exist_ok=True)

        plot_tokenize(data, save_path)

    if args.plot_file:

        plot_file(args.data_path, folder=args.folder_path, files=args.file_path)

    if args.get_metainfos:

        assert args.folder_path is not None, f'`folder_path` should not be None!'
        get_metainfos(args.folder_path)

    if args.plot_comparison:

        methods = ['ours', 'smart']
        gt_folders = [
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
        ]
        rollouts_paths = [
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_ours0',
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_smart',
            '/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_cslft',
        ]
        save_path = f'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/{args.comparison_folder}/'
        os.makedirs(save_path, exist_ok=True)

        scenario_ids = ['72ff3e1540b28431','a16c927b1a1cca74','a504d55ea6658de7','639949ea1d16125b']
        plot_comparison(methods, rollouts_paths, gt_folders,
                        save_path=save_path,
                        scenario_ids=scenario_ids)