"""A module to simulate optical transfer functions and point spread functions.

If this file is run as a script (python -m pyotf.otf) it will compare
the HanserPSF to the SheppardPSF in a plot.

https://en.wikipedia.org/wiki/Optical_transfer_function
https://en.wikipedia.org/wiki/Point_spread_function

Copyright (c) 2020, David Hoffman
"""

import copy
import logging
from functools import cached_property

import numpy as np
from numpy.fft import fftfreq, fftshift, ifftn
from numpy.linalg import norm

from .utils import NumericProperty, cart2pol, easy_fft, easy_ifft, psqrt, slice_maker
from .display import psf_plot, otf_plot
from .zernike import name2noll, zernike

logger = logging.getLogger(__name__)


class BasePSF(object):
    """A base class for objects that can calculate OTF's and PSF's.
    
    It is not intended to be used alone

    To fully describe a PSF or OTF of an objective lens, assuming no
    abberation, we generally need a few parameters:
    - The wavelength of operation (assume monochromatic light)
    - the numerical aperature of the objective
    - the index of refraction of the medium

    For numerical calculations we'll also want to know the x/y resolution
    and number of points. Note that it is assumed that z is the optical
    axis of the objective lens
    """

    # Define all the numeric properties of the base class
    wl = NumericProperty(attr="_wl", vartype=(float, int), doc="Wavelength of emission, in nm")
    na = NumericProperty(attr="_na", vartype=(float, int), doc="Numerical Aperature")
    ni = NumericProperty(attr="_ni", vartype=(float, int), doc="Refractive index")
    size = NumericProperty(attr="_size", vartype=int, doc="x/y size")
    zsize = NumericProperty(attr="_zsize", vartype=int, doc="z size")

    def __init__(
        self, wl, na, ni, res, size, zres=None, zsize=None, vec_corr="none", condition="sine"
    ):
        """Generate a PSF object.

        Parameters
        ----------
        wl : numeric
            Emission wavelength of the simulation
        na : numeric
            Numerical aperature of the simulation
        ni : numeric
            index of refraction for the media
        res : numeric
            x/y resolution of the simulation, must have same units as wl
        size : int
            x/y size of the simulation

        Optional Parameters
        -------------------
        zres : numeric
            z resolution of simuation, must have same units a wl
        zsize : int
            z size of simulation
        vec_corr : str
            keyword to indicate whether to include vectorial effects
                Valid options are: "none", "x", "y", "z", "total"
                Default is: "none"
        condition : str
            keyword to indicate whether to model the sine or herschel conditions
            **Herschel's Condition** invariance of axial magnification
            **Abbe's Sine Condition** invariance of lateral magnification
            conditions
                Valid options are: "none", "sine", "herschel"
                Default is: "sine"
                Note: "none" is not a physical solution
        """
        self.wl = wl
        self.na = na
        self.ni = ni
        self.res = res
        self.size = size
        # if zres is not passed, set it to res
        if zres is None:
            zres = res
        self.zres = zres
        # if zsize isn't passed set it to size
        if zsize is None:
            zsize = size
        self.zsize = zsize
        self.vec_corr = vec_corr
        self.condition = condition

    def __repr__(self):
        """Return representation of PSF object."""
        return (
            f"{self.__class__.__name__}(wl={self.wl}, na={self.na}, ni={self.ni},"
            + f" res={self.res}, size={self.size}, zres={self.zres}, zsize={self.zsize},"
            + f" vec_corr='{self.vec_corr}', condition='{self.condition}')"
        )

    def _attribute_changed(self):
        """Attribute has changed.

        Sets internal state variables to None so that when the
        user asks for them they are recalculated
        """
        for attr in ("PSFa", "OTFa", "PSFi", "OTFi"):
            try:
                delattr(self, attr)
            except AttributeError:
                logger.debug(f"{attr} wasn't available to delete")

    @property
    def zres(self):
        """Z resolution."""
        return self._zres

    @zres.setter
    def zres(self, value):
        # make sure z res is positive
        if not value > 0:
            raise ValueError("zres must be positive")
        self._zres = value
        self._attribute_changed()

    @property
    def res(self):
        """X/Y resolution."""
        return self._res

    @res.setter
    def res(self, value):
        # max_val is the abbe limit, but for an accurate simulation
        # the pixel size must be smaller than half this number
        # thinking in terms of the convolution that is implicitly
        # performed when generating the OTFi we also don't want
        # any wrapping effects. However, allowing the number
        # to be the Abbe limit can allow phase retrieval for
        # larger pixels
        abbe_limit = 1 / (2 * self.na / self.wl)
        if value >= abbe_limit:
            raise ValueError(
                f"{value} is larger than the Abbe Limit, try a number smaller than {abbe_limit}"
            )
        if value >= abbe_limit / 2:
            logger.info(
                f"res has been set to {value} which is greater than the Nyquist limit of {abbe_limit / 2}"
            )
        self._res = value
        self._attribute_changed()

    @property
    def vec_corr(self):
        """Take into account the vectorial nature of light.
        
        Valid values are: "none", "x", "y", "z", "total"
        """
        return self._vec_corr

    @vec_corr.setter
    def vec_corr(self, value):
        valid_values = {"none", "x", "y", "z", "total"}
        if value not in valid_values:
            raise ValueError("Vector correction must be one of {}".format(", ".join(valid_values)))
        self._vec_corr = value
        self._attribute_changed()

    @property
    def condition(self):
        """Imaging condition to simulate."""
        return self._condition

    @condition.setter
    def condition(self, value):
        valid_values = {"none", "sine", "herschel"}
        if value not in valid_values:
            raise ValueError(("Condition must be one of {}").format(", ".join(valid_values)))
        self._condition = value
        self._attribute_changed()

    @cached_property
    def OTFa(self):
        """Amplitude OTF (coherent transfer function), complex array."""
        raise NotImplementedError

    @cached_property
    def PSFa(self):
        """Amplitude PSF, complex array."""
        raise NotImplementedError

    @cached_property
    def PSFi(self):
        """Intensity PSF, real array."""
        return (abs(self.PSFa) ** 2).sum(axis=0)

    @cached_property
    def OTFi(self):
        """Intensity OTF, complex array."""
        return easy_fft(self.PSFi)

    def _validate_zrange(self):
        """Check zrange for uniform step size."""
        try:
            # make sure there's only one size of z-step and that there's more than one
            zsteps = np.diff(self.zrange)
            if len(zsteps) < 2 or not np.allclose(zsteps, zsteps.mean()):
                raise RuntimeError(f"{self} doesn't have uniform z-steps ---> {zsteps}")
        except AttributeError:
            pass

    def plot_psf(self, **kwargs):
        """Plot the intensity PSF.
        
        See `pyotf.display.psf_plot` for details and possible kwargs
        """
        self._validate_zrange()
        # smart cropping
        # nice lateral extent
        lateral_extent = self.wl / 2 / self.na / self.res * 32
        axial_extent = self.wl / (self.ni - np.sqrt(self.ni ** 2 - self.na ** 2)) / self.zres * 16

        max_loc = np.unravel_index(self.PSFi.argmax(), self.PSFi.shape)
        crop = slice_maker(max_loc, (axial_extent, lateral_extent, lateral_extent))
        return psf_plot(self.PSFi[crop], zres=self.zres, res=self.res, **kwargs)

    def plot_otf(self, **kwargs):
        """Plot the intensity OTF.
        
        See `pyotf.display.otf_plot` for details and possible kwargs
        """
        self._validate_zrange()

        # normalize OTF and make sure it's real
        otf = abs(self.OTFi)
        otf = np.fmax(otf / otf.max(), np.finfo(float).eps)

        # nice default plotting kwargs
        dkwargs = dict(vmin=1e-4)
        dkwargs.update(kwargs)

        return otf_plot(
            otf, na=self.na, ni=self.ni, wl=self.wl, zres=self.zres, res=self.res, **dkwargs,
        )


