# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Define tools to plot Gridded PSF models.
"""

import astropy
import numpy as np
from astropy.utils import minversion
from astropy.visualization import simple_norm

__all__ = ['ModelGridPlotMixin']


class ModelGridPlotMixin:
    """
    Mixin class to plot a grid of ePSF models.
    """

    def _reshape_grid(self, data):
        """
        Reshape the 3D ePSF grid as a 2D array of horizontally and
        vertically stacked ePSFs.

        Parameters
        ----------
        data : `numpy.ndarray`
            The 3D array of ePSF data.

        Returns
        -------
        reshaped_data : `numpy.ndarray`
            The 2D array of ePSF data.
        """
        nypsfs = self._ygrid.shape[0]
        nxpsfs = self._xgrid.shape[0]
        ny, nx = self.data.shape[1:]
        data.shape = (nypsfs, nxpsfs, ny, nx)

        return data.transpose([0, 2, 1, 3]).reshape(nypsfs * ny, nxpsfs * nx)

    def plot_grid(self, *, ax=None, vmax_scale=None, peak_norm=False,
                  deltas=False, cmap='viridis', dividers=True,
                  divider_color='darkgray', divider_ls='-', figsize=None):
        """
        Plot the grid of ePSF models.

        Parameters
        ----------
        ax : `matplotlib.axes.Axes` or `None`, optional
            The matplotlib axes on which to plot. If `None`, then the
            current `~matplotlib.axes.Axes` instance is used.

        vmax_scale : float, optional
            Scale factor to apply to the image stretch limits. This
            value is multiplied by the peak ePSF value to determine the
            plotting ``vmax``. The defaults are 1.0 for plotting the
            ePSF data and 0.03 for plotting the ePSF difference data
            (``deltas=True``). If ``deltas=True``, the ``vmin`` is set
            to ``-vmax``. If ``deltas=False`` the ``vmin`` is set to
            ``vmax`` / 1e4.

        peak_norm : bool, optional
            Whether to normalize the ePSF data by the peak value. The
            default shows the ePSF flux per pixel.

        deltas : bool, optional
            Set to `True` to show the differences between each ePSF
            and the average ePSF.

        cmap : str or `matplotlib.colors.Colormap`, optional
            The colormap to use. The default is 'viridis'.

        dividers : bool, optional
            Whether to show divider lines between the ePSFs.

        divider_color, divider_ls : str, optional
            Matplotlib color and linestyle options for the divider
            lines between ePSFs. These keywords have no effect unless
            ``show_dividers=True``.

        figsize : (float, float), optional
            The figure (width, height) in inches.

        Returns
        -------
        fig : `matplotlib.figure.Figure`
            The matplotlib figure object. This will be the current
            figure if ``ax=None``. Use ``fig.savefig()`` to save the
            figure to a file.

        Notes
        -----
        This method returns a figure object. If you are using this
        method in a script, you will need to call ``plt.show()`` to
        display the figure. If you are using this method in a Jupyter
        notebook, the figure will be displayed automatically.

        When in a notebook, if you do not store the return value of this
        function, the figure will be displayed twice due to the REPL
        automatically displaying the return value of the last function
        call. Alternatively, you can append a semicolon to the end of
        the function call to suppress the display of the return value.
        """
        import matplotlib.pyplot as plt
        from mpl_toolkits.axes_grid1 import make_axes_locatable

        data = self.data.copy()
        if deltas:
            # Compute mean ignoring any blank (all zeros) ePSFs.
            # This is the case for MIRI with its non-square FOV.
            mask = np.zeros(data.shape[0], dtype=bool)
            for i, arr in enumerate(data):
                if np.count_nonzero(arr) == 0:
                    mask[i] = True
            data -= np.mean(data[~mask], axis=0)
            data[mask] = 0.0

        data = self._reshape_grid(data)

        if ax is None:
            if figsize is None and self.meta.get('detector', '') == 'NRCSW':
                figsize = (20, 8)
            fig, ax = plt.subplots(figsize=figsize)
        else:
            fig = plt.gcf()

        if peak_norm and data.max() != 0:
            # normalize relative to peak
            data /= data.max()

        if deltas:
            if vmax_scale is None:
                vmax_scale = 0.03
            vmax = data.max() * vmax_scale
            vmin = -vmax
            if minversion(astropy, '6.1'):
                norm = simple_norm(data, 'linear', vmin=vmin, vmax=vmax)
            else:
                norm = simple_norm(data, 'linear', min_cut=vmin, max_cut=vmax)
        else:
            if vmax_scale is None:
                vmax_scale = 1.0
            vmax = data.max() * vmax_scale
            vmin = vmax / 1.0e4
            if minversion(astropy, '6.1'):
                norm = simple_norm(data, 'log', vmin=vmin, vmax=vmax,
                                   log_a=1.0e4)
            else:
                norm = simple_norm(data, 'log', min_cut=vmin, max_cut=vmax,
                                   log_a=1.0e4)

        # Set up the coordinate axes to later set tick labels based on
        # detector ePSF coordinates. This sets up axes to have, behind the
        # scenes, the ePSFs centered at integer coords 0, 1, 2, 3 etc.
        # extent order: left, right, bottom, top
        nypsfs = self._ygrid.shape[0]
        nxpsfs = self._xgrid.shape[0]
        extent = [-0.5, nxpsfs - 0.5, -0.5, nypsfs - 0.5]

        axim = ax.imshow(data, extent=extent, norm=norm, cmap=cmap,
                         origin='lower')

        # Use the axes set up above to set appropriate tick labels
        xticklabels = self._xgrid.astype(int)
        yticklabels = self._ygrid.astype(int)
        if self.meta.get('detector', '') == 'NRCSW':
            xticklabels = list(xticklabels[0:5]) * 4
            yticklabels = list(yticklabels[0:5]) * 2
        ax.set_xticks(np.arange(nxpsfs))
        ax.set_xticklabels(xticklabels)
        ax.set_xlabel('ePSF location in detector X pixels')
        ax.set_yticks(np.arange(nypsfs))
        ax.set_yticklabels(yticklabels)
        ax.set_ylabel('ePSF location in detector Y pixels')

        if dividers:
            for ix in range(nxpsfs - 1):
                ax.axvline(ix + 0.5, color=divider_color, ls=divider_ls)
            for iy in range(nypsfs - 1):
                ax.axhline(iy + 0.5, color=divider_color, ls=divider_ls)

        instrument = self.meta.get('instrument', '')
        if not instrument:
            # WebbPSF output
            instrument = self.meta.get('instrume', '')
        detector = self.meta.get('detector', '')
        filtername = self.meta.get('filter', '')

        # WebbPSF outputs a tuple with the comment in the second element
        if isinstance(instrument, (tuple, list, np.ndarray)):
            instrument = instrument[0]
        if isinstance(detector, (tuple, list, np.ndarray)):
            detector = detector[0]
        if isinstance(filtername, (tuple, list, np.ndarray)):
            filtername = filtername[0]

        title = f'{instrument} {detector} {filtername}'
        if title != '':
            # add extra space at end
            title += ' '

        if deltas:
            minus = '\u2212'
            ax.set_title(f'{title}(ePSFs {minus} <ePSF>)')
            if peak_norm:
                label = 'Difference relative to average ePSF peak'
            else:
                label = 'Difference relative to average ePSF values'
        else:
            ax.set_title(f'{title}ePSFs')
            if peak_norm:
                label = 'Scale relative to ePSF peak pixel'
            else:
                label = 'ePSF flux per pixel'

        divider = make_axes_locatable(ax)
        cax_cbar = divider.append_axes('right', size='3%', pad='3%')
        cbar = fig.colorbar(axim, cax=cax_cbar, label=label)

        if not deltas:
            cbar.ax.set_yscale('log')

        if self.meta.get('detector', '') == 'NRCSW':
            # NIRCam NRCSW STDPSF files contain all detectors.
            # The plot gets extra divider lines and SCA name labels.
            nxpsfs = len(self._xgrid)
            nypsfs = len(self._ygrid)
            plt.axhline(nypsfs / 2 - 0.5, color='orange')
            for i in range(1, 4):
                ax.axvline(nxpsfs / 4 * i - 0.5, color='orange')

            det_labels = [['A1', 'A3', 'B4', 'B2'], ['A2', 'A4', 'B3', 'B1']]
            for i in range(2):
                for j in range(4):
                    ax.text(j * nxpsfs / 4 - 0.45,
                            (i + 1) * nypsfs / 2 - 0.55,
                            det_labels[i][j], color='orange',
                            verticalalignment='top', fontsize=12)

        fig.tight_layout()

        return fig
