import torch


def fft_rec_loss2(tgt, src, mask):
    """Calculate FFT loss between pair of images. The loss is masked by `mask`. FFT loss is calculated on GPU and, thus, very fast. We calculate phase and amplitude L1 losses.

    Args:
        tgt (torch.Tensor): Target image. Shape [B,C,H,W]
        src (torch.Tensor): Source image. Shape [B,C,H,W]
        mask (torch.Tensor): Boolean mask for valid points

    Returns:
        torch.Tensor: Scalar sum of phase and amplitude losses.
    """

    # Apply 2D FFT on the last two dimensions, e.g., [H,W] channels.
    fft_tgt = torch.fft.rfftn(
        tgt, s=tgt.shape[-2:], dim=[-2, -1], norm="forward")  # [B,C,H,W,2]
    fft_src = torch.fft.rfftn(
        src, s=tgt.shape[-2:], dim=[-2, -1], norm="forward")  # [B,C,H,W,2]
    # fft_diff = torch.fft.rfftn(tgt-src, s=tgt.shape[-2:], dim=[-2,-1], norm="ortho") # [B,C,H,W,2]
    # fft_diff = fft_tgt - fft_src

    # fft_diff = torch.view_as_real(fft_diff)
    # mag_diff = fft_diff[...,0].abs().sum() #20*torch.log10(fft_diff[...,0]) # mag2db
    # pha_diff = fft_diff[...,1].abs().sum()

    # Convolution over pixels is FFT on frequencies.
    # We may find a more clever way.
    # fft_conv = fft_tgt*fft_src
    # inv_fft_conv = torch.fft.irfftn(fft_conv, s=tgt.shape[-2:], dim=[-2,-1], norm="forward") # [B,C,H,W,2]
    # mask_diff = fft_diff #fft_tgt-fft_src
    # print(mask_diff.shape)
    # print(mask.shape)
    # mask_diff = mask_diff * mask

    # l = 20*torch.log10(fft_diff.abs()) # mag2db 20*log10(abs(complex))

    # Derivative for angle is not implemented yet.
    # pha_diff = torch.abs(fft_tgt.angle() - fft_src.angle())
    mag_diff = torch.abs(fft_tgt.abs() - fft_src.abs())

    fft_tgt = torch.view_as_real(fft_tgt)
    fft_src = torch.view_as_real(fft_src)
    pha_tgt = torch.atan2(fft_tgt[..., 1], fft_tgt[..., 0])
    pha_src = torch.atan2(fft_src[..., 1], fft_src[..., 0])
    # mag_tgt = torch.sqrt(fft_tgt[...,1]**2 + fft_tgt[...,0]**2)
    # mag_src = torch.sqrt(fft_src[...,1]**2 + fft_src[...,0]**2)

    pha_diff = torch.abs(pha_tgt-pha_src)
    # mag_diff = torch.abs(mag_tgt - mag_src)

    # print(pha_diff.sum())
    # print(mag_diff.sum())

    l = 1e-4*mag_diff.sum() + pha_diff.sum()

    return l