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

__all__ = ['array_split_divpoints_ntotal', 'fftshift_checkerboard', 'cartesian_aberrations', 'memory_mb', 'memory_gb',
           'fourier_coordinates_2D', 'array_split_divpoints', 'R_factor', 'distance', 'rel_dist', 'Param',
           'validate_standard_param', 'format_standard_param', 'asParam', 'make_default', 'PARAM_PREFIX',
           'single_sideband_reconstruction', 'single_sideband_kernel', 'sector_mask', 'wavelength', 'DOF',
           'dense_to_sparse_kernel', 'advanced_raster_scan', 'advanced_raster_scan', 'get_qx_qy_1D', 'get_qx_qy_2D',
           'scatter_add_patches', 'gather_patches', 'array_split_divpoints', 'HSV_to_RGB', 'P1A_to_HSV', 'imsave',
           'plot_complex_multi', 'plot', 'zplot', 'plotAbsAngle', 'center_of_mass_kernel',
           'sparse_to_dense_datacube_kernel_crop', 'sparse_to_dense_datacube_crop',
           'sparse_to_dense_datacube_crop_gain_mask_kernel', 'sparse_to_dense_datacube_crop_gain_mask',
           'fftshift_kernel', 'fftshift_pad_kernel', 'virtual_annular_image_kernel',
           'crop_symmetric_around_center_kernel', 'crop_symmetric_around_center', 'rotate_kernel', 'rotate',
           'sum_kernel', 'sum_frames', 'rebin_kernel', 'disk_overlap_function', 'disk_overlap_kernel',
           'double_overlap_intensitities_in_range', 'find_rotation_angle_with_double_disk_overlap',
           'cartesian_aberrations_single', 'cartesian_aberrations', 'aperture', 'ZernikeProbe2', 'ZernikeProbeSingle',
           'ZernikeProbe', 'sdebug', 'h5write', 'h5append', 'h5read', 'h5info', 'h5options', 'prox_D_gaussian_kernel',
           'prox_D_gaussian', 'gradz_poisson_sparse_kernel', 'gradz_poisson_sparse', 'gradz_gaussian_sparse_kernel',
           'gradz_gaussian_sparse', 'sparse_amplitude_loss_kernel', 'sparse_amplitude_loss',
           'sparse_smooth_truncated_amplitude_loss_kernel', 'sparse_smooth_truncated_amplitude_loss',
           'sparse_smooth_truncated_amplitude_prox_kernel', 'sparse_smooth_truncated_amplitude_prox',
           'sparse_amplitude_prox_kernel', 'sparse_amplitude_prox']

# Cell
from numpy.fft import fftfreq
import numpy as np
from .torch_imports import *
from numba import cuda
import math as m
import cmath as cm
import sigpy as sp

# Cell

def array_split_divpoints_ntotal(Ntotal, indices_or_sections):
    """
    Split an array into multiple sub-arrays.
    Please refer to the ``split`` documentation.  The only difference
    between these functions is that ``array_split`` allows
    `indices_or_sections` to be an integer that does *not* equally
    divide the axis. For an array of length l that should be split
    into n sections, it returns l % n sub-arrays of size l//n + 1
    and the rest of size l//n.
    See Also
    --------
    split : Split array into multiple sub-arrays of equal size.
    Examples
    --------
    >>> x = np.arange(8.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.,  5.]), array([6.,  7.])]
    >>> x = np.arange(7.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.]), array([5.,  6.])]
    """
    try:
        # handle array case.
        Nsections = len(indices_or_sections) + 1
        div_points = [0] + list(indices_or_sections) + [Ntotal]
    except TypeError:
        # indices_or_sections is a scalar, not an array.
        Nsections = int(indices_or_sections)
        if Nsections <= 0:
            raise ValueError('number sections must be larger than 0.')
        Neach_section, extras = divmod(Ntotal, Nsections)
        section_sizes = ([0] +
                         extras * [Neach_section + 1] +
                         (Nsections - extras) * [Neach_section])
        div_points = np.array(section_sizes, dtype=np.intp).cumsum()

    return div_points

# Cell
def fftshift_checkerboard(w, h):
    re = np.r_[w * [-1, 1]]  # even-numbered rows
    ro = np.r_[w * [1, -1]]  # odd-numbered rows
    return np.row_stack(h * (re, ro))

# Cell
def cartesian_aberrations(qx, qy, lam, C):
    """
    Zernike polynomials in the cartesian coordinate system
    :param qx:
    :param qy:
    :param lam: wavelength in Angstrom
    :param C:   12 x D
    :return:
    """

    u = qx * lam
    v = qy * lam
    u2 = u ** 2
    u3 = u ** 3
    u4 = u ** 4
    # u5 = u ** 5

    v2 = v ** 2
    v3 = v ** 3
    v4 = v ** 4
    # v5 = v ** 5

    aberr = Param()
    aberr.C1 = C[0].unsqueeze(1).unsqueeze(1)
    aberr.C12a = C[1].unsqueeze(1).unsqueeze(1)
    aberr.C12b = C[2].unsqueeze(1).unsqueeze(1)
    aberr.C21a = C[3].unsqueeze(1).unsqueeze(1)
    aberr.C21b = C[4].unsqueeze(1).unsqueeze(1)
    aberr.C23a = C[5].unsqueeze(1).unsqueeze(1)
    aberr.C23b = C[6].unsqueeze(1).unsqueeze(1)
    aberr.C3 = C[7].unsqueeze(1).unsqueeze(1)
    aberr.C32a = C[8].unsqueeze(1).unsqueeze(1)
    aberr.C32b = C[9].unsqueeze(1).unsqueeze(1)
    aberr.C34a = C[10].unsqueeze(1).unsqueeze(1)
    aberr.C34b = C[11].unsqueeze(1).unsqueeze(1)

    chi = 0

    # r-2 = x-2 +y-2.
    chi += 1 / 2 * aberr.C1 * (u2 + v2) # r^2
    #r-2 cos(2*phi) = x"2 -y-2.
    # r-2 sin(2*phi) = 2*x*y.
    chi += 1 / 2 * (aberr.C12a * (u2 - v2) + 2 * aberr.C12b * u * v) # r^2 cos(2 phi) + r^2 sin(2 phi)
    # r-3 cos(3*phi) = x-3 -3*x*y'2. r"3 sin(3*phi) = 3*y*x-2 -y-3.
    chi += 1 / 3 * (aberr.C23a * (u3 - 3 * u * v2) + aberr.C23b * (3 * u2 * v - v3))# r^3 cos(3phi) + r^3 sin(3 phi)
    # r-3 cos(phi) = x-3 +x*y-2.
    # r-3 sin(phi) = y*x-2 +y-3.
    chi += 1 / 3 * (aberr.C21a * (u3 + u * v2) + aberr.C21b * (v3 + u2 * v))# r^3 cos(phi) + r^3 sin(phi)
    # r-4 = x-4 +2*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C3 * (u4 + v4 + 2 * u2 * v2)# r^4
    # r-4 cos(4*phi) = x-4 -6*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C34a * (u4 - 6 * u2 * v2 + v4)# r^4 cos(4 phi)
    # r-4 sin(4*phi) = 4*x-3*y -4*x*y-3.
    chi += 1 / 4 * aberr.C34b * (4 * u3 * v - 4 * u * v3) # r^4 sin(4 phi)
    # r-4 cos(2*phi) = x-4 -y-4.
    chi += 1 / 4 * aberr.C32a * (u4 - v4)
    # r-4 sin(2*phi) = 2*x-3*y +2*x*y-3.
    chi += 1 / 4 * aberr.C32b * (2 * u3 * v + 2 * u * v3)
    # r-5 cos(phi) = x-5 +2*x-3*y-2 +x*y-4.
    # r-5 sin(phi) = y*x"4 +2*x-2*y-3 +y-5.
    # r-5 cos(3*phi) = x-5 -2*x-3*y-2 -3*x*y-4.
    # r-5 sin(3*phi) = 3*y*x-4 +2*x-2*y-3 -y-5.
    # r-5 cos(5*phi) = x-5 -10*x-3*y-2 +5*x*y-4.
    # r-5 sin(5*phi) = 5*y*x-4 -10*x-2*y-3 +y-5.

    chi *= 2 * np.pi / lam

    return chi

# Cell
def memory_mb(x, dtype=None):
    if isinstance(x, th.Tensor):
        return x.nelement() * x.element_size() / 2 ** 20
    elif isinstance(x, tuple):
        assert dtype is not None, 'memory_mb: dtype must not be None'
        element_size = th.zeros(1, dtype=dtype).element_size()
        nelement = np.prod(np.asarray(x))
        return nelement * element_size / 2 ** 20


def memory_gb(x, dtype=None):
    if isinstance(x, th.Tensor):
        return x.nelement() * x.element_size() / 2 ** 30
    elif isinstance(x, tuple):
        assert dtype is not None, 'memory_mb: dtype must not be None'
        element_size = th.zeros(1, dtype=dtype).element_size()
        nelement = np.prod(np.asarray(x))
        return nelement * element_size / 2 ** 30

# Cell
def fourier_coordinates_2D(N, dx=[1.0, 1.0], centered=True):
    qxx = fftfreq(N[1], dx[1])
    qyy = fftfreq(N[0], dx[0])
    if centered:
        qxx += 0.5 / N[1] / dx[1]
        qyy += 0.5 / N[0] / dx[0]
    qx, qy = np.meshgrid(qxx, qyy)
    q = np.array([qy, qx]).astype(np.float32)
    return q

# Cell
def array_split_divpoints(ary, indices_or_sections, axis=0):
    """
    Split an array into multiple sub-arrays.
    Please refer to the ``split`` documentation.  The only difference
    between these functions is that ``array_split`` allows
    `indices_or_sections` to be an integer that does *not* equally
    divide the axis. For an array of length l that should be split
    into n sections, it returns l % n sub-arrays of size l//n + 1
    and the rest of size l//n.
    See Also
    --------
    split : Split array into multiple sub-arrays of equal size.
    Examples
    --------
    >>> x = np.arange(8.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.,  5.]), array([6.,  7.])]
    >>> x = np.arange(7.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.]), array([5.,  6.])]
    """
    try:
        Ntotal = ary.shape[axis]
    except AttributeError:
        Ntotal = len(ary)
    try:
        # handle array case.
        Nsections = len(indices_or_sections) + 1
        div_points = [0] + list(indices_or_sections) + [Ntotal]
    except TypeError:
        # indices_or_sections is a scalar, not an array.
        Nsections = int(indices_or_sections)
        if Nsections <= 0:
            raise ValueError('number sections must be larger than 0.')
        Neach_section, extras = divmod(Ntotal, Nsections)
        section_sizes = ([0] +
                         extras * [Neach_section + 1] +
                         (Nsections - extras) * [Neach_section])
        div_points = np.array(section_sizes, dtype=np.intp).cumsum()

    return div_points

# Cell
def R_factor(z, a, world_size=1):
    """Calculate R error metric = sum(|z-a|)/sum(|a|)."""
    num = th.norm(th.abs(z) - a,p=1)
    denom = th.norm(a,p=1)
    if world_size > 1:
        dist.all_reduce(num)
        dist.all_reduce(denom)
    return num / denom

# Cell
def distance(z, x):
    """
    Distance of two complex vectors
    :param z: tensor
    :param x: tensor
    :return:
    """
    c = th.vdot(z.ravel(), x.ravel())
    phi = -th.angle(c)
    exp_minus_phi = th.exp(1j * phi)
    p = exp_minus_phi.to(x.device)
    x_hat = x * p
    res = th.norm(z - x_hat,2)
    return res

# Cell
def rel_dist(z, x):
    """
    Distance of two complex vectors
    :param z: tensor
    :param x: tensor
    :return:
    """
    d = distance(z, x)
    x_norm = th.norm(x,2)
    return d / x_norm

# Cell
PARAM_PREFIX = 'pars'
class Param(dict):
    """
    Convenience class: a dictionary that gives access to its keys
    through attributes.

    Note: dictionaries stored in this class are also automatically converted
    to Param objects:
    >>> p = Param()
    >>> p.x = {}
    >>> p
    Param({})

    While dict(p) returns a dictionary, it is not recursive, so it is better in this case
    to use p.todict(). However, p.todict does not check for infinite recursion. So please
    don't store a dictionary (or a Param) inside itself.

    BE: Please note also that the recursive behavior of the update function will create
    new references. This will lead inconsistency if other objects refer to dicts or Params
    in the updated Param instance.
    """
    _display_items_as_attributes = True
    _PREFIX = PARAM_PREFIX

    def __init__(self, __d__=None, **kwargs):
        """
        A Dictionary that enables access to its keys as attributes.
        Same constructor as dict.
        """
        dict.__init__(self)
        if __d__ is not None: self.update(__d__)
        self.update(kwargs)

    def __getstate__(self):
        return self.__dict__.items()

    def __setstate__(self, items):
        for key, val in items:
            self.__dict__[key] = val

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, dict.__repr__(self))

    # def __str__(self):
    #     from .verbose import report
    #     return report(self,depth=7,noheader=True)

    def __setitem__(self, key, value):
        # BE: original behavior modified as implicit conversion may destroy references
        # Use update(value,Convert=True) instead
        # return super(Param, self).__setitem__(key, Param(value) if type(value) == dict else value)
        return super(Param, self).__setitem__(key, value)

    def __getitem__(self, name):
        # item = super(Param, self).__getitem__(name)
        # return Param(item) if type(item) == dict else item
        return super(Param, self).__getitem__(name)

    def __delitem__(self, name):
        return super(Param, self).__delitem__(name)

    def __delattr__(self, name):
        return super(Param, self).__delitem__(name)

    # __getattr__ = __getitem__
    def __getattr__(self, name):
        try:
            return self.__getitem__(name)
        except KeyError as ke:
            raise AttributeError(ke)

    __setattr__ = __setitem__

    def copy(self, depth=0):
        """
        :returns Param: A (recursive) copy of P with depth `depth`
        """
        d = Param(self)
        if depth > 0:
            for k, v in d.iteritems():
                if isinstance(v, self.__class__): d[k] = v.copy(depth - 1)
        return d

    def __dir__(self):
        """
        Defined to include the keys when using dir(). Useful for
        tab completion in e.g. ipython.
        If you do not wish the dict key's be displayed as attributes
        (although they are still accessible as such) set the class
        attribute `_display_items_as_attributes` to False. Default is
        True.
        """
        if self._display_items_as_attributes:
            return self.keys()
            # return [item.__dict__.get('name',str(key)) for key,item in self.iteritems()]
        else:
            return []

    def update(self, __d__=None, in_place_depth=0, Convert=False, **kwargs):
        """
        Update Param - almost same behavior as dict.update, except
        that all dictionaries are converted to Param if `Convert` is set
        to True, and update may occur in-place recursively for other Param
        instances that self refers to.

        Parameters
        ----------
        Convert : bool
                  If True, convert all dict-like values in self also to Param.
                  *WARNING*
                  This mey result in misdirected references in your environment
        in_place_depth : int
                  Counter for recursive in-place updates
                  If the counter reaches zero, the Param to a key is
                  replaced instead of updated
        """

        def _k_v_update(k, v):
            # If an element is itself a dict, convert it to Param
            if Convert and hasattr(v, 'keys'):
                # print 'converting'
                v = Param(v)
            # new key
            if not k in self:
                self[k] = v
            # If this key already exists and is already dict-like, update it
            elif in_place_depth > 0 and hasattr(v, 'keys') and isinstance(self[k], self.__class__):
                self[k].update(v, in_place_depth - 1)
                """
                if isinstance(self[k],self.__class__):
                    # Param gets recursive in_place updates
                    self[k].update(v, in_place_depth - 1)
                else:
                    # dicts are only updated in-place once
                    self[k].update(v)
                """
            # Otherwise just replace it
            else:
                self[k] = v

        if __d__ is not None:
            if hasattr(__d__, 'keys'):
                # Iterate through dict-like argument
                for k, v in __d__.items():
                    _k_v_update(k, v)

            else:
                # here we assume a (key,value) list.
                for (k, v) in __d__:
                    _k_v_update(k, v)

        for k, v in kwargs.items():
            _k_v_update(k, v)

        return None

    def _to_dict(self, Recursive=False):
        """
        Convert to dictionary (recursively if needed).
        """
        if not Recursive:
            return dict(self)
        else:
            d = dict(self)
            for k, v in d.items():
                if isinstance(v, self.__class__): d[k] = v._to_dict(Recursive)
        return d

    @classmethod
    def _from_dict(cls, dct):
        """
        Make Param from dict. This is similar to the __init__ call
        """
        # p=Param()
        # p.update(dct.copy())
        return Param(dct.copy())


