# -*- coding: utf-8 -*-
"""
A module to plot SOSS data

Authors: Joe Filippazzo
"""
import copy

from astropy.io import fits
from bokeh.plotting import figure, show
from bokeh.models import Tabs, TabPanel, ColumnDataSource, HoverTool, CustomJSHover, LogColorMapper, FixedTicker, BasicTickFormatter, CustomJSTickFormatter, BasicTicker, LogTicker, LinearColorMapper, ColorBar, Span, CustomJS, Slider, Range1d
from bokeh.layouts import gridplot, column
from bokeh.palettes import viridis
import numpy as np


def quicklook(file, nint=0, ngrp=0, **kwargs):
    """
    Function to quickly plot a single frame of data from a JWST data file

    Parameters
    ----------
    file: str
        The filepath
    nint: int
        The integration to look at
    ngrp: int
        The group to look at
    """
    # Get the data
    header = fits.getheader(file)
    data = fits.getdata(file, extname='SCI')
    try:
        udata = fits.getdata(file, extname='ERR')
    except KeyError:
        udata = np.sqrt(data)
        print("No ERR extension")

    # Check the data shape
    if data.ndim == 4:
        nints, ngrps, rows, cols = data.shape
        if nint > nints:
            print("{} > {}: Choose a different nint".format(nint, nints))
        if ngrp > ngrps:
            print("{} > {}: Choose a different ngrp".format(ngrp, ngrps))
        frame = data[nint, ngrp]
        uframe = udata[nint, ngrp]
    elif data.ndim == 3:
        nints, rows, cols = data.shape
        if nint > nints:
            print("{} > {}: Choose a different nint".format(nint, nints))
        frame = data[nint]
        uframe = udata[nint]
    elif data.ndim == 2:
        frame = data
        uframe = udata
    else:
        raise ValueError("{}: Not sure what to do with that shape data.".format(data.shape))

    # Pass it to plot_frame
    fig = plot_frame(frame, uframe=uframe, **kwargs)
    show(fig)

    return header


def compare_frames(frames, titles, uframes=None, units='ADU/s', scale='linear', wavecal=None, fullWell=65536., saturation=0.8, color_map='Viridis256', title="SOSS Frame Comparison"):
    """
    Plot several frames simultaneously with a shared cross-section plot

    Parameters
    ----------
    frames: sequence
        A list of the frames to plot

    Returns
    -------

    """
    # Set shared plot params
    nrows, ncols = frames[0].shape
    x_range = (0, ncols)
    y_range = (0, nrows)
    height = int(nrows/2.)+120

    # Add a plot for each frame
    img_plots = []
    line_data = []
    for n, frame in enumerate(frames):

        # Link all image axis ranges
        if n > 0:
            x_range = img_plots[0][0].x_range
            y_range = img_plots[0][0].y_range

        # Make the figure
        fig = figure(x_range=x_range, y_range=y_range, width=800, height=height, title=titles[n])

        # Get data, snr, and saturation for plotting
        dat = frame
        snr = np.sqrt(frame)
        if uframes is not None:
            snr = frame / uframes[n]

        sat = dat > saturation * fullWell
        sat = sat.astype(int)
        dh, dw = dat.shape

        # Fix if log scale
        if scale == 'log':
            dat[dat < 1.] = 1.

        # Set the source data
        source = ColumnDataSource(data=dict(img=[dat], snr=[snr], saturation=[sat]))
        line_data.append(np.nansum(dat, axis=0))

        # Set the tooltips
        tooltips = [("(x,y)", "($x{int}, $y{int})"), (units, "@data"), ("SNR", "@snr"), ('Saturation', '@saturation')]

        # Add wavelength calibration if possible
        if isinstance(wavecal, np.ndarray):
            if wavecal.shape == frame.shape:
                source['wave1'] = [wavecal]
                tooltips.append(("Wavelength", "@wave1"))
            if wavecal.ndim == 3 and wavecal.shape[0] == 3:
                source['wave1'] = [wavecal[0]]
                source['wave2'] = [wavecal[1]]
                source['wave3'] = [wavecal[2]]
                tooltips.append(("Wave 1", "@wave1"))
                tooltips.append(("Wave 2", "@wave2"))
                tooltips.append(("Wave 3", "@wave3"))

        vmin = int(np.nanmin(dat[dat >= 0]))
        vmax = int(np.nanmax(dat[dat < np.inf]))

        # Set the plot scale
        if scale == 'log':
            mapper = LogColorMapper(palette=color_map, low=vmin, high=vmax)
        else:
            mapper = LinearColorMapper(palette=color_map, low=vmin, high=vmax)

        # Add a line renderer for each series
        line_renderer = fig.image(source=source, image='img', x=0, y=0, dw=dw, dh=dh, color_mapper=mapper)

        # Create a HoverTool for each series with unique tooltips
        hover_tool = HoverTool(tooltips=tooltips, renderers=[line_renderer])
        fig.add_tools(hover_tool)

        img_plots.append([fig])

    # Add the column plot
    colplt = figure(x_range=x_range, y_range=y_range, width=800, height=height)
    for color, ldata, titl in zip(['blue','red','green'], line_data, titles):
        colplt.line(np.arange(2048), ldata, color=color, legend_label=titl)

    show(gridplot(img_plots + [[colplt]]))


