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

__all__ = ['calc_psi', 'AtF2', 'dS', 'A_realspace', 'calc_psi_denom', 'Qoverlap_real2', 'split', 'A', 'complex_matmul',
           'A_fast_full2', 'A_fast_full3', 'A_fast_full5', 'A_fast_full4', 'AH_S', 'smatrix_phase_factorsBDK',
           'smatrix_phase_factorsKB']

# Cell
from .kernels import *
import torch as th
import numpy as np
import math as m
import numba.cuda as cuda

# Cell
def calc_psi(r, t, z, out):
    out[:] = 0
    K = r.shape[0]
    MY, MX = out.shape
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(K * MY * MX / threadsperblock)
    psi_kernel[blockspergrid, threadsperblock](r, t, z, out)
    return out

# Cell


# Cell
def AtF2(z, psi, r, out):
    """

    :param z:   K x MY x MX
    :param psi: B x K x MY x MX
    :param r:   K x 2
    :param out: B x NY x NX
    :return:
    """
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(z.shape) / threadsperblock)
    AtF2_kernel[blockspergrid, threadsperblock](z, psi, r, out)
    return out

# Cell
def dS(z, z_old, psi, psi_int, psi_int_max, r, out, alpha=0.1):
    """
    :param z:           K x MY x MX
    :param z_old:       K x MY x MX
    :param psi:         B x K x MY x MX
    :param psi_int:     B x K x MY x MX
    :param psi_int_max: B x K
    :param r:           K x 2
    :param out:         B x NY x NX
    :param alpha:       float
    :return:
    """

    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(z.shape) / threadsperblock)
    dS_kernel[blockspergrid, threadsperblock](z, z_old, psi, psi_int, psi_int_max, alpha, r, th.view_as_real(out))
    return out

# Cell
def A_realspace(r, t, psi, out):
    """

    :param r:   K x 2
    :param t:   BB x NY x NX
    :param psi: B x K x MY x MX
    :param out: K x MY x MX
    :return:
    """
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(out.shape) / threadsperblock)
    # print(r.shape,t.shape,psi.shape,out.shape)
    A_realspace_kernel[blockspergrid, threadsperblock](r, th.view_as_real(t), th.view_as_real(psi), th.view_as_real(out))
    return out

# Cell

def calc_psi_denom(r, t, out):
    """

    :param r: K x 2
    :param t: BB x NY x NX
    :param out: BB x MY x MX
    :return:
    """

    out[:] = 0
    K = r.shape[0]
    BB, MY, MX = out.shape
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(BB * K * MY * MX / threadsperblock)
    psi_denom_kernel[blockspergrid, threadsperblock](r, t, out)
    return out

# Cell
def Qoverlap_real2(r, z, out):
    """

    :param r: K x 2
    :param z: BB x K x MY x MX
    :param out: BB x NY x NX
    :return: out
    """
    BB = out.shape[0]
    K = r.shape[0]
    out[:] = 1
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(BB * K * np.prod(z.shape) / threadsperblock)
    overlap_kernel_real2[blockspergrid, threadsperblock](r, z, out)
    return out

# Cell
def split(S, r, MY, MX):
    """

    :param S:  B x NY x NX
    :param r:  K x 2
    :param MY: int
    :param MX: int
    :return:
    """
    B = S.shape[0]
    K, _ = r.shape
    out = th.zeros((B, K, MY, MX), dtype=th.complex64, device=S.device)
    gpu = cuda.get_current_device()
    threadsperblock = gpu.MAX_THREADS_PER_BLOCK // 2
    blockspergrid = m.ceil(np.prod(out.shape) / threadsperblock)
    stream = th.cuda.current_stream().cuda_stream
    split_kernel4[blockspergrid, threadsperblock, stream](S, r, out)
    return out


# Cell

def A(S, Psi, r, r_min, out=None, Mx=0, My=0):
    """ Fastest version, takes precomputed phase factors, assumes S-matrix with beam tilt included

    :param S:               B x NY x NX
    :param phase_factors:   B x D x K x 2
    :param r:               D x K x 2
    :param out:             D x K x MY x MX
    :return: exit waves in out
    """
    B = S.shape[0]
    if out is None and My > 0 and Mx > 0:
        D, K, _ = r.shape
        out = th.zeros((D, K, My, Mx, 2), dtype=th.float32, device=S.device)
    else:
        out[:] = 0
    D, K, MY, MX, _ = out.shape
    gpu = cuda.get_current_device()
    threadsperblock = 128#gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array((B, MX, MY))) / threadsperblock)
    smatrix_forward_kernel[blockspergrid, threadsperblock, th.cuda.current_stream().cuda_stream]\
        (th.view_as_real(S), Psi.phase_factors, r, r_min, out)
    return th.view_as_complex(out)

@th.jit.script
def complex_matmul(a, b):
    """
    Complex matrix multiplication of tensors a and b.

    Pass conjugate = True to conjugate tensor b in the multiplication.
    """
    are, aim = th.unbind(a, -1)
    bre, bim = th.unbind(b, -1)
    real = are @ bre - aim @ bim
    imag = are @ bim + aim @ bre
    return th.stack([real, imag], -1)