def validate_standard_param(sp, p=None, prefix=None):
    """\
    validate_standard_param(sp) checks if sp follows the standard parameter convention.
    validate_standard_param(sp, p) attemps to check if p is a valid implementation of sp.

    NOT VERY SOPHISTICATED FOR NOW!
    """
    if p is None:
        good = True
        for k, v in sp.iteritems():
            if k.startswith('_'): continue
            if type(v) == type(sp):
                pref = k if prefix is None else '.'.join([prefix, k])
                good &= validate_standard_param(v, prefix=pref)
                continue
            else:
                try:
                    a, b, c = v
                    if prefix is not None:
                        print('    %s.%s = %s' % (prefix, k, str(v)))
                    else:
                        print('    %s = %s' % (k, str(v)))
                except:
                    good = False
                    if prefix is not None:
                        print('!!! %s.%s = %s <--- Incorrect' % (prefix, k, str(v)))
                    else:
                        print('!!! %s = %s <--- Incorrect' % (k, str(v)))

        return good
    else:
        raise RuntimeError('Checking if a param fits with a standard is not yet implemented')


def format_standard_param(p):
    """\
    Pretty-print a Standard Param class.
    """
    lines = []
    if not validate_standard_param(p):
        print('Standard parameter does not')
    for k, v in p.iteritems():
        if k.startswith('_'): continue
        if type(v) == type(p):
            sublines = format_standard_param(v)
            lines += [k + '.' + s for s in sublines]
        else:
            lines += ['%s = %s #[%s] %s' % (k, str(v[1]), v[0], v[2])]
    return lines


def asParam(obj):
    """
    Convert the input to a Param.

    Parameters
    ----------
    a : dict_like
        Input structure, in any format that can be converted to a Param.

    Returns:
    out : Param
        The Param structure built from a. No copy is done if the input
        is already a Param.
    """
    return obj if isinstance(obj, Param) else Param(obj)


def make_default(default_dict_or_file):
    """
    convert description dict to a module dict using a possibly verbose Q & A game
    """
    pass

# Cell

def single_sideband_reconstruction(G, Qx_all, Qy_all, Kx_all, Ky_all, aberrations, theta_rot, alpha_rad,
                                   Ψ_Qp, Ψ_Qp_left_sb, Ψ_Qp_right_sb, eps, lam):
    xp = sp.backend.get_array_module(G)
    threadsperblock = 2 ** 8
    blockspergrid = m.ceil(np.prod(G.shape) / threadsperblock)
    strides = xp.array((np.array(G.strides) / (G.nbytes / G.size)).astype(np.int))
    scale = 1
    single_sideband_kernel[blockspergrid, threadsperblock](G, strides, Qx_all, Qy_all, Kx_all, Ky_all, aberrations,
                                                           theta_rot, alpha_rad, Ψ_Qp, Ψ_Qp_left_sb,
                                                           Ψ_Qp_right_sb, eps, lam, scale)
    xp.cuda.Device(Ψ_Qp.device).synchronize()

@cuda.jit
def single_sideband_kernel(G, strides, Qx_all, Qy_all, Kx_all, Ky_all, aberrations, theta_rot, alpha,
                           Ψ_Qp, Ψ_Qp_left_sb, Ψ_Qp_right_sb, eps, lam, scale):
    def aperture2(qx, qy, lam, alpha_max, scale):
        qx2 = qx ** 2
        qy2 = qy ** 2
        q = m.sqrt(qx2 + qy2)
        ktheta = m.asin(q * lam)
        return (ktheta < alpha_max) * scale

    def chi3(qy, qx, lam, C):
        """
        Zernike polynomials in the cartesian coordinate system
        :param qx:
        :param qy:
        :param lam: wavelength in Angstrom
        :param C:   (12 ,)
        :return:
        """

        u = qx * lam
        v = qy * lam
        u2 = u ** 2
        u3 = u ** 3
        u4 = u ** 4
        # u5 = u ** 5

        v2 = v ** 2
        v3 = v ** 3
        v4 = v ** 4
        # v5 = v ** 5

        # aberr = Param()
        # aberr.C1 = C[0]
        # aberr.C12a = C[1]
        # aberr.C12b = C[2]
        # aberr.C21a = C[3]
        # aberr.C21b = C[4]
        # aberr.C23a = C[5]
        # aberr.C23b = C[6]
        # aberr.C3 = C[7]
        # aberr.C32a = C[8]
        # aberr.C32b = C[9]
        # aberr.C34a = C[10]
        # aberr.C34b = C[11]

        chi = 0

        # r-2 = x-2 +y-2.
        chi += 1 / 2 * C[0] * (u2 + v2)  # r^2
        # r-2 cos(2*phi) = x"2 -y-2.
        # r-2 sin(2*phi) = 2*x*y.
        chi += 1 / 2 * (C[1] * (u2 - v2) + 2 * C[2] * u * v)  # r^2 cos(2 phi) + r^2 sin(2 phi)
        # r-3 cos(3*phi) = x-3 -3*x*y'2. r"3 sin(3*phi) = 3*y*x-2 -y-3.
        chi += 1 / 3 * (C[5] * (u3 - 3 * u * v2) + C[6] * (3 * u2 * v - v3))  # r^3 cos(3phi) + r^3 sin(3 phi)
        # r-3 cos(phi) = x-3 +x*y-2.
        # r-3 sin(phi) = y*x-2 +y-3.
        chi += 1 / 3 * (C[3] * (u3 + u * v2) + C[4] * (v3 + u2 * v))  # r^3 cos(phi) + r^3 sin(phi)
        # r-4 = x-4 +2*x-2*y-2 +y-4.
        chi += 1 / 4 * C[7] * (u4 + v4 + 2 * u2 * v2)  # r^4
        # r-4 cos(4*phi) = x-4 -6*x-2*y-2 +y-4.
        chi += 1 / 4 * C[10] * (u4 - 6 * u2 * v2 + v4)  # r^4 cos(4 phi)
        # r-4 sin(4*phi) = 4*x-3*y -4*x*y-3.
        chi += 1 / 4 * C[11] * (4 * u3 * v - 4 * u * v3)  # r^4 sin(4 phi)
        # r-4 cos(2*phi) = x-4 -y-4.
        chi += 1 / 4 * C[8] * (u4 - v4)
        # r-4 sin(2*phi) = 2*x-3*y +2*x*y-3.
        chi += 1 / 4 * C[9] * (2 * u3 * v + 2 * u * v3)
        # r-5 cos(phi) = x-5 +2*x-3*y-2 +x*y-4.
        # r-5 sin(phi) = y*x"4 +2*x-2*y-3 +y-5.
        # r-5 cos(3*phi) = x-5 -2*x-3*y-2 -3*x*y-4.
        # r-5 sin(3*phi) = 3*y*x-4 +2*x-2*y-3 -y-5.
        # r-5 cos(5*phi) = x-5 -10*x-3*y-2 +5*x*y-4.
        # r-5 sin(5*phi) = 5*y*x-4 -10*x-2*y-3 +y-5.

        chi *= 2 * np.pi / lam

        return chi

    gs = G.shape
    N = gs[0] * gs[1] * gs[2] * gs[3]
    n = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
    iqy = n // strides[0]
    iqx = (n - iqy * strides[0]) // strides[1]
    iky = (n - (iqy * strides[0] + iqx * strides[1])) // strides[2]
    ikx = (n - (iqy * strides[0] + iqx * strides[1] + iky * strides[2])) // strides[3]

    if n < N:

        Qx = Qx_all[iqx]
        Qy = Qy_all[iqy]
        Kx = Kx_all[ikx]
        Ky = Ky_all[iky]

        Qx_rot = Qx * m.cos(theta_rot) - Qy * m.sin(theta_rot)
        Qy_rot = Qx * m.sin(theta_rot) + Qy * m.cos(theta_rot)

        Qx = Qx_rot
        Qy = Qy_rot

        A = aperture2(Ky, Kx, lam, alpha, scale) * cm.exp(-1j * chi3(Ky, Kx, lam, aberrations))
        chi_KplusQ = chi3(Ky + Qy, Kx + Qx, lam, aberrations)
        A_KplusQ = aperture2(Ky + Qy, Kx + Qx, lam, alpha, scale) * cm.exp(-1j * chi_KplusQ)
        chi_KminusQ = chi3(Ky - Qy, Kx - Qx, lam, aberrations)
        A_KminusQ = aperture2(Ky - Qy, Kx - Qx, lam, alpha, scale) * cm.exp(-1j * chi_KminusQ)

        Γ = A.conjugate() * A_KminusQ - A * A_KplusQ.conjugate()

        Kplus = sqrt((Kx + Qx) ** 2 + (Ky + Qy) ** 2)
        Kminus = sqrt((Kx - Qx) ** 2 + (Ky - Qy) ** 2)
        K = sqrt(Kx ** 2 + Ky ** 2)
        bright_field = K < alpha / lam
        double_overlap1 = (Kplus < alpha / lam) * bright_field * (Kminus > alpha / lam)
        double_overlap2 = (Kplus > alpha / lam) * bright_field * (Kminus < alpha / lam)

        Γ_abs = abs(Γ)
        take = Γ_abs > eps and bright_field
        if take:
            val = G[iqy, iqx, iky, ikx] * Γ.conjugate()
            cuda.atomic.add(Ψ_Qp.real, (iqy, iqx), val.real)
            cuda.atomic.add(Ψ_Qp.imag, (iqy, iqx), val.imag)
        if double_overlap1:
            val = G[iqy, iqx, iky, ikx] * Γ.conjugate()
            cuda.atomic.add(Ψ_Qp_left_sb.real, (iqy, iqx), val.real)
            cuda.atomic.add(Ψ_Qp_left_sb.imag, (iqy, iqx), val.imag)
        if double_overlap2:
            val = G[iqy, iqx, iky, ikx] * Γ.conjugate()
            cuda.atomic.add(Ψ_Qp_right_sb.real, (iqy, iqx), val.real)
            cuda.atomic.add(Ψ_Qp_right_sb.imag, (iqy, iqx), val.imag)
        if iqx == 0 and iqy == 0:
            val = abs(G[iqy, iqx, iky, ikx]) + 1j * 0
            cuda.atomic.add(Ψ_Qp.real, (iqy, iqx), val.real)
            cuda.atomic.add(Ψ_Qp_left_sb.real, (iqy, iqx), val.real)
            cuda.atomic.add(Ψ_Qp_right_sb.real, (iqy, iqx), val.real)

# Cell

import matplotlib
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
import matplotlib.font_manager as fm
import numpy as np
from PIL import Image

def sector_mask(shape, centre, radius, angle_range=(0,360)):
    """
    Return a boolean mask for a circular sector. The start/stop angles in
    `angle_range` should be given in clockwise order.
    """

    x, y = np.ogrid[:shape[0], :shape[1]]
    cx, cy = centre
    tmin, tmax = np.deg2rad(angle_range)

    # ensure stop angle > start angle
    if tmax < tmin:
        tmax += 2 * np.pi

    # convert cartesian --> polar coordinates
    r2 = (x - cx) * (x - cx) + (y - cy) * (y - cy)
    theta = np.arctan2(x - cx, y - cy) - tmin

    # wrap angles between 0 and 2*pi
    theta %= (2 * np.pi)

    # circular mask
    circmask = r2 <= radius * radius

    # angular mask
    anglemask = theta <= (tmax - tmin)

    return circmask * anglemask

# Cell
from math import sqrt
def wavelength(E_eV):
    emass = 510.99906;  # electron rest mass in keV
    hc = 12.3984244;  # h*c
    lam = hc / m.sqrt(E_eV * 1e-3 * (2 * emass + E_eV * 1e-3))  # in Angstrom
    return lam


def DOF(alpha, E_eV):
    E0 = E_eV

    # Calculate wavelength and electron interaction parameter
    m = 9.109383 * 10 ** -31
    e = 1.602177 * 10 ** -19
    c = 299792458
    h = 6.62607 * 10 ** -34

    lam = h / sqrt(2 * m * e * E0) / sqrt(1 + e * E0 / 2 / m / c ** 2) * 10 ** 10
    DOF = 2 * lam / alpha ** 2
    return DOF

# Cell
@cuda.jit
def dense_to_sparse_kernel(dense, indices, counts, frame_dimensions):
    ny, nx = cuda.grid(2)
    NY, NX, MYBIN, MXBIN = dense.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        k = 0
        for mx in range(MX):
            for my in range(MY):
                idx1d = my * MX + mx
                if dense[ny,nx,my,mx] > 0:
                    indices[ny,nx,k] = idx1d
                    counts[ny,nx,k] = dense[ny,nx,my,mx]
                    k += 1

def advanced_raster_scan(ny=10, nx=10, fast_axis=1, mirror=[1, 1], theta=0, dy=1, dx=1):
    """
    Generates as raster scan.

    Parameters
    ----------
    ny, nx : int
        Number of steps in *y* (vertical) and *x* (horizontal) direction
        *x* is the fast axis

    dy, dx : float
        Step size (grid spacinf) in *y* and *x*
        2
    Returns
    -------
    pos : ndarray
        A (N,2)-array of positions.

    Examples
    --------
    """
    iix, iiy = np.indices((nx, ny))
    if fast_axis != 1:
        tmp = iix
        iix = iiy
        iiy = tmp

    # print iix.shape, iiy.shape
    positions = np.array([(dx * i, dy * j) for i, j in zip(iix.ravel(), iiy.ravel())]).astype(np.float32)

    mins = np.array([positions[:, 0].min(), positions[:, 1].min()])
    maxs = np.array([positions[:, 0].max(), positions[:, 1].max()])

    center = mins + (maxs - mins) / 2.0
    positions -= center

    positions[:, 0] *= mirror[0]
    positions[:, 1] *= mirror[1]

    theta_rad = theta / 180.0 * np.pi
    R = np.array([[np.cos(theta_rad), -np.sin(theta_rad)],
                  [np.sin(theta_rad), np.cos(theta_rad)]])
    # rotate counterclockwise by theta
    positions = positions.dot(R)
    mins = np.array([positions[:, 0].min(), positions[:, 1].min()])
    positions -= mins
    return positions.astype(np.float32)