def plot(frame, scale='linear', title="My Plot", histogram=False, x_range=None, y_range=None, low=None, high=None, width=400, height=400, epsilon=1e-5, palette="Viridis256", bins=100):
    """
    Generic function to plot 2D data with tooltips

    Parameters
    ----------
    frame: sequence
        The 2D array to plot
    scale: str
        The scaling of the data
    title: str
        The title for the plot
    histogram: bool
        Also plot a histogram of the data
    x_range: sequence
        The (start, end) values for the x-axis
    y_range: sequence
        The (start, end) values for the y_axis
    low: float
        The low for the color range
    high: float
        The high for the color range
    width: int
        The plot width
    height: int
        The plot height
    epsilon: float
        The value to set for n<0 if log scale
    palette: str
        The bokeh color palette to use

    Returns
    -------
    bokeh.figure.Figure
        The bokeh plot
    """
    # Validate scale input
    if scale not in ['linear', 'log']:
        raise ValueError("Scale must be either 'linear' or 'log'")

    # Adjust array for log scale to avoid zero or negative values
    if scale == 'log':
        frame = np.where(frame <= 0, epsilon, frame)

    # Get the dimensions of the array
    nrows, ncols = frame.shape

    # Create 2D data source
    source = ColumnDataSource(data=dict(image=[frame]))

    # Create a figure
    p = figure(title=title, width=width, height=height, x_range=x_range or (0, ncols), y_range=y_range or (0, nrows), tools='pan,box_zoom,reset,box_select,save', tooltips=[("x", "$x{0}"), ("y", "$y{0}"), ("Value", "@image")])

    # Define color mappers for linear and logarithmic scales
    mapper = LinearColorMapper if scale == 'linear' else LogColorMapper
    color_mapper = mapper(palette=palette, low=low or np.nanmin(frame), high=high or np.nanmax(frame))

    # Add the image glyph to the figure
    p.image(image="image", x=0, y=0, dw=ncols, dh=nrows, color_mapper=color_mapper, source=source)
    color_bar = ColorBar(color_mapper=color_mapper, location=(0, 0))
    p.add_layout(color_bar, 'right')

    if histogram:

        # Remove NaN values
        array = frame[~np.isnan(frame)].flatten()

        # Create histogram data source
        hist_source = ColumnDataSource(data=dict(top=[], left=[], right=[]))

        # Create the histogram
        hist_fig = figure(title="Histogram", width=width, height=height, tools="pan,box_zoom,reset,save")
        hist_fig.quad(top='top', bottom=0, left='left', right='right', line_color="white", fill_color="navy",
                      fill_alpha=0.5, source=hist_source)

        # Create CustomJS callback for the histogram update
        callback = CustomJS(args=dict(source=source, hist_source=hist_source, array=array, hist_fig=hist_fig), code="""
            const indices = source.selected.indices;
            if (indices.length === 0) {
                return;
            }
            const data = array;
            const selected_data = indices.map(i => data[Math.floor(i / ncols)][i % ncols]);
    
            // Compute the histogram
            const bins = 30;
            const min = Math.min(...selected_data);
            const max = Math.max(...selected_data);
            const step = (max - min) / bins;
    
            const hist = Array(bins).fill(0);
            for (let i = 0; i < selected_data.length; i++) {
                const bin = Math.floor((selected_data[i] - min) / step);
                hist[Math.min(bin, bins - 1)]++;
            }
    
            const left = [];
            const right = [];
            const top = hist;
            for (let i = 0; i < bins; i++) {
                left.push(min + i * step);
                right.push(min + (i + 1) * step);
            }
    
            hist_source.data = {top: top, left: left, right: right};
            hist_source.change.emit();
    
            // Update the range of the histogram plot
            hist_fig.x_range.start = min;
            hist_fig.x_range.end = max;
        """)

        # Attach the callback to the source's selected attribute
        source.selected.js_on_change('indices', callback)

        return gridplot([[p, hist_fig]])

    else:

        return p


