import torch


def fft_frame(img):
    """Calculate FFT of a given image.

    Args:
        img (torch.Tensor): Image tensor. Shape [B,C,H,W]

    Returns:
        torch.Tensor: Amplitude and phase of FFT result. Shape [B,C,H,W]
    """

    fft_im = torch.fft.rfftn(
        img, s=img.shape[-2:], dim=[-2, -1], norm="forward")  # [B,C,H,W,2]
    fft_amp = fft_im[..., 0]**2 + fft_im[..., 1]**2
    fft_amp = torch.sqrt(fft_amp)  # this is the amplitude [B,C,H,W]
    # this is the phase [B,C,H,W]
    fft_pha = torch.atan2(fft_im[..., 1], fft_im[..., 0])

    return fft_amp, fft_pha