import numpy as np
import scipy.ndimage.interpolation

import colorcet as cc
import xarray as xr
from arpes.plotting.interactive_utils import BokehInteractiveTool
from arpes.typing import DataType
from arpes.utilities.funcutils import Debounce
from arpes.utilities.normalize import normalize_to_spectrum

__all__ = ['ComparisonTool', 'compare']


class ComparisonTool(BokehInteractiveTool):
    auto_zero_nans = False
    auto_rebin = False
    other = None
    compared = None

    def __init__(self, other, **kwargs):
        super().__init__()

        self.load_settings(**kwargs)

        self.app_main_size = self.settings.get('app_main_size', 600)
        self.other = other

    def tool_handler(self, doc):
        from bokeh.layouts import row, column, widgetbox
        from bokeh.models import widgets
        from bokeh.models.mappers import LinearColorMapper
        from bokeh.plotting import figure

        default_palette = self.default_palette
        difference_palette = cc.coolwarm

        intensity_slider = widgets.Slider(
            title='Relative Intensity Scaling', start=0.5, end=1.5,
            step=0.005, value=1)

        self.app_context.update({
            'A': self.arr,
            'B': self.other,
            'compared': self.compared,
            'plots': {},
            'figures': {},
            'widgets': {},
            'data_range': self.arr.T.range(),
            'color_maps': {},
        })

        self.color_maps['main'] = LinearColorMapper(
            default_palette, low=np.min(self.arr.values), high=np.max(self.arr.values), nan_color='black')

        figure_kwargs = {
            'tools': ['reset', 'wheel_zoom'],
            'plot_width': self.app_main_size,
            'plot_height': self.app_main_size,
            'min_border': 10,
            'toolbar_location': 'left',
            'x_range': self.data_range['x'],
            'y_range': self.data_range['y'],
            'x_axis_location': 'below',
            'y_axis_location': 'right',
        }

        self.figures['A'] = figure(title='Spectrum A', **figure_kwargs)
        self.figures['B'] = figure(title='Spectrum B', **figure_kwargs)
        self.figures['compared'] = figure(title='Comparison', **figure_kwargs)

        self.compared = self.arr - self.other
        diff_low, diff_high = np.min(self.arr.values), np.max(self.arr.values)
        diff_range = np.sqrt((abs(diff_low) + 1) * (abs(diff_high) + 1)) * 1.5
        self.color_maps['difference'] = LinearColorMapper(
            difference_palette, low=-diff_range, high=diff_range, nan_color='white')

        self.plots['A'] = self.figures['A'].image(
            [self.arr.values], x=self.data_range['x'][0], y=self.data_range['y'][0],
            dw=self.data_range['x'][1] - self.data_range['x'][0],
            dh=self.data_range['y'][1] - self.data_range['y'][0],
            color_mapper=self.color_maps['main'],
        )

        self.plots['B'] = self.figures['B'].image(
            [self.other.values], x=self.data_range['x'][0], y=self.data_range['y'][0],
            dw=self.data_range['x'][1] - self.data_range['x'][0],
            dh=self.data_range['y'][1] - self.data_range['y'][0],
            color_mapper=self.color_maps['main']
        )

        self.plots['compared'] = self.figures['compared'].image(
            [self.compared.values], x=self.data_range['x'][0], y=self.data_range['y'][0],
            dw=self.data_range['x'][1] - self.data_range['x'][0],
            dh=self.data_range['y'][1] - self.data_range['y'][0],
            color_mapper=self.color_maps['difference']
        )

        x_axis_name = self.arr.dims[0]
        y_axis_name = self.arr.dims[1]

        stride = self.arr.T.stride()
        delta_x_axis = stride['x']
        delta_y_axis = stride['y']

        delta_x_slider = widgets.Slider(
            title='{} Shift'.format(x_axis_name), start=-20 * delta_x_axis,
            step=delta_x_axis / 2, end=20 * delta_x_axis, value=0)

        delta_y_slider = widgets.Slider(
            title='{} Shift'.format(y_axis_name), start=-20 * delta_y_axis,
            step=delta_y_axis / 2, end=20 * delta_y_axis, value=0)

        @Debounce(0.5)
        def update_summed_figure(attr, old, new):
            # we don't actually use the args because we need to pull all the data out
            shifted = (intensity_slider.value) * scipy.ndimage.interpolation.shift(self.other.values, [
                delta_x_slider.value / delta_x_axis,
                delta_y_slider.value / delta_y_axis,
            ], order=1, prefilter=False, cval=np.nan)
            self.compared = self.arr - xr.DataArray(
                shifted,
                coords=self.arr.coords,
                dims=self.arr.dims)

            self.compared.attrs.update(**self.arr.attrs)
            try:
                del self.compared.attrs['id']
            except KeyError:
                pass

            self.app_context['compared'] = self.compared
            self.plots['compared'].data_source.data = {
                'image': [self.compared.values]
            }

        layout = column(
            row(
                column(self.app_context['figures']['A']),
                column(self.app_context['figures']['B']),
            ),
            row(
                column(self.app_context['figures']['compared']),
                widgetbox(
                    intensity_slider,
                    delta_x_slider,
                    delta_y_slider,
                ),
            )
        )

        update_summed_figure(None, None, None)

        delta_x_slider.on_change('value', update_summed_figure)
        delta_y_slider.on_change('value', update_summed_figure)
        intensity_slider.on_change('value', update_summed_figure)

        doc.add_root(layout)
        doc.title = 'Comparison Tool'


def compare(A: DataType, B: DataType):
    A = normalize_to_spectrum(A)
    attrs = A.attrs
    B = normalize_to_spectrum(B)

    # normalize total intensity
    TOTAL_INTENSITY = 1000000
    A = A / (A.sum(A.dims) / TOTAL_INTENSITY)
    B = B / (B.sum(B.dims) / TOTAL_INTENSITY)
    A.attrs.update(**attrs)

    tool = ComparisonTool(other=B)
    return tool.make_tool(A)