def plot_frame(frame, cols=None, uframe=None, units='ADU/s', scale='log', trace_coeffs=None, saturation=0.8, plot_width=900, plot_height=None, low=None, high=None, title=None, wavecal=None, color_map=None, tabs=False):
    """
    Plot a SOSS frame

    Parameters
    ----------
    frame: sequence
        The 2D frame to plot
    cols: int, list, tuple
        The 1D column(s) to plot
    uframe: sequence (optional)
        The uncertainty frame to plot
    units: str
        The units of the data
    scale: str
        Plot scale, ['linear', 'log']
    trace_coeffs: sequence
        Plot the traces for the given coefficients
    saturation: float
        The full-well fraction to be considered saturated, (0, 1)
    plot_width: int
        The width of the plot in pixels
    title: str
        A title for the plot
    wavecal: np.ndarray
        A wavelength calibration map for each pixel
    """
    # Determine subarray
    nrows, ncols = frame.shape

    # Get data, snr, and saturation for plotting
    dat = frame
    snr = np.sqrt(frame)
    if uframe is not None:
        if uframe.shape == frame.shape:
            snr = frame / uframe
    fullWell = 65536.0
    sat = dat > saturation * fullWell
    sat = sat.astype(int)
    dh, dw = dat.shape

    # Fix if log scale
    if scale == 'log':
        dat[dat < 1.] = 1.

    # Set the source data
    source = dict(data=[dat], snr=[snr], saturation=[sat])

    # Set the tooltips
    tooltips = [("(x,y)", "($x{int}, $y{int})"), (units, "@data"), ("SNR", "@snr"), ('Saturation', '@saturation')]

    # Add wavelength calibration if possible
    if isinstance(wavecal, np.ndarray):
        if wavecal.shape == frame.shape:
            source['wave1'] = [wavecal]
            tooltips.append(("Wavelength", "@wave1"))
        if wavecal.ndim == 3 and wavecal.shape[0] == 3:
            source['wave1'] = [wavecal[0]]
            source['wave2'] = [wavecal[1]]
            source['wave3'] = [wavecal[2]]
            tooltips.append(("Wave 1", "@wave1"))
            tooltips.append(("Wave 2", "@wave2"))
            tooltips.append(("Wave 3", "@wave3"))

    # Set shared plot params
    x_range = (0, dat.shape[1])
    y_range = (0, dat.shape[0])
    height = int(nrows/2.)+160
    toolbar = 'above'

    # Draw the figures
    plot_tabs = []
    color_map = color_map or viridis(256)
    for pname, ptype in zip([units, 'SNR', 'Saturation ({}% Full Well)'.format(saturation*100)], ['data', 'snr', 'saturation']):

        # Make the figure
        fig_title = '{} - {}'.format(title, pname)
        fig = figure(x_range=x_range, y_range=y_range, tooltips=tooltips, width=plot_width, height=plot_height or height, title=fig_title, toolbar_location=toolbar, toolbar_sticky=True)

        # Get the data
        vals = source[ptype][0]

        # Saturation plot is different
        if ptype == 'saturation':
            vmin = 0
            vmax = 1
            formatter = CustomJSTickFormatter(code="""return {0: 'Unsaturated', 1: 'Saturated'}[tick]""")
            color_map = ['#404387', '#FDE724']
            ticker = FixedTicker(ticks=[vmin, vmax])

        # Counts and SNR are similar plots
        else:
            # vmin = low or int(np.nanmin(vals[vals >= 0]))
            # vmax = high or int(np.nanmax(vals[vals < np.inf]))
            vmin = np.nanmin(vals)
            vmax = np.nanmax(vals)
            formatter = BasicTickFormatter()
            ticker = BasicTicker()

        # Set the plot scale
        if scale == 'log':
            mapper = LogColorMapper(palette=color_map, low=vmin, high=vmax)
        else:
            mapper = LinearColorMapper(palette=color_map, low=vmin, high=vmax)

        # Plot the frame
        fig.image(source=source, image=ptype, x=0, y=0, dw=dw, dh=dh, color_mapper=mapper)
        color_bar = ColorBar(color_mapper=mapper, ticker=ticker, formatter=formatter, orientation="horizontal", location=(0, 0))

        # Plot the trace polynomials
        if trace_coeffs is not None:
            X = np.linspace(0, 2048, 2048)

            # Check number of traces
            if np.array(trace_coeffs).ndim == 1:
                trace_coeffs = [trace_coeffs]

            for coeffs in trace_coeffs:
                if len(coeffs) == 2048:
                    Y = coeffs
                else:
                    Y = np.polyval(coeffs, X)
                fig.line(X, Y, color='red', line_dash='dashed')

        # Add the colorbar
        fig.add_layout(color_bar, 'below')

        # Plot the column data
        col_fig = None
        col_colors = ['blue', 'green', 'purple', 'cyan', 'orange']
        if cols is not None:

            # Make sure it's iterable
            if isinstance(cols, int):
                cols = [cols]

            # Initialize column figure
            col_fig = figure(width=plot_width, height=300, toolbar_location=toolbar, toolbar_sticky=True)

            for n, col in enumerate(cols):
                col_color = col_colors[n]
                col_fig.step(np.arange(nrows), vals[:, col], color=col_color, legend_label='Col {}'.format(col))
                col_fig.y_range = Range1d(vmin * 0.9, vmax * 1.1)
                col_fig.x_range = Range1d(*y_range)

                # Add line to image
                fig.line([col, col], [0, nrows], color=col_color, line_width=3)

            # Update click policy
            col_fig.legend.click_policy = 'hide'

        if col_fig is not None:

            # Add the figure to the tab list
            if tabs:
                plot_tabs.append(TabPanel(child=column([fig, col_fig]), title=pname))
            else:
                plot_tabs.append(fig)

        else:

            # No column object
            if tabs:
                plot_tabs.append(TabPanel(child=fig, title=pname))
            else:
                plot_tabs.append(fig)

    # Make the final tabbed figure
    if tabs:
        final = TabPanel(tabs=plot_tabs)
    else:
        final = plot_tabs[0]

    return final