def A_fast_full2(S, phase_factors, r, r_min, MY, MX):
    """ Fastest version, takes precomputed phase factors, assumes S-matrix with beam tilt included

    :param S:               B x NY x NX
    :param phase_factors:   K x B x 2
    :param r:               K x 2
    :param out:             K x MY x MX x 2
    :return: exit waves in out
    """
    B = S.shape[0]
    K, _ = r.shape
    out = th.zeros((K, MY, MX, B, 2), dtype=th.float32, device=S.device)
    K, MYMX, _, _, _ = out.shape
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = 256  # gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array((K, MX, MY, B))) / threadsperblock)

    phase_factors2 = phase_factors.unsqueeze(2)
    # 1 - get crops from S-matrix
    split_kernel[blockspergrid, threadsperblock, stream](th.view_as_real(S), r, out)
    out = out.view((K, MY * MX, B, 2))
    # 2 - complex batched matmul: K x MY*MX x B x 2 @ K x B x 1 x 2
    # print(out.shape)
    # print(phase_factors2.shape)
    exitwaves = complex_matmul(out, phase_factors2)
    # 3 - reshape
    exitwaves = exitwaves.view((K, MY, MX, 2))
    return exitwaves

def A_fast_full3(S, phase_factors, r, r_min, MY, MX):
    """ Fastest version, takes precomputed phase factors, assumes S-matrix with beam tilt included

    :param S:               B x NY x NX
    :param phase_factors:   K x B
    :param r:               K x 2
    :param out:             K x MY x MX
    :return: exit waves in out
    """
    B = S.shape[0]
    K, _ = r.shape
    out = th.zeros((K, MY, MX, B, 2), dtype=th.float32, device=S.device)
    K, MYMX, _, _, _ = out.shape
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream

    threadsperblock = 128  # gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array((K, MX, MY, B))) / threadsperblock)
    # 1 - get crops from S-matrix
    split_kernel[blockspergrid, threadsperblock, stream](th.view_as_real(S), r, out)

    # threadsperblock = 128  # gpu.MAX_THREADS_PER_BLOCK
    # blockspergrid = m.ceil(np.prod(np.array((K, B))) / threadsperblock)
    # # 1 - get crops from S-matrix
    # split_kernel2[blockspergrid, threadsperblock, stream](th.view_as_real(S), r, out)

    out = out.view((K, MY * MX, B, 2))
    out = th.view_as_complex(out)
    # 1.5 - convert to cupy
    # 2 - complex batched matmul: K x MY*MX x B @ K x B x 1
    # print(out.shape)
    # print(phase_factors2.shape)
    # print(out.dtype)
    # print(phase_factors2.dtype)
    phase_factors2 = phase_factors.unsqueeze(2)
    exitwaves = out @ phase_factors2
    # 3 - reshape
    exitwaves = exitwaves.view((K, MY, MX))
    #4 convert to pytorch
    return exitwaves

def A_fast_full5(S, phase_factors, r, r_min, MY, MX):
    """ Fastest version, takes precomputed phase factors, assumes S-matrix with beam tilt included

    :param S:               B x NY x NX
    :param phase_factors:   K x B
    :param r:               K x 2
    :param out:             K x MY x MX
    :return: exit waves in out
    """
    B = S.shape[0]
    K, _ = r.shape
    out = th.zeros((K, B, MY, MX), dtype=th.complex64, device=S.device)
    K, B, MY, MX = out.shape
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream

    threadsperblock = gpu.MAX_THREADS_PER_BLOCK // 2
    blockspergrid = m.ceil(np.prod(np.array((K, B, MY, MX))) / threadsperblock)
    # 1 - get crops from S-matrix
    split_kernel4[blockspergrid, threadsperblock, stream](th.view_as_real(S), r, th.view_as_real(out))

    # threadsperblock = 128  # gpu.MAX_THREADS_PER_BLOCK
    # blockspergrid = m.ceil(np.prod(np.array((K, B))) / threadsperblock)
    # # 1 - get crops from S-matrix
    # split_kernel2[blockspergrid, threadsperblock, stream](th.view_as_real(S), r, out)

    out = out.view((K, B, MY * MX))
    # 1.5 - convert to cupy
    # 2 - complex batched matmul: K x 1 x B @ K x B x MY*MX --> K x 1 x MY * MX
    # print(out.shape)
    # print(phase_factors2.shape)
    # print(out.dtype)
    # print(phase_factors2.dtype)
    phase_factors2 = phase_factors.unsqueeze(1)
    exitwaves = phase_factors2 @ out
    # 3 - reshape
    exitwaves = exitwaves.view((K, MY, MX))
    #4 convert to pytorch
    return exitwaves

