from typing import Callable
import torch
from einops import rearrange


def video_to_image(func: Callable) -> Callable:
    def wrapper(self, x: torch.Tensor, *args, **kwargs):
        if x.dim() == 5:
            t = x.shape[2]
            x = rearrange(x, "b c t h w -> (b t) c h w")
            x = func(self, x, *args, **kwargs)
            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
        else:
            x = func(self, x, *args, **kwargs)
        return x

    return wrapper


def nonlinearity(x):
    return x * torch.sigmoid(x)


def cast_tuple(t, length=1):
    return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length)