def plot_frames(data, unc=None, idx=0, col=0, units='Counts', scale='linear', trace_coeffs=None, saturation=0.8, width=1000, title=None, wavecal=None):
    """
    Plot a SOSS frame

    Parameters
    ----------
    data: sequence
        The 3D data to plot
    unc: sequence
        The 3D uncertainty to plot
    units: str
        The units of the data
    scale: str
        Plot scale, ['linear', 'log']
    trace_coeffs: sequence
        Plot the traces for the given coefficients
    saturation: float
        The full-well fraction to be considered saturated, (0, 1)
    title: str
        A title for the plot
    wavecal: np.ndarray
        A wavelength calibration map for each pixel
    """
    # Determine subarray
    nframes, nrows, ncols = data.shape

    # Remove infs
    data[data == np.inf] = np.nan

    # Get data, snr, and saturation for plotting
    dat = data
    if unc is None:
        snr = np.sqrt(data)
    else:
        snr = data / unc
    fullWell = 65536.0
    sat = dat > saturation * fullWell
    sat = sat.astype(int)

    # Fix log scale plot values
    if scale == 'log':
        dat[dat < 1.] = 1.
        snr[snr < 1.] = 1.

    # Broadcast the data
    frames = np.arange(nframes)
    columns = np.arange(ncols)
    rows = np.arange(nrows)
    verticals = np.tile(np.arange(ncols), (nrows, 1)).T

    # Wrap the data in ColumnDataSources
    source_available = ColumnDataSource(data=dict(**{'counts{}'.format(n): dat[n] for n in frames}, **{'snr{}'.format(n): snr[n] for n in frames}, **{'saturation{}'.format(n): sat[n] for n in frames}))
    source_visible = ColumnDataSource(data=dict(counts=[dat[idx]], snr=[snr[idx]], saturation=[sat[idx]]))
    vertical_available = ColumnDataSource(data=dict(**{'vertical{}'.format(n): vert for n, vert in enumerate(verticals)}))
    vertical_visible = ColumnDataSource(data=dict(column=rows, vertical=verticals[col]))
    col_visible = ColumnDataSource(data=dict(columns=rows, counts=dat[0, :, col], snr=snr[0, :, col], saturation=sat[0, :, col]))
    col_dict = {}
    for fnum in frames:
        for cnum in columns:
            for datacube, pname in zip([dat, snr, sat], ['counts', 'snr', 'saturation']):
                col_dict['{}{}_{}'.format(pname, cnum, fnum)] = datacube[fnum, :, cnum]
    col_available = ColumnDataSource(data=col_dict)

    # Set the tooltips
    tooltips = [("(x,y)", "($x{int}, $y{int})"), (units, "@counts"), ("SNR", "@snr"), ('Saturation', '@saturation')]
    col_color = 'blue'

    # Add wavelength calibration if possible
    if isinstance(wavecal, np.ndarray):
        if wavecal.shape == dat[idx].shape:
            source_visible.data['wave1'] = [wavecal]
            tooltips.append(("Wavelength", "@wave1"))
        if wavecal.ndim == 3 and wavecal.shape[0] == 3:
            source_visible.data['wave1'] = [wavecal[0]]
            source_visible.data['wave2'] = [wavecal[1]]
            source_visible.data['wave3'] = [wavecal[2]]
            tooltips.append(("Wave 1", "@wave1"))
            tooltips.append(("Wave 2", "@wave2"))
            tooltips.append(("Wave 3", "@wave3"))

    # Set shared plot params
    x_range = (0, ncols)
    y_range = (0, nrows)
    height = int(nrows/2.)+160
    toolbar = 'right'

    # Draw the figures
    tabs = []
    for pdata, pname, ptype, ylabel in zip([dat, snr, sat], [units, 'SNR', 'Saturation ({}% Full Well)'.format(saturation*100)], ['counts', 'snr', 'saturation'], ['Count Rate [ADU/s]', 'SNR', 'Saturated']):

        # Make the figure
        fig_title = '{} - {}'.format(title, pname)

        # Get min and max
        vmin = np.nanmin(pdata)
        vmax = np.nanmax(pdata)

        # Saturation plot is different
        if ptype == 'saturation':
            formatter = CustomJSTickFormatter(code="""return {0: 'Unsaturated', 1: 'Saturated'}[tick]""")
            color_map = ['#404387', '#FDE724']
            ticker = FixedTicker(ticks=[vmin, vmax])

        # Counts and SNR are similar plots
        else:
            formatter = BasicTickFormatter()
            color_map = 'Viridis256'
            ticker = BasicTicker()

        # Set the plot scale
        if scale == 'log':
            mapper = LogColorMapper(palette=color_map, low=vmin, high=vmax)
        else:
            mapper = LinearColorMapper(palette=color_map, low=vmin, high=vmax)

        # Plot the image data
        fig = figure(x_range=x_range, y_range=y_range, tooltips=tooltips, width=width, height=height, title=fig_title, toolbar_location=toolbar, toolbar_sticky=True)
        fig.image(source=source_visible, image=ptype, x=0, y=0, dw=ncols, dh=nrows, color_mapper=mapper)

        # Plot the line indicating the column
        fig.line('vertical', 'column', source=vertical_visible, color=col_color, line_width=3)

        # Add the colorbar
        color_bar = ColorBar(color_mapper=mapper, ticker=ticker, formatter=formatter, orientation="horizontal", location=(0, 0))
        fig.add_layout(color_bar, 'below')

        # Plot the column data
        col_fig = figure(width=width, height=300, toolbar_location=toolbar, toolbar_sticky=True)
        col_fig.step('columns', ptype, source=col_visible, color=col_color)
        col_fig.xaxis.axis_label = 'Row'
        col_fig.yaxis.axis_label = ylabel
        col_fig.y_range = Range1d(vmin*0.9, vmax*1.1)
        col_fig.x_range = Range1d(*y_range)

        # Plot the trace polynomials
        if trace_coeffs is not None:
            for coeffs in trace_coeffs:
                Y = np.polyval(coeffs, columns)
                fig.line(columns, Y, color='red')

        # Add the figure to the tab list
        tabs.append(TabPanel(child=column(fig, col_fig), title=pname))

    # Make the final tabbed figure
    final = Tabs(tabs=tabs)

    # Write JS code
    code ="""
        var vis = visible.data;
        var avail = available.data;
        var frame = fr_slide.value.toString(10);

        var viscol = col_vis.data;
        var availcol = col_avail.data;
        var col = col_slide.value.toString(10);

        var visvert = vert_vis.data;
        var availvert = vert_avail.data;

        vis['counts'] = [avail['counts'.concat(frame)]];
        vis['snr'] = [avail['snr'.concat(frame)]];
        vis['saturation'] = [avail['saturation'.concat(frame)]];

        viscol['counts'] = availcol['counts'.concat(col, '_', frame)];
        viscol['snr'] = availcol['snr'.concat(col, '_', frame)];
        viscol['saturation'] = availcol['saturation'.concat(col, '_', frame)];

        visvert['vertical'] = availvert['vertical'.concat(col)];

        visible.change.emit();
        col_vis.change.emit();
        vert_vis.change.emit();
    """

    # Make the column slider
    column_slider = Slider(title='Column', value=col, start=0, end=ncols-1, step=1)

    # Make the frame slider
    if nframes-1 > 0:
        frame_slider = Slider(title='Frame', value=idx, start=0, end=nframes-1, step=1)
    else:
        frame_slider = None
        code = code.replace('fr_slide.value.toString(10);', '0')

    # CustomJS callback to update the three plots on slider changes
    callback = CustomJS(args=dict(visible=source_visible, available=source_available, col_vis=col_visible, col_avail=col_available, vert_vis=vertical_visible, vert_avail=vertical_available, fr_slide=frame_slider, col_slide=column_slider), code=code)

    # Add callback to column slider
    column_slider.js_on_change('value', callback)

    # Add callback to frame slider
    if frame_slider is not None:
        frame_slider.js_on_change('value', callback)
        return column(final, frame_slider, column_slider)

    else:
        return column(final, column_slider)


