planner/test_worker.py

#######################################################################
# Name: test_worker.py
#
# - Runs robot in environment using RL Planner
#######################################################################

from .test_parameter import *

import imageio
import os
import copy
import numpy as np
import torch
from time import time
from pathlib import Path
from skimage.transform import resize
from taxabind_avs.satbind.kmeans_clustering import CombinedSilhouetteInertiaClusterer
from .env import Env
from .robot import Robot

np.seterr(invalid='raise', divide='raise')


class TestWorker:
    def __init__(self, meta_agent_id, n_agent, policy_net, global_step, device='cuda', greedy=False, save_image=False, clip_seg_tta=None):
        self.device = device
        self.greedy = greedy
        self.n_agent = n_agent
        self.metaAgentID = meta_agent_id
        self.global_step = global_step
        self.k_size = K_SIZE
        self.save_image = save_image
        self.clip_seg_tta = clip_seg_tta
        self.execute_tta = EXECUTE_TTA      # Added to interface with app.py

        self.env = Env(map_index=self.global_step, n_agent=n_agent, k_size=self.k_size, plot=save_image, test=True)
        self.local_policy_net = policy_net

        self.robot_list = []
        self.all_robot_positions = []
        for i in range(self.n_agent):
            robot_position = self.env.start_positions[i]
            robot = Robot(robot_id=i, position=robot_position, plot=save_image)
            self.robot_list.append(robot)
            self.all_robot_positions.append(robot_position)

        self.perf_metrics = dict()
        self.bad_mask_init = False

        # NOTE: Option to override gifs_path to interface with app.py
        self.gifs_path = GIFS_PATH

        # NOTE: updated due to app.py (hf does not allow heatmap to persist)
        if LOAD_AVS_BENCH: 
            if clip_seg_tta is not None: 
                heatmap, heatmap_unnormalized, heatmap_unnormalized_initial, patch_embeds = self.clip_seg_tta.reset(sample_idx=self.global_step)
                self.clip_seg_tta.heatmap = heatmap
                self.clip_seg_tta.heatmap_unnormalized = heatmap_unnormalized
                self.clip_seg_tta.heatmap_unnormalized_initial = heatmap_unnormalized_initial
                self.clip_seg_tta.patch_embeds = patch_embeds   

            # Override target positions in env
            self.env.target_positions = [(pose[1], pose[0]) for pose in self.clip_seg_tta.target_positions] 

            # Override segmentation mask
            if not USE_CLIP_PREDS and OVERRIDE_MASK_DIR != "":
                score_mask_path = os.path.join(OVERRIDE_MASK_DIR, self.clip_seg_tta.gt_mask_name)
                print("score_mask_path: ", score_mask_path)
                if os.path.exists(score_mask_path):
                    self.env.segmentation_mask = self.env.import_segmentation_mask(score_mask_path)
                    self.env.begin(self.env.map_start_position)
                else:
                    print(f"\n\n{RED}ERROR: Trying to override, but score mask not found at path:{NC} ", score_mask_path)
                    self.bad_mask_init = True

            # Save clustered embeds from sat encoder
            if USE_CLIP_PREDS:
                self.kmeans_clusterer = CombinedSilhouetteInertiaClusterer(
                    k_min=1,
                    k_max=8,
                    k_avg_max=4,
                    silhouette_threshold=0.15,
                    relative_threshold=0.15,
                    random_state=0,
                    min_patch_size=5,   
                    n_smooth_iter=2,    
                    ignore_label=-1,
                    plot=self.save_image,
                    gifs_dir = GIFS_PATH      
                )    
                # Generate kmeans clusters
                self.kmeans_sat_embeds_clusters = self.kmeans_clusterer.fit_predict(
                    patch_embeds=self.clip_seg_tta.patch_embeds,
                    map_shape=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1]),
                )
                print("Chosen k:", self.kmeans_clusterer.final_k)

                # if EXECUTE_TTA:
                #     print("Will execute TTA...")

        # Define Poisson TTA params
        self.step_since_tta = 0
        self.steps_to_first_tgt = None
        self.steps_to_mid_tgt = None
        self.steps_to_last_tgt = None


    def run_episode(self, curr_episode):

        # Return all metrics as None if faulty mask init
        if self.bad_mask_init:
            self.perf_metrics['tax'] = None
            self.perf_metrics['travel_dist'] = None
            self.perf_metrics['travel_steps'] = None
            self.perf_metrics['steps_to_first_tgt'] = None
            self.perf_metrics['steps_to_mid_tgt'] = None
            self.perf_metrics['steps_to_last_tgt'] = None
            self.perf_metrics['explored_rate'] = None
            self.perf_metrics['targets_found'] = None
            self.perf_metrics['targets_total'] = None
            self.perf_metrics['kmeans_k'] = None
            self.perf_metrics['tgts_gt_score'] = None
            self.perf_metrics['clip_inference_time'] = None
            self.perf_metrics['tta_time'] = None
            self.perf_metrics['success_rate'] = None
            return

        eps_start = time()
        done = False
        for robot_id, deciding_robot in enumerate(self.robot_list):
            deciding_robot.observations = self.get_observations(deciding_robot.robot_position)
            if LOAD_AVS_BENCH and USE_CLIP_PREDS:
                if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
                    heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
                    self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
                    unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
                    self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
                    print("Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
                else:   
                    self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
                    self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)

        ### Run episode ###
        for step in range(NUM_EPS_STEPS):

            next_position_list = []
            dist_list = []
            travel_dist_list = []
            dist_array = np.zeros((self.n_agent, 1))
            for robot_id, deciding_robot in enumerate(self.robot_list):
                observations = deciding_robot.observations

                ### Forward pass through policy to get next position ###
                next_position, action_index = self.select_node(observations)
                dist = np.linalg.norm(next_position - deciding_robot.robot_position)

                ### Log results of action (e.g. distance travelled) ###
                dist_array[robot_id] = dist
                dist_list.append(dist)
                travel_dist_list.append(deciding_robot.travel_dist)
                next_position_list.append(next_position)
                self.all_robot_positions[robot_id] = next_position

            arriving_sequence = np.argsort(dist_list)
            next_position_list = np.array(next_position_list)
            dist_list = np.array(dist_list)
            travel_dist_list = np.array(travel_dist_list)
            next_position_list = next_position_list[arriving_sequence]
            dist_list = dist_list[arriving_sequence]
            travel_dist_list = travel_dist_list[arriving_sequence]

            ### Take Action (Deconflict if 2 agents choose the same target position) ###
            next_position_list, dist_list = self.solve_conflict(arriving_sequence, next_position_list, dist_list)
            reward_list, done = self.env.multi_robot_step(next_position_list, dist_list, travel_dist_list)

            ### Update observations + rewards from action ###
            for reward, robot_id in zip(reward_list, arriving_sequence):
                robot = self.robot_list[robot_id]
                robot.save_trajectory_coords(self.env.find_index_from_coords(robot.robot_position), self.env.num_new_targets_found)

                # # TTA Update via Poisson Test (with KMeans clustering stats)
                if LOAD_AVS_BENCH and USE_CLIP_PREDS and self.execute_tta:  
                    self.poisson_tta_update(robot, self.global_step, step)

                robot.observations = self.get_observations(robot.robot_position)
                robot.save_reward_done(reward, done)

                # Update metrics
                self.log_metrics(step=step) 

            ### Save a frame to generate gif of robot trajectories ###
            if self.save_image:
                robots_route = []
                for robot in self.robot_list:
                    robots_route.append([robot.xPoints, robot.yPoints])
                if not os.path.exists(self.gifs_path):
                    os.makedirs(self.gifs_path)
                if LOAD_AVS_BENCH:
                    # NOTE: Replaced since using app.py
                    self.env.plot_heatmap(self.gifs_path, step, max(travel_dist_list), robots_route)

            if done:
                break

        if LOAD_AVS_BENCH:
            tax = Path(self.clip_seg_tta.gt_mask_name).stem
            self.perf_metrics['tax'] = " ".join(tax.split("_")[1:])
        else:
            self.perf_metrics['tax'] = None
        self.perf_metrics['travel_dist'] = max(travel_dist_list)
        self.perf_metrics['travel_steps'] = step + 1
        self.perf_metrics['steps_to_first_tgt'] = self.steps_to_first_tgt
        self.perf_metrics['steps_to_mid_tgt'] = self.steps_to_mid_tgt
        self.perf_metrics['steps_to_last_tgt'] = self.steps_to_last_tgt
        self.perf_metrics['explored_rate'] = self.env.explored_rate
        self.perf_metrics['targets_found'] = self.env.targets_found_rate
        self.perf_metrics['targets_total'] = len(self.env.target_positions)
        if USE_CLIP_PREDS:
            self.perf_metrics['kmeans_k'] = self.kmeans_clusterer.final_k
            self.perf_metrics['tgts_gt_score'] = self.clip_seg_tta.tgts_gt_score
            self.perf_metrics['clip_inference_time'] = self.clip_seg_tta.clip_inference_time
            self.perf_metrics['tta_time'] = self.clip_seg_tta.tta_time
        else:
            self.perf_metrics['kmeans_k'] = None
            self.perf_metrics['tgts_gt_score'] = None
            self.perf_metrics['clip_inference_time'] = None
            self.perf_metrics['tta_time'] = None
        if FORCE_LOGGING_DONE_TGTS_FOUND and self.env.targets_found_rate == 1.0:
            self.perf_metrics['success_rate'] = True
        else:
            self.perf_metrics['success_rate'] = done

        # save gif
        if self.save_image:
            path = self.gifs_path   # NOTE: Set to self.gifs_path since using app.py
            self.make_gif(path, curr_episode)

        print(YELLOW, f"[Eps {curr_episode} Completed] Time Taken: {time()-eps_start:.2f}s, Steps: {step+1}", NC)

    def get_observations(self, robot_position):
        """ Get robot's sensor observation of environment given position """
        current_node_index = self.env.find_index_from_coords(robot_position)
        current_index = torch.tensor([current_node_index]).unsqueeze(0).unsqueeze(0).to(self.device)  # (1,1,1)
    
        node_coords = copy.deepcopy(self.env.node_coords)
        graph = copy.deepcopy(self.env.graph)
        node_utility = copy.deepcopy(self.env.node_utility)
        guidepost = copy.deepcopy(self.env.guidepost)
        segmentation_info_mask = copy.deepcopy(self.env.filtered_seg_info_mask)

        n_nodes = node_coords.shape[0]
        node_coords = node_coords / 640
        node_utility = node_utility / 50
        node_utility_inputs = node_utility.reshape((n_nodes, 1))

        occupied_node = np.zeros((n_nodes, 1))
        for position in self.all_robot_positions:
            index = self.env.find_index_from_coords(position)
            if index == current_index.item():
                occupied_node[index] = -1
            else:
                occupied_node[index] = 1

        node_inputs = np.concatenate((node_coords, segmentation_info_mask, guidepost), axis=1)
        node_inputs = torch.FloatTensor(node_inputs).unsqueeze(0).to(self.device)  
        node_padding_mask = None

        graph = list(graph.values())
        edge_inputs = []
        for node in graph:
            node_edges = list(map(int, node))
            edge_inputs.append(node_edges)

        bias_matrix = self.calculate_edge_mask(edge_inputs)
        edge_mask = torch.from_numpy(bias_matrix).float().unsqueeze(0).to(self.device)

        for edges in edge_inputs:
            while len(edges) < self.k_size:
                edges.append(0)        

        edge_inputs = torch.tensor(edge_inputs).unsqueeze(0).to(self.device)  
        edge_padding_mask = torch.zeros((1, len(edge_inputs), K_SIZE), dtype=torch.int64).to(self.device)
        one = torch.ones_like(edge_padding_mask, dtype=torch.int64).to(self.device)
        edge_padding_mask = torch.where(edge_inputs == 0, one, edge_padding_mask)

        observations = node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask
        return observations
    

    def select_node(self, observations):
        """ Forward pass through policy to get next position to go to on map """
        node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask = observations
        with torch.no_grad():
            logp_list = self.local_policy_net(node_inputs, edge_inputs, current_index, node_padding_mask, edge_padding_mask, edge_mask)

        if self.greedy:
            action_index = torch.argmax(logp_list, dim=1).long()
        else:
            action_index = torch.multinomial(logp_list.exp(), 1).long().squeeze(1)

        next_node_index = edge_inputs[:, current_index.item(), action_index.item()]

        next_position = self.env.node_coords[next_node_index]

        return next_position, action_index

    def solve_conflict(self, arriving_sequence, next_position_list, dist_list):
        """ Deconflict if 2 agents choose the same target position """
        for j, [robot_id, next_position] in enumerate(zip(arriving_sequence, next_position_list)):
            moving_robot = self.robot_list[robot_id]
            # if next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
            #     dist_to_next_position = np.argsort(np.linalg.norm(self.env.node_coords - next_position, axis=1))
            #     k = 0
            #     while next_position[0] + next_position[1] * 1j in (next_position_list[:, 0] + next_position_list[:, 1] * 1j)[:j]:
            #         k += 1
            #         next_position = self.env.node_coords[dist_to_next_position[k]]

            dist = np.linalg.norm(next_position - moving_robot.robot_position)
            next_position_list[j] = next_position
            dist_list[j] = dist
            moving_robot.travel_dist += dist
            moving_robot.robot_position = next_position

        return next_position_list, dist_list

    def work(self, currEpisode):
        '''
        Interacts with the environment. The agent gets either gradients or experience buffer
        '''
        self.run_episode(currEpisode)

    def calculate_edge_mask(self, edge_inputs):
        size = len(edge_inputs)
        bias_matrix = np.ones((size, size))
        for i in range(size):
            for j in range(size):
                if j in edge_inputs[i]:
                    bias_matrix[i][j] = 0
        return bias_matrix

    def make_gif(self, path, n):
        """ Generate a gif given list of images """
        with imageio.get_writer('{}/{}_target_rate_{:.2f}.gif'.format(path, n, self.env.targets_found_rate), mode='I',
                                fps=5) as writer:
            for frame in self.env.frame_files:
                image = imageio.imread(frame)
                writer.append_data(image)
        print('gif complete\n')

        # Remove files
        for filename in self.env.frame_files[:-1]:
            os.remove(filename)

        # For gif during TTA
        if LOAD_AVS_BENCH:
            with imageio.get_writer('{}/{}_kmeans_stats.gif'.format(path, n), mode='I',
                                    fps=5) as writer:
                for frame in self.kmeans_clusterer.kmeans_frame_files:
                    image = imageio.imread(frame)
                    writer.append_data(image)
            print('Kmeans Clusterer gif complete\n')

            # Remove files
            for filename in self.kmeans_clusterer.kmeans_frame_files[:-1]:
                os.remove(filename)

    ################################################################################
    # SPPP Related Fns
    ################################################################################

    def log_metrics(self, step):
        # Update tgt found metrics
        if self.steps_to_first_tgt is None and self.env.num_targets_found == 1:
            self.steps_to_first_tgt = step + 1
        if self.steps_to_mid_tgt is None and self.env.num_targets_found == int(len(self.env.target_positions) / 2):
            self.steps_to_mid_tgt = step + 1
        if self.steps_to_last_tgt is None and self.env.num_targets_found == len(self.env.target_positions):
            self.steps_to_last_tgt = step + 1


    def transpose_flat_idx(self, idx, H=NUM_COORDS_HEIGHT, W=NUM_COORDS_WIDTH):
        """
        Transpose a flat index from an ``H×W`` grid to the equivalent
        position in the ``W×H`` transposed grid while **keeping the result
        in 1-D**.
        """
        # --- Safety check to catch out-of-range indices ---
        assert 0 <= idx < H * W, f"idx {idx} out of bounds for shape ({H}, {W})"

        # Original (row, col)
        row, col = divmod(idx, W)
        # After transpose these coordinates swap
        row_T, col_T = col, row

        # Flatten back into 1-D (row-major) for the W×H grid
        return row_T * H + col_T
    

    def poisson_tta_update(self, robot, episode, step):

        # Generate Kmeans Clusters Stats
        # Scale index back to CLIP_GRIDS_DIMS to be compatible with CLIP patch size
        if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:
            # High-res remap via pixel coordinates preserves exact neighbourhood
            filt_traj_coords, filt_targets_found_on_path = self.scale_trajectory(
                robot.trajectory_coords,
                self.env.target_positions,
                old_dims=(NUM_COORDS_HEIGHT, NUM_COORDS_WIDTH),
                full_dims=(512, 512),
                new_dims=(CLIP_GRIDS_DIMS[0], CLIP_GRIDS_DIMS[1])
            )
        else:
            filt_traj_coords = [self.transpose_flat_idx(idx) for idx in robot.trajectory_coords]
            filt_targets_found_on_path = robot.targets_found_on_path

        region_stats_dict = self.kmeans_clusterer.compute_region_statistics(
            self.kmeans_sat_embeds_clusters,
            self.clip_seg_tta.heatmap_unnormalized,
            filt_traj_coords,
            episode_num=episode,
            step_num=step
        )

        # Prep & execute TTA
        self.step_since_tta += 1
        if robot.targets_found_on_path[-1] or self.step_since_tta % STEPS_PER_TTA == 0:    

            # NOTE: integration with app.py on hf
            self.clip_seg_tta.executing_tta = True

            num_cells = self.clip_seg_tta.heatmap.shape[0] * self.clip_seg_tta.heatmap.shape[1]
            pos_sample_weight_scale, neg_sample_weight_scale = [], []

            for i, sample_loc in enumerate(filt_traj_coords):
                label = self.kmeans_clusterer.get_label_id(self.kmeans_sat_embeds_clusters, sample_loc)
                num_patches = region_stats_dict[label]['num_patches']
                patches_visited = region_stats_dict[label]['patches_visited']
                expectation = region_stats_dict[label]['expectation']

                # Exponent like focal loss to wait for more samples before confidently decreasing
                pos_weight = 4.0 
                neg_weight = min(1.0, (patches_visited/(3*num_patches))**GAMMA_EXPONENT)  
                pos_sample_weight_scale.append(pos_weight)
                neg_sample_weight_scale.append(neg_weight)

            # # # Adaptative LR (as samples increase, increase LR to fit more datapoints)
            adaptive_lr = MIN_LR + (MAX_LR - MIN_LR) * (step / num_cells)
                
            # TTA Update
            # NOTE: updated due to app.py (hf does not allow heatmap to persist)
            heatmap = self.clip_seg_tta.execute_tta(
                filt_traj_coords, 
                filt_targets_found_on_path, 
                tta_steps=NUM_TTA_STEPS, 
                lr=adaptive_lr,
                pos_sample_weight=pos_sample_weight_scale,
                neg_sample_weight=neg_sample_weight_scale,
                reset_weights=RESET_WEIGHTS
            )
            self.clip_seg_tta.heatmap = heatmap

            if NUM_COORDS_WIDTH != CLIP_GRIDS_DIMS[0] or NUM_COORDS_HEIGHT != CLIP_GRIDS_DIMS[1]:   # If heatmap is resized from clip original dims
                heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
                self.env.segmentation_info_mask = np.expand_dims(heatmap.T.flatten(), axis=1)
                unnormalized_heatmap = self.convert_heatmap_resolution(self.clip_seg_tta.heatmap_unnormalized, full_dims=(512, 512), new_dims=(NUM_COORDS_WIDTH, NUM_COORDS_HEIGHT))
                self.env.segmentation_info_mask_unnormalized = np.expand_dims(unnormalized_heatmap.T.flatten(), axis=1)
                print("~Resized heatmap to", NUM_COORDS_WIDTH, "x", NUM_COORDS_HEIGHT)
            else:  
                self.env.segmentation_info_mask = np.expand_dims(self.clip_seg_tta.heatmap.T.flatten(), axis=1)
                self.env.segmentation_info_mask_unnormalized = np.expand_dims(self.clip_seg_tta.heatmap_unnormalized.T.flatten(), axis=1)

            self.step_since_tta = 0

            # NOTE: integration with app.py on hf
            self.clip_seg_tta.executing_tta = False
    

    def convert_heatmap_resolution(self, heatmap, full_dims=(512, 512), new_dims=(24, 24)):
        heatmap_large = resize(heatmap, full_dims, order=1,  # order=1 → bilinear
             mode='reflect', anti_aliasing=True)
        
        coords = self.env.graph_generator.grid_coords   # (N, N, 2)
        rows, cols = coords[...,1], coords[...,0]
        heatmap_resized = heatmap_large[rows, cols]
        heatmap_resized = heatmap_resized.reshape(new_dims[1], new_dims[0])
        return heatmap_resized


    def convert_labelmap_resolution(self, labelmap, full_dims=(512, 512), new_dims=(24, 24)):
        """
        1) Upsample via nearest‐neighbor to full_dims
        2) Sample back down to your graph grid using grid_coords
        """
        # 1) Upsample with nearest‐neighbor, preserving integer labels
        up = resize(
            labelmap,
            full_dims,
            order=0,                # nearest‐neighbor
            mode='edge',            # padding mode
            preserve_range=True,    # don't normalize labels
            anti_aliasing=False     # must be False for labels
        ).astype(labelmap.dtype)    # back to original integer dtype

        # 2) Downsample via your precomputed grid coords
        coords = self.env.graph_generator.grid_coords   # shape (N, N, 2)
        rows = coords[...,1].astype(int)
        cols = coords[...,0].astype(int)

        small = up[rows, cols]      # shape (N, N)
        small = small.reshape(new_dims[0], new_dims[1])
        return small
    

    def scale_trajectory(self,
                        flat_indices,
                        targets,
                        old_dims=(17, 17),
                        full_dims=(512, 512),
                        new_dims=(24, 24)):
        """
        Args:
            flat_indices: list of ints in [0..old_H*old_W-1]
            targets:      list of (y_pix, x_pix) in [0..full_H-1]
            old_dims:     (old_H, old_W)
            full_dims:    (full_H, full_W)
            new_dims:     (new_H, new_W)

        Returns:
            new_flat_traj: list of unique flattened indices in new_H×new_W
            counts:        list of ints, same length as new_flat_traj
        """
        old_H, old_W = old_dims
        full_H, full_W = full_dims
        new_H,  new_W  = new_dims

        # 1) bin targets into new grid
        cell_h_new = full_H / new_H
        cell_w_new = full_W / new_W
        grid_counts = [[0]*new_W for _ in range(new_H)]
        for x_pix, y_pix in targets:  # note (x, y) order as in original implementation
            i_t = min(int(y_pix / cell_h_new), new_H - 1)
            j_t = min(int(x_pix / cell_w_new), new_W - 1)
            grid_counts[i_t][j_t] += 1

        # 2) Walk the trajectory indices and project each old cell's *entire
        #    pixel footprint* onto the finer 24×24 grid. 
        cell_h_full = full_H / old_H
        cell_w_full = full_W / old_W

        seen = set()
        new_flat_traj = []

        for node_idx in flat_indices:
            if node_idx < 0 or node_idx >= len(self.env.graph_generator.node_coords):
                continue

            coord_xy = self.env.graph_generator.node_coords[node_idx]
            try:
                row_old, col_old = self.env.graph_generator.find_index_from_grid_coords_2d(coord_xy)
            except Exception:
                continue

            # Bounding box of the old cell in full-resolution pixel space
            y0 = row_old * cell_h_full
            y1 = (row_old + 1) * cell_h_full
            x0 = col_old * cell_w_full
            x1 = (col_old + 1) * cell_w_full

            # Which new-grid rows & cols overlap? (inclusive ranges)
            i_start = max(0, min(int(y0 / cell_h_new), new_H - 1))
            i_end   = max(0, min(int((y1 - 1) / cell_h_new), new_H - 1))
            j_start = max(0, min(int(x0 / cell_w_new), new_W - 1))
            j_end   = max(0, min(int((x1 - 1) / cell_w_new), new_W - 1))

            for ii in range(i_start, i_end + 1):
                for jj in range(j_start, j_end + 1):
                    f_new = ii * new_W + jj
                    if f_new not in seen:
                        seen.add(f_new)
                        new_flat_traj.append(f_new)

        # 3) annotate counts
        counts = []
        for f in new_flat_traj:
            i_new, j_new = divmod(f, new_W)
            counts.append(grid_counts[i_new][j_new])

        return new_flat_traj, counts

    ################################################################################