# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/50_functional.ipynb (unless otherwise specified).

__all__ = ['SMatrixSubpix', 'SparseSmoothTruncatedAmplitudeLoss', 'SparseSmoothTruncatedAmplitudeProx',
           'SparseAmplitudeLoss']

# Cell
from .torch_imports import *
from .util import cartesian_aberrations, fftshift_checkerboard
import numpy as np
import torch as th
from .operators import AtF2, A_realspace, dS
from numpy.fft import fftshift, fftfreq
from torch.autograd import Function
from .operators import split
from .util import sparse_amplitude_loss, plotAbsAngle, sparse_smooth_truncated_amplitude_loss, sparse_smooth_truncated_amplitude_prox

class SMatrixSubpix(th.autograd.Function):
    @staticmethod
    def forward(ctx, S: th.Tensor, psi: th.Tensor, r: th.Tensor) -> th.Tensor:
        """
        :param S: BB x N1 x N2       tensor
        :param psi: BB X M1 x M2     tensor in fourier space
        :param pos: K x 2 real                  tensor

        :return: D x K x M1 x M2                measured amplitudes in the far field
        """
        B, MY, MX = psi.shape
        K, _ = r.shape
        qy, qx = np.meshgrid(fftfreq(MY), fftfreq(MX), indexing='ij')
        # 2 x M1 x M2
        q = th.stack([th.as_tensor(qy), th.as_tensor(qx)]).float().cuda()
        rs = r - r.int()
        r_int = r.long()
        # K x MY x MX
        ramp = th.exp(-2j * np.pi *
            (q[0][None, None, ...] * rs[:, 0][:, None, None] + q[1][None, None, ...] * rs[:, 1][:, None, None]))

        # B x K x M1 x M2
        frames_exit = split(S.detach(), r_int, MY, MX)
        # plotAbsAngle(frames_exit[0, 0].detach().cpu(),'frames_exit')
        # B x K x M x M
        #subpixel shift the whole batch
        # print(f'psi norm: {th.norm(psi)**2}')
        Psi = th.fft.fft2(psi, norm='ortho')
        Psi = Psi.unsqueeze(1)
        Psi = Psi.repeat(1, ramp.shape[0], 1, 1)
        psi = th.fft.ifft2(Psi * ramp, norm='ortho')
        # print(f'psi norm: {th.norm(psi[0,0])**2}')
        # psi = th.ifft(Psi, 2, True)
        # plotAbsAngle(psi[0,0].detach().cpu(),'psi input')
        psi_exit = th.zeros((K, MY, MX), dtype=th.complex64, device=S.device)
        psi_exit = A_realspace(r_int, S.detach(), psi, psi_exit)
        # print(f'psi_exit norm: {th.norm(psi_exit[0])**2}')
        # plotAbsAngle(psi_exit[0].detach().cpu(), 'psi_exit')
        Psi_exit = th.fft.fft2(psi_exit, norm='ortho')
        # plotAbsAngle(Psi_exit[0].detach().cpu(), 'Psi_exit')
        # K x M x M
        a_model = th.abs(Psi_exit)
        # plot(a_model[0].cpu().detach().numpy(),'a_model')
        ctx.save_for_backward(frames_exit, Psi_exit, psi_exit, psi, r_int, ramp, th.Tensor([S.shape])[0].int())
        return a_model

    def backward(ctx, grad_output):
        frames_exit, Psi_exit, psi_exit, psi, r_int, ramp, ss = ctx.saved_tensors
        # (K x M x M)  * (K x M x M)
        grad_Psi = Psi_exit * grad_output
        # print(f'grad_Psi.shape {grad_output.shape}')
        # plotAbsAngle(fftshift(grad_Psi[0].cpu().numpy()), 'grad_Psi')
        # K x M x M
        grad_Psi = th.fft.ifft2(grad_Psi, norm='ortho')

        # print(f'grad_Psi.shape {grad_output.shape}')
        # plotAbsAngle(grad_Psi[0].cpu().numpy(),'grad_Psi')

        if ctx.needs_input_grad[0]:
            grad_S = th.zeros(tuple(ss.numpy()), device=Psi_exit.device, dtype=th.complex64)
            # print(Psi_exit.shape, grad_Psi.shape, psi.shape, r.shape, grad_S.shape)

            psi_int = th.abs(psi)**2
            psi_int_max = th.max(th.max(psi_int,-1).values,-1).values
            grad_S = dS(grad_Psi, psi_exit, psi, psi_int, psi_int_max, r_int, grad_S, alpha=0.1)
            # AtF2(grad_Psi, psi, r_int, grad_S)
        else:
            grad_S = None

        if ctx.needs_input_grad[1]:
            # (B x K x MY x MX) * (B x K x MY x MX)
            frames_exit_int_max = th.max(th.abs(frames_exit)**2)
            grad_psi = grad_Psi * frames_exit.conj() / frames_exit_int_max
            grad_psi = th.fft.fft2(grad_psi, norm='ortho')
            ramp = th.view_as_real(ramp)
            ramp[..., 1] *= -1
            ramp = th.view_as_complex(ramp)
            #(B x K x MY x MX) * (K x MY x MX)
            grad_psi = th.fft.ifft2(grad_psi * ramp, norm='ortho')
            #(B x MY x MX)
            grad_psi = th.mean(grad_psi, 1)
            # plotAbsAngle(grad_psi[0].cpu().numpy(), 'grad_psi')
            # grad_Psi = th.fft.ifft2(grad_psi, norm='ortho')
        else:
            grad_psi = None

        if ctx.needs_input_grad[2]:
            grad_pos = None
        else:
            grad_pos = None

        return grad_S, grad_psi, grad_pos

class SparseSmoothTruncatedAmplitudeLoss(Function):

    @staticmethod
    def forward(ctx, a_model, indices_target, counts_target):
        frame_dimensions = th.as_tensor(a_model.shape[1:], device=a_model.device)
        loss, grad = sparse_smooth_truncated_amplitude_loss(a_model, indices_target, counts_target, frame_dimensions)
        loss.requires_grad = True
        ctx.save_for_backward(grad)
        return loss

    @staticmethod
    def backward(ctx, *grad_outputs):
        grad_input, = ctx.saved_tensors
        grad_indices = None
        grad_counts = None
        return grad_input, grad_indices, grad_counts

class SparseSmoothTruncatedAmplitudeProx(Function):

    @staticmethod
    def forward(ctx, a_model, indices_target, counts_target):
        frame_dimensions = th.as_tensor(a_model.shape[1:], device=a_model.device)
        loss, grad = sparse_smooth_truncated_amplitude_prox(a_model, indices_target, counts_target, frame_dimensions)
        loss.requires_grad = True
        ctx.save_for_backward(grad)
        return loss

    @staticmethod
    def backward(ctx, *grad_outputs):
        grad_input, = ctx.saved_tensors
        grad_indices = None
        grad_counts = None
        return grad_input, grad_indices, grad_counts

class SparseAmplitudeLoss(Function):

    @staticmethod
    def forward(ctx, a_model, indices_target, counts_target):
        frame_dimensions = th.as_tensor(a_model.shape[1:], device=a_model.device)
        loss, grad = sparse_amplitude_loss(a_model, indices_target, counts_target, frame_dimensions)
        loss.requires_grad = True
        ctx.save_for_backward(grad)
        return loss

    @staticmethod
    def backward(ctx, *grad_outputs):
        grad_input, = ctx.saved_tensors
        grad_indices = None
        grad_counts = None
        return grad_input, grad_indices, grad_counts