# Cell
def advanced_raster_scan(ny=10, nx=10, fast_axis=1, mirror=[1, 1], theta=0, dy=1, dx=1):
    """
    Generates as raster scan.

    Parameters
    ----------
    ny, nx : int
        Number of steps in *y* (vertical) and *x* (horizontal) direction
        *x* is the fast axis

    dy, dx : float
        Step size (grid spacinf) in *y* and *x*

    Returns
    -------
    pos : ndarray
        A (N,2)-array of positions.

    Examples
    --------
    """
    iiy, iix = np.indices((ny, nx))
    if fast_axis != 1:
        tmp = iix
        iix = iiy
        iiy = tmp

    # print iix.shape, iiy.shape
    positions = np.array([(dy * i, dx * j) for i, j in zip(iiy.ravel(), iix.ravel())]).astype(np.float32)

    mins = np.array([positions[:, 0].min(), positions[:, 1].min()])
    maxs = np.array([positions[:, 0].max(), positions[:, 1].max()])

    center = mins + (maxs - mins) / 2.0
    positions -= center

    positions[:, 0] *= mirror[0]
    positions[:, 1] *= mirror[1]

    theta_rad = theta / 180.0 * np.pi
    R = np.array([[np.cos(theta_rad), -np.sin(theta_rad)],
                  [np.sin(theta_rad), np.cos(theta_rad)]])
    # rotate counterclockwise by theta
    positions = positions.dot(R)
    mins = np.array([positions[:, 0].min(), positions[:, 1].min()])
    positions -= mins
    return positions.astype(np.float32)

# Cell

def get_qx_qy_1D(M, dx, dtype, fft_shifted=False):
    xp = sp.backend.get_array_module(dx)
    qxa = xp.fft.fftfreq(M[0], dx[0]).astype(dtype)
    qya = xp.fft.fftfreq(M[1], dx[1]).astype(dtype)
    if fft_shifted:
        qxa = xp.fft.fftshift(qxa)
        qya = xp.fft.fftshift(qya)
    return qxa, qya


def get_qx_qy_2D(M, dx, dtype, fft_shifted=False):
    xp = sp.backend.get_array_module(dx)
    qxa = xp.fft.fftfreq(M[0], dx[0]).astype(dtype)
    qya = xp.fft.fftfreq(M[1], dx[1]).astype(dtype)
    [qxn, qyn] = xp.meshgrid(qxa, qya)
    if fft_shifted:
        qxn = xp.fft.fftshift(qxn)
        qyn = xp.fft.fftshift(qyn)
    return qxn, qyn

# Cell
import numpy as np
import torch as th
from math import sqrt

def scatter_add_patches(input: th.Tensor, out: th.Tensor, axes, positions, patch_size, reduce_dim=None) -> th.Tensor:
    """
    Scatter_adds K patches of size :patch_size: at axes [ax1, ax2] into the output tensor. The patches are added at
    positions :positions:. Additionally, several dimensions of the input tensor can be summed, specified by reduce_dims.

    :param input:   K x M1 x M2 at least 3-dimensional tensor of patches,
    :param out: at least two-dimensional tensor that is the scatter_add target
    :param axes: (2,) axes at which to scatter the input
    :param positions: K x 2 LongTensor
    :param patch_size: (2,) LongTensor
    :param reduce_dims: (N,) LongTensor
    :return: out, the target of scatter_add_
    """
    if reduce_dim is not None:
        other1 = th.split(input, 1, dim=reduce_dim)
        other = tuple()
        for one in other1:
            other += (one.squeeze_().contiguous(),)
        # now we have D tensors of shape K B M1 M2 2
    else:
        other = [input]

    r = positions
    s = patch_size
    K = r.shape[0]

    # patches has dimension  K x M1 x M2 x 2

    index0 = th.arange(s[0], device=input.device, dtype=th.long).view(s[0], 1).expand(s[0], s[1])
    index1 = th.arange(s[1], device=input.device, dtype=th.long).view(1, s[1]).expand(s[0], s[1])

    # size is patch_size
    # print(f"strides at axes: {input.stride(axes[0]), input.stride(axes[1])}")
    index = out.stride(axes[0]) * (index0 + r[:, 0].view(K, 1, 1)) + out.stride(axes[1]) * (
            index1 + r[:, 1].view(K, 1, 1))

    # print(f"new index shape: {index.shape}")
    higher_dim_offsets = th.arange(out.stride(axes[1]), device=input.device).view(1, 1, 1, out.stride(axes[1]))
    index = index.view(index.shape[0], index.shape[1], index.shape[2], 1).expand(
        (index.shape[0], index.shape[1], index.shape[2], out.stride(axes[1])))

    # print(higher_dim_offsets, higher_dim_offsets.shape, index.shape)
    index = index + higher_dim_offsets
    # now we have the K x M1 x M2 x 2 indices into the N1 x N2 x 2 array

    # print(f"new index shape: {index.shape}")
    # index = index.view(index.shape[0], 1, index.shape[1], index.shape[2], index.shape[3]).expand(
    #     (index.shape[0], B, index.shape[1], index.shape[2], index.shape[3]))

    # print(f"max index   : {th.max(index.view(-1))}")
    # print(f"len   out   : {out.view(-1).shape[0]}")
    # print(f"index shape : {index.view(-1).shape[0]}")
    # print(f"others shape: {other[0].view(-1).shape[0]}")

    for i, one in enumerate(other):
        # print(f"others [{i}] shape: {one.view(-1).shape[0]}")
        # print(one.shape)
        # print(index.shape)
        out.view(-1).scatter_add_(0, index.view(-1), one.view(-1))
    return out

# Cell
def gather_patches(input, axes, positions, patch_size, out=None) -> th.Tensor:
    """
    Gathers K patches of size :patch_size: at axes [ax1, ax2] of the input tensor. The patches are collected started at
    K positions pos.

    if :input: is an n-dimensional tensor with size (x_0, x_1, x_2, ..., x_a, x_ax1, x_ax2, x_b, ..., x_{n-1})
    then :out: is an n-dimensional tensor with size  (K, x_0, x_1, x_2, ..., x_a, patch_size[0], patch_size[1], x_3, ..., x_{n-1})

    :param input: at least two-dimensional tensor
    :param axes: axes at which to gather the patches
    :param positions: K x 2 LongTensor
    :param patch_size: (2,) LongTensor
    :param out: n-dimensional tensor with size  (K, x_0, x_1, x_2, ..., x_a, patch_size[0], patch_size[1], x_3, ..., x_{n-1})
    :return:
    """
    # print(f"input shape: {input.shape}")
    # print(f"positions.dtype {positions.dtype}")
    r = positions
    s = patch_size
    K = positions.shape[0]

    # condense all dimensions x_0 ... x_a
    dim0size = th.prod(th.Tensor([input.shape[:axes[0]]])).int().item() if axes[0] > 0 else 1
    view = [dim0size]
    for d in input.shape[axes[0]:]:
        view.append(d)
    y = input.view(th.Size(view)).squeeze()

    index0 = th.arange(s[0], device=input.device, dtype=th.long).view(s[0], 1).expand(s[0], s[1])
    index1 = th.arange(s[1], device=input.device, dtype=th.long).view(1, s[1]).expand(s[0], s[1])

    # size is patch_size
    # print(f"strides at axes: {input.stride(axes[0]), input.stride(axes[1])}")
    # print(index0.dtype, r.dtype)
    index = input.stride(axes[0]) * (index0 + r[:, 0].view(K, 1, 1)) + input.stride(axes[1]) * (
            index1 + r[:, 1].view(K, 1, 1))
    # print(f"new index shape: {index.shape}")
    higher_dim_offsets = th.arange(input.stride(axes[1]), device=input.device).view(1, 1, 1, input.stride(axes[1]))
    index = index.view(index.shape[0], index.shape[1], index.shape[2], 1).expand(
        (index.shape[0], index.shape[1], index.shape[2], input.stride(axes[1])))
    # print(higher_dim_offsets, higher_dim_offsets.shape, index.shape)
    index = index + higher_dim_offsets
    # print(f"new index shape: {index.shape}")
    index = index.view(index.shape[0], 1, index.shape[1], index.shape[2], index.shape[3]).expand(
        (index.shape[0], dim0size, index.shape[1], index.shape[2], index.shape[3]))
    # print(f"new index shape: {index.shape}")
    lower_dim_offset = th.arange(dim0size, device=input.device) * y.stride(0)
    lower_dim_offset = lower_dim_offset.view(1, dim0size, 1, 1, 1).long()
    index = index + lower_dim_offset
    # print(f"new index shape: {index.shape}")
    index = index.contiguous().view(-1)
    out = th.index_select(y.view(-1), 0, index, out=out)

    out_view = (K,)
    for ax in input.shape[:axes[0]]:
        out_view += (ax,)
    out_view += (patch_size[0].item(),)
    out_view += (patch_size[1].item(),)
    for ax in input.shape[axes[1] + 1:]:
        out_view += (ax,)

    out = out.view(out_view)
    return out



# Cell
def array_split_divpoints(ary, indices_or_sections, axis=0):
    """
    Split an array into multiple sub-arrays.
    Please refer to the ``split`` documentation.  The only difference
    between these functions is that ``array_split`` allows
    `indices_or_sections` to be an integer that does *not* equally
    divide the axis. For an array of length l that should be split
    into n sections, it returns l % n sub-arrays of size l//n + 1
    and the rest of size l//n.
    See Also
    --------
    split : Split array into multiple sub-arrays of equal size.
    Examples
    --------
    >>> x = np.arange(8.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.,  5.]), array([6.,  7.])]
    >>> x = np.arange(7.0)
    >>> np.array_split(x, 3)
        [array([0.,  1.,  2.]), array([3.,  4.]), array([5.,  6.])]
    """
    try:
        Ntotal = ary.shape[axis]
    except AttributeError:
        Ntotal = len(ary)
    try:
        # handle array case.
        Nsections = len(indices_or_sections) + 1
        div_points = [0] + list(indices_or_sections) + [Ntotal]
    except TypeError:
        # indices_or_sections is a scalar, not an array.
        Nsections = int(indices_or_sections)
        if Nsections <= 0:
            raise ValueError('number sections must be larger than 0.')
        Neach_section, extras = divmod(Ntotal, Nsections)
        section_sizes = ([0] +
                         extras * [Neach_section + 1] +
                         (Nsections - extras) * [Neach_section])
        div_points = np.array(section_sizes, dtype=np.intp).cumsum()

    return div_points

# Cell
import matplotlib
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
import matplotlib.font_manager as fm
import numpy as np
from PIL import Image

# Cell
def HSV_to_RGB(cin):
    """\
    HSV to RGB transformation.
    """

    # HSV channels
    h, s, v = cin

    i = (6. * h).astype(int)
    f = (6. * h) - i
    p = v * (1. - s)
    q = v * (1. - s * f)
    t = v * (1. - s * (1. - f))
    i0 = (i % 6 == 0)
    i1 = (i == 1)
    i2 = (i == 2)
    i3 = (i == 3)
    i4 = (i == 4)
    i5 = (i == 5)

    imout = np.zeros(h.shape + (3,), dtype=h.dtype)
    imout[:, :, 0] = 255 * (i0 * v + i1 * q + i2 * p + i3 * p + i4 * t + i5 * v)
    imout[:, :, 1] = 255 * (i0 * t + i1 * v + i2 * v + i3 * q + i4 * p + i5 * p)
    imout[:, :, 2] = 255 * (i0 * p + i1 * p + i2 * t + i3 * v + i4 * v + i5 * q)

    return imout
def P1A_to_HSV(cin, vmin=None, vmax=None):
    """\
    Transform a complex array into an RGB image,
    mapping phase to hue, amplitude to value and
    keeping maximum saturation.
    """
    # HSV channels
    h = .5 * np.angle(cin) / np.pi + .5
    s = np.ones(cin.shape)

    v = abs(cin)
    if vmin is None: vmin = 0.
    if vmax is None: vmax = v.max()
    assert vmin < vmax
    v = (v.clip(vmin, vmax) - vmin) / (vmax - vmin)

    return HSV_to_RGB((h, s, v))


def imsave(a, filename=None, vmin=None, vmax=None, cmap=None):
    """
    imsave(a) converts array a into, and returns a PIL image
    imsave(a, filename) returns the image and also saves it to filename
    imsave(a, ..., vmin=vmin, vmax=vmax) clips the array to values between vmin and vmax.
    imsave(a, ..., cmap=cmap) uses a matplotlib colormap.
    """

    if a.dtype.kind == 'c':
        # Image is complex
        if cmap is not None:
            print('imsave: Ignoring provided cmap - input array is complex')
        i = P1A_to_HSV(a, vmin, vmax)
        im = Image.fromarray(np.uint8(i), mode='RGB')

    else:
        if vmin is None:
            vmin = a.min()
        if vmax is None:
            vmax = a.max()
        im = Image.fromarray((255 * (a.clip(vmin, vmax) - vmin) / (vmax - vmin)).astype('uint8'))
        if cmap is not None:
            r = im.point(lambda x: cmap(x / 255.0)[0] * 255)
            g = im.point(lambda x: cmap(x / 255.0)[1] * 255)
            b = im.point(lambda x: cmap(x / 255.0)[2] * 255)
            im = Image.merge("RGB", (r, g, b))
            # b = (255*(a.clip(vmin,vmax)-vmin)/(vmax-vmin)).astype('uint8')
            # im = Image.fromstring('L', a.shape[-1::-1], b.tostring())

    if filename is not None:
        im.save(filename)
    return im

# Cell
def plot_complex_multi(x, title='_', figsize=(10, 10), savePath=None, scale=None, show=True):
    n, h, w = x.shape
    rows = int(np.floor(np.sqrt(n)))
    cols = n // rows + 1
    fontprops = fm.FontProperties(size=18)
    fig = plt.figure(dpi=300, constrained_layout=True)
    gs1 = gridspec.GridSpec(rows, cols)
    gs1.update(wspace=0.1, hspace=0.1)
    for r in range(rows):
        for c in range(cols):
            i = cols * r + c
            ax = plt.subplot(gs1[i])
            if i < n:
                imax1 = ax.imshow(imsave(x[i]), interpolation='nearest')
                if scale is not None and i == 0:
                    scalebar = AnchoredSizeBar(ax.transData,
                                               scale[0], scale[1], 'lower right',
                                               pad=0.1,
                                               color='white',
                                               frameon=False,
                                               size_vertical=x.shape[0] / 40,
                                               fontproperties=fontprops)
                    ax.add_artist(scalebar)
            else:
                ax.axis('off')
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            ax.set_xticks([])
            ax.set_yticks([])

    plt.suptitle(title)
    plt.grid(False)
    fig.set_constrained_layout_pads(w_pad=0.1, h_pad=0.1)
    if savePath is not None:
        fig.savefig(savePath + '.png')
        fig.savefig(savePath + '.pdf')
    if show:
        plt.show()
    return fig

# Cell
def plot(img, title='Image', savePath=None, cmap='inferno', show=True, vmax=None, figsize=(10, 10), scale=None):
    fig, ax = plt.subplots(figsize=figsize)
    im = ax.imshow(img, interpolation='nearest', cmap=plt.cm.get_cmap(cmap), vmax=vmax)
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    ax.set_title(title)
    fontprops = fm.FontProperties(size=18)
    if scale is not None:
        scalebar = AnchoredSizeBar(ax.transData,
                                   scale[0], scale[1], 'lower right',
                                   pad=0.1,
                                   color='white',
                                   frameon=False,
                                   size_vertical=img.shape[0] / 40,
                                   fontproperties=fontprops)

        ax.add_artist(scalebar)
    ax.grid(False)
    if savePath is not None:
        fig.savefig(savePath + '.pdf', dpi=600)
        # fig.savefig(savePath + '.eps', dpi=600)
    if show:
        plt.show()
    return fig