class HanserPSF(BasePSF):
    """A class defining the pupil function and its closely related methods.

    Based on the following work

    [(1) Hanser, B. M.; Gustafsson, M. G. L.; Agard, D. A.; Sedat, J. W.
    Phase-Retrieved Pupil Functions in Wide-Field Fluorescence Microscopy.
    Journal of Microscopy 2004, 216 (1), 32–48.](dx.doi.org/10.1111/j.0022-2720.2004.01393.x)
    [(2) Hanser, B. M.; Gustafsson, M. G. L.; Agard, D. A.; Sedat, J. W.
    Phase Retrieval for High-Numerical-Aperture Optical Systems.
    Optics Letters 2003, 28 (10), 801.](dx.doi.org/10.1364/OL.28.000801)
    """

    def __init__(self, *args, zrange=None, **kwargs):  # noqa: D205,D208,D400,D403
        """zrange : array-like
            An alternate way to specify the z range for the calculation
            must be expressed in the same units as wavelength
        """
        super().__init__(*args, **kwargs)
        if zrange is None:
            self._gen_zrange()
        else:
            self.zrange = zrange

    # include parent documentation
    __init__.__doc__ = BasePSF.__init__.__doc__ + __init__.__doc__

    def __repr__(self):
        """Represent HanserPSF."""
        return super().__repr__()[:-1] + f", zrange={self.zrange!r})"

    def _gen_zrange(self):
        """Generate the zrange from zsize and zres."""
        self.zrange = (np.arange(self.zsize) - (self.zsize + 1) // 2) * self.zres

    @BasePSF.zsize.setter
    def zsize(self, value):
        """Set zsize."""
        # we need override this setter so that the zrange is recalculated
        BasePSF.zsize.fset(self, value)
        # try and except is necessary for initialization
        try:
            self._gen_zrange()
        except AttributeError:
            pass

    @BasePSF.zres.setter
    def zres(self, value):
        """Set zres."""
        # same as for zsize
        BasePSF.zres.fset(self, value)
        try:
            self._gen_zrange()
        except AttributeError:
            pass

    @property
    def zrange(self):
        """Return range overwhich to calculate the psf."""
        return self._zrange

    @zrange.setter
    def zrange(self, value):
        self._zrange = np.asarray(value)
        # check if passed value is scalar
        if not self._zrange.shape:
            # convert to array for later multiplications
            self._zrange.shape = (1,)
        self._attribute_changed()

    def _gen_kr(self):
        """Generate coordinate system and other internal parameters."""
        k = self._k = fftfreq(self.size, self.res)
        kxx, kyy = np.meshgrid(k, k)
        self._kr, self._phi = cart2pol(kyy, kxx)
        # kmag is the radius of the spherical shell of the OTF
        self._kmag = self.ni / self.wl
        # because the OTF only exists on a spherical shell we can calculate
        # a kz value for any pair of kx and ky values
        self._kz = psqrt(self._kmag ** 2 - self._kr ** 2)

    def _gen_pupil(self):
        """Generate an ideal pupil."""
        kr = self._kr
        # define the diffraction limit
        # remember we"re working with _coherent_ data _not_ intensity,
        # so drop the factor of 2
        diff_limit = self._na / self._wl
        # return a circle of intensity 1 over the ideal passband of the
        # objective make sure data is complex
        return (kr < diff_limit).astype(complex)

    def _calc_defocus(self):
        """Calculate the defocus to apply to the base pupil."""
        kz = self._kz
        return np.exp(2 * np.pi * 1j * kz * self.zrange[:, np.newaxis, np.newaxis])

    def _gen_psf(self, pupil_base=None):
        """Generate the PSF.

        kwargs
        ------
        pupil_base : ndarray
            provided so that phase retrieval algorithms can hook into this
            method.

        NOTE: that the internal state is created with fftfreq, which creates
        _unshifted_ frequences
        """
        # clear internal state
        self._attribute_changed()
        # generate internal coordinates
        self._gen_kr()
        # generate the pupil
        if pupil_base is None:
            pupil_base = self._gen_pupil()
        else:
            assert pupil_base.ndim == 2, f"`pupil_base` is wrong shape: {pupil_base.shape}"
            # Maybe we should do ifftshift here so user doesn't have too
        # pull relevant internal state variables
        kr = self._kr
        phi = self._phi
        kmag = self._kmag
        # apply the defocus to the base_pupil
        pupil = pupil_base * self._calc_defocus()
        # calculate theta, this is possible because we know that the
        # OTF is only non-zero on a spherical shell
        theta = np.arcsin((kr < kmag) * kr / kmag)
        # The authors claim that the following code is unecessary as the
        # sine condition is already taken into account in the definition
        # of the pupil, but I call bullshit
        if self.condition != "none":
            if self.condition == "sine":
                a = 1.0 / np.sqrt(np.cos(theta))
            elif self.condition == "herschel":
                a = 1.0 / np.cos(theta)
            else:
                raise RuntimeError("You should never see this")
            pupil *= a
        # apply the vectorial corrections, if requested
        if self.vec_corr != "none":
            plist = []
            if self.vec_corr == "z" or self.vec_corr == "total":
                plist.append(np.sin(theta) * np.cos(phi))  # Pzx
                plist.append(np.sin(theta) * np.sin(phi))  # Pzy
            if self.vec_corr == "y" or self.vec_corr == "total":
                plist.append((np.cos(theta) - 1) * np.sin(phi) * np.cos(phi))  # Pyx
                plist.append(np.cos(theta) * np.sin(phi) ** 2 + np.cos(phi) ** 2)  # Pyy
            if self.vec_corr == "x" or self.vec_corr == "total":
                plist.append(np.cos(theta) * np.cos(phi) ** 2 + np.sin(phi) ** 2)  # Pxx
                plist.append((np.cos(theta) - 1) * np.sin(phi) * np.cos(phi))  # Pxy
            # apply the corrections to the base pupil
            pupils = pupil * np.array(plist)[:, np.newaxis]
        else:
            # if no correction we still need one more axis for the following
            # code to work generally
            pupils = pupil[np.newaxis]
        # save the pupil for inspection, not necessary
        # self._pupils = pupils
        # because the internal state is created with fftfreq, no initial shift
        # is necessary.
        PSFa = fftshift(ifftn(pupils, axes=(2, 3)), axes=(2, 3))
        # save the PSF internally
        return PSFa

    def apply_pupil(self, pupil):
        """Apply a pupil function to the model."""
        self._attribute_changed()
        self.PSFa = self._gen_psf(pupil)

    @cached_property
    def OTFa(self):
        """Amplitude OTF."""
        return easy_fft(self.PSFa, axes=(1, 2, 3))

    @cached_property
    def PSFa(self):
        """Amplitude PSF."""
        return self._gen_psf()


class SheppardPSF(BasePSF):
    """A class defining the 3D pupil function and its closely related methods.
    
    Based on the following work:

    [(1) Arnison, M. R.; Sheppard, C. J. R. A 3D Vectorial Optical Transfer
    Function Suitable for Arbitrary Pupil Functions. Optics Communications
    2002, 211 (1–6), 53–63.](dx.doi.org/10.1016/S0030-4018(02)01857-6)
    """

    dual = NumericProperty(attr="_dual", vartype=bool, doc="Simulate dual objectives")

    def __init__(self, *args, dual=False, **kwargs):  # noqa: D205,D208,D400,D403
        """dual : bool
            Simulate dual objectives
        """
        super().__init__(*args, **kwargs)
        self.dual = dual

    # include parent documentation
    __init__.__doc__ = BasePSF.__init__.__doc__ + __init__.__doc__

    def __repr__(self):
        """Represent SheppardPSF."""
        return super().__repr__()[:-1] + f", dual={self.dual})"

    @property
    def dual(self):
        """Simulate opposing objectives."""
        return self._dual

    @dual.setter
    def dual(self, value):
        if not isinstance(value, bool):
            raise TypeError("`dual` must be a boolean")
        self._dual = value
        self._attribute_changed()

    @BasePSF.zres.setter
    def zres(self, value):
        """Set zres."""
        # this checks the nyquist limit for z
        # remember that because we create a spherical shell for
        # The amplitude OTF not nyquist for the final intensity OTF ...
        max_val = self.wl / 2 / self.ni
        if value >= max_val:
            # this will cause a fftconvolution error when calculating the
            # intensity OTF
            raise ValueError(f"{value} is too large try a number smaller than {max_val}")
        BasePSF.zres.fset(self, value)

    def _gen_kr(self):
        """Generate internal state."""
        # generate internal kspace coordinates
        k = fftfreq(self.size, self.res)
        kz = fftfreq(self.zsize, self.zres)
        k_tot = np.meshgrid(kz, k, k, indexing="ij")
        # calculate r
        kr = norm(k_tot, axis=0)
        # calculate the radius of the spherical shell in k-space
        self.kmag = kmag = self.ni / self.wl
        # determine k-space pixel size
        dk, dkz = k[1] - k[0], kz[1] - kz[0]
        # save output for user
        self.dk, self.dkz = dk, dkz
        # determine the min value for kz given the NA and wavelength
        kz_min = np.sqrt(kmag ** 2 - (self.na / self.wl) ** 2)
        # make sure we're not crazy
        assert kz_min >= 0, "Something went horribly wrong"
        # if the user gave us different z and x/y res we need to calculate
        # the positional "error" in k-space to draw the spherical shell
        if dk != dkz:
            with np.errstate(invalid="ignore"):
                dd = np.array((dkz, dk, dk)).reshape(3, 1, 1, 1)
                dkr = norm(np.array(k_tot) * dd, axis=0) / kr
            # we know the origin is zero so replace it
            dkr[0, 0, 0] = 0.0
        else:
            dkr = dk
        if self.dual:
            # if we want dual objectives we need two spherical shells
            kzz = abs(k_tot[0])
        else:
            kzz = k_tot[0]
        # calculate the points on the spherical shell, save them and the
        # corresponding kz, ky and kx coordinates
        self.valid_points = np.logical_and(abs(kr - kmag) < dkr, kzz > kz_min + dkr)
        self.kzz, self.kyy, self.kxx = [k[self.valid_points] for k in k_tot]

    def _gen_otf(self):
        """Generate the OTFs."""
        # clear internal state
        self._attribute_changed()
        # generate coordinate space
        self._gen_kr()
        kxx, kyy, kzz = self.kxx, self.kyy, self.kzz
        # generate direction cosines
        m, n, s = np.array((kxx, kyy, kzz)) / norm((kxx, kyy, kzz), axis=0)
        # apply a given imaging condition
        if self.condition == "sine":
            a = 1.0 / np.sqrt(s)
        elif self.condition == "herschel":
            a = 1.0 / s
        elif self.condition == "none":
            a = 1.0
        else:
            raise RuntimeError("You should never see this")
        # apply the vectorial corrections if requested
        if self.vec_corr != "none":
            plist = []
            if self.vec_corr == "z" or self.vec_corr == "total":
                plist.append(-m)  # Pzx
                plist.append(-n)  # Pzy
            if self.vec_corr == "y" or self.vec_corr == "total":
                plist.append(-n * m / (1 + s))  # Pyx
                plist.append(1 - n ** 2 / (1 + s))  # Pyy
            if self.vec_corr == "x" or self.vec_corr == "total":
                plist.append(1 - m ** 2 / (1 + s))  # Pxx
                plist.append(-m * n / (1 + s))  # Pxy
            # generate empty otf
            otf = np.zeros((len(plist), self.zsize, self.size, self.size), dtype="D")
            # fill in the valid poins
            for o, p in zip(otf, plist):
                o[self.valid_points] = p * a
        else:
            # TODO: we can actually do a LOT better here.
            # if the vectorial correction is None then we can
            # calculate a 2D (kz, kr) OTF and interpolate it out to
            # the full 3D size.
            # otf_sub = self._gen_radsym_otf()
            # otf = otf_sub[np.newaxis]
            otf_sub = np.zeros((self.zsize, self.size, self.size), dtype="D")
            otf_sub[self.valid_points] = 1.0
            otf = otf_sub[np.newaxis]
        # we're already calculating the OTF, so we just need to shift it into
        # the right place.
        return fftshift(otf, axes=(1, 2, 3))

    @cached_property
    def OTFa(self):
        """Amplitude OTF."""
        return self._gen_otf()

    @cached_property
    def PSFa(self):
        """Amplitude PSF."""
        return easy_ifft(self.OTFa, axes=(1, 2, 3))


def apply_aberration(model, mcoefs, pcoefs):
    """Apply a set of abberations to a model PSF.

    Parameters
    ----------
    model : HanserPSF
        The model PSF to which to apply the aberrations
    mcoefs : ndarray (n, )
        The magnitude coefficiencts
    pcoefs : ndarray (n, )
        The phase coefficients
    
    Note: this function assumes the mcoefs and pcoefs are Noll ordered
    """
    # sanity checks
    assert isinstance(model, HanserPSF), "Model must be a HanserPSF"

    model = copy.copy(model)

    if mcoefs is None and pcoefs is None:
        logger.warning("No abberation applied")
        return model

    if mcoefs is None:
        mcoefs = np.zeros_like(pcoefs)

    if pcoefs is None:
        pcoefs = np.zeros_like(mcoefs)

    assert len(mcoefs) == len(pcoefs), "Coefficient lengths don't match"

    # extract kr
    model._gen_kr()
    kr = model._kr
    theta = model._phi
    # make zernikes (need to convert kr to r where r = 1 when kr is at
    # diffraction limit)
    r = kr * model.wl / model.na
    zerns = zernike(r, theta, np.arange(len(mcoefs)) + 1)

    pupil_phase = (zerns * pcoefs[:, None, None]).sum(0)
    pupil_mag = (zerns * mcoefs[:, None, None]).sum(0)

    # apply aberrations to unaberrated pupil (N.B. the unaberrated phase is 0)
    pupil_mag += abs(model._gen_pupil())

    # generate the PSF, assign to attribute
    pupil_total = pupil_mag * np.exp(1j * pupil_phase)
    model.apply_pupil(pupil_total)

    return model


def apply_named_aberration(model, aberration, magnitude):
    """Apply a specific named aberration to the PSF. This will only effect the phase."""
    pcoefs = named_aberration_to_pcoefs(aberration, magnitude)
    return apply_aberration(model, None, pcoefs)


def named_aberration_to_pcoefs(aberration, magnitude):
    """Convert named aberration into phase coefficients.

    Parameters
    ----------
    aberration: str
        Name of aberration
    magnitude: int
        Magnitude of aberration

    Returns
    -------
    np.ndarray
        Corresponding phase coefficients
    """
    try:
        noll = name2noll[aberration]
    except KeyError as e:
        raise KeyError(
            f"Aberration '{aberration}' unknown, choose from: '"
            + "', '".join(name2noll.keys())
            + "'"
        )
    pcoefs = np.zeros(len(name2noll))
    pcoefs[noll - 1] = magnitude
    return pcoefs


def apply_named_aberrations(model, aberrations):
    """Use to apply multiple named aberration to the PSF. This will only affect the phase.

    Parameters
    ----------
    model: PSF
        PSF model onto which aberration will be applied
    aberrations: dict()
        Dictionary of aberration-magnitude pairs

    Returns
    -------
    PSF
        Aberrated model

    """
    pcoefs = np.zeros(len(name2noll))
    for aberration, magnitude in aberrations.items():
        # Sum phase coefficients
        pcoefs = np.add(pcoefs, named_aberration_to_pcoefs(aberration, magnitude))
    return apply_aberration(model, None, pcoefs)


if __name__ == "__main__":
    # import plotting
    from matplotlib import pyplot as plt

    # generate a comparison
    kwargs = dict(
        wl=520e-3,
        na=1.27,
        ni=1.33,
        res=90e-3,
        size=256,
        zres=190e-3,
        zsize=128,
        vec_corr="none",
        condition="none",
    )
    psfs = HanserPSF(**kwargs), SheppardPSF(**kwargs)

    with plt.style.context("dark_background"):

        fig, axs = plt.subplots(2, 2, figsize=(9, 6), gridspec_kw=dict(width_ratios=(1, 2)))

        for psf, ax_sub in zip(psfs, axs):
            psf.plot_otf()
            psf.plot_psf(interpolation="bicubic")
            # make coordinates
            ax_yx, ax_zx = ax_sub
            # get magnitude
            otf = abs(psf.OTFi)
            # normalize
            otf /= otf.max()
            otf /= otf.mean()
            otf = np.log(otf + np.finfo(float).eps)

            # plot
            style = dict(vmin=-3, vmax=5, cmap="inferno", interpolation="bicubic")
            ax_yx.matshow(otf[otf.shape[0] // 2], **style)
            ax_yx.set_title("{} $k_y k_x$ plane".format(psf.__class__.__name__))
            ax_zx.matshow(otf[..., otf.shape[1] // 2], **style)
            ax_zx.set_title("{} $k_z k_x$ plane".format(psf.__class__.__name__))

            for ax in ax_sub:
                ax.xaxis.set_major_locator(plt.NullLocator())
                ax.yaxis.set_major_locator(plt.NullLocator())
        fig.tight_layout()

    # NOTE: the results are _very_ close on a qualitative scale, but they do not match exactly
    # as theory says they should (they're mathematically identical to one another)

    model_kwargs = dict(
        wl=525, na=1.27, ni=1.33, res=70, size=256, zrange=[0], vec_corr="none", condition="none",
    )
    model = HanserPSF(**model_kwargs)

    with plt.style.context("dark_background"):
        fig, axs = plt.subplots(3, 5, figsize=(12, 8))
        # fill out plot
        for ax, name in zip(axs.ravel(), name2noll.keys()):
            model2 = apply_named_aberration(model, name, 1)
            ax.imshow(
                model2.PSFi.squeeze()[104:-104, 104:-104], cmap="inferno", interpolation="bicubic"
            )
            ax.set_xlabel(name.replace(" ", "\n", 1).title())
            ax.xaxis.set_major_locator(plt.NullLocator())
            ax.yaxis.set_major_locator(plt.NullLocator())

        # fig.tight_layout()
    plt.show()
