import torch
from typing import Callable


def batch_elementwise(input: torch.Tensor, param: torch.Tensor,
                      op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
                      input_batch_dim: int = 0, pndim: int = 1) -> torch.Tensor:
    """
    Do elementwise operation in groups.

    :param input: input, any shape, [..., Ci, Cj, ...]
    :param param: the parameter, shape [N, Ci, Cj....], in which case B % N == 0, or [Ci, Cj....]
    :param input_batch_dim: which dimension is the batch in the input
    :param op: the operation to perform
    :param pndim: number of parameter dimensions without batch
    :return: input with the op performed, the same shape as input
    """

    if param.ndim == pndim+1:
        # If the param is batched, check if it the batch size is 1
        param = param.squeeze(0)

    if param.ndim == pndim:
        # If the param has no batch, do the normal op
        return op(input, param)

    assert param.ndim == pndim + 1
    assert input.shape[input_batch_dim] % param.shape[0] == 0

    input_r = input.view(*input.shape[:input_batch_dim], param.shape[0], -1, *input.shape[input_batch_dim+1:])

    param_r = param.view(*([1]*input_batch_dim), param.shape[0], *([1]*(input_r.ndim - input_batch_dim - param.ndim)),
                         *param.shape[1:])

    return op(input_r, param_r).view_as(input)