from __future__ import with_statement
# must succeed, min requirement py 2.5
import sys
from collections import defaultdict
from visir_utils import log_invocation, Future
import types
log_invocation()

# import the needed modules
try:
    import numpy
    try:
        import pyfits
    except ImportError:
        import astropy.io.fits as pyfits
    # needs pyfits
    import fitsloader
    import matplotlib
    from matplotlib.widgets import RadioButtons, Button
    matplotlib.use('WXAgg')
    # import before wx to ensure we get 2.8 even if 2.6 is also installed
    import matplotlib.backends.backend_wxagg
    from matplotlib.widgets import SpanSelector

    from plot_utils import LinePlot, ImagePlot, ScaleNormalize, freedman_bins
    from pipeline_display import ScatterDisplay
    from reflexlib import VisirParameter
    from reflexgui import ReflexBaseGui, PlotDialog, BeamDialog

    import wx

    import_success = True
except (ImportError, NameError), e:
    import_success = False
    print "Error importing modules: ", str(e)


def myspanignore(self, event):
    # ignore in pan and zoom mode
    # https://github.com/matplotlib/matplotlib/issues/1368
    if not self.canvas.widgetlock.available(self):
        return True
    return SpanSelector.ignore(self, event)


def get_keys_wb(keyword, fl, files):
    cor = fl.getval_all(keyword, files=files, skip_first=True)
    bounds = numpy.cumsum([x.size for x in cor])
    cor = numpy.concatenate(cor)
    return cor, [(0, 0)] + [(x, y) for x, y in zip(bounds, files)]


def get_slice(fits, y, x):
    vals = []
    for ext in fits[1:]:
        vals.append(float(ext.data[int(y), int(x)]))
    return vals


def get_beam_coords(fl, files):
    res = []
    for fn in files:
        bid = fl.getval(fn, 'ESO QC BEAMID', 0)
        x = fl.getval(fn, 'ESO QC BEAMX', 0)
        y = fl.getval(fn, 'ESO QC BEAMY', 0)
        res.append((bid, x, y))
    return res


