"""
This script allows the user to manually select the sky background regions
"""

import os, sys
import copy
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, Slider
import matplotlib.transforms as mtransforms
from scipy.interpolate import RectBivariateSpline

from pypeit import specobjs
from pypeit import msgs
from pypeit.io import write_to_fits
from pypeit.core import pixels

operations = dict({'cursor': "Add sky region (LMB drag)\n" +
                   "         Remove sky region (RMB drag)\n" +
                   "         Note: If you would like to pan or zoom, you need to activate\n" +
                   "         the pan/zoom tool with the 'p' key, or by selecting the\n" +
                   "         pan/zoom tool on the Matplotlib navigation tool menu. You\n" +
                   "         can also zoom using the magnifying glass (select this option\n" +
                   "         from the Matplotlib navigation tool menu). While the pan/zoom\n" +
                   "         feature is enabled, you will not be able to update sky regions.\n",
                   'c': "Center the window at the location of the mouse",
                   'd': "Delete all sky regions and start again",
                   'h/r': "Return zoom to the original plotting limits",
                   'p': "Toggle pan/zoom with the cursor",
                   '?': "Display the available options",
                   })


class SkySubGUI(object):
    """
    GUI to interactively define the sky regions. The GUI can be run within
    PypeIt during data reduction, or as a standalone script outside of
    PypeIt. To initialize the GUI, call the initialize() function in this
    file.
    """

    def __init__(self, canvas, image, frame, outname, det, tslits_dict, axes,
                 printout=False, runtime=False, resolution=1000):
        """Controls for the interactive sky regions definition tasks in PypeIt.

        The main goal of this routine is to interactively select sky background
        regions.

        Args:
            canvas : Matploltib figure canvas
                The canvas on which all axes are contained
            image : AxesImage
                The image plotted to screen
            frame : ndarray
                The image data
            outname : str
                The output filename to save the sky regions mask
            det : int
                Detector to add a slit on
            tslits_dict : dict
                Dictionary containing information about the slit traces
            axes : dict
                Dictionary of four Matplotlib axes instances (Main spectrum panel, two for residuals, one for information)
            printout : bool
                Should the results be printed to screen
            runtime : bool
                Is the GUI being launched during data reduction?
            resolution : int
                The resolution of the skysub definitions. It is the number of pixels to divide the
                 slit width by (i.e. 1000 pixels means a resolution of 0.1% of the slit width).
        """
        # Store the axes
        self._det = det
        self.image = image
        self.frame = frame
        self._outname = outname
        self.nspec, self.nspat = frame.shape[0], frame.shape[1]
        self._spectrace = np.arange(self.nspec)
        self._printout = printout
        self._runtime = runtime
        self.axes = axes
        self._currslit = -1
        self.tslits_dict = tslits_dict
        self._nslits = tslits_dict['slit_left'].shape[1]
        self._resolution = int(resolution)
        self._allreg = np.zeros(int(resolution), dtype=np.bool)
        self._specx = np.arange(int(resolution))

        # Unset some of the matplotlib keymaps
        matplotlib.pyplot.rcParams['keymap.fullscreen'] = ''        # toggling fullscreen (Default: f, ctrl+f)
        #matplotlib.pyplot.rcParams['keymap.home'] = ''              # home or reset mnemonic (Default: h, r, home)
        matplotlib.pyplot.rcParams['keymap.back'] = ''              # forward / backward keys to enable (Default: left, c, backspace)
        matplotlib.pyplot.rcParams['keymap.forward'] = ''           # left handed quick navigation (Default: right, v)
        #matplotlib.pyplot.rcParams['keymap.pan'] = ''              # pan mnemonic (Default: p)
        matplotlib.pyplot.rcParams['keymap.zoom'] = ''              # zoom mnemonic (Default: o)
        matplotlib.pyplot.rcParams['keymap.save'] = ''              # saving current figure (Default: s)
        matplotlib.pyplot.rcParams['keymap.quit'] = ''              # close the current figure (Default: ctrl+w, cmd+w)
        matplotlib.pyplot.rcParams['keymap.grid'] = ''              # switching on/off a grid in current axes (Default: g)
        matplotlib.pyplot.rcParams['keymap.yscale'] = ''            # toggle scaling of y-axes ('log'/'linear') (Default: l)
        matplotlib.pyplot.rcParams['keymap.xscale'] = ''            # toggle scaling of x-axes ('log'/'linear') (Default: L, k)
        matplotlib.pyplot.rcParams['keymap.all_axes'] = ''          # enable all axes (Default: a)

        # Initialise the main canvas tools
        canvas.mpl_connect('draw_event', self.draw_callback)
        canvas.mpl_connect('button_press_event', self.button_press_callback)
        canvas.mpl_connect('key_press_event', self.key_press_callback)
        canvas.mpl_connect('button_release_event', self.button_release_callback)
        canvas.mpl_connect('motion_notify_event',  self.mouse_move_callback)
        self.canvas = canvas

        # Interaction variables
        self._respreq = [False, None]  # Does the user need to provide a response before any other operation will be permitted? Once the user responds, the second element of this array provides the action to be performed.
        self._qconf = False  # Confirm quit message
        self._changes = False
        self._use_updates = True
        self._inslit = -1  # Which slit is the mouse in
        self.mmx, self.mmy = 0, 0
        self._fitr = []  # Matplotlib shaded fit region
        self._fita = None

        # Draw the spectrum
        self.canvas.draw()

        self.initialize_menu()
        self.reset_regions()

    @classmethod
    def initialize(cls, det, frame, tslits_dict, outname="skyregions.fits", runtime=False, printout=False):
        """Initialize the 'ObjFindGUI' window for interactive object tracing

            Args:
                frame : ndarray
                    Sky subtracted science image
                tslits_dict : dict, None
                    Dictionary containing slit trace information
                det : int
                    Detector index
                printout : bool
                    Should the results be printed to screen
                runtime : bool
                    Is this GUI being launched during a data reduction?

            Returns:
                ObjFindGUI: Returns an instance of the ObjFindGUI class
        """
        # This allows the input lord and rord to either be (nspec, nslit) arrays or a single
        # vectors of size (nspec)
        if tslits_dict['slit_left'].ndim == 2:
            nslit = tslits_dict['slit_left'].shape[1]
        else:
            nslit = 1
            tslits_dict['slit_left'] = tslits_dict['slit_left'].reshape((tslits_dict['slit_left'].size, 1))
            tslits_dict['slit_righ'] = tslits_dict['slit_righ'].reshape((tslits_dict['slit_righ'].size, 1))
        lordloc = tslits_dict['slit_left']
        rordloc = tslits_dict['slit_righ']

        # Determine the scale of the image
        med = np.median(frame)
        mad = np.median(np.abs(frame - med))
        vmin = med - 3 * mad
        vmax = med + 3 * mad

        # Add the main figure axis
        fig, ax = plt.subplots(figsize=(16, 9), facecolor="white")
        plt.subplots_adjust(bottom=0.05, top=0.85, left=0.05, right=0.8)
        image = ax.imshow(frame, aspect='auto', cmap='Greys', vmin=vmin, vmax=vmax)

        # Overplot the slit traces
        specarr = np.arange(lordloc.shape[0])
        for sl in range(nslit):
            ax.plot(lordloc[:, sl], specarr, 'g-')
            ax.plot(rordloc[:, sl], specarr, 'b-')

        # Add an information GUI axis
        axinfo = fig.add_axes([0.15, .92, .7, 0.07])
        axinfo.get_xaxis().set_visible(False)
        axinfo.get_yaxis().set_visible(False)
        axinfo.text(0.5, 0.5, "Press '?' to list the available options", transform=axinfo.transAxes,
                    horizontalalignment='center', verticalalignment='center')
        axinfo.set_xlim((0, 1))
        axinfo.set_ylim((0, 1))

        axes = dict(main=ax, info=axinfo)
        # Initialise the object finding window and display to screen
        fig.canvas.set_window_title('PypeIt - Sky regions')
        srgui = SkySubGUI(fig.canvas, image, frame, outname, det, tslits_dict, axes,
                          printout=printout, runtime=runtime)
        plt.show()

        return srgui

    def print_help(self):
        """Print the keys and descriptions that can be used for Identification
        """
        keys = operations.keys()
        print("===============================================================")
        print("Define the sky background regions in each slit by using the left")
        print("mouse button to click and drag over the sky background region.")
        print("Use the right mouse button (click and drag) to delete a region.")
        print("If you click 'Continue (and save changes)' the sky background")
        print("regions file will be printed to the terminal, where you can")
        print("copy them into your .pypeit file.")
        print("")
        print("thin green/blue lines  = slit edges")
        print("shaded red regions     = selected sky regions")
        print("===============================================================")
        print("       OTHER OPERATIONS")
        for key in keys:
            print("{0:6s} : {1:s}".format(key, operations[key]))
        print("---------------------------------------------------------------")

    def initialize_menu(self):
        """Initialize the menu buttons
        """
        axcolor = 'lightgoldenrodyellow'
        # Continue with reduction (using updated specobjs)
        ax_cont = plt.axes([0.82, 0.85, 0.15, 0.05])
        self._ax_cont = Button(ax_cont, "Continue (and save changes)", color=axcolor, hovercolor='y')
        self._ax_cont.on_clicked(self.button_cont)
        # Continue with reduction (using original specobjs)
        ax_exit = plt.axes([0.82, 0.79, 0.15, 0.05])
        self._ax_exit = Button(ax_exit, "Continue (don't save changes)", color=axcolor, hovercolor='y')
        self._ax_exit.on_clicked(self.button_exit)
        # Frame for the sliders
        self.axes['allslitreg'] = plt.axes([0.82, 0.68, 0.15, 0.04], facecolor='black', title="Assign sky regions to all slits")
        self.axes['allslitreg'].get_xaxis().set_ticks([])
        self.axes['allslitreg'].get_yaxis().set_ticks([])
        self.axes['allslitreg'].axvspan(0, self._resolution-1, color='lightgrey')
        self.axes['allslitreg'].set_xlim(-self._resolution/10, self._resolution+self._resolution/10)
        self.axes['allslitreg'].set_ylim(0, 1)

    def button_cont(self, event):
        """What to do when the 'exit and save' button is clicked

        Args:
            event : Event
                A matplotlib event instance
        """
        self._respreq = [True, "exit_update"]
        self.update_infobox(message="Are you sure you want to exit and save the newly defined sky regions?", yesno=True)

    def button_exit(self, event):
        """What to do when the 'exit and do not save changes' button is clicked

        Args:
            event : Event
                A matplotlib event instance
        """
        self._respreq = [True, "exit_restore"]
        self.update_infobox(message="Are you sure you want to exit without saving the  sky regions?", yesno=True)

    def replot(self):
        """Redraw the entire canvas
        """
        self.canvas.restore_region(self.background)
        self.draw_regions()
        self.canvas.draw()

    def draw_regions(self):
        """Refresh the fit regions
        """
        # Remove the regions and reset the patches
        for rr in range(len(self._fitr)):
            self._fitr[rr].remove()
        if self._fita is not None:
            self._fita.remove()
        self._fitr = []
        # Loop through all slits:
        for sl in range(self._nslits):
            # Fill fraction of the slit
            diff = self.tslits_dict['slit_righ'][:,sl] - self.tslits_dict['slit_left'][:,sl]
            tmp = np.zeros(self._resolution+2)
            tmp[1:-1] = self._skyreg[sl]
            wl = np.where(tmp[1:] > tmp[:-1])[0]
            wr = np.where(tmp[1:] < tmp[:-1])[0]
            for rr in range(wl.size):
                left = self.tslits_dict['slit_left'][:, sl] + wl[rr]*diff/(self._resolution-1.0)
                righ = self.tslits_dict['slit_left'][:, sl] + wr[rr]*diff/(self._resolution-1.0)
                self._fitr.append(self.axes['main'].fill_betweenx(self._spectrace, left, righ, facecolor='red',
                                                                  alpha=0.5))
        # Plot the region on top of the "all slits" panel
        trans = mtransforms.blended_transform_factory(self.axes['allslitreg'].transData,
                                                      self.axes['allslitreg'].transAxes)
        self._fita = self.axes['allslitreg'].fill_between(self._specx, 0, 1, transform=trans,
                                                          where=self._allreg, facecolor='red', alpha=0.5, zorder=10)

    def draw_callback(self, event):
        """Draw callback (i.e. everytime the canvas is being drawn/updated)

        Args:
            event : Event
                A matplotlib event instance
        """
        # Get the background
        self.background = self.canvas.copy_from_bbox(self.axes['main'].bbox)
        self.draw_regions()

    def get_current_slit(self, event):
        """Get the index of the slit closest to the cursor

        Args:
            event : Event
                Matplotlib event instance containing information about the event
        """
        # Find the current slit
        self._currslit = -1
        yv = np.argmin(np.abs(event.ydata-self._spectrace))
        wsl = np.where((event.xdata > self.tslits_dict['slit_left'][yv, :]) &
                       (event.xdata < self.tslits_dict['slit_righ'][yv, :]))[0]
        # Double check there's only one solution
        if wsl.size == 1:
            self._currslit = int(wsl[0])
        return

    def get_axisID(self, event):
        """Get the ID of the axis where an event has occurred

        Args:
            event : Event
                Matplotlib event instance containing information about the event

        Returns:
            int, None: Axis where the event has occurred
        """
        if event.inaxes == self.axes['main']:
            return 0
        elif event.inaxes == self.axes['info']:
            return 1
        elif event.inaxes == self.axes['allslitreg']:
            return 2
        return None

    def mouse_move_callback(self, event):
        """Store the locations of mouse as it moves across the canvas
        """
        if event.inaxes is None:
            return
        if event.inaxes == self.axes['main']:
            self.mmx, self.mmy = event.xdata, event.ydata

    def button_press_callback(self, event):
        """What to do when the mouse button is pressed

        Args:
            event : Event
                Matplotlib event instance containing information about the event
        """
        if event.inaxes is None:
            return
        if self.canvas.toolbar.mode != "":
            return
        if event.button == 1:
            self._addsub = 1
            self.get_current_slit(event)
        elif event.button == 3:
            self._addsub = 0
            self.get_current_slit(event)
        if event.inaxes == self.axes['main']:
            self._start = [event.xdata, event.ydata]
        elif event.inaxes == self.axes['allslitreg']:
            self._start = [int(round(event.xdata)), event.ydata]

    def button_release_callback(self, event):
        """What to do when the mouse button is released

        Args:
            event : Event
                Matplotlib event instance containing information about the event
        """
        if event.inaxes is None:
            return
        if event.inaxes == self.axes['info']:
            if (event.xdata > 0.8) and (event.xdata < 0.9):
                answer = "y"
            elif event.xdata >= 0.9:
                answer = "n"
            else:
                return
            self.operations(answer, -1)
            self.update_infobox(default=True)
            return
        elif self._respreq[0]:
            # The user is trying to do something before they have responded to a question
            return
        if self.canvas.toolbar.mode != "":
            return
        # Draw an actor
        axisID = self.get_axisID(event)
        if axisID is not None:
            if axisID == 2:
                # Set the slide l value
                self._end = [int(round(event.xdata)), event.ydata]
                self.add_region_all()
            else:
                self._end = [event.xdata, event.ydata]
                if (self._end[0] == self._start[0]) and (self._end[1] == self._start[1]):
                    # The mouse button was pressed (not dragged)
                    pass
                elif self._end != self._start:
                    # The mouse button was dragged
                    if axisID == 0:
                        if self._start[0] > self._end[0]:
                            self._start[0], self._end[0] = self._end[0], self._start[0]
                        # Now do something
                        self.add_region()
        self.replot()

    def key_press_callback(self, event):
        """What to do when a key is pressed

        Args:
            event : Event
                Matplotlib event instance containing information about the event
        """
        # Check that the event is in an axis...
        if not event.inaxes:
            return
        # ... but not the information box!
        if event.inaxes == self.axes['info']:
            return
        axisID = self.get_axisID(event)
        self.operations(event.key, axisID)

    def operations(self, key, axisID):
        """Canvas operations

        Args:
            key : str
                Which key has been pressed
            axisID : int
                The index of the axis where the key has been pressed (see get_axisID)
        """
        # Check if the user really wants to quit
        if key == 'q' and self._qconf:
            if self._changes:
                self.update_infobox(message="WARNING: There are unsaved changes!!\nPress q again to exit", yesno=False)
                self._qconf = True
            else:
                msgs.bug("Need to change this to kill and return the results to PypeIt")
                plt.close()
        elif self._qconf:
            self.update_infobox(default=True)
            self._qconf = False

        # Manage responses from questions posed to the user.
        if self._respreq[0]:
            if key != "y" and key != "n":
                return
            else:
                # Switch off the required response
                self._respreq[0] = False
                # Deal with the response
                if self._respreq[1] == "exit_update" and key == "y":
                    self._use_updates = True
                    self.operations("qu", None)
                elif self._respreq[1] == "exit_restore" and key == "y":
                    self._use_updates = False
                    self.operations("qr", None)
                else:
                    return
            # Reset the info box
            self.update_infobox(default=True)
            return

        if key == '?':
            self.print_help()
        elif key == 'd':
            if axisID == 0:
                # If this is pressed on the main window
                self.reset_regions()
        elif key == 'c':
            if axisID == 0:
                # If this is pressed on the main window
                self.recenter()
        elif key == 'qu' or key == 'qr':
            if self._changes:
                self.update_infobox(message="WARNING: There are unsaved changes!!\nPress q again to exit", yesno=False)
                self._qconf = True
            else:
                plt.close()
        self.replot()

    def get_result(self):
        """Save a mask containing the skysub regions, and print information
        for what the user should include in their .pypeit file
        """
        # Only do this if the user wishes to save the result
        if self._use_updates:
            # Generate the mask
            inmask = self.generate_mask()
            # Save the mask
            write_to_fits(inmask, self._outname, name="SKYREG")
            # Print the output to screen
            msgs.info("Include the following info in your .pypeit file:\n")
            print("STILL WORKING ON THIS!!!")
        return

    def generate_mask(self):
        """Generate the mask of sky regions

        Returns:
            ndarray : Boolean mask containing sky regions
        """
        nreg = 0
        left_edg, righ_edg = np.zeros((self.nspec, 0)), np.zeros((self.nspec, 0))
        spec_min, spec_max = np.array([]), np.array([])
        for sl in range(self._nslits):
            diff = self.tslits_dict['slit_righ'][:, sl] - self.tslits_dict['slit_left'][:, sl]
            tmp = np.zeros(self._resolution+2)
            tmp[1:-1] = self._skyreg[sl]
            wl = np.where(tmp[1:] > tmp[:-1])[0]
            wr = np.where(tmp[1:] < tmp[:-1])[0]
            for rr in range(wl.size):
                left = self.tslits_dict['slit_left'][:, sl] + wl[rr]*diff/(self._resolution-1.0)
                righ = self.tslits_dict['slit_left'][:, sl] + wr[rr]*diff/(self._resolution-1.0)
                left_edg = np.append(left_edg, left[:, np.newaxis], axis=1)
                righ_edg = np.append(righ_edg, righ[:, np.newaxis], axis=1)
                nreg += 1
                spec_min = np.append(spec_min, self.tslits_dict['spec_min'][sl])
                spec_max = np.append(spec_max, self.tslits_dict['spec_max'][sl])
        reg_dict = dict(pad=0, slit_left=left_edg, slit_righ=righ_edg, nslits=nreg,
                        nspec=self.nspec, nspat=self.nspat, spec_min=spec_min, spec_max=spec_max)
        return (pixels.tslits2mask(reg_dict) >= 0).astype(np.int)

    def recenter(self):
        xlim = self.axes['main'].get_xlim()
        ylim = self.axes['main'].get_ylim()
        xmin = self.mmx - 0.5*(xlim[1]-xlim[0])
        xmax = self.mmx + 0.5*(xlim[1]-xlim[0])
        ymin = self.mmy - 0.5*(ylim[1]-ylim[0])
        ymax = self.mmy + 0.5*(ylim[1]-ylim[0])
        self.axes['main'].set_xlim([xmin, xmax])
        self.axes['main'].set_ylim([ymin, ymax])

    def update_infobox(self, message="Press '?' to list the available options",
                       yesno=True, default=False):
        """Send a new message to the information window at the top of the canvas

        Args:
            message : str
                Message to be displayed
            yesno : bool
                Is a yes/no option desired?
            default : bool
                Would you like to refresh the info box and just display the default message
        """
        self.axes['info'].clear()
        if default:
            self.axes['info'].text(0.5, 0.5, "Press '?' to list the available options", transform=self.axes['info'].transAxes,
                          horizontalalignment='center', verticalalignment='center')
            self.canvas.draw()
            return
        # Display the message
        self.axes['info'].text(0.5, 0.5, message, transform=self.axes['info'].transAxes,
                      horizontalalignment='center', verticalalignment='center')
        if yesno:
            self.axes['info'].fill_between([0.8, 0.9], 0, 1, facecolor='green', alpha=0.5, transform=self.axes['info'].transAxes)
            self.axes['info'].fill_between([0.9, 1.0], 0, 1, facecolor='red', alpha=0.5, transform=self.axes['info'].transAxes)
            self.axes['info'].text(0.85, 0.5, "YES", transform=self.axes['info'].transAxes,
                          horizontalalignment='center', verticalalignment='center')
            self.axes['info'].text(0.95, 0.5, "NO", transform=self.axes['info'].transAxes,
                          horizontalalignment='center', verticalalignment='center')
        self.axes['info'].set_xlim((0, 1))
        self.axes['info'].set_ylim((0, 1))
        self.canvas.draw()

    def add_region(self):
        """ Add/subtract a defined region
        """
        # Figure out the locations of the start values
        ys = np.argmin(np.abs(self._start[1]-self._spectrace))
        difs = self.tslits_dict['slit_righ'][ys, self._currslit] - self.tslits_dict['slit_left'][ys, self._currslit]
        sval = (self._start[0]-self.tslits_dict['slit_left'][ys, self._currslit]) / difs
        sidx = int(round(self._resolution*sval))
        # Figure out the locations of the start values
        yf = np.argmin(np.abs(self._end[1]-self._spectrace))
        diff = self.tslits_dict['slit_righ'][yf, self._currslit] - self.tslits_dict['slit_left'][yf, self._currslit]
        fval = (self._end[0]-self.tslits_dict['slit_left'][yf, self._currslit]) / diff
        fidx = int(round(self._resolution*fval))
        # Switch the indices if needed
        if sidx > fidx:
            sidx, fidx = fidx, sidx
        # Check that we are within bounds
        if sidx < 0:
            sidx = 0
        if fidx > self._resolution:
            fidx = self._resolution
        # Assign the sky regions
        self._skyreg[self._currslit][sidx:fidx] = self._addsub
        # If some regions are removed, remove this from the "all slits" regions, as well
        if self._addsub == 0:
            self._allreg[sidx:fidx] = 0

    def add_region_all(self):
        """ Set the sky regions for all slits simultaneously
        """
        # Do some checks
        xmin, xmax = self._start[0], self._end[0]
        if xmax < xmin:
            xmin, xmax = xmax, xmin
        if xmin < 0:
            xmin = 0
        if xmax > self._resolution:
            xmax = self._resolution
        # Apply to all slits
        for sl in range(self._nslits):
            self._skyreg[sl][xmin:xmax] = self._addsub
        # Set the all regions parameter
        self._allreg[xmin:xmax] = self._addsub

    def reset_regions(self):
        """ Reset the sky regions for all slits simultaneously
        """
        self._skyreg = [np.zeros(self._resolution, dtype=np.bool) for all in range(self._nslits)]
        self._allreg[:] = False
