"""
Provides utilities for generating and stacking regular grid positional encodings
into simulation input tensors. Includes functions to build coordinate grids and
concatenate them with existing data to add spatial position information as channels.
"""

import torch


def stack_regular_grid(x: torch.Tensor, d: int, start=0, end=1):
    """
    Stack a regular coordinate grid onto input feature tensors as additional channels.
    """
    grid = get_regular_grid(x, d, start, end)

    if x.ndim == d:
        x = x.unsqueeze(0).unsqueeze(-1)  # add batch and channel dims
    elif x.ndim == d + 1:
        x = x.unsqueeze(-1)  # add channel dim

    return torch.cat([grid, x], dim=-1)


def get_regular_grid(x: torch.Tensor, d: int, start=0, end=1):
    """
    Generate a regular d-dimensional coordinate grid matching the spatial shape of x.
    """
    is_batch = x.ndim == d + 1

    # Spatial sizes for each dimension
    sizes = x.shape[1:] if is_batch else x.shape

    # Create 1D grids for each axis
    grids_1d = [
        torch.linspace(start, end, steps=sizes[i], device=x.device) for i in range(d)
    ]

    coord_list = torch.meshgrid(*grids_1d, indexing="ij")
    grid = torch.stack(coord_list, dim=-1)

    if is_batch:
        grid = grid.unsqueeze(0).expand(x.shape[0], *grid.shape)

    return grid