def plot_ramp(data):
    """
    Plot the total flux on each frame to display the ramp
    """
    fig = figure()
    x = range(len(data))
    y = np.sum(data, axis=(-1, -2))
    fig.circle(x, y, size=8)
    fig.xaxis.axis_label = 'Group'
    fig.yaxis.axis_label = 'Count Rate [ADU/s]'

    return fig


def plot_spectrum(wavelength, flux, fig=None, scale='log', legend=None, ylabel='Flux Density', xlabel='Wavelength [um]', width=1024, height=500, **kwargs):
    """
    Plot a generic spectrum

    Parameters
    ----------
    wavelength: array-like
        The 1D wavelength array
    flux: array-like
        The 1D counts or flux
    fig: bokeh.plotting.figure.Figure (optional)
        The plot to add the spectrum to 
    scale: str
        Plot scale, ['linear', 'log']
    legend: str
        The text for the legend
    ylabel: str
        The text for the y-axis
    xlabel: str
        The text for the x-axis
    width: int
        The width of the plot
    height: int
        The height of the plot

    Returns
    -------
    bokeh.plotting.figure.Figure
        The figure
    """
    # Make the figure
    if fig is None:
        fig = figure(x_axis_type=scale, y_axis_type=scale, width=width, height=height)

    # Add the spectrum plot
    if legend is None:
        fig.step(wavelength, flux, mode='center', **kwargs)
    else:
        fig.step(wavelength, flux, mode='center', legend_label=legend, **kwargs)
    fig.yaxis.axis_label = ylabel
    fig.xaxis.axis_label = xlabel

    return fig