class DataPlotterManager:
    stats = dict(
                 stdev=dict(key='ESO QC STDEV',
                            name='stdev', title='Standard deviation',
                            label='Std. deviation', unit='ADU / s'),
                 mad=dict(key='ESO QC MAD', name='mad',
                          title='Median absolute deviation',
                          label='MAD', unit='ADU / s'),
                 median=dict(key='ESO QC MEDIAN', name='median',
                             title='Median', label='Median', unit='ADU / s'),
                 fwhmx=dict(key='ESO QC FWHMX', name='fwhmx', unit='pixel',
                            title='Full width half maximum X', label='FWHM X'),
                 fwhmy=dict(key='ESO QC FWHMY', name='fwhmy', unit='pixel',
                            title='Full width half maximum Y', label='FWHM Y'),
                 correl=dict(key='ESO QC CORR', name='correl',
                             title='Cross-correlation', label='Correlation'),
                )
    dflt_stat = 'correl'
    SHOW_STATS = 1
    SHOW_SLICE = 2

    def __init__(self, fitsloader):
        self.fl = fitsloader
        self.dlg = None
        self.beams = None
        self.undef_beams = False

    def readFitsData(self, fitsFiles):
        #Read all the products
        fits = defaultdict(list)
        for frame in fitsFiles:
            if frame == '':
                continue
            category = frame.category
            fits[category].append(frame.name)
        self.fits = fits

        self.templates = []
        for f in fits['CORRELATION_TEMPLATE']:
            # workaround pyfits <= 2.1.1 unicode issue (CNSHD750437)
            try:
                self.templates.extend(pyfits.open(f))
            except AttributeError:
                self.templates.extend(pyfits.open(str(f)))

        try:
            self.mean = pyfits.getdata(fits['IMG_REPACKED_MEAN'][0], 0)
        except:
            self.mean = None

    def addSubplots(self, figure):
        if len(self.fits['BEAM_DETECTED']) > 0:
            self.figure = figure
            self.plt = dict()
            self.plt['av'] = ImagePlot(figure, figure.add_subplot(334))
            self.plt['img'] = ImagePlot(figure, figure.add_subplot(335))
            self.plt['tmpl'] = ImagePlot(figure, figure.add_subplot(336))
            self.plt['shifts'] = ImagePlot(figure, figure.add_subplot(313))
            self.plt['stats'] = LinePlot(figure, figure.add_axes([0.30, 0.70,
                                                                 0.675, 0.25]))
            self.rax = figure.add_axes([0.02, 0.75, 0.20, 0.13])
            self.rax.tooltip = 'Select statistics to plot'
            self.statskeys = sorted(self.stats.keys())
            self.statslabels = [self.stats[x]['label'] for x in self.statskeys]
            self.rad = RadioButtons(self.rax, self.statslabels,
                                    active=self.statskeys.index(self.dflt_stat))
            self.bax = figure.add_axes([0.08, 0.65, 0.13, 0.05])
            self.beamselect = Button(self.bax, 'Select beams', color='gray')
            self.beamselect.on_clicked(self.select_beams)
        else:
            self.nodata = figure.add_subplot(111)

    def get_image(self, idx):
        bounds = [(0, 0), (1000, self.fits['BEAM_DETECTED'][0])]
        if idx != 0:
            bounds = self.statsdata[self.dflt_stat].get()[1]
        idx = int(idx) + 1
        for i, v in enumerate(bounds):
            if idx <= v[0]:
                return self.fl.getdata(v[1], idx - bounds[i - 1][0])
        return numpy.zeros((200, 200))

    def get_slice(self, files, x, y):
        res = []
        for fn in files:
            res.append(self.fl.custom(fn, get_slice, (x, y)))
        return  numpy.concatenate(res)

    def rad_clicked(self, label):
        self.plt['stats'].set_loading()
        # add_done_callback would be useful here
        label = self.statskeys[self.statslabels.index(label)]
        Future(lambda: self.statsdata[label].get(),
               callback=lambda f: wx.CallAfter(self.plot_stats, f.get(), label))
        self.showing = self.SHOW_STATS

    def change_clim(self, vmin, vmax, plt=None):
        if plt is None or not plt.ax.images:
            return
        plt.set_clim([vmin, vmax])
        plt.figure.canvas.draw_idle()

    def detach(self, plt):
        if self.figure.canvas.HasCapture():
            self.figure.canvas.ReleaseMouse()
        dlg = PlotDialog(self.frame, title=plt.ax.get_title(),
                         size=wx.Size(700, 700))
        dlg.plt = ImagePlot(dlg.figure, dlg.figure.add_subplot(111))
        dlg.plt.set_data(plt.data)
        dlg.plt.draw(interpolation='nearest', origin='lower')
        dlg.plt.figure.colorbar(dlg.plt.ax.images[0])
        dlg.plt.detached = True
        dlg.on_click = lambda e: wx.CallAfter(self.show_popup, dlg.plt) \
                                 if e.button == 3 else None
        dlg.figure.canvas.mpl_connect('button_press_event', dlg.on_click)

        # close scale dialog
        def dlgclose(event):
            dlg.Destroy()
            if self.dlg:
                self.dlg.Destroy()
                self.dlg = None
        dlg.Bind(wx.EVT_CLOSE, dlgclose)
        dlg.Bind(wx.EVT_BUTTON, dlgclose)

        dlg.Show()

    def scale_dialog(self, plt):
        if self.figure.canvas.HasCapture():
            self.figure.canvas.ReleaseMouse()
        if self.dlg is None:
            self.dlg = PlotDialog(self.frame, title='Colorscale')
        else:
            self.dlg.figure.clf()
        ax = self.dlg.figure.add_subplot(111)

        ax.hist(plt.data.ravel(), histtype='stepfilled',
                bins=freedman_bins(plt.data))
        self.dlg.span = SpanSelector(ax,
                                     lambda x, y : self.change_clim(x, y, plt),
                                     'horizontal', useblit=True,
                                     rectprops=dict(alpha=0.5, facecolor='red'))
        self.dlg.span.ignore = types.MethodType(myspanignore, self.dlg.span)

        self.dlg.figure.canvas.draw_idle()
        def dlgclose(event):
            self.dlg.Destroy()
            self.dlg = None
        self.dlg.Bind(wx.EVT_CLOSE, dlgclose)
        self.dlg.Bind(wx.EVT_BUTTON, dlgclose)
        self.dlg.Show()

    def select_beams(self, event):
        if self.figure.canvas.HasCapture():
            self.figure.canvas.ReleaseMouse()
        try:
            mheader = pyfits.getheader(self.fits['IMG_REPACKED_MEAN'][0], 0)
            pthrow = mheader['ESO TEL CHOP THROW'] / \
                     float(mheader['ESO INS PFOV'])
        except:
            pthrow = 100
        dlg = BeamDialog(self.frame, data=self.mean, dflt_size=pthrow / 2)
        ret = dlg.ShowModal()
        bp = ''
        for s in dlg.get_selections():
            x, y, w = s[1]
            sign = s[0]
            bp += '%s:%d,%d,%d;' % (sign, x, y, w)

        print bp, ret, wx.ID_OK
        if bp and ret == wx.ID_OK:
            self.parWidgets['beampos'].set_value(bp)
        dlg.Destroy()

    def show_popup(self, plt):
        try:
            if not hasattr(plt, 'popupmenu'):
                plt.popupmenu = wx.Menu()
                from functools import partial
                f = partial(self.OnPopupItemSelected, plt=plt)
                cmenu = wx.Menu()
                plt.popupmenu.AppendSubMenu(cmenu,'Color')
                pmenu = wx.Menu()
                plt.popupmenu.AppendSubMenu(pmenu, 'Colorscale')
                item = plt.popupmenu.Append(-1, 'Scale parameters')
                plt.figure.canvas.Bind(wx.EVT_MENU, f, item)

                if not hasattr(plt, 'detached'):
                    item = plt.popupmenu.Append(-1, 'Detach')
                    plt.figure.canvas.Bind(wx.EVT_MENU, f, item)

                for text in ('jet', 'gray', 'hot'):
                    item = cmenu.AppendRadioItem(-1, text)
                    plt.figure.canvas.Bind(wx.EVT_MENU, f, item)
                for text in ('linear', 'arcsinh', 'square', 'sqrt', 'log'):
                    item = pmenu.AppendRadioItem(-1, text)
                    plt.figure.canvas.Bind(wx.EVT_MENU, f, item)

            plt.figure.canvas.PopupMenu(plt.popupmenu)
            #self.popupmenu.Destroy()

            plt.draw()
        finally:
            # something does not release capture with popups, so force it
            if plt.figure.canvas.HasCapture():
                plt.figure.canvas.ReleaseMouse()

    def OnPopupItemSelected(self, event, plt=None):
        item = plt.popupmenu.FindItemById(event.GetId())
        v = item.GetText()
        if v == 'arcsinh':
            plt.set_norm(ScaleNormalize(numpy.arcsinh))
        elif v == 'square':
            plt.set_norm(ScaleNormalize(numpy.square))
        elif v == 'sqrt':
            plt.set_norm(ScaleNormalize(numpy.sqrt))
        elif v == 'log':
            plt.set_norm(ScaleNormalize(numpy.log))
        elif v == 'linear':
            plt.set_norm(None)
        elif v == 'Scale parameters':
            self.scale_dialog(plt)
        elif v == 'Detach':
            self.detach(plt)
        else:
            for im in plt.ax.images:
                im.set_cmap(getattr(matplotlib.cm, v, matplotlib.cm.jet))

    def onclick(self, event):
        if not self.figure.canvas.widgetlock.available(self):
            return
        if event.button == 3:
            plt = None
            for v in self.plt.values():
                if v.ax == event.inaxes and isinstance(v, ImagePlot) and \
                   v != self.plt['shifts']:
                    plt = v
                    break
            if plt:
                wx.CallAfter(self.show_popup, plt)

        elif self.showing == self.SHOW_STATS and \
            self.plt['stats'].inaxes(event) and \
            hasattr(self, 'cor') and \
            event.xdata >= 1 and round(event.xdata) < len(self.cor) + 1:
            x = int(round(event.xdata)) - 1
            img = self.get_image(x)
            try:
                ax = self.dlg.figure.axes[0]
                ax.cla()
                ax.hist(img.ravel(), histtype='stepfilled',
                        bins=freedman_bins(img))
                self.dlg.figure.canvas.draw()
            except:
                pass
            self.plot_image(img, x)

        elif self.plt['av'].inaxes(event):
            self.showing = self.SHOW_SLICE
            for c in self.rad.circles:
                c.set_facecolor(self.plt['stats'].ax.get_axis_bgcolor())
            x, y = int(event.ydata), int(event.xdata)
            self.plt['stats'].set_loading()
            Future(self.get_slice, args=(self.fits['BKG_CORRECTED'], x, y),
                   callback=lambda f: wx.CallAfter(self.plot_slice,
                                                   f.get(), x, y))

    def report_pixel(self, plt, x, y):
        xorig, yorig = x, y
        if hasattr(plt, 'scale_coord'):
            x, y = plt.scale_coord(x, y)
        if numpy.array(plt.data).size and \
           y > 0 and x > 0 and \
           x < plt.data.shape[1] and y < plt.data.shape[0]:
            return "x=%s y=%s value=%s" % (plt.ax.format_xdata(xorig),
                                           plt.ax.format_ydata(yorig),
                                           plt.ax.format_ydata(plt.data[int(y), int(x)]))
        else:
            return "x=%s y=%s" % (plt.ax.format_xdata(xorig),
                                  plt.ax.format_ydata(yorig))

    def plotProductsGraphics(self, figure, canvas):
        if len(self.fits['BEAM_DETECTED']) == 0:
            #Data not found info
            self.nodata.set_axis_off()
            self.text_nodata = \
            """Data not found. Possibly all images have been rejected"""
            self.nodata.text(0.1, 0.6, self.text_nodata,
                                     color='#11557c', fontsize=18, ha='left',
                                     va='center', alpha=1.0)
            self.nodata.tooltip = 'beam detection not found in the products'
            return

        from functools import partial
        get_bkeys = partial(get_keys_wb, fl=self.fl,
                            files=self.fits['BEAM_DETECTED'])

        plt_stats = self.plt['stats']
        plt_av = self.plt['av']
        plt_tmpl = self.plt['tmpl']
        plt_img = self.plt['img']

        plt_av.tooltip = 'Input image'
        try:
            self.plot_average(self.mean)
        except:
            plt_av.set_loading()
            Future(self.get_average,
                   callback=lambda f: wx.CallAfter(self.plot_average, f.get()))

        self.statsdata = dict()
        self.statsdata[self.dflt_stat] = \
            Future(get_bkeys, args=(self.stats[self.dflt_stat]['key'],),
                   callback=lambda f: wx.CallAfter(self.plot_stats,
                                                   f.get(), self.dflt_stat))
        for k in self.stats.keys():
            if k in self.statsdata:
                continue
            self.statsdata[k] = Future(get_bkeys, args=(self.stats[k]['key'],))

        plt_stats.set_loading()
        self.showing = self.SHOW_STATS

        plt_tmpl.set_data(self.templates[0].data)
        plt_tmpl.draw(interpolation='nearest', origin='lower')

        plt_img.tooltip = 'Plane'
        plt_img.set_loading()
        Future(lambda: self.get_image(0),
               callback=lambda f: wx.CallAfter(self.plot_image, f.get(), 0))

        plt_img.ax.format_coord = types.MethodType(self.report_pixel, plt_img)
        plt_av.ax.format_coord = types.MethodType(self.report_pixel, plt_av)
        plt_tmpl.ax.format_coord = types.MethodType(self.report_pixel, plt_tmpl)

        plt_tmpl.ax.set_title('Correlation template')
        plt_tmpl.ax.tooltip = 'Correlation template'

        self.plt['shifts'].set_loading()
        Future(lambda: (get_bkeys('CRPIX1')[0], get_bkeys('CRPIX2')[0]),
               callback=lambda f: wx.CallAfter(self.plot_shifts, f.get()))

        self.rad.on_clicked(self.rad_clicked)
        figure.canvas.mpl_connect('button_press_event', self.onclick)

    def add_beams(self, data):
        self.beams = data
        from matplotlib.patches import RegularPolygon
        for c in data:
            if c[0].upper() == 'UNDEFINED':
                self.undef_beams = True
                self.beamselect.color = 'red'
                self.bax.figure.canvas.draw_idle()
            p = RegularPolygon(c[1:], numVertices=4, alpha=0.2,
                               radius=self.plt['av'].data.shape[0] / 40,
                               facecolor='white')

            self.plt['av'].ax.add_patch(p)
        self.plt['av'].figure.canvas.draw_idle()

    def get_average(self):
        nimages = self.fl.custom(self.fits['BKG_CORRECTED'][0], len) - 1
        imginput = self.fl.getdata(self.fits['BKG_CORRECTED'][0],
                                   range(1, min(nimages, 100)))
        return numpy.average(numpy.array(imginput), axis=0)

    def plot_image(self, data, idx):
        plt_img = self.plt['img']
        plt_img.set_data(data)
        plt_img.ax.set_title('Image plane #%d' % (idx + 1))
        plt_img.draw(interpolation='nearest', origin='lower')

    def plot_average(self, data):
        plt_av = self.plt['av']
        plt_av.set_data(data)
        plt_av.ax.set_title('Input image')
        plt_av.draw(origin='lower', interpolation='nearest', rescale=False)

        Future(get_beam_coords, args=(self.fl, self.fits['BEAM_DETECTED']),
               callback=lambda f: wx.CallAfter(self.add_beams, f.get()))

    def plot_stats(self, data_, label):
        data, bounds = data_
        plt_stats = self.plt['stats']
        plt_stats.ax.set_title(self.stats[label]['title'] + ' (Click to view image)')
        plt_stats.ax.tooltip = 'Statistics of all planes'
        plt_stats.ax.set_xlabel('Plane number')
        if 'unit' in self.stats[label]:
            plt_stats.ax.set_ylabel('[%s]' % self.stats[label]['unit'])
        plt_stats.set_data(numpy.array((data, (numpy.arange(1, len(data) + 1)))))
        # ugly, better: disable click handles while loading
        self.cor = data
        label = self.statslabels[self.statskeys.index(label)]
        for i, l in enumerate(self.rad.labels):
            if l.get_text() == label:
                self.rad.circles[i].set_facecolor(self.rad.activecolor)
            else:
                self.rad.circles[i].set_facecolor(plt_stats.ax.get_axis_bgcolor())

        plt_stats.draw()

    def plot_slice(self, data, x, y):
        self.plt['stats'].set_data(numpy.array((data, (numpy.arange(1, len(data) + 1)))))
        self.plt['stats'].ax.set_title('Values over x=%d y=%d' % (x, y))
        self.plt['stats'].draw()

    def plot_shifts(self, scat):
        title_scatter = 'Shifts'
        tooltip_scatter = 'Shifts of the beams'
        scatx, scaty = scat
        # avoid crash on empty input
        if scatx.size == 0 or scaty.size == 0:
            scatx, scaty = [0], [0]
        scatterdisp = ScatterDisplay()
        scatterdisp.setLabels('X [pixel]', 'Y [pixel]')
        scatterdisp.display(self.plt['shifts'].ax, title_scatter,
                            tooltip_scatter, scatx, scaty)
        # add a histogram for many points
        if (len(scatx) > 50):
            self.plt['shifts'].data, binsx, binsy, aximg = \
                self.plt['shifts'].ax.hist2d(scatx, scaty, bins=25, alpha=0.75)
            aximg.set_cmap('gray')
            def shift_scale(x, y):
                # x, y transposed in data array
                return ((y - binsy[0]) / numpy.diff(binsy[:2]),
                        (x - binsx[0]) / numpy.diff(binsx[:2]))
            self.plt['shifts'].scale_coord = shift_scale
            self.plt['shifts'].ax.format_coord = types.MethodType(self.report_pixel, self.plt['shifts'])
        self.plt['shifts'].figure.canvas.draw()

    def setInteractiveParameters(self):
        rec = 'visir_util_detect_shift'
        paramList = [
            VisirParameter(rec, name='template', type='file',
                           description='Correlation template used to detect shifts and reject bad images. If none given an averaged image is used.'),
            VisirParameter(rec, name='naverage', type='int',
                           description='Number of planes to average before attempting to detect the shifts.'),
            VisirParameter(rec, name='no-reject', type='bool',
                           description='If true images with bad statistics will not be rejected and no shift correction is done.'),
            VisirParameter(rec, name='no-shift', type='bool',
                           description='Whether to apply the determined shifts in following recipes. Set to TRUE of shifts appear too large'),
            VisirParameter(rec, name='min-correlation', type='float',
                           description='Minimal correlation to the template required to be accepted.'),
            VisirParameter(rec, name='max-shift', type='float',
                           description='Maximal allowed object shift from the template for an image to be accepted.'),
            VisirParameter(rec, name='max-mad', type='float',
                           description='Maximal allowed median absolute deviation for an image to be accepted. <= 0 equals no limit.'),
            VisirParameter(rec, name='beampos', type='string',
                           description='Set via the "Select beams" button.'),
            VisirParameter(rec, name='method', type='list',
                           options=['xcorrelate', 'brightest'], description="""\
Method to determine beam shifts:
xcorrelate: use the cross correlation shift.
brightest: use the position of the brightest pixel"""),
        ]
        return paramList

    def verifyParameters(self, params):
        if params['no-reject']['value'].upper() == 'TRUE':
            if any((not params['max-mad']['isdefault'],
                    not params['min-correlation']['isdefault'],
                    not params['max-shift']['isdefault'])):
                    return 'Rejection parameters changed but "no-reject" is ' \
                      'still set. In order to reject frames this must be unset.'
        return None


#This is the 'main' function
if __name__ == '__main__':
    # for debugging

    # Create commandline interface
    from reflex_interactive_app import PipelineInteractiveApp
    app = PipelineInteractiveApp()
    app.parse_args()

    #Open the interactive window if enabled
    try:
        if app.isGUIEnabled() and import_success:
            # hard dependency on wx, has own io parsing
            #Get the specific functions for this window
            with fitsloader.get_fitsloader() as fl:
                dataPlotManager = DataPlotterManager(fl)
                dataPlotManager.in_pars = app.inputs.in_sop
                ReflexBaseGui.show(app, dataPlotManager)
        else:
            app.passProductsThrough()
    except (ImportError, NameError) as e:
        print 'Error:', e
        app.passProductsThrough()

    app.print_outputs()
    # prevents hang when using the threading module in fitsloader
    sys.exit(0)