# Cell
def zplot(imgs, suptitle='Image', savePath=None, cmap=['hot', 'hsv'], title=['Abs', 'Phase'], show=True,
          figsize=(9, 5), scale=None):
    im1, im2 = imgs
    fig = plt.figure(figsize=figsize, dpi=300)
    fig.suptitle(suptitle, fontsize=15, y=0.8)
    gs1 = gridspec.GridSpec(1, 2)
    gs1.update(wspace=0, hspace=0)  # set the spacing between axes.
    ax1 = plt.subplot(gs1[0])
    ax2 = plt.subplot(gs1[1])
    div1 = make_axes_locatable(ax1)
    div2 = make_axes_locatable(ax2)

    imax1 = ax1.imshow(im1, interpolation='nearest', cmap=plt.cm.get_cmap(cmap[0]))
    imax2 = ax2.imshow(im2, interpolation='nearest', cmap=plt.cm.get_cmap(cmap[1]))

    cax1 = div1.append_axes("left", size="10%", pad=0.4)
    cax2 = div2.append_axes("right", size="10%", pad=0.4)

    cbar1 = plt.colorbar(imax1, cax=cax1)
    cbar2 = plt.colorbar(imax2, cax=cax2)

    cax1.yaxis.set_ticks_position('left')
    ax2.yaxis.set_ticks_position('right')

    ax1.set_title(title[0])
    ax2.set_title(title[1])

    if scale is not None:
        fontprops = fm.FontProperties(size=18)
        scalebar = AnchoredSizeBar(ax1.transData,
                                   scale[0], scale[1], 'lower right',
                                   pad=0.1,
                                   color='white',
                                   frameon=False,
                                   size_vertical=im1.shape[0] / 40,
                                   fontproperties=fontprops)

        ax1.add_artist(scalebar)

    ax1.grid(False)
    ax2.grid(False)

    if show:
        plt.show()
    if savePath is not None:
        # print 'saving'
        fig.savefig(savePath + '.png')


def plotAbsAngle(img, suptitle='Image', savePath=None, cmap=['gray', 'gray'], title=['Abs', 'Phase'], show=True,
                 figsize=(10, 10), scale=None):
    zplot([np.abs(img), np.angle(img)], suptitle, savePath, cmap, title, show, figsize, scale)

# Cell
import numpy as np
import torch as th
import sigpy as sp
import numpy as np
from .torch_imports import *
import numba.cuda as cuda

@cuda.jit
def center_of_mass_kernel(comx, comy, indices, counts, frame_dimensions, no_count_indicator, qx, qy):
    ny, nx = cuda.grid(2)
    NY, NX, _ = indices.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            if idx1d != no_count_indicator:
                cuda.atomic.add(comy, (ny, nx), counts[ny, nx, i] * qy[my, mx])
                cuda.atomic.add(comx, (ny, nx), counts[ny, nx, i] * qx[my, mx])

# Cell

from numba import cuda
@cuda.jit
def sparse_to_dense_datacube_kernel_crop(dc, indices, counts, frame_dimensions, bin, start, end, no_count_indicator):
    ny, nx = cuda.grid(2)
    NY, NX, MYBIN, MXBIN = dc.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            if my >= start[0] and mx >= start[1] and my < end[0] and mx < end[1]:
                mybin = (my - start[0]) // bin
                mxbin = (mx - start[1]) // bin
                if idx1d != no_count_indicator:
                    cuda.atomic.add(dc, (ny, nx, mybin, mxbin), counts[ny, nx, i])

# Cell
def sparse_to_dense_datacube_crop(indices, counts, scan_dimensions, frame_dimensions, center, radius, bin=1):
    xp = sp.backend.get_array_module(indices)
    radius = int(np.ceil(radius / bin) * bin)
    start = center - radius
    end = center + radius
    frame_size = 2 * radius // bin

    dc = xp.zeros((*scan_dimensions, frame_size, frame_size), dtype=indices.dtype)

    threadsperblock = (16, 16)
    blockspergrid = tuple(np.ceil(np.array(indices.shape[:2]) / threadsperblock).astype(np.int))

    no_count_indicator = np.iinfo(indices.dtype).max

    sparse_to_dense_datacube_kernel_crop[blockspergrid, threadsperblock](dc, indices, counts, xp.array(frame_dimensions), bin,
                                                                         start, end, no_count_indicator)
    return dc

