# Copyright (c) OpenMMLab. All rights reserved.
from typing import List

import torch
from torch import Tensor, nn

from mmdet3d.registry import MODELS


@MODELS.register_module()
class PointPillarsScatter(nn.Module):
    """Point Pillar's Scatter.

    Converts learned features from dense tensor to sparse pseudo image.

    Args:
        in_channels (int): Channels of input features.
        output_shape (list[int]): Required output shape of features.
    """

    def __init__(self, in_channels: int, output_shape: List[int]):
        super().__init__()
        self.output_shape = output_shape
        self.ny = output_shape[0]
        self.nx = output_shape[1]
        self.in_channels = in_channels

    def forward(self,
                voxel_features: Tensor,
                coors: Tensor,
                batch_size: int = None) -> Tensor:
        """Foraward function to scatter features."""
        # TODO: rewrite the function in a batch manner
        # no need to deal with different batch cases
        if batch_size is not None:
            return self.forward_batch(voxel_features, coors, batch_size)
        else:
            return self.forward_single(voxel_features, coors)

    def forward_single(self, voxel_features: Tensor, coors: Tensor) -> Tensor:
        """Scatter features of single sample.

        Args:
            voxel_features (torch.Tensor): Voxel features in shape (N, M, C).
            coors (torch.Tensor): Coordinates of each voxel.
                The first column indicates the sample ID.
        """
        # Create the canvas for this sample
        canvas = torch.zeros(
            self.in_channels,
            self.nx * self.ny,
            dtype=voxel_features.dtype,
            device=voxel_features.device)

        indices = coors[:, 2] * self.nx + coors[:, 3]
        indices = indices.long()
        voxels = voxel_features.t()
        # Now scatter the blob back to the canvas.
        canvas[:, indices] = voxels
        # Undo the column stacking to final 4-dim tensor
        canvas = canvas.view(1, self.in_channels, self.ny, self.nx)
        return canvas

    # def forward_batch(self, voxel_features: Tensor, coors: Tensor,
    #                   batch_size: int) -> Tensor:
    #     """Scatter features of single sample.

    #     Args:
    #         voxel_features (torch.Tensor): Voxel features in shape (N, M, C).
    #         coors (torch.Tensor): Coordinates of each voxel in shape (N, 4).
    #             The first column indicates the sample ID.
    #         batch_size (int): Number of samples in the current batch.
    #     """
    #     # batch_canvas will be the final output.
    #     batch_canvas = []
    #     for batch_itt in range(batch_size):
    #         # Create the canvas for this sample
    #         canvas = torch.zeros(
    #             self.in_channels,
    #             self.nx * self.ny,
    #             dtype=voxel_features.dtype,
    #             device=voxel_features.device)

    #         # Only include non-empty pillars
    #         batch_mask = coors[:, 0] == batch_itt
    #         this_coors = coors[batch_mask, :]
    #         indices = this_coors[:, 2] * self.nx + this_coors[:, 3]
    #         indices = indices.type(torch.long)
    #         voxels = voxel_features[batch_mask, :]
    #         voxels = voxels.t()

    #         # Now scatter the blob back to the canvas.
    #         canvas[:, indices] = voxels

    #         # Append to a list for later stacking.
    #         batch_canvas.append(canvas)

    #     # Stack to 3-dim tensor (batch-size, in_channels, nrows*ncols)
    #     batch_canvas = torch.stack(batch_canvas, 0)

    #     # Undo the column stacking to final 4-dim tensor
    #     batch_canvas = batch_canvas.view(batch_size, self.in_channels, self.ny,
    #                                      self.nx)

    #     return batch_canvas

    def forward_batch(self, voxel_features: Tensor, coors: Tensor, batch_size: int) -> Tensor:
        """
        voxel_features: (N, C)
        coors         : (N, 4)  -> [batch_idx, z, y, x]
        return        : (B, C, H, W)
        """
        N, C = voxel_features.shape
        H, W = self.ny, self.nx
        device = voxel_features.device
        spatial_size = H * W  # 缓存

        # 1. 提取坐标
        b = coors[:, 0]        # (N,)
        y = coors[:, 2]        # (N,)
        x = coors[:, 3]        # (N,)
        spatial_offset = y * W + x  # (N,)

        # 2. 构造全局索引 (N, C) → view(-1)
        c_range = torch.arange(C, device=device)          # (C,)
        bc_offset = b.unsqueeze(1) * C + c_range          # (N, C)
        # 利用广播避免 repeat_interleave
        global_indices = (bc_offset * spatial_size + spatial_offset.unsqueeze(1)).view(-1)  # (N*C,)

        # 3. 初始化 canvas
        canvas = torch.zeros(batch_size * C * spatial_size, dtype=voxel_features.dtype, device=device)

        # 4. 散射累加（可用 index_add_ 替代，实测选择）
        features_flat = voxel_features.view(-1)  # (N*C,)
        canvas.scatter_add_(0, global_indices, features_flat)  # 或 scatter_add_

        # 5. reshape
        canvas = canvas.view(-1, C, H, W)
        return canvas