def plot_time_series_spectra(flux, wavelength=None, time=None, xlabel='Column', ylabel='Frame', width=1024, height=300, title=None, **kwargs):
    """
    Plot time series 1D spectra as an image

    Parameters
    ----------
    flux: array-like
        The 2D counts or flux
    wavelength: sequence (optional)
        The 1D wavelength array
    time: sequence (optional)
        The 1D time array
    xlabel: str
        The label for the data x-axis
    ylabel: str
        The label for the data y-axis
    width: int
        The width of the plot
    height: int
        The height of the plot
    title: str
        A title for the plot

    Returns
    -------
    bokeh.plotting.figure.Figure
        The figure
    """
    # Check that flux is 2D
    if not flux.ndim == 2:
        raise ValueError("{}: 'flux' must be a 2D array.".format(flux.shape))

    # Copy flux array
    flx = copy.copy(flux)

    # Get plot params
    dh, dw = flx.shape
    fmin = np.nanmin(flx)
    fmax = np.nanmax(flx)
    wmin, wmax = 0, dw
    tmin, tmax = 0, dh
    lightcurves = flx.T

    # Make sure time array is correct length
    if time is not None:
        if len(time) != dh:
            raise ValueError("{} != {}: 'time' array must be the same length as 'flux' cube.".format(len(time), flux.shape[0]))

    # Make sure wavelength array is correct length
    if wavelength is not None:
        if len(wavelength) != dw:
            raise ValueError("{} != {}: 'wavelength' array must be the same depth as 'flux' cube.".format(len(wavelength), flux.shape[1]))

    # Major tick labels
    waxis = np.arange(wmin, wmax)
    wstart, wskip = 128, 256
    taxis = np.arange(tmin, tmax)
    tstart, tskip = 1, max(4, min(tmax//2, 10))

    # Set the source data
    wave_data = np.array([wavelength if wavelength is not None else np.arange(dw)]*flx.shape[0])
    time_data = np.array([time if time is not None else np.arange(dh)]*flx.shape[1]).T
    sourceW = ColumnDataSource(data=dict(flux=[flx], wavelength=[wave_data], time=[time_data]))
    sourceX = ColumnDataSource(data=dict(wavelength=waxis, flux=flx[0], **{'flux{}'.format(n): fx for n, fx in enumerate(flx)}))
    sourceY = ColumnDataSource(data=dict(time=taxis, lightcurve=lightcurves[0], **{'lightcurve{}'.format(n): lc for n, lc in enumerate(lightcurves)}))
    sourceZ = ColumnDataSource(data=dict(x=[0.5], y=[0.5]))

    # ====================================================================

    # Make the 2D spectra figure
    spec_fig = figure(x_range=(wmin, wmax), y_range=(tmin, tmax), x_axis_label=xlabel, y_axis_label=ylabel,
                      width=width, height=height, title=title, tools="")

    # Plot the image
    color_mapper = LogColorMapper(palette="Viridis256", low=fmin, high=fmax)
    spec_fig.image(image='flux', x=wmin, y=tmin, dw=wmax, dh=tmax, color_mapper=color_mapper, alpha=0.8, name='img', source=sourceW)
    color_bar = ColorBar(color_mapper=color_mapper, ticker=LogTicker(), orientation="horizontal", label_standoff=12, border_line_color=None, location=(0, 0))

    # Add current lightcurve line to plot
    spec_fig.vbar(x='x', width=1, top=tmax, source=sourceZ, color='red', alpha=0.3)

    # Add current spectrum line to plot
    spec_fig.hbar(y='y', height=1, right=wmax, source=sourceZ, color='blue', alpha=0.3)

    # Set the tooltips
    spec_tt = HoverTool(name="img", tooltips=[("(x,y)", "($x{int}, $y{int})"), ("Flux", "@flux"), ('Wavelength', '@wavelength'), ('Time', '@time{0.00000}')])
    spec_fig.add_tools(spec_tt)

    # Change y tick labels
    if time is not None:
        spec_fig.yaxis.ticker = taxis[tstart::tskip]
        spec_fig.yaxis.major_label_overrides = {int(n): '{:.2f}'.format(t) for n, t in zip(taxis[tstart::tskip], time[tstart::tskip])}

    # Change x tick labels
    if wavelength is not None:
        spec_fig.xaxis.ticker = waxis[wstart::wskip]
        spec_fig.xaxis.major_label_overrides = {int(n): '{:.3f}'.format(w) for n, w in zip(waxis[wstart::wskip], wavelength[wstart::wskip])}

    # Remove grid lines
    spec_fig.xgrid.grid_line_color = None
    spec_fig.ygrid.grid_line_color = None

    # ====================================================================

    # Make the 1D spectrum figure
    sp_fig = figure(x_range=(wmin, wmax), y_range=(fmin, fmax), width=width, height=height,
                    x_axis_label=xlabel, y_axis_label='Flux Density', title='Spectrum', tools="")

    # Draw the spectrum
    sp_fig.step('wavelength', 'flux', source=sourceX, color='blue', line_width=3, line_alpha=0.6, mode='center', name='wf')

    # Set the tooltips
    sp_tt = HoverTool(name='wf', tooltips=[("Flux", "@flux"), ('Wavelength', '@wavelength')], mode='vline')
    sp_fig.add_tools(sp_tt)

    # Change x tick labels
    if wavelength is not None:
        sp_fig.xaxis.ticker = waxis[wstart::wskip]
        sp_fig.xaxis.major_label_overrides = {int(n): '{:.3f}'.format(w) for n, w in zip(waxis[wstart::wskip], wavelength[wstart::wskip])}

    # Make the spectrum slider
    sp_slider = Slider(value=0, start=tmin, end=tmax-1, step=1, width=30, title=ylabel, orientation='vertical', direction='rtl', bar_color='blue')

    # ====================================================================

    # Make the 1D lightcurve figure
    lc_fig = figure(x_range=(tmin, tmax), y_range=(fmin, fmax), width=width, height=height,
                    x_axis_label=ylabel, y_axis_label='Flux Density', title='Lightcurve', tools="")

    # Draw the lightcurve
    lc_fig.step('time', 'lightcurve', source=sourceY, color='red', line_width=3, line_alpha=0.6, mode='center', name='tl')

    # Set the tooltips
    lc_tt = HoverTool(name='tl', tooltips=[("Time", "@time"), ('Flux', '@lightcurve')], mode='vline')
    lc_fig.tools.append(lc_tt)

    # Change x tick labels
    if time is not None:
        lc_fig.xaxis.ticker = taxis[tstart::tskip]
        lc_fig.xaxis.major_label_overrides = {int(n): '{:.2f}'.format(t) for n, t in zip(taxis[tstart::tskip], time[tstart::tskip])}

    # Make the lightcurve slider
    lc_slider = Slider(value=0, start=wmin, end=wmax-1, step=1, width=width, title=xlabel, bar_color='red')

    # ====================================================================

    # CustomJS callback to update the three plots on slider changes
    callback = CustomJS(args=dict(sourcex=sourceX, sourcey=sourceY, sourcez=sourceZ, sp_slide=sp_slider, lc_slide=lc_slider), code="""
        var datax = sourcex.data;
        var datay = sourcey.data;
        var dataz = sourcez.data;
        var sp = sp_slide.value;
        var lc = lc_slide.value;
        var wavelength = datax['wavelength'];
        var time = datay['time'];
        datax['flux'] = datax['flux'.concat(sp.toString(10))];
        datay['lightcurve'] = datay['lightcurve'.concat(lc.toString(10))];
        dataz['x'] = [lc+0.5];
        dataz['y'] = [sp+0.5];
        sourcex.change.emit();
        sourcey.change.emit();
        sourcez.change.emit();
    """)

    # Add callback to spectrum slider
    sp_slider.js_on_change('value', callback)

    # Add callback to lightcurve slider
    lc_slider.js_on_change('value', callback)

    return gridplot([[sp_fig, None],[spec_fig, sp_slider], [lc_slider, None], [lc_fig, None]])