def A_fast_full4(S, phase_factors, r, r_min, out=None, Mx=0, My=0):
    """ Fastest version, takes precomputed phase factors, assumes S-matrix with beam tilt included

    :param S:               B x NY x NX
    :param phase_factors:   B x D x K x 2
    :param r:               D x K x 2
    :param out:             D x K x MY x MX
    :return: exit waves in out
    """
    B = S.shape[0]
    if out is None and My > 0 and Mx > 0:
        D, K, _ = r.shape
        out = th.zeros((D, K, My, Mx, 2), dtype=th.float32, device=S.device)
    D, K, MY, MX, _ = out.shape
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    tbp = int(gpu.MAX_THREADS_PER_BLOCK**(1/3))

    # max dim of thread block (1024,1024,64), with a total of 1024 max
    # max dim of grid (2^32-1 , 2^16-1, 2^16-1)
    #                 (2^32-1 , 65535, 65535)

    threadsperblock = (tbp, tbp, tbp)
    blockspergrid = tuple(np.ceil((K/tbp, MY/tbp, MX/tbp)).astype(np.int))# m.ceil(np.prod(np.array((B, MX, MY))) / threadsperblock)
    smatrix_forward_kernel_fast_full4[blockspergrid, threadsperblock, stream](th.view_as_real(S), phase_factors, r, r_min, out)
    return out

# Cell
from .kernels import smatrix_backward_kernel_S
def AH_S(z, Psi, r, r_min, out=None, tau = th.tensor([1.0]), Ny=-1, Nx=-1):
    """
    Adjoint S-matrix operator for the full S-matrix. Expects pre-computed phase-factors as inputs.

    :param z:               D x K x My x Mx
    :param Psi:             D x My x Mx
    :param r:               D x K x 2
    :param r_min:           2
    :param out:             B x NY x NX x 2
    :param Ny:              optional, int
    :param Nx:              optional, int
    :return: result of adjoint S-matrix operator, shape (B x NY x NX x 2)
    """
    D, K, MY, MX = z.shape
    B = Psi.phase_factors.shape[0]

    if out is None and Ny > 0 and Nx > 0:
        out_is_gradient = True
        out = th.zeros((B, Ny, Nx, 2), dtype=th.float32, device=z.device)
    else:
        out_is_gradient = False
        tau /= (K * D)

    # shape D
    mean_probe_intensities = th.norm(Psi, p=2, dim=(1, 2))
    mean_probe_intensities /= MX * MY

    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = 128#gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array((B, MY, MX))) / threadsperblock)
    smatrix_backward_kernel_S[blockspergrid, threadsperblock, stream] \
        (th.view_as_real(z), Psi.phase_factors, mean_probe_intensities, r, r_min, out, tau)

    if out_is_gradient:
        out /= (K * D)

    return th.view_as_complex(out)

# Cell
from .kernels import phase_factor_kernelDBK
def smatrix_phase_factorsBDK(Psi, r, take_beams, q, B, out=None):
    """
    Abbreviations:
    B: number of (input) beams in S-matrix
    D: number of scans/ aperture functions
    K: number of scan positions
    MY/MX: detector shape
    NY/NX: S-matrix shape

    :param Psi: q           D x B
    :param r:               D x K x 2
    :param take_beams:      MY x MX
    :param q:               2 x MY x MX
    :param out:             B x D x K x 2
    :return:
    """
    if out is None:
        D, K, _ = r.shape
        out = th.zeros((B, D, K, 2), dtype=th.float32, device=Psi.device)
    else:
        out[:] = 0
        _, D, K, c = out.shape

    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = 128#gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array(out.shape[:-1])) / threadsperblock)
    tb = take_beams[None, ...].expand(*Psi.shape)
    Psi_DB = Psi[tb].reshape(D, B)
    tb = take_beams[None, ...].expand(*q.shape)
    qB = q[tb].reshape(2, B)
    phase_factor_kernelDBK[blockspergrid, threadsperblock, stream](th.view_as_real(Psi_DB), r, qB, out)
    return out

# Cell
from .kernels import phase_factor_kernelKB
def smatrix_phase_factorsKB(Psi, r, take_beams, q, q_indices, B, out=None):
    """
    Abbreviations:
    B: number of (input) beams in S-matrix
    K: number of scan positions
    MY/MX: detector shape
    NY/NX: S-matrix shape

    :param Psi: q           MY x MX
    :param r:               K x 2
    :param take_beams:      MY x MX
    :param q:               2 x MY x MX
    :param q_indices:       2 x MY x MX
    :param out:             K x B x 2
    :return:
    """
    if out is None:
        K, _ = r.shape
        out = th.zeros((K, B, 2), dtype=th.float32, device=Psi.device)
    else:
        out[:] = 0
        K, B, c = out.shape

    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = 128#gpu.MAX_THREADS_PER_BLOCK
    blockspergrid = m.ceil(np.prod(np.array(out.shape[:-1])) / threadsperblock)
    tb = take_beams.expand(*Psi.shape)
    Psi_B = Psi[tb].reshape(B)
    tb = take_beams[None, ...].expand(*q_indices.shape)
    qB = q[tb].reshape(2, B)
    phase_factor_kernelKB[blockspergrid, threadsperblock, stream](th.view_as_real(Psi_B), r, qB, out)
    return out