# Cell
@cuda.jit
def sparse_to_dense_datacube_crop_gain_mask_kernel(dc, frames, counts,
                                                   frame_dimensions,
                                                   center_frame, center_data,
                                                   radius_data_int, binning,
                                                   fftshift):
    ny, nx = cuda.grid(2)
    NY, NX, MYBIN, MXBIN = dc.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        for i in range(frames[ny, nx].shape[0]):
            idx1d = frames[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            my_center = my - center_data[0]
            mx_center = mx - center_data[1]
            dist_center = m.sqrt(my_center ** 2 + mx_center ** 2)
            if dist_center < radius_data_int:
                mybin = int(center_frame[0] + my_center // binning)
                mxbin = int(center_frame[1] + mx_center // binning)
                if fftshift:
                    mybin = (mybin - center_frame[0]) % (center_frame[0] * 2)
                    mxbin = (mxbin - center_frame[1]) % (center_frame[1] * 2)
                if (mxbin >= 0 and mybin >= 0):
                    cuda.atomic.add(dc, (ny, nx, mybin, mxbin), counts[ny, nx, i])

# Cell

def sparse_to_dense_datacube_crop_gain_mask(indices, counts, scan_dimensions, frame_dimensions, center_data,
                                            radius_data,
                                            radius_max, n_batches=4, binning=1, fftshift=False):
    radius_data_int = int(np.ceil(radius_data / binning) * binning)
    radius_max_int = int(np.ceil(radius_max / binning) * binning)
    frame_size = 2 * radius_max_int // binning
    print(f'radius_data_int : {radius_data_int} ')
    print(f'radius_max_int  : {radius_max_int} ')
    print(f'Dense frame size: {frame_size}x {frame_size}')
    cuda.select_device(1)
    dev = th.device('cuda:1')
    stream = th.cuda.current_stream().cuda_stream

    K = indices.shape[0]
    divpts = array_split_divpoints_ntotal(K, n_batches)
    K_batch = [divpts[i] - divpts[i - 1] for i in range(1, n_batches + 1)]
    dc0 = np.zeros((scan_dimensions[0], scan_dimensions[1], frame_size, frame_size), dtype=np.uint8)

    for b in range(n_batches):
        dc = th.zeros((K_batch[b], scan_dimensions[1], frame_size, frame_size), dtype=th.float32, device=dev)

        center_frame = th.tensor([frame_size // 2, frame_size // 2], device=dev)
        fd = th.as_tensor(frame_dimensions, device=dev)
        center = th.as_tensor(center_data, device=dev)
        inds = th.as_tensor(indices[divpts[b]:divpts[b + 1], ...], device=dev)
        cts = th.as_tensor(counts[divpts[b]:divpts[b + 1], ...].astype(np.float32), dtype=th.float32, device=dev)

        threadsperblock = (16, 16)
        blockspergrid = tuple(np.ceil(np.array(indices.shape[:2]) / threadsperblock).astype(np.int))
        print(f'batch {b}: dtypes:', dc.dtype, inds.dtype, cts.dtype, frame_dimensions.dtype)
        sparse_to_dense_datacube_crop_gain_mask_kernel[blockspergrid, threadsperblock, stream](dc, inds, cts, fd,
                                                                                               center_frame, center,
                                                                                               radius_data_int, binning,
                                                                                               fftshift)
        dc0[divpts[b]:divpts[b + 1], ...] = dc.cpu().type(th.uint8).numpy()

    cuda.select_device(0)
    return dc0

# Cell
from numba import cuda
@cuda.jit
def fftshift_kernel(indices, center_frame, scan_dimensions, no_count_indicator):
    ny, nx = cuda.grid(2)
    NY, NX = scan_dimensions
    MY = center_frame[0] * 2
    MX = center_frame[1] * 2
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            mysh = (my - center_frame[0]) % (center_frame[0] * 2)
            mxsh = (mx - center_frame[1]) % (center_frame[1] * 2)
            if idx1d != no_count_indicator:
                indices[ny, nx, i] = mysh * MX + mxsh

# Cell
from numba import cuda
@cuda.jit
def fftshift_pad_kernel(indices, center_frame, scan_dimensions, new_frame_dimensions, no_count_indicator_old, no_count_indicator_new):
    ny, nx = cuda.grid(2)
    NY, NX = scan_dimensions
    MX = center_frame[1] * 2
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            mysh = (my - center_frame[0]) % (new_frame_dimensions[0])
            mxsh = (mx - center_frame[1]) % (new_frame_dimensions[1])
            if idx1d != no_count_indicator_old:
                indices[ny, nx, i] = mysh * new_frame_dimensions[1] + mxsh
            else:
                indices[ny, nx, i] = no_count_indicator_new

# Cell
from numba import cuda
@cuda.jit
def virtual_annular_image_kernel(img, indices, counts, radius_inner, radius_outer, center_frame, frame_dimensions, no_count_indicator):
    ny, nx = cuda.grid(2)
    NY, NX, _ = indices.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            my_center = my - center_frame[0]
            mx_center = mx - center_frame[1]
            radius = m.sqrt(my_center ** 2 + mx_center ** 2)
            if radius < radius_outer and radius >= radius_inner and idx1d != no_count_indicator:
                cuda.atomic.add(img, (ny,nx), counts[ny, nx, i])

# Cell
@cuda.jit
def crop_symmetric_around_center_kernel(new_frames, old_frames, center_frame, old_frame_dimensions, center_data, radius_data_int):
    ny, nx = cuda.grid(2)
    NY, NX, _ = old_frames.shape
    MY, MX = old_frame_dimensions
    MXnew = center_frame[1] * 2
    if ny < NY and nx < NX:
        k = 0
        for i in range(old_frames[ny, nx].shape[0]):
            idx1d = old_frames[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            my_center = my - center_data[0]
            mx_center = mx - center_data[1]
            dist_center = m.sqrt(my_center ** 2 + mx_center ** 2)
            if dist_center < radius_data_int:
                mybin = int(center_frame[0] + my_center)
                mxbin = int(center_frame[1] + mx_center)
                new_frames[ny, nx, k] = mybin * MXnew + mxbin
                k += 1

# Cell
def crop_symmetric_around_center(old_frames, old_frame_dimensions, center_data, max_radius):
    xp = sp.backend.get_array_module(old_frames)
    max_radius_int = int(max_radius)
    frame_size = 2 * max_radius_int
    center_frame = xp.array([frame_size // 2, frame_size // 2])
    new_frame_dimensions = np.array([frame_size,frame_size])

    threadsperblock = (16, 16)
    blockspergrid = tuple(np.ceil(np.array(old_frames.shape[:2]) / threadsperblock).astype(np.int))

    new_frames = xp.zeros_like(old_frames)
    new_frames[:] = xp.iinfo(new_frames.dtype).max

    crop_symmetric_around_center_kernel[blockspergrid, threadsperblock](new_frames, old_frames, center_frame,
                                                                        xp.array(old_frame_dimensions),
                                                                        xp.array(center_data), max_radius_int)

    max_counts = xp.max(xp.sum(new_frames > 0, 2).ravel())
    # print(f'max counts: {max_counts}')
    res = new_frames[:,:,:max_counts].get()
    return res, new_frame_dimensions

# Cell
from numba import cuda
import math as m

@cuda.jit
def rotate_kernel(frames, center_frame, old_frame_dimensions, center_data, no_count_indicator, angle_rad):
    ny, nx = cuda.grid(2)
    NY, NX, _ = frames.shape
    MY, MX = old_frame_dimensions
    MXnew = center_frame[1] * 2
    if ny < NY and nx < NX:
        for i in range(frames[ny, nx].shape[0]):
            idx1d = frames[ny, nx, i]
            if idx1d != no_count_indicator:
                my = idx1d // MX
                mx = idx1d - my * MX
                my_center = my - center_data[0]
                mx_center = mx - center_data[1]
                #rotate
                mx_center_rot = round(mx_center * m.cos(angle_rad) - my_center * m.sin(angle_rad))
                my_center_rot = round(mx_center * m.sin(angle_rad) + my_center * m.cos(angle_rad))
                mybin = int(center_frame[0] + my_center_rot)
                mxbin = int(center_frame[1] + mx_center_rot)
                frames[ny, nx, i] = mybin * MXnew + mxbin

# Cell

def rotate(frames, old_frame_dimensions, center, angle_rad):
    threadsperblock = (16, 16)
    blockspergrid = tuple(np.ceil(np.array(frames.shape[:2]) / threadsperblock).astype(np.int))
    no_count_indicator = th.iinfo(frames.dtype).max
    new_frames= th.tensor(frames, device=frames.device)
    rotate_kernel[blockspergrid, threadsperblock](new_frames, center, th.tensor(old_frame_dimensions, device=frames.device),
        th.tensor(center, device=frames.device), no_count_indicator, angle_rad)
    return new_frames.get()

# Cell

@cuda.jit
def sum_kernel(indices, counts, frame_dimensions, sum, no_count_indicator):
    ny, nx = cuda.grid(2)
    NY, NX, _ = indices.shape
    MY, MX = frame_dimensions
    if ny < NY and nx < NX:
        for i in range(indices[ny, nx].shape[0]):
            idx1d = indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            if idx1d != no_count_indicator:
                cuda.atomic.add(sum, (my, mx), counts[ny, nx, i])

# Cell
def sum_frames(frames, counts, frame_dimensions):
    threadsperblock = (16, 16)
    blockspergrid = tuple(np.ceil(np.array(frames.shape[:2]) / threadsperblock).astype(np.int))

    frames1 = sp.to_device(frames,0)
    xp = sp.backend.get_array_module(frames1)
    sum = xp.zeros(frame_dimensions)
    counts1 = xp.array(counts)
    no_count_indicator = xp.iinfo(frames.dtype).max

    sum_kernel[blockspergrid, threadsperblock](frames1, counts1, xp.array(frame_dimensions), sum, no_count_indicator)
    return sum.get()

# Cell


@cuda.jit
def rebin_kernel(indices, counts, new_frame_center, old_indices, old_counts, old_frame_center, no_count_indicator,
                 bin_factor):
    ny, nx = cuda.grid(2)
    NY, NX, _ = indices.shape
    MY = old_frame_center[0] * 2
    MX = old_frame_center[1] * 2
    MXnew = new_frame_center[1] * 2
    if ny < NY and nx < NX:
        k = 0
        for i in range(old_indices[ny, nx].shape[0]):
            idx1d = old_indices[ny, nx, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            my_center = my - old_frame_center[0]
            mx_center = mx - old_frame_center[1]
            if idx1d != no_count_indicator:
                mybin = int(new_frame_center[0] + my_center // bin_factor)
                mxbin = int(new_frame_center[1] + mx_center // bin_factor)
                indices[ny, nx, k] = mybin * MXnew + mxbin
                counts[ny, nx, k] = old_counts[ny, nx, i]
                k += 1

# Cell
from .torch_imports import *
import numpy as np
import time
from numba import cuda
from math import *
import math as m
import cmath as cm
import sigpy as sp

# Cell

def disk_overlap_function(Qx_all, Qy_all, Kx_all, Ky_all, aberrations, theta_rot, alpha, lam):
    n_batch = Qx_all.shape[0]
    xp = sp.backend.get_array_module(aberrations)
    Gamma = xp.zeros((n_batch,) + (Ky_all.shape[0], Kx_all.shape[0]), dtype=xp.complex64)
    gs = Gamma.shape
    threadsperblock = 2 ** 8
    blockspergrid = m.ceil(np.prod(gs) / threadsperblock)
    strides = xp.array((np.array(Gamma.strides) / (Gamma.nbytes / Gamma.size)).astype(np.int))
    disk_overlap_kernel[blockspergrid, threadsperblock](Gamma, strides, Qx_all, Qy_all, Kx_all, Ky_all, aberrations,
                                                        theta_rot, alpha, lam)
    return Gamma


@cuda.jit
def disk_overlap_kernel(Γ, strides, Qx_all, Qy_all, Kx_all, Ky_all, aberrations, theta_rot, alpha, lam):
    def aperture2(qx, qy, lam, alpha_max):
        qx2 = qx ** 2
        qy2 = qy ** 2
        q = m.sqrt(qx2 + qy2)
        ktheta = m.asin(q * lam)
        return ktheta < alpha_max

    def chi3(qy, qx, lam, C):
        """
        Zernike polynomials in the cartesian coordinate system
        :param qx:
        :param qy:
        :param lam: wavelength in Angstrom
        :param C:   (12 ,)
        :return:
        """

        u = qx * lam
        v = qy * lam
        u2 = u ** 2
        u3 = u ** 3
        u4 = u ** 4
        # u5 = u ** 5

        v2 = v ** 2
        v3 = v ** 3
        v4 = v ** 4
        # v5 = v ** 5

        # aberr = Param()
        # aberr.C1 = C[0]
        # aberr.C12a = C[1]
        # aberr.C12b = C[2]
        # aberr.C21a = C[3]
        # aberr.C21b = C[4]
        # aberr.C23a = C[5]
        # aberr.C23b = C[6]
        # aberr.C3 = C[7]
        # aberr.C32a = C[8]
        # aberr.C32b = C[9]
        # aberr.C34a = C[10]
        # aberr.C34b = C[11]

        chi = 0

        # r-2 = x-2 +y-2.
        chi += 1 / 2 * C[0] * (u2 + v2)  # r^2
        # r-2 cos(2*phi) = x"2 -y-2.
        # r-2 sin(2*phi) = 2*x*y.
        chi += 1 / 2 * (C[1] * (u2 - v2) + 2 * C[2] * u * v)  # r^2 cos(2 phi) + r^2 sin(2 phi)
        # r-3 cos(3*phi) = x-3 -3*x*y'2. r"3 sin(3*phi) = 3*y*x-2 -y-3.
        chi += 1 / 3 * (C[5] * (u3 - 3 * u * v2) + C[6] * (3 * u2 * v - v3))  # r^3 cos(3phi) + r^3 sin(3 phi)
        # r-3 cos(phi) = x-3 +x*y-2.
        # r-3 sin(phi) = y*x-2 +y-3.
        chi += 1 / 3 * (C[3] * (u3 + u * v2) + C[4] * (v3 + u2 * v))  # r^3 cos(phi) + r^3 sin(phi)
        # r-4 = x-4 +2*x-2*y-2 +y-4.
        chi += 1 / 4 * C[7] * (u4 + v4 + 2 * u2 * v2)  # r^4
        # r-4 cos(4*phi) = x-4 -6*x-2*y-2 +y-4.
        chi += 1 / 4 * C[10] * (u4 - 6 * u2 * v2 + v4)  # r^4 cos(4 phi)
        # r-4 sin(4*phi) = 4*x-3*y -4*x*y-3.
        chi += 1 / 4 * C[11] * (4 * u3 * v - 4 * u * v3)  # r^4 sin(4 phi)
        # r-4 cos(2*phi) = x-4 -y-4.
        chi += 1 / 4 * C[8] * (u4 - v4)
        # r-4 sin(2*phi) = 2*x-3*y +2*x*y-3.
        chi += 1 / 4 * C[9] * (2 * u3 * v + 2 * u * v3)
        # r-5 cos(phi) = x-5 +2*x-3*y-2 +x*y-4.
        # r-5 sin(phi) = y*x"4 +2*x-2*y-3 +y-5.
        # r-5 cos(3*phi) = x-5 -2*x-3*y-2 -3*x*y-4.
        # r-5 sin(3*phi) = 3*y*x-4 +2*x-2*y-3 -y-5.
        # r-5 cos(5*phi) = x-5 -10*x-3*y-2 +5*x*y-4.
        # r-5 sin(5*phi) = 5*y*x-4 -10*x-2*y-3 +y-5.

        chi *= 2 * np.pi / lam

        return chi

    gs = Γ.shape
    N = gs[0] * gs[1] * gs[2]
    n = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
    j = n // strides[0]
    iky = (n - j * strides[0]) // strides[1]
    ikx = (n - (j * strides[0] + iky * strides[1])) // strides[2]

    if n < N:
        Qx = Qx_all[j]
        Qy = Qy_all[j]
        Kx = Kx_all[ikx]
        Ky = Ky_all[iky]

        Qx_rot = Qx * cos(theta_rot) - Qy * sin(theta_rot)
        Qy_rot = Qx * sin(theta_rot) + Qy * cos(theta_rot)

        Qx = Qx_rot
        Qy = Qy_rot

        chi = chi3(Ky, Kx, lam, aberrations)
        A = aperture2(Ky, Kx, lam, alpha) * cm.exp(-1j * chi)
        chi = chi3(Ky + Qy, Kx + Qx, lam, aberrations)
        Ap = aperture2(Ky + Qy, Kx + Qx, lam, alpha) * cm.exp(-1j * chi)
        chi = chi3(Ky - Qy, Kx - Qx, lam, aberrations)
        Am = aperture2(Ky - Qy, Kx - Qx, lam, alpha) * cm.exp(-1j * chi)

        Γ[j, iky, ikx] = A.conjugate() * Am - A * Ap.conjugate()

# Cell
def double_overlap_intensitities_in_range(G_max, thetas, Qx_max, Qy_max, Kx, Ky, aberrations,
                                          aberration_angles, alpha_rad, lam, do_plot=False):
    xp = sp.backend.get_array_module(G_max)
    intensities = np.zeros((len(thetas)))
    for i, theta_rot in enumerate(thetas):
        if th.cuda.is_available():
            Gamma = disk_overlap_function(Qx_max, Qy_max, Kx, Ky, aberrations, theta_rot, alpha_rad,lam)

        intensities[i] = xp.sum(xp.abs(G_max * Gamma.conj()))

    if do_plot:
        f, ax = plt.subplots()
        ax.scatter(np.rad2deg(thetas), intensities)
        plt.show()

    return intensities

# Cell
def find_rotation_angle_with_double_disk_overlap(G, lam, k_max, dxy, alpha_rad, mask=None, n_fit=6, ranges=[360, 30],
                                                 partitions=[144, 120], verbose=False, manual_frequencies=None, aberrations=None):
    """
        Finds the best rotation angle by maximizing the double disk overlap intensity of the 4D dataset. Only valid
        for datasets where the scan step size is roughly on the same length scale as the illumination half-angle alpha.

    :param G: G function. 4DSTEM dataset Fourier transformed along the scan coordinates
    :param lam:
    :param k_max:
    :param dxy:
    :param alpha_rad:
    :param n_fit: number of object spatial frequencies to fit
    :param ranges:
    :param verbose:
    :return: the best rotation angle in radians.
    """
    ny, nx, nky, nkx = G.shape
    xp = sp.backend.get_array_module(G)

    def get_qx_qy_1D(M, dx, dtype, fft_shifted=False):
        qxa = xp.fft.fftfreq(M[0], dx[0]).astype(dtype)
        qya = xp.fft.fftfreq(M[1], dx[1]).astype(dtype)
        if fft_shifted:
            qxa = xp.fft.fftshift(qxa)
            qya = xp.fft.fftshift(qya)
        return qxa, qya

    Kx, Ky = get_qx_qy_1D([nkx, nky], k_max, G[0, 0, 0, 0].real.dtype, fft_shifted=True)
    Qx, Qy = get_qx_qy_1D([nx, ny], dxy, G[0, 0, 0, 0].real.dtype, fft_shifted=False)

    if aberrations is None:
        aberrations = xp.zeros((12))
    aberration_angles = xp.zeros((12))

    if manual_frequencies is None:
        Gabs = xp.sum(xp.abs(G), (2, 3))
        if mask is not None:
            gg = Gabs * mask
#             plot(gg.get(), 'Gabs * mask')
            inds = xp.argsort((gg).ravel()).get()
        else:
            inds = xp.argsort(Gabs.ravel()).get()
        strongest_object_frequencies = np.unravel_index(inds[-1 - n_fit:-1], G.shape[:2])

        G_max = G[strongest_object_frequencies]
        Qy_max = Qy[strongest_object_frequencies[0]]
        Qx_max = Qx[strongest_object_frequencies[1]]
    else:
        strongest_object_frequencies = manual_frequencies
        G_max = G[strongest_object_frequencies]
        Qy_max = Qy[strongest_object_frequencies[0]]
        Qx_max = Qx[strongest_object_frequencies[1]]

    if verbose:
        print(f"strongest_object_frequencies: {strongest_object_frequencies}")

    best_angle = 0

    for j, (range, parts) in enumerate(zip(ranges, partitions)):
        thetas = np.linspace(best_angle - np.deg2rad(range / 2), best_angle + np.deg2rad(range / 2), parts)
        intensities = double_overlap_intensitities_in_range(G_max, thetas, Qx_max, Qy_max, Kx, Ky, aberrations,
                                                            aberration_angles, alpha_rad, lam, do_plot=False)

        sortind = np.argsort(intensities)
        max_ind0 = sortind[-1]
        max_ind1 = sortind[0]
        best_angle = thetas[max_ind0]
        best_angle1 = thetas[max_ind1]
        if verbose:
            A = xp.zeros(G_max.shape[1:], dtype=xp.complex64)
            Ap = xp.zeros(G_max.shape[1:], dtype=xp.complex64)
            Am = xp.zeros(G_max.shape[1:], dtype=xp.complex64)
            print(f"Iteration {j}: current best rotation angle: {np.rad2deg(best_angle)}")
            Gamma = disk_overlap_function(Qx_max, Qy_max, Kx, Ky, aberrations, best_angle, alpha_rad,lam)
#             plotcxmosaic(Gamma.get() * G_max.get(), f"best theta_rot = {np.rad2deg(best_angle)} (maximum)")

            Gamma = disk_overlap_function(Qx_max, Qy_max, Kx, Ky, aberrations, best_angle1, alpha_rad,lam)
#             plotcxmosaic(Gamma.get() * G_max.get(), f"best theta_rot = {np.rad2deg(best_angle1)} (minimum)")

    max_ind = np.argsort(intensities)[-1]

    return max_ind, thetas, intensities

# Cell
import torch as th
import torch.nn as nn
from .util import fftshift_checkerboard, Param
import numpy as np

# def zernike_aberrations(qx, qy, lam, C):
#     """
#     Calculates the geometrical aberrations up to the order specified in C. C is a complex array specifying the
#     aberration coefficients.
#
#     :param qx: x coordinate system, can be (i_batch, qxx, qxy)
#     :param qy: y coordinate system, can be (i_batch, qyx, qyy)
#     :param lam: wavelength in A
#     :param C: (n_batch, m_max, n_max) complex array
#     :return: chi, the real aberration function
#     """
#     if isinstance(qx, np.ndarray):
#         xp = np
#
#         if C.ndim == 3:
#             n_batch, m_max, n_max = C.shape
#         elif C.ndim == 2:
#             m_max, n_max = C.shape
#             n_batch = 1
#             C = C[None, ...]
#         else:
#             raise RuntimeError("Aberration coefficient array has wrong dimension")
#
#         sh = (n_batch, *qx.shape)
#         chi = xp.zeros(sh, dtype=np.float32) + 1j * xp.zeros(sh, dtype=np.float32)
#         ω = lam * (qx + 1j * qy)
#
#         for i_batch in range(n_batch):
#             for m in range(m_max):
#                 if m <= 3:
#                     n_max = m + 1
#                 elif m <= 5:
#                     n_max = m // 2
#                 else:
#                     n_max = 1
#                 for n in range(n_max):
#                     if m > 0 or n > 0:
#                         # print(f"C[{m}, {n}] = {C[i_batch, m, n]}")
#                         ab = C[i_batch, m, n] * (ω ** m) * (ω.conj() ** n)
#                         ab /= (m + n)
#                         chi[i_batch] += ab
#         chi = 2 * np.pi / lam * chi.real
#         chi = xp.squeeze(chi)
#     elif isinstance(qx, th.Tensor):
#         if C.ndimension() == 4:
#             n_batch, m_max, n_max, _ = C.shape
#         elif C.ndimension() == 3:
#             m_max, n_max, _ = C.shape
#             n_batch = 1
#             C = C.unsqueeze(0)
#         else:
#             raise RuntimeError("Aberration coefficient array has wrong dimension")
#
#         sh = (n_batch, *qx.shape)
#
#         chi = th.zeros(sh, dtype=th.float32)
#
#         # print(chi.shape)
#
#         ω = th.zeros_like(qx)
#         ω[re] = lam * qx
#         ω[im] = lam * qy
#         C_mn = make_real(th.zeros(qx.shape))
#         for i_batch in range(n_batch):
#             for m in range(m_max):
#                 if m <= 3:
#                     n_max = m + 1
#                 elif m <= 5:
#                     n_max = m // 2
#                 else:
#                     n_max = 1
#                 for n in range(n_max):
#                     if m > 0 or n > 0:
#                         # print(f"C[{m}, {n}] = {C[i_batch, m, n]}")
#                         C_mn[re] = C[i_batch, m, n][re]
#                         C_mn[im] = C[i_batch, m, n][im]
#                         ab = complex_mul(C_mn, cpow(ω, n))
#                         ab = complex_mul(ab, cpow(conj(ω), n))
#                         ab /= (m + n)
#                         # print(chi.shape)
#                         chi[i_batch] += ab
#         chi = 2 * np.pi / lam * chi[re]
#         chi = chi.squeeze()
#     return chi

def cartesian_aberrations_single(qx, qy, lam, C):
    """
    Zernike polynomials in the cartesian coordinate system
    :param qx:
    :param qy:
    :param lam: wavelength in Angstrom
    :param C:   (12 ,)
    :return:
    """

    u = qx * lam
    v = qy * lam
    u2 = u ** 2
    u3 = u ** 3
    u4 = u ** 4
    # u5 = u ** 5

    v2 = v ** 2
    v3 = v ** 3
    v4 = v ** 4
    # v5 = v ** 5

    aberr = Param()
    aberr.C1 = C[0]
    aberr.C12a = C[1]
    aberr.C12b = C[2]
    aberr.C21a = C[3]
    aberr.C21b = C[4]
    aberr.C23a = C[5]
    aberr.C23b = C[6]
    aberr.C3 = C[7]
    aberr.C32a = C[8]
    aberr.C32b = C[9]
    aberr.C34a = C[10]
    aberr.C34b = C[11]

    chi = 0

    # r-2 = x-2 +y-2.
    chi += 1 / 2 * aberr.C1 * (u2 + v2) # r^2
    #r-2 cos(2*phi) = x"2 -y-2.
    # r-2 sin(2*phi) = 2*x*y.
    chi += 1 / 2 * (aberr.C12a * (u2 - v2) + 2 * aberr.C12b * u * v) # r^2 cos(2 phi) + r^2 sin(2 phi)
    # r-3 cos(3*phi) = x-3 -3*x*y'2. r"3 sin(3*phi) = 3*y*x-2 -y-3.
    chi += 1 / 3 * (aberr.C23a * (u3 - 3 * u * v2) + aberr.C23b * (3 * u2 * v - v3))# r^3 cos(3phi) + r^3 sin(3 phi)
    # r-3 cos(phi) = x-3 +x*y-2.
    # r-3 sin(phi) = y*x-2 +y-3.
    chi += 1 / 3 * (aberr.C21a * (u3 + u * v2) + aberr.C21b * (v3 + u2 * v))# r^3 cos(phi) + r^3 sin(phi)
    # r-4 = x-4 +2*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C3 * (u4 + v4 + 2 * u2 * v2)# r^4
    # r-4 cos(4*phi) = x-4 -6*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C34a * (u4 - 6 * u2 * v2 + v4)# r^4 cos(4 phi)
    # r-4 sin(4*phi) = 4*x-3*y -4*x*y-3.
    chi += 1 / 4 * aberr.C34b * (4 * u3 * v - 4 * u * v3) # r^4 sin(4 phi)
    # r-4 cos(2*phi) = x-4 -y-4.
    chi += 1 / 4 * aberr.C32a * (u4 - v4)
    # r-4 sin(2*phi) = 2*x-3*y +2*x*y-3.
    chi += 1 / 4 * aberr.C32b * (2 * u3 * v + 2 * u * v3)
    # r-5 cos(phi) = x-5 +2*x-3*y-2 +x*y-4.
    # r-5 sin(phi) = y*x"4 +2*x-2*y-3 +y-5.
    # r-5 cos(3*phi) = x-5 -2*x-3*y-2 -3*x*y-4.
    # r-5 sin(3*phi) = 3*y*x-4 +2*x-2*y-3 -y-5.
    # r-5 cos(5*phi) = x-5 -10*x-3*y-2 +5*x*y-4.
    # r-5 sin(5*phi) = 5*y*x-4 -10*x-2*y-3 +y-5.

    chi *= 2 * np.pi / lam

    return chi

def cartesian_aberrations(qx, qy, lam, C):
    """
    Zernike polynomials in the cartesian coordinate system
    :param qx:
    :param qy:
    :param lam: wavelength in Angstrom
    :param C:   12 x D
    :return:
    """

    u = qx * lam
    v = qy * lam
    u2 = u ** 2
    u3 = u ** 3
    u4 = u ** 4
    # u5 = u ** 5

    v2 = v ** 2
    v3 = v ** 3
    v4 = v ** 4
    # v5 = v ** 5

    aberr = Param()
    aberr.C1 = C[0].unsqueeze(1).unsqueeze(1)
    aberr.C12a = C[1].unsqueeze(1).unsqueeze(1)
    aberr.C12b = C[2].unsqueeze(1).unsqueeze(1)
    aberr.C21a = C[3].unsqueeze(1).unsqueeze(1)
    aberr.C21b = C[4].unsqueeze(1).unsqueeze(1)
    aberr.C23a = C[5].unsqueeze(1).unsqueeze(1)
    aberr.C23b = C[6].unsqueeze(1).unsqueeze(1)
    aberr.C3 = C[7].unsqueeze(1).unsqueeze(1)
    aberr.C32a = C[8].unsqueeze(1).unsqueeze(1)
    aberr.C32b = C[9].unsqueeze(1).unsqueeze(1)
    aberr.C34a = C[10].unsqueeze(1).unsqueeze(1)
    aberr.C34b = C[11].unsqueeze(1).unsqueeze(1)

    chi = 0

    # r-2 = x-2 +y-2.
    chi += 1 / 2 * aberr.C1 * (u2 + v2) # r^2
    #r-2 cos(2*phi) = x"2 -y-2.
    # r-2 sin(2*phi) = 2*x*y.
    chi += 1 / 2 * (aberr.C12a * (u2 - v2) + 2 * aberr.C12b * u * v) # r^2 cos(2 phi) + r^2 sin(2 phi)
    # r-3 cos(3*phi) = x-3 -3*x*y'2. r"3 sin(3*phi) = 3*y*x-2 -y-3.
    chi += 1 / 3 * (aberr.C23a * (u3 - 3 * u * v2) + aberr.C23b * (3 * u2 * v - v3))# r^3 cos(3phi) + r^3 sin(3 phi)
    # r-3 cos(phi) = x-3 +x*y-2.
    # r-3 sin(phi) = y*x-2 +y-3.
    chi += 1 / 3 * (aberr.C21a * (u3 + u * v2) + aberr.C21b * (v3 + u2 * v))# r^3 cos(phi) + r^3 sin(phi)
    # r-4 = x-4 +2*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C3 * (u4 + v4 + 2 * u2 * v2)# r^4
    # r-4 cos(4*phi) = x-4 -6*x-2*y-2 +y-4.
    chi += 1 / 4 * aberr.C34a * (u4 - 6 * u2 * v2 + v4)# r^4 cos(4 phi)
    # r-4 sin(4*phi) = 4*x-3*y -4*x*y-3.
    chi += 1 / 4 * aberr.C34b * (4 * u3 * v - 4 * u * v3) # r^4 sin(4 phi)
    # r-4 cos(2*phi) = x-4 -y-4.
    chi += 1 / 4 * aberr.C32a * (u4 - v4)
    # r-4 sin(2*phi) = 2*x-3*y +2*x*y-3.
    chi += 1 / 4 * aberr.C32b * (2 * u3 * v + 2 * u * v3)
    # r-5 cos(phi) = x-5 +2*x-3*y-2 +x*y-4.
    # r-5 sin(phi) = y*x"4 +2*x-2*y-3 +y-5.
    # r-5 cos(3*phi) = x-5 -2*x-3*y-2 -3*x*y-4.
    # r-5 sin(3*phi) = 3*y*x-4 +2*x-2*y-3 -y-5.
    # r-5 cos(5*phi) = x-5 -10*x-3*y-2 +5*x*y-4.
    # r-5 sin(5*phi) = 5*y*x-4 -10*x-2*y-3 +y-5.

    chi *= 2 * np.pi / lam

    return chi


def aperture(q: th.Tensor, lam, alpha_max, edge=2):
    ktheta = th.asin(q.norm(dim=0) * lam)
    qmax = alpha_max / lam
    dk = q[0][1][0]

    arr = th.zeros_like(q[1])
    arr[ktheta < alpha_max] = 1
    if edge > 0:
        dEdge = edge / (qmax / dk)  # fraction of aperture radius that will be smoothed
        # some fancy indexing: pull out array elements that are within
        #    our smoothing edges
        ind = (ktheta / alpha_max > (1 - dEdge)) * (ktheta / alpha_max < (1 + dEdge))
        arr[ind] = 0.5 * (1 - th.sin(np.pi / (2 * dEdge) * (ktheta[ind] / alpha_max - 1)))
    return arr


class ZernikeProbe2(nn.Module):
    def __init__(self, q: th.Tensor, lam, fft_shifted=True):
        """
        Creates an aberration surface from aberration coefficients. The output is backpropable

        :param q: 2 x M1 x M2 tensor of x coefficients of reciprocal space
        :param lam: wavelength in Angstrom
        :param C: aberration coefficients
        :return: (maximum size of all aberration tensors) x MY x MX
        """

        super(ZernikeProbe2, self).__init__()
        self.q = q
        self.lam = lam
        self.fft_shifted = fft_shifted

        if self.fft_shifted:
            cb = fftshift_checkerboard(self.q.shape[1] // 2, self.q.shape[2] // 2)
            self.cb = th.from_numpy(cb).float().to(q.device)

    def forward(self, C, A):
        chi = cartesian_aberrations(self.q[1], self.q[0], self.lam, C)
        Psi = th.exp(-1j*chi) * A.expand_as(chi)

        if self.fft_shifted:
            Psi = Psi * self.cb

        return Psi

class ZernikeProbeSingle(nn.Module):
    def __init__(self, q: th.Tensor, lam, fft_shifted=True):
        """
        Creates an aberration surface from aberration coefficients. The output is backpropable

        :param q: 2 x M1 x M2 tensor of x coefficients of reciprocal space
        :param lam: wavelength in Angstrom
        :param C: aberration coefficients
        :return: (maximum size of all aberration tensors) x MY x MX
        """

        super(ZernikeProbeSingle, self).__init__()
        self.q = q
        self.lam = lam
        self.fft_shifted = fft_shifted

        if self.fft_shifted:
            cb = fftshift_checkerboard(self.q.shape[1] // 2, self.q.shape[2] // 2).T
            self.cb = th.from_numpy(cb).float().to(q.device)

    def forward(self, C, A):
        chi = cartesian_aberrations_single(self.q[1], self.q[0], self.lam, C)
        Psi = th.exp(-1j*chi) * A.expand_as(chi)

        if self.fft_shifted:
            Psi = Psi * self.cb

        return Psi

class ZernikeProbe(nn.Module):
    def __init__(self, q: th.Tensor, lam, A_init, A_requires_grad, fft_shifted=True,
                 C1=th.zeros(1), C12a=th.zeros(1), C12b=th.zeros(1), C21a=th.zeros(1), C21b=th.zeros(1),
                 C23a=th.zeros(1), C23b=th.zeros(1), C3=th.zeros(1), C32a=th.zeros(1), C32b=th.zeros(1),
                 C34a=th.zeros(1), C34b=th.zeros(1)):
        """
        Creates an aberration surface from aberration coefficients. The output is backpropable

        :param qx: M1 x M2 tensor of x coefficients of reciprocal space
        :param qy: M1 x M2 tensor of y coefficients of reciprocal space
        :param lam: wavelength in Angstrom
        :param C: shift                               in Angstrom
        :return: (maximum size of all aberration tensors) x MY x MX
        """

        super(ZernikeProbe, self).__init__()
        self.qx = q[1]
        self.qy = q[0]
        self.lam = lam
        self.fft_shifted = fft_shifted
        self.dtype = q.dtype
        d = self.dtype
        # defocus
        self.C1 = nn.Parameter(data=C1.type(d), requires_grad=th.max(th.abs(C1)).item() > 1e-10)
        # two-fold stig
        self.C12a = nn.Parameter(data=C12a.type(d), requires_grad=th.max(th.abs(C12a)).item() > 1e-10).to(C1.device)
        self.C12b = nn.Parameter(data=C12b.type(d), requires_grad=th.max(th.abs(C12b)).item() > 1e-10).to(C1.device)

        # axial coma
        self.C21a = nn.Parameter(data=C21a.type(d), requires_grad=th.max(th.abs(C21a)).item() > 1e-10).to(C1.device)
        self.C21b = nn.Parameter(data=C21b.type(d), requires_grad=th.max(th.abs(C21b)).item() > 1e-10).to(C1.device)
        # 3-fold stig
        self.C23a = nn.Parameter(data=C23a.type(d), requires_grad=th.max(th.abs(C23a)).item() > 1e-10).to(C1.device)
        self.C23b = nn.Parameter(data=C23b.type(d), requires_grad=th.max(th.abs(C23b)).item() > 1e-10).to(C1.device)
        # spherical
        self.C3 = nn.Parameter(data=C3.type(d), requires_grad=th.max(th.abs(C3)).item() > 1e-10).to(C1.device)
        # star
        self.C32a = nn.Parameter(data=C32a.type(d), requires_grad=th.max(th.abs(C32a)).item() > 1e-10).to(C1.device)
        self.C32b = nn.Parameter(data=C32b.type(d), requires_grad=th.max(th.abs(C32b)).item() > 1e-10).to(C1.device)
        # 4-fold stig
        self.C34a = nn.Parameter(data=C34a.type(d), requires_grad=th.max(th.abs(C34a)).item() > 1e-10).to(C1.device)
        self.C34b = nn.Parameter(data=C34b.type(d), requires_grad=th.max(th.abs(C34b)).item() > 1e-10).to(C1.device)

        self.A = nn.Parameter(data=A_init.type(d), requires_grad=A_requires_grad).to(C1.device)

        self.optimized_parameters_ = []

        if self.fft_shifted:
            self.cb = th.from_numpy(fftshift_checkerboard(self.qx.shape[0] // 2, self.qx.shape[1] // 2)).float().to(
                q.device).type(d)

        for p in self.parameters():
            if p.requires_grad:
                self.optimized_parameters_.append(p)

    def optimized_parameters(self):
        return self.optimized_parameters_

    def forward(self):
        max_size = np.max(
            [a.shape[0] if len(a.shape) > 0 else 1 for a in
             [self.C1, self.C12a, self.C12b, self.C21a, self.C21b, self.C23a, self.C23b, self.C3,
              self.C32a, self.C32b, self.C34a, self.C34b]])
        C = th.stack((self.C1.expand(max_size), self.C12a.expand(max_size), self.C12b.expand(max_size),
                      self.C21a.expand(max_size), self.C21b.expand(max_size), self.C23a.expand(max_size),
                      self.C23b.expand(max_size), self.C3.expand(max_size), self.C32a.expand(max_size),
                      self.C32b.expand(max_size), self.C34a.expand(max_size), self.C34b.expand(max_size)), 0)

        chi = cartesian_aberrations(self.qx, self.qy, self.lam, C).type(self.dtype)
        Psi = th.exp(-1j*chi) * self.A.expand_as(chi)

        if self.fft_shifted:
            Psi = Psi * self.cb

        return Psi



# Cell
import h5py
import numpy as np
# import cPickle
import time
import os.path
import glob

h5options = dict(
    H5RW_VERSION='0.1',
    H5PY_VERSION=h5py.version.version,
    # UNSUPPORTED = 'pickle'
    # UNSUPPORTED = 'ignore'
    UNSUPPORTED='fail',
    SLASH_ESCAPE='_SLASH_')


def sdebug(f):
    """
    debugging decorator for _store functions
    """

    def newf(*args, **kwds):
        print('{0:20} {1:20}'.format(f.func_name, args[2]))
        return f(*args, **kwds)

    newf.__doc__ = f.__doc__
    return newf


def _h5write(filename, mode, *args, **kwargs):
    """\
    _h5write(filename, mode, {'var1'=..., 'var2'=..., ...})
    _h5write(filename, mode, var1=..., var2=..., ...)
    _h5write(filename, mode, dict, var1=..., var2=...)

    Writes variables var1, var2, ... to file filename. The file mode
    can be chosen according to the h5py documentation. The key-value
    arguments have precedence on the provided dictionnary.

    supported variable types are:
    * scalars
    * numpy arrays
    * strings
    * lists
    * dictionaries

    (if the option UNSUPPORTED is equal to 'pickle', any other type
    is pickled and saved. UNSUPPORTED = 'ignore' silently eliminates
    unsupported types. Default is 'fail', which raises an error.)

    The file mode can be chosen according to the h5py documentation.
    It defaults to overwriting an existing file.
    """

    filename = os.path.abspath(os.path.expanduser(filename))

    ctime = time.asctime()
    mtime = ctime

    # Update input dictionnary
    if args:
        d = args[0].copy()  # shallow copy
    else:
        d = {}
    d.update(kwargs)

    # List of object ids to make sure we are not saving something twice.
    ids = []

    # This is needed to store strings
    dt = h5py.new_vlen(str)

    def check_id(id):
        if id in ids:
            raise RuntimeError('Circular reference detected! Aborting save.')
        else:
            ids.append(id)

    def pop_id(id):
        ids[:] = [x for x in ids if x != id]

    # @sdebug
    def _store_numpy(group, a, name, compress=True):
        if compress:
            dset = group.create_dataset(name, data=a, compression='gzip')
        else:
            dset = group.create_dataset(name, data=a)
        dset.attrs['type'] = 'array'
        return dset

    # @sdebug
    def _store_string(group, s, name):
        dset = group.create_dataset(name, data=np.asarray(s), dtype=dt)
        dset.attrs['type'] = 'string'
        return dset

    # @sdebug
    def _store_unicode(group, s, name):
        dset = group.create_dataset(name, data=np.asarray(s.encode('utf8')), dtype=dt)
        dset.attrs['type'] = 'unicode'
        return dset

    # @sdebug
    def _store_list(group, l, name):
        check_id(id(l))
        arrayOK = len(set([type(x) for x in l])) == 1
        if arrayOK:
            try:
                # Try conversion to a numpy array
                la = np.array(l)
                if la.dtype.type is np.string_:
                    arrayOK = False
                else:
                    dset = _store_numpy(group, la, name)
                    dset.attrs['type'] = 'arraylist'
            except:
                arrayOK = False
        if not arrayOK:
            # inhomogenous list. Store all elements individually
            dset = group.create_group(name)
            for i, v in enumerate(l):
                _store(dset, v, '%05d' % i)
            dset.attrs['type'] = 'list'
        pop_id(id(l))
        return dset

    # @sdebug
    def _store_tuple(group, t, name):
        dset = _store_list(group, list(t), name)
        dset_type = dset.attrs['type']
        dset.attrs['type'] = 'arraytuple' if dset_type == 'arraylist' else 'tuple'
        return dset

    # @sdebug
    def _store_dict(group, d, name):
        check_id(id(d))
        if any([type(k) not in [str,] for k in d.keys()]):
            raise RuntimeError('Only dictionaries with string keys are supported.')
        dset = group.create_group(name)
        dset.attrs['type'] = 'dict'
        for k, v in d.items():
            if k.find('/') > -1:
                k = k.replace('/', h5options['SLASH_ESCAPE'])
                ndset = _store(dset, v, k)
                if ndset is not None:
                    ndset.attrs['escaped'] = '1'
            else:
                _store(dset, v, k)
        pop_id(id(d))
        return dset

    def _store_dict_new(group, d, name):
        check_id(id(d))
        dset = group.create_group(name)
        dset.attrs['type'] = 'dict'
        for i, kv in enumerate(d.iteritems()):
            _store(dset, kv, '%05d' % i)
        pop_id(id(d))
        return dset

    # @sdebug
    def _store_None(group, a, name):
        dset = group.create_dataset(name, data=np.zeros((1,)))
        dset.attrs['type'] = 'None'
        return dset

    # @sdebug
    # def _store_pickle(group, a, name):
    #     apic = cPickle.dumps(a)
    #     dset = group.create_dataset(name, data=np.asarray(apic), dtype=dt)
    #     dset.attrs['type'] = 'pickle'
    #     return dset

    # @sdebug
    def _store(group, a, name):
        if type(a) is str:
            dset = _store_string(group, a, name)
        # elif type(a) is unicode:
        #     dset = _store_unicode(group, a, name)
        elif type(a) is dict:
            dset = _store_dict(group, a, name)
        elif type(a) is list:
            dset = _store_list(group, a, name)
        elif type(a) is tuple:
            dset = _store_tuple(group, a, name)
        elif type(a) is np.ndarray:
            dset = _store_numpy(group, a, name)
        elif np.isscalar(a):
            dset = _store_numpy(group, np.asarray(a), name, compress=False)
            dset.attrs['type'] = 'scalar'
        elif a is None:
            dset = _store_None(group, a, name)
        else:
            if h5options['UNSUPPORTED'] == 'fail':
                raise RuntimeError('Unsupported data type : %s' % type(a))
            elif h5options['UNSUPPORTED'] == 'pickle':
                dset = _store_pickle(group, a, name)
            else:
                dset = None
        return dset

    # Open the file and save everything
    with h5py.File(filename, mode) as f:
        f.attrs['h5rw_version'] = h5options['H5RW_VERSION']
        f.attrs['ctime'] = ctime
        f.attrs['mtime'] = mtime
        for k, v in d.items():
            _store(f, v, k)

    return


def h5write(filename, *args, **kwargs):
    """\
    h5write(filename, {'var1'=..., 'var2'=..., ...})
    h5write(filename, var1=..., var2=..., ...)
    h5write(filename, dict, var1=..., var2=...)

    Writes variables var1, var2, ... to file filename. The key-value
    arguments have precedence on the provided dictionnary.

    supported variable types are:
    * scalars
    * numpy arrays
    * strings
    * lists
    * dictionaries

    (if the option UNSUPPORTED is equal to 'pickle', any other type
    is pickled and saved. UNSUPPORTED = 'ignore' silently eliminates
    unsupported types. Default is 'fail', which raises an error.)

    The file mode can be chosen according to the h5py documentation.
    It defaults to overwriting an existing file.
    """

    _h5write(filename, 'w', *args, **kwargs)
    return


def h5append(filename, *args, **kwargs):
    """\
    h5append(filename, {'var1'=..., 'var2'=..., ...})
    h5append(filename, var1=..., var2=..., ...)
    h5append(filename, dict, var1=..., var2=...)

    Appends variables var1, var2, ... to file filename. The
    key-value arguments have precedence on the provided dictionnary.

    supported variable types are:
    * scalars
    * numpy arrays
    * strings
    * lists
    * dictionaries

    (if the option UNSUPPORTED is equal to 'pickle', any other type
    is pickled and saved. UNSUPPORTED = 'ignore' silently eliminates
    unsupported types. Default is 'fail', which raises an error.)

    The file mode can be chosen according to the h5py documentation.
    It defaults to overwriting an existing file.
    """

    _h5write(filename, 'a', *args, **kwargs)
    return


def h5read(filename, *args, **kwargs):
    """\
    h5read(filename)
    h5read(filename, s1, s2, ...)
    h5read(filename, (s1,s2, ...))

    Read variables from a hdf5 file created with h5write and returns them as
    a dictionary.

    If specified, only variable named s1, s2, ... are loaded.

    Variable names support slicing and group access. For instance, provided
    that the file contains the appropriate objects, the following syntax is
    valid:

    a = h5read('file.h5', 'myarray[2:4]')
    a = h5read('file.h5', 'adict.thekeyIwant')

    h5read(filename_with_wildcard, ... , doglob=True)
    Reads sequentially all globbed filenames.

    """
    doglob = kwargs.get('doglob', None)

    # Used if we read a list of files
    fnames = []
    if not isinstance(filename, str):
        # We have a list
        fnames = filename
    else:
        if doglob is None:
            # glob only if there is a wildcard in the filename
            doglob = glob.has_magic(filename)
        if doglob:
            fnames = sorted(glob.glob(filename))
            if not fnames:
                raise IOError('%s : no match.' % filename)

    if fnames:
        # We are here only if globbing was allowed.
        dl = []
        # Loop over file names
        for f in fnames:
            # Call again, but this time without globbing.
            d = h5read(f, *args, doglob=False, **kwargs)
            dl.append(d)
        return dl

    # We are here only if there was no globbing (fnames is empty)
    filename = os.path.abspath(os.path.expanduser(filename))

    def _load_dict_new(dset):
        d = {}
        keys = dset.keys()
        keys.sort()
        for k in keys:
            dk, dv = _load(dset[k])
            d[dk] = dv
        return d

    def _load_dict(dset):
        d = {}
        for k, v in dset.items():
            if v.attrs.get('escaped', None) is not None:
                k = k.replace(h5options['SLASH_ESCAPE'], '/')
            d[k] = _load(v)
        return d

    def _load_list(dset):
        l = []
        keys = dset.keys()
        keys.sort()
        for k in keys:
            l.append(_load(dset[k]))
        return l

    def _load_numpy(dset, sl=None):
        if sl is not None:
            return dset[sl]
        else:
            return dset[...]

    def _load_scalar(dset):
        try:
            return dset[...].item()
        except:
            return dset[...]

    def _load_str(dset):
        return dset.value

    def _load_unicode(dset):
        return dset.value.decode('utf8')

    # def _load_pickle(dset):
    #     return cPickle.loads(dset[...])

    def _load(dset, sl=None):
        dset_type = dset.attrs.get('type', None)

        # Treat groups as dicts
        if (dset_type is None) and (type(dset) is h5py.Group):
            dset_type = 'dict'

        if dset_type == b'dict' or dset_type == 'dict':
            if sl is not None:
                raise RuntimeError('Dictionaries do not support slicing')
            val = _load_dict(dset)
        elif dset_type == b'list' or dset_type == 'list':
            val = _load_list(dset)
            if sl is not None:
                val = val[sl]
        elif dset_type == b'array' or dset_type == 'array':
            val = _load_numpy(dset, sl)
        elif dset_type == b'arraylist' or dset_type == 'arraylist':
            val = [x for x in _load_numpy(dset)]
            if sl is not None:
                val = val[sl]
        elif dset_type == b'tuple' or dset_type == 'tuple':
            val = tuple(_load_list(dset))
            if sl is not None:
                val = val[sl]
        elif dset_type == b'arraytuple' or dset_type == 'arraytuple':
            val = tuple(_load_numpy(dset).tolist())
            if sl is not None:
                val = val[sl]
        elif dset_type == b'string' or dset_type == 'string':
            val = _load_str(dset)
            if sl is not None:
                val = val[sl]
        elif dset_type == b'unicode' or dset_type == 'unicode':
            val = _load_str(dset)
            if sl is not None:
                val = val[sl]
        elif dset_type == b'scalar' or dset_type == 'scalar':
            val = _load_scalar(dset)
        elif dset_type == 'None':
            # 24.4.13 : B.E. commented due to hr5read not being able to return None type
            # try:
            #   val = _load_numpy(dset)
            # except:
            #    val = None
            val = None
        # elif dset_type == 'pickle':
        #     val = _load_pickle(dset)
        elif dset_type is None:
            val = _load_numpy(dset)
        else:
            raise RuntimeError('Unsupported data type : %s' % dset_type)
        return val

    outdict = {}
    with h5py.File(filename, 'r') as f:
        h5rw_version = f.attrs.get('h5rw_version', None)
        #        if h5rw_version is None:
        #            print('Warning: this file does not seem to follow h5read format.')
        ctime = f.attrs.get('ctime', None)
        if ctime is not None:
            print(f"file created : {ctime}")
        if len(args) == 0:
            # no input arguments - load everything
            key_list = f.keys()
        else:
            if (len(args) == 1) and (type(args[0]) is list):
                # input argument is a list of object names
                key_list = args[0]
            else:
                # arguments form a list
                key_list = list(args)
        for k in key_list:
            # detect slicing
            if '[' in k:
                k, slice_string = k.split('[')
                slice_string = slice_string.split(']')[0]
                sldims = slice_string.split(',')
                sl = tuple(
                    [slice(*[None if x.strip() == '' else int(x) for x in (s.split(':') + ['', '', ''])[:3]]) for s in
                     sldims])
            else:
                sl = None

            # detect group access
            if '.' in k:
                glist = k.split('.')
                k = glist[-1]
                gr = f[glist[0]]
                for gname in glist[1:-1]:
                    gr = gr[gname]
                outdict[k] = _load(gr[k], sl)
            else:
                outdict[k] = _load(f[k], sl)

    return outdict


def h5info(filename, output=None):
    """\
    h5info(filename)

    Prints out a tree structure of given h5 file.

    [17/01/2012 guillaume potdevin]
    added optional argument output:
    	if output is set to 1, then the printed string is returned
    """

    indent = 4
    filename = os.path.abspath(os.path.expanduser(filename))

    def _format_dict(key, dset):
        stringout = ' ' * key[0] + ' * %s [dict]:\n' % key[1]
        for k, v in dset.items():
            if v.attrs.get('escaped', None) is not None:
                k = k.replace(h5options['SLASH_ESCAPE'], '/')
            stringout += _format((key[0] + indent, k), v)
        return stringout

    def _format_list(key, dset):
        stringout = ' ' * key[0] + ' * %s [list]:\n' % key[1]
        keys = dset.keys()
        keys.sort()
        for k in keys:
            stringout += _format((key[0] + indent, ''), dset[k])
        return stringout

    def _format_tuple(key, dset):
        stringout = ' ' * key[0] + ' * %s [tuple]:\n' % key[1]
        keys = dset.keys()
        keys.sort()
        for k in keys:
            stringout += _format((key[0] + indent, ''), dset[k])
        return stringout

    def _format_arraytuple(key, dset):
        a = dset[...]
        if len(a) < 5:
            stringout = ' ' * key[0] + ' * ' + key[1] + ' [tuple = ' + str(tuple(a.ravel())) + ']\n'
        else:
            try:
                float(a.ravel()[0])
                stringout = ' ' * key[0] + ' * ' + key[1] + ' [tuple = (' + (
                            ('%f, ' * 4) % tuple(a.ravel()[:4])) + ' ...)]\n'
            except ValueError:
                stringout = ' ' * key[0] + ' * ' + key[1] + ' [tuple = (%d x %s objects)]\n' % (a.size, str(a.dtype))
        return stringout

    def _format_arraylist(key, dset):
        a = dset[...]
        if len(a) < 5:
            stringout = ' ' * key[0] + ' * ' + key[1] + ' [list = ' + str(a.tolist()) + ']\n'
        else:
            try:
                float(a.ravel()[0])
                stringout = ' ' * key[0] + ' * ' + key[1] + ' [list = [' + (
                            ('%f, ' * 4) % tuple(a.ravel()[:4])) + ' ...]]\n'
            except ValueError:
                stringout = ' ' * key[0] + ' * ' + key[1] + ' [list = [%d x %s objects]]\n' % (a.size, str(a.dtype))
        return stringout

    def _format_numpy(key, dset):
        a = dset[...]
        if len(a) < 5 and a.ndim == 1:
            stringout = ' ' * key[0] + ' * ' + key[1] + ' [array = ' + str(a.ravel()) + ']\n'
        else:
            stringout = ' ' * key[0] + ' * ' + key[1] + ' [' + (('%dx' * (a.ndim - 1) + '%d') % a.shape) + ' ' + str(
                a.dtype) + ' array]\n'
        return stringout

    def _format_scalar(key, dset):
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [scalar = ' + str(dset[...]) + ']\n'
        return stringout

    def _format_str(key, dset):
        s = str(dset[...])
        if len(s) > 40:
            s = s[:40] + '...'
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [string = "' + s + '"]\n'
        return stringout

    def _format_unicode(key, dset):
        s = str(dset[...]).decode('utf8')
        if len(s) > 40:
            s = s[:40] + '...'
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [unicode = "' + s + '"]\n'
        return stringout

    def _format_pickle(key, dset):
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [pickled object]\n'
        return stringout

    def _format_None(key, dset):
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [None]\n'
        return stringout

    def _format_unknown(key, dset):
        stringout = ' ' * key[0] + ' * ' + key[1] + ' [unknown]\n'
        return stringout

    def _format(key, dset):
        dset_type = dset.attrs.get('type', None)

        # Treat groups as dicts
        if (dset_type is None) and (type(dset) is h5py.Group):
            dset_type = 'dict'

        if dset_type == 'dict':
            stringout = _format_dict(key, dset)
        elif dset_type == 'list':
            stringout = _format_list(key, dset)
        elif dset_type == 'array':
            stringout = _format_numpy(key, dset)
        elif dset_type == 'arraylist':
            stringout = _format_arraylist(key, dset)
        elif dset_type == 'tuple':
            stringout = _format_tuple(key, dset)
        elif dset_type == 'arraytuple':
            stringout = _format_arraytuple(key, dset)
        elif dset_type == 'string':
            stringout = _format_str(key, dset)
        elif dset_type == 'unicode':
            stringout = _format_unicode(key, dset)
        elif dset_type == 'scalar':
            stringout = _format_scalar(key, dset)
        elif dset_type == 'None':
            try:
                stringout = _format_numpy(key, dset)
            except:
                stringout = _format_None(key, dset)
        elif dset_type == 'pickle':
            stringout = _format_pickle(dset)
        elif dset_type is None:
            stringout = _format_numpy(key, dset)
        else:
            stringout = _format_unknown(key, dset)
        return stringout

    with h5py.File(filename, 'r') as f:
        h5rw_version = f.attrs.get('h5rw_version', None)
        #        if h5rw_version is None:
        #            print('Warning: this file does not seem to follow h5read format.')
        ctime = f.attrs.get('ctime', None)
        #        if ctime is not None:
        #            print('File created : ' + ctime)
        key_list = f.keys()
        outstring = ''
        for k in key_list:
            outstring += _format((0, k), f[k])

    print(outstring)

    # return string if output variable passed as option
    if output != None:
        return outstring



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

@cuda.jit
def prox_D_gaussian_kernel(z, z_hat, a, beta, a_strides):
    """

    :param z:           D x K x My x Mx x 2
    :param z_hat:       D x K x My x Mx x 2
    :param a:           D x K x My x Mx
    :param beta:        1
    :param a_strides:   (4,)
    :return:
    """
    n = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
    D, K, MY, MX, _ = z.shape
    N1 = D * K * MY * MX

    d = n // a_strides[0]
    k = (n - d * a_strides[0]) // a_strides[1]
    my = (n - d * a_strides[0] - k * a_strides[1]) // a_strides[2]
    mx = (n - d * a_strides[0] - k * a_strides[1] - my * a_strides[2]) // a_strides[3]

    if n < N1:
        z_hatc = z_hat[d, k, my, mx, 0] + 1j * z_hat[d, k, my, mx, 1]
        abs_zhat_c = abs(z_hatc)
        if abs_zhat_c != 0:
            sgn_zhat = z_hatc / abs_zhat_c
            fac = (a[d, k, my, mx] + beta * abs_zhat_c) / (1.0 + beta)
            zc = fac * sgn_zhat
            z[d, k, my, mx, 0] = zc.real
            z[d, k, my, mx, 1] = zc.imag

# Cell
def prox_D_gaussian(z, z_hat, a, beta):
    """
    Proximal operator of the Gaussian log-likelihood.

    :param z:           D x K x My x Mx x 2, updated exit waves
    :param z_hat:       D x K x My x Mx x 2, current model exit waves
    :param a:           D x K x My x Mx,     measured amplitudes
    :param beta:        float                hyperparameter

    :return: z
    """
    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(a.shape)) / threadsperblock)
    strides = th.tensor(a.stride()).to(z.device)
    prox_D_gaussian_kernel[blockspergrid, threadsperblock, stream](z, z_hat, a, beta, strides)
    return z

# Cell
@cuda.jit
def gradz_poisson_sparse_kernel(out, z, a_indices, a_counts, no_count_indicator, total_cts):
    """

    :param z:           D x K x My x Mx x 2
    :param z_hat:       D x K x My x Mx x 2
    :param a_indices:   D x K x counts
    :param a_counts:    D x K x counts
    :param beta:        float
    :param frame_dimensions: (2,)
    :param no_count_indicator: float or int
    :return:
    """

    d, k = cuda.grid(2)
    D, K, MY, MX, _ = z.shape
    if d < D and k < K:
        if total_cts[d,k] > 0:
            for i in range(a_indices[d, k].shape[0]):
                idx1d = a_indices[d, k, i]
                if idx1d != no_count_indicator:
                    my = idx1d // MX
                    mx = idx1d - my * MX
                    zc = z[d, k, my, mx, 0] + 1j * z[d, k, my, mx, 1]
                    abs_zc = abs(zc)
                    # if abs_zc != 0:
                    fac = 1 - (a_counts[d, k, i] / (abs_zc**2+1e-2))
                    # else:
                    #     fac = 1 - (a_counts[d, k, i] / 1e-3)
                    zc *= fac
                    out[d, k, my, mx, 0] = zc.real
                    out[d, k, my, mx, 1] = zc.imag

# Cell
def gradz_poisson_sparse(out, z, a_indices, a_counts):
    """
    Proximal operator of the Gaussian log-likelihood. Sparse version

    :param z:           D x K x My x Mx x 2, updated exit waves
    :param z_hat:       D x K x My x Mx x 2, current model exit waves
    :param a_indices:   D x K x cts,     measured amplitude indices
    :param a_counts:    D x K x cts,     measured amplitude counts
    :param beta:        float                hyperparameter

    :return: z
    """
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = (2, 32)
    blockspergrid = tuple(np.ceil(np.array(z.shape[:2]) / threadsperblock).astype(np.int))
    no_count_indicator = th.iinfo(a_indices.dtype).max
    total_cts = th.sum(a_counts,2)
    gradz_poisson_sparse_kernel[blockspergrid, threadsperblock, stream](out, z, a_indices, a_counts, no_count_indicator,
                                                                        total_cts)
    return z

# Cell
@cuda.jit
def gradz_gaussian_sparse_kernel(out, z, a_indices, a_counts, no_count_indicator):
    """

    :param z:           D x K x My x Mx x 2
    :param z_hat:       D x K x My x Mx x 2
    :param a_indices:   D x K x counts
    :param a_counts:    D x K x counts
    :param beta:        float
    :param frame_dimensions: (2,)
    :param no_count_indicator: float or int
    :return:
    """
    d, k = cuda.grid(2)
    D, K, MY, MX, _ = z.shape
    if d < D and k < K:
        for i in range(a_indices[d, k].shape[0]):
            idx1d = a_indices[d, k, i]
            if idx1d != no_count_indicator:
                my = idx1d // MX
                mx = idx1d - my * MX
                zc = z[d, k, my, mx, 0] + 1j * z[d, k, my, mx, 1]
                abs_zc = abs(zc)
                if abs_zc != 0:
                    fac = 1 - (float(a_counts[d, k, i]) / abs_zc)
                else:
                    fac = 1 - (float(a_counts[d, k, i]) / 1e-3)
                zc *= fac
                out[d, k, my, mx, 0] = zc.real
                out[d, k, my, mx, 1] = zc.imag

# Cell

def gradz_gaussian_sparse(out, z, a_indices, a_counts):
    """
    Proximal operator of the Gaussian log-likelihood. Sparse version

    :param z:           D x K x My x Mx x 2, updated exit waves
    :param z_hat:       D x K x My x Mx x 2, current model exit waves
    :param a_indices:   D x K x cts,     measured amplitude indices
    :param a_counts:    D x K x cts,     measured amplitude counts
    :param beta:        float                hyperparameter

    :return: z
    """
    gpu = cuda.get_current_device()
    stream = th.cuda.current_stream().cuda_stream
    threadsperblock = (2, 32)
    blockspergrid = tuple(np.ceil(np.array(z.shape[:2]) / threadsperblock).astype(np.int))
    no_count_indicator = th.iinfo(a_indices.dtype).max
    gradz_gaussian_sparse_kernel[blockspergrid, threadsperblock, stream](out, z, a_indices, a_counts, no_count_indicator)
    return z

# Cell
import torch as th
import numpy as np

from numba import cuda


@cuda.jit
def sparse_amplitude_loss_kernel(a_model, indices_target, counts_target, loss, grad, frame_dimensions,
                                 no_count_indicator):
    k = cuda.grid(1)
    K, _ = indices_target.shape
    MY, MX = frame_dimensions
    if k < K:
        for i in range(indices_target[k].shape[0]):
            idx1d = indices_target[k, i]
            my = idx1d // MX
            mx = idx1d - my * MX
            if idx1d != no_count_indicator:
                grad[k, my, mx] = 1 - (counts_target[k, i] / (a_model[k, my, mx] + 1e-2))
                cuda.atomic.add(loss, (0), (a_model[k, my, mx] - counts_target[k, i]) ** 2)

def sparse_amplitude_loss(a_model, indices_target, counts_target, frame_dimensions):
    """

    :param a_model:             K x M1 x M2
    :param indices_target:      K x num_max_counts
    :param counts_target:       K x num_max_counts
    :param frame_dimensions:    2
    :return: loss (1,), grad (K x M1 x M2)
    """
    threadsperblock = (256,)
    blockspergrid = tuple(np.ceil(np.array(indices_target.shape[0]) / threadsperblock).astype(np.int))

    loss = th.zeros((1,), device=a_model.device, dtype=th.float32)
    grad = th.ones_like(a_model)
    no_count_indicator = th.iinfo(indices_target.dtype).max
    sparse_amplitude_loss_kernel[blockspergrid, threadsperblock](a_model.detach(), indices_target.detach(),
                                                                 counts_target.detach(), loss.detach(), grad.detach(),
                                                                 frame_dimensions, no_count_indicator)
    return loss, grad

import math as m
import cmath as cm
@cuda.jit
def sparse_smooth_truncated_amplitude_loss_kernel(a_model, indices_target, counts_target, loss, grad, no_count_indicator, eps):
    n = cuda.grid(1)
    K, MY, MX = a_model.shape
    N = K * MY * MX

    k = n // (MY * MX)
    my = (n - k * (MY * MX)) // MX
    mx = (n - k * (MY * MX) - my * MX)

    if n < N:
        idx1d = mx + my * MX
        a_measure = 0
        a_model = a_model[k, my, mx]
        for i in range(indices_target.shape[1]):
            if indices_target[k, i] == idx1d and no_count_indicator != indices_target[k, i]:
                a_measure = m.sqrt(counts_target[k, i])

        if a_model < eps * a_measure:
            grad[k, my, mx] =  (1-1/eps)
            loss_k = (1-eps)/2*(a_measure**2-(1/eps)*a_model**2)
        else:
            grad[k, my, mx] =  1 - (a_measure/a_model)
            loss_k = .5 * abs(a_model - a_measure)**2
        loss[k] = loss_k
        # cuda.atomic.add(loss, (k), loss_k)


def sparse_smooth_truncated_amplitude_loss(a_model, indices_target, counts_target, frame_dimensions, eps=0.1):
    """
    Smooth truncated amplitude loss from Chang et al., Overlapping Domain Decomposition Methods for Ptychographic Imaging, (2020)

    :param a_model:             K x M1 x M2
    :param indices_target:      K x num_max_counts
    :param counts_target:       K x num_max_counts
    :param frame_dimensions:    2
    :return: loss (K,), grad (K x M1 x M2)
    """

    threadsperblock = (256,)
    blockspergrid = tuple(np.ceil(np.array(np.prod(a_model.shape)) / threadsperblock).astype(np.int))

    loss = th.zeros((a_model.shape[0],), device=a_model.device, dtype=th.float32)
    grad = th.ones_like(a_model)
    no_count_indicator = th.iinfo(indices_target.dtype).max

    sparse_smooth_truncated_amplitude_loss_kernel[blockspergrid, threadsperblock](a_model.detach(), indices_target.detach(),
                                                                 counts_target.detach(), loss.detach(), grad.detach(),
                                                                 no_count_indicator, eps)
    return loss, grad

@cuda.jit
def sparse_smooth_truncated_amplitude_prox_kernel(a_model, indices_target, counts_target, loss, grad, no_count_indicator, eps, lam):
    n = cuda.grid(1)
    K, MY, MX = a_model.shape
    N = K * MY * MX

    k = n // (MY * MX)
    my = (n - k * (MY * MX)) // MX
    mx = (n - k * (MY * MX) - my * MX)

    if n < N:
        idx1d = mx + my * MX
        a_measure = 0
        a_model = a_model[k, my, mx]
        for i in range(indices_target.shape[1]):
            if indices_target[k, i] == idx1d and no_count_indicator != indices_target[k, i]:
                a_measure = m.sqrt(counts_target[k, i])

        # if a_model < eps * a_measure:
        #     grad[k, my, mx] =  max(0, lam * a_measure/(lam-((1-eps)/eps)))/ a_model
        #     loss_k = (1-eps)/2*(a_measure**2-(1/eps)*a_model**2)
        # else:
        #     grad[k, my, mx] =  ((a_measure+lam * a_model)/(1+lam))/ a_model
        #     loss_k = .5 * abs(a_model - a_measure)**2

        grad[k, my, mx] =  a_measure/ a_model
        loss_k = .5 * abs(a_model - a_measure)**2

        loss[k] = loss_k


def sparse_smooth_truncated_amplitude_prox(a_model, indices_target, counts_target, frame_dimensions, eps=0.5, lam=6e-1):
    """
    Smooth truncated amplitude loss from Chang et al., Overlapping Domain Decomposition Methods for Ptychographic Imaging, (2020)

    :param a_model:             K x M1 x M2
    :param indices_target:      K x num_max_counts
    :param counts_target:       K x num_max_counts
    :param frame_dimensions:    2
    :return: loss (K,), grad (K x M1 x M2)
    """

    threadsperblock = (256,)
    blockspergrid = tuple(np.ceil(np.array(np.prod(a_model.shape)) / threadsperblock).astype(np.int))

    loss = th.zeros((a_model.shape[0],), device=a_model.device, dtype=th.float32)
    grad = th.ones_like(a_model)
    no_count_indicator = th.iinfo(indices_target.dtype).max

    sparse_smooth_truncated_amplitude_prox_kernel[blockspergrid, threadsperblock](a_model.detach(), indices_target.detach(),
                                                                 counts_target.detach(), loss.detach(), grad.detach(),
                                                                 no_count_indicator, eps, lam)
    return loss, grad

@cuda.jit
def sparse_amplitude_prox_kernel(a_model, indices_target, counts_target, loss, grad, no_count_indicator, eps, lam):
    n = cuda.grid(1)
    K, MY, MX = a_model.shape
    N = K * MY * MX

    k = n // (MY * MX)
    my = (n - k * (MY * MX)) // MX
    mx = (n - k * (MY * MX) - my * MX)

    if n < N:
        idx1d = mx + my * MX
        a_measure = 0
        a_model = a_model[k, my, mx]
        for i in range(indices_target.shape[1]):
            if indices_target[k, i] == idx1d and no_count_indicator != indices_target[k, i]:
                a_measure = m.sqrt(counts_target[k, i])

        grad[k, my, mx] =  a_measure/ a_model
        loss_k = .5 * abs(a_model - a_measure)**2

        loss[k] = loss_k


def sparse_amplitude_prox(a_model, indices_target, counts_target, frame_dimensions, eps=0.5, lam=6e-1):
    """
    Smooth truncated amplitude loss from Chang et al., Overlapping Domain Decomposition Methods for Ptychographic Imaging, (2020)

    :param a_model:             K x M1 x M2
    :param indices_target:      K x num_max_counts
    :param counts_target:       K x num_max_counts
    :param frame_dimensions:    2
    :return: loss (K,), grad (K x M1 x M2)
    """

    threadsperblock = (256,)
    blockspergrid = tuple(np.ceil(np.array(np.prod(a_model.shape)) / threadsperblock).astype(np.int))

    loss = th.zeros((a_model.shape[0],), device=a_model.device, dtype=th.float32)
    grad = th.ones_like(a_model)
    no_count_indicator = th.iinfo(indices_target.dtype).max

    sparse_amplitude_prox_kernel[blockspergrid, threadsperblock](a_model.detach(), indices_target.detach(),
                                                                 counts_target.detach(), loss.detach(), grad.detach(),
                                                                 no_count_indicator, eps, lam)
    return loss, grad