from pymodaq.daq_utils.gui_utils.file_io import select_file
from qtpy import QtGui, QtWidgets, QtCore
from qtpy.QtCore import Qt, QObject, Slot, Signal

import sys
import os
import numpy as np
import tables

from pyqtgraph.parametertree import Parameter, ParameterTree
from pymodaq.daq_utils.plotting.data_viewers.viewer2D_basic import Viewer2DBasic
from pymodaq.daq_utils.plotting.items.image import UniformImageItem, SpreadImageItem
from pymodaq.daq_utils import daq_utils as utils
from pymodaq.daq_utils.h5modules import H5Saver, browse_data, H5BrowserUtil
from pymodaq.daq_utils import gui_utils as gutils

from pymodaq.daq_utils.config import get_set_local_dir, Config

local_path = get_set_local_dir()
navigator_path = os.path.join(local_path, 'navigator_temp_files')
if not os.path.isdir(navigator_path):
    os.makedirs(navigator_path)

logger = utils.set_logger(utils.get_module_name(__file__))

Ntick = 128
colors_red = np.array([(int(r), 0, 0) for r in np.linspace(0, 255, Ntick)])
colors_green = np.array([(0, int(g), 0) for g in np.linspace(0, 255, Ntick)])
colors_blue = np.array([(0, 0, int(b)) for b in np.linspace(0, 255, Ntick)])
config = Config()


class Navigator(QObject):
    log_signal = Signal(str)
    sig_double_clicked = Signal(float, float)

    def __init__(self, parent=None, h5file_path=None):
        super().__init__(parent)

        if parent is None:
            parent = QtWidgets.QWidget()
        self.parent = parent
        self.title = 'Navigator'
        self.status_time = 1000
        self.x_range = []
        self.y_range = []
        self.filters = tables.Filters(complevel=5)
        self.next_scan_index = 0
        self.viewer = None
        self.overlays = []  # %list of imageItem items displaying 2D scans info
        self.h5module_path = h5file_path
        self.h5module = H5BrowserUtil()

        if h5file_path is not None:
            self.h5module.open_file(h5file_path, 'a')
            self.settings.child('settings', 'filepath').setValue(h5file_path)
            self.settings.child('settings', 'Load h5').hide()
            self.show_overlay()

        self.setupUI()

    def create_toolbar(self):
        iconload = QtGui.QIcon()
        iconload.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/NewLayer.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        self.loadaction = QtWidgets.QAction(iconload, "Load scan file (.h5)", None)
        self.toolbar.addAction(self.loadaction)
        self.loadaction.triggered.connect(self.load_data)

        iconloadim = QtGui.QIcon()
        iconloadim.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/Open_File_32.png"), QtGui.QIcon.Normal,
                             QtGui.QIcon.Off)
        self.loadactionim = QtWidgets.QAction(iconloadim, "Load image file (.h5)", None)
        self.toolbar.addAction(self.loadactionim)
        self.loadactionim.triggered.connect(self.load_image)

        icon_ratio = QtGui.QIcon()
        icon_ratio.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/Zoom_1_1.png"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        self.action_ratio = QtWidgets.QAction(icon_ratio, "Set viewbox aspect ratio to 1", None)
        self.action_ratio.setCheckable(True)
        self.toolbar.addAction(self.action_ratio)
        self.action_ratio.triggered.connect(self.set_aspect_ratio)

        icon_moveat = QtGui.QIcon()
        icon_moveat.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/move_contour.png"), QtGui.QIcon.Normal,
                              QtGui.QIcon.Off)
        self.moveat_action = QtWidgets.QAction(icon_moveat,
                                               "When selected, double clicking on viewbox will move DAQ_Move modules",
                                               None)
        self.moveat_action.setCheckable(True)
        self.toolbar.addAction(self.moveat_action)

        icon_sel_all = QtGui.QIcon()
        icon_sel_all.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/select_all2.png"), QtGui.QIcon.Normal,
                               QtGui.QIcon.Off)
        self.sel_all_action = QtWidgets.QAction(icon_sel_all, "Select (show) all 2D scans on the viewer", None)
        self.toolbar.addAction(self.sel_all_action)
        self.sel_all_action.triggered.connect(self.show_all)

        icon_sel_none = QtGui.QIcon()
        icon_sel_none.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/select_none.png"), QtGui.QIcon.Normal,
                                QtGui.QIcon.Off)
        self.sel_none_action = QtWidgets.QAction(icon_sel_none, "Unselect (hide) all 2D scans on the viewer", None)
        self.toolbar.addAction(self.sel_none_action)
        self.sel_none_action.triggered.connect(self.show_none)

        icon_histo = QtGui.QIcon()
        icon_histo.addPixmap(QtGui.QPixmap(":/icons/Icon_Library/Histogram.png"), QtGui.QIcon.Normal,
                             QtGui.QIcon.Off)
        self.histo_action = QtWidgets.QAction(icon_histo, "Show (hide) histograms", None)
        self.toolbar.addAction(self.histo_action)
        self.histo_action.setCheckable(True)
        self.histo_action.triggered.connect(self.show_histo)

    @Slot(float, float)
    def move_at(self, posx, posy):
        if self.moveat_action.isChecked():
            self.sig_double_clicked.emit(posx, posy)

    def show_histo(self):
        show_state = self.histo_action.isChecked()
        self.viewer.histo_widget.setVisible(show_state)

    def show_all(self):
        self.show_scans()

    def show_none(self):
        self.show_scans(False)

    def show_scans(self, show=True):
        for child in self.settings.child('scans'):
            val = child.value()
            val['checked'] = show
            child.setValue(val)
            child.sigValueChanged.emit(child, val)

    def list_2Dscans(self):
        try:
            scans = self.h5module.get_h5file_scans()
            # settings=[dict(scan_name=node._v_name,path=node._v_pathname, pixmap=nparray2Qpixmap(node.read()))),...]
            params = []
            for child in self.settings.child(('scans')).children():
                if 'Scan' in child.name():
                    self.settings.child(('scans')).removeChild(child)
            for scan in scans:
                params.append({'name': scan['scan_name'], 'type': 'pixmap_check',
                               'value': dict(data=scan['data'], checked=False, path=scan['path'],
                                             info=scan['scan_name'])})
            self.settings.child(('scans')).addChildren(params)

            for child in self.settings.child(('scans')).children():
                val = child.value()
                val.update(dict(checked=True))
                child.setValue(val)
                child.sigValueChanged.emit(child, child.value())

        except Exception as e:
            logger.exception(str(e))

    def load_image(self):
        # image_filepath = str(utils.select_file(start_path=None, save=False, ext='h5'))
        data, fname, node_path = browse_data(ret_all=True)
        if data is not None and fname != '':
            self.h5module_image = H5BrowserUtil()
            self.h5module_image.open_file(fname, 'a')
            node = self.h5module_image.get_node(node_path)
            pixmaps = self.h5module_image.get_h5file_scans(node.parent_node)

            self.settings.child('settings', 'imagepath').setValue(fname)
            other_child = [child for child in self.settings.child(('scans')).children() if 'Scan' not in child.name()]
            if len(other_child) >= 1:
                for child in other_child:
                    self.settings.child(('scans')).removeChild(child)
            params = []
            for pixmap in pixmaps:
                params.append({'name': pixmap['scan_name'], 'type': 'pixmap_check',
                               'value': dict(data=pixmap['data'], checked=False, path=pixmap['path'])})
            self.settings.child(('scans')).addChildren(params)

            val = self.settings.child('scans', pixmaps[0]['scan_name']).value()
            val.update(dict(checked=True))
            self.settings.child('scans', pixmaps[0]['scan_name']).setValue(val)
            self.settings.child('scans', pixmaps[0]['scan_name']).sigValueChanged.emit(
                self.settings.child('scans', pixmaps[0]['scan_name']),
                self.settings.child('scans', pixmaps[0]['scan_name']).value())

    def load_data(self):
        self.h5module_path = str(
            select_file(start_path=config('data_saving', 'h5file', 'save_path'),
                                                            save=False, ext='h5'))
        if self.h5module_path != '':
            self.settings.child('settings', 'filepath').setValue(self.h5module_path)
            if self.h5module is not None:
                if self.h5module.isopen():
                    self.h5module.close_file()
            self.h5module.open_file(self.h5module_path, 'a')
            self.list_2Dscans()

    def set_aspect_ratio(self):
        if self.action_ratio.isChecked():
            self.viewer.image_widget.plotitem.vb.setAspectLocked(lock=True, ratio=1)
        else:
            self.viewer.image_widget.plotitem.vb.setAspectLocked(lock=False, ratio=1)

    def settings_changed(self, param, changes):
        """
            Check for changes in the given (parameter,change,information) tuple list.
            In case of value changed, update the DAQscan_settings tree consequently.

            =============== ============================================ ==============================
            **Parameters**    **Type**                                     **Description**
            *param*           instance of pyqtgraph parameter              the parameter to be checked
            *changes*         (parameter,change,information) tuple list    the current changes state
            =============== ============================================ ==============================
        """
        for param, change, data in changes:
            path = self.settings.childPath(param)
            if path is not None:
                childName = '.'.join(path)
            else:
                childName = param.name()
            if change == 'childAdded':
                pass

            elif change == 'value':
                if param.parent().name() == 'scans':
                    if data['checked']:
                        try:
                            if 'Scan' in param.name():
                                h5module = self.h5module
                                nodes = [node for node in h5module.walk_nodes(data['path'])]
                            else:
                                h5module = self.h5module_image
                                nodes = [h5module.get_node(data['path'])]
                            ind = 0
                            for node in nodes:
                                flag = False
                                if 'type' in node.attrs.attrs_name and 'data_dimension' in node.attrs.attrs_name:
                                    if 'scan_type' in node.attrs.attrs_name:
                                        if node.attrs['scan_type'] == 'scan2D' and node.attrs['data_dimension'] == '0D':
                                            # 2d scan of 0D data
                                            flag = True
                                        elif node.attrs['scan_type'] == '' and node.attrs['data_dimension'] == '2D':
                                            # image data (2D) with no scan
                                            flag = True

                                if flag:
                                    isadaptive = 'adaptive' in node.attrs['scan_subtype'].lower()
                                    if isadaptive:
                                        im = SpreadImageItem()
                                    else:
                                        im = UniformImageItem()
                                    im.setOpacity(1)
                                    # im.setOpts(axisOrder='row-major')
                                    self.viewer.image_widget.plotitem.addItem(im)
                                    im.setCompositionMode(QtGui.QPainter.CompositionMode_Plus)

                                    if 'Scan' in param.name():
                                        if isadaptive:
                                            x_axis = h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                                       utils.capitalize('scan_x_axis')).read()
                                            y_axis = h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                                       utils.capitalize('scan_y_axis')).read()
                                        else:
                                            x_axis = np.unique(
                                                h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                                  utils.capitalize('scan_x_axis')).read())
                                            y_axis = np.unique(
                                                h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                                  utils.capitalize('scan_y_axis')).read())
                                    else:
                                        x_axis = np.unique(
                                            h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                              utils.capitalize('x_axis')).read())
                                        y_axis = np.unique(
                                            h5module.get_node(h5module.get_node(data['path']).parent_node,
                                                              utils.capitalize('y_axis')).read())
                                    if not isadaptive:
                                        rect = QtCore.QRectF(np.min(x_axis), np.min(y_axis),
                                                             (np.max(x_axis) - np.min(x_axis)),
                                                             (np.max(y_axis) - np.min(y_axis)))
                                        im.setOpts(rescale=rect)
                                        im.setImage(node.read())
                                    else:
                                        im.setImage(np.vstack((x_axis, y_axis, node.read())).T)

                                    if ind == 0:
                                        # im.setLookupTable(colors_red)
                                        self.viewer.histogram_red.setImageItem(im)
                                        if not self.viewer.histogram_red.isVisible():
                                            self.viewer.histogram_red.setVisible(True)
                                    elif ind == 1:
                                        # im.setLookupTable(colors_green)
                                        self.viewer.histogram_green.setImageItem(im)
                                        if not self.viewer.histogram_green.isVisible():
                                            self.viewer.histogram_green.setVisible(True)
                                    else:
                                        # im.setLookupTable(colors_blue)
                                        self.viewer.histogram_blue.setImageItem(im)
                                        if not self.viewer.histogram_blue.isVisible():
                                            self.viewer.histogram_blue.setVisible(True)

                                    self.overlays.append(dict(name='{:s}_{:03d}'.format(param.name(), ind), image=im))

                                    ind += 1
                            # self.viewer.image_widget.view.autoRange()
                        except Exception as e:
                            logger.exception(str(e))

                    else:
                        for overlay in self.overlays[:]:
                            if param.name() in overlay['name']:
                                ind = self.overlays.index(overlay)
                                self.viewer.image_widget.plotitem.removeItem(overlay['image'])
                                self.overlays.pop(ind)

            elif change == 'parent':
                for overlay in self.overlays[:]:
                    if param.name() in overlay['name']:
                        ind = self.overlays.index(overlay)
                        self.viewer.image_widget.plotitem.removeItem(overlay['image'])
                        self.overlays.pop(ind)

    def setupUI(self):
        self.ui = QObject()
        layout = QtWidgets.QVBoxLayout()

        self.parent.setLayout(layout)
        sett_widget = QtWidgets.QWidget()
        self.sett_layout = QtWidgets.QHBoxLayout()
        sett_widget.setLayout(self.sett_layout)
        # creating a toolbar
        self.toolbar = QtWidgets.QToolBar()
        self.create_toolbar()
        layout.addWidget(self.toolbar)

        splitter = QtWidgets.QSplitter(Qt.Horizontal)
        layout.addWidget(splitter)

        # set viewer area
        widg = QtWidgets.QWidget()
        # self.viewer = Viewer2D(widg)
        self.viewer = Viewer2DBasic(widg)
        self.viewer.histogram_red.setVisible(False)
        self.viewer.histogram_green.setVisible(False)
        self.viewer.histogram_blue.setVisible(False)
        self.viewer.sig_double_clicked.connect(self.move_at)

        # displaying the scan list tree
        self.settings_tree = ParameterTree()
        # self.settings_tree.setMaximumWidth(300)
        self.settings_tree.setMinimumWidth(300)
        # self.settings_tree.setVisible(False)
        params_scan = [
            {'title': 'Settings', 'name': 'settings', 'type': 'group', 'children': [
                {'title': 'Load h5:', 'name': 'Load h5', 'type': 'action'},
                {'title': 'h5 path:', 'name': 'filepath', 'type': 'str', 'value': '', 'readonly': True},
                {'title': 'Load Image:', 'name': 'Load Image', 'type': 'action'},
                {'title': 'Image path:', 'name': 'imagepath', 'type': 'str', 'value': '', 'readonly': True},
            ]},
            {'title': 'Scans', 'name': 'scans', 'type': 'group', 'children': []},
        ]
        self.settings = Parameter.create(name='settings', type='group', children=params_scan)
        self.settings_tree.setParameters(self.settings, showTop=False)
        self.settings.sigTreeStateChanged.connect(self.settings_changed)

        self.settings.child('settings', 'Load h5').sigActivated.connect(self.load_data)
        self.settings.child('settings', 'Load Image').sigActivated.connect(self.load_image)

        self.ui.statusbar = QtWidgets.QStatusBar()
        self.ui.statusbar.setMaximumHeight(25)
        layout.addWidget(self.ui.statusbar)
        self.ui.log_message = QtWidgets.QLabel('Initializing')
        self.ui.statusbar.addPermanentWidget(self.ui.log_message)

        self.sett_layout.addWidget(self.settings_tree)
        splitter.addWidget(sett_widget)
        splitter.addWidget(self.viewer.parent)

    def show_image(self, data):
        """

        Parameters
        ----------
        data: (dict) with keys 'names', 'data', 'x_axis', 'y_axis', 'pixmap2D'
        """

        if self.h5module is None:
            scan_path, current_filename, dataset_path = H5Saver.set_current_scan_path(navigator_path, base_name='Scan',
                                                                                      update_h5=True,
                                                                                      next_scan_index=self.next_scan_index,
                                                                                      create_scan_folder=False)
            self.h5module = H5BrowserUtil()
            self.h5module.open_file(str(dataset_path.joinpath(dataset_path.name + ".h5")), 'w')

        else:
            scan_path, current_filename, dataset_path = H5Saver.set_current_scan_path(navigator_path, base_name='Scan',
                                                                                      update_h5=False,
                                                                                      next_scan_index=self.next_scan_index,
                                                                                      create_scan_folder=False)
            if not self.h5module.isopen():
                self.h5module.open_file(str(dataset_path.joinpath(dataset_path.name + ".h5")), 'a')

        h5group = self.h5module.root()
        data2D_group = self.h5module.get_set_group(h5group, 'Data2D')
        data2D_group.attrs.type = 'data2D'

        self.next_scan_index += 1
        curr_group = self.h5module.get_set_group('/Data2D', current_filename)
        live_group = self.h5module.get_set_group(curr_group, 'Live_scan_2D')
        live_group.attrs['pixmap2D'] = data['pixmap2D']

        xdata = data['x_axis']
        if isinstance(xdata, dict):
            xdata = xdata['data']
        xarray = self.h5module.create_carray(curr_group, "Scan_x_axis", obj=xdata,
                                             title=current_filename)
        xarray.attrs['type'] = 'navigation_axis'
        xarray.attrs['data_dimension'] = '1D'
        xarray.attrs['nav_index'] = 0

        ydata = data['y_axis']
        if isinstance(ydata, dict):
            ydata = ydata['data']
        yarray = self.h5module.create_carray(curr_group, "Scan_y_axis", obj=ydata,
                                             title=current_filename)
        yarray.attrs['type'] = 'navigation_axis'
        yarray.attrs['data_dimension'] = '1D'
        yarray.attrs['nav_index'] = 1

        for ind_channel, name in enumerate(data['names']):
            try:
                channel_group = self.h5module.get_set_group(live_group, name)
                channel_group.attrs.Channel_name = name
                array = self.h5module.create_carray(channel_group, current_filename + '_' + name,
                                                    obj=data['data'][ind_channel],
                                                    title='data', )
                array.attrs['type'] = 'data'
                array.attrs['data_dimension'] = '0D'
                array.attrs['data_name'] = name
                array.attrs['scan_type'] = 'scan2D'
                array.attrs['scan_subtype'] = ''
            except Exception as e:
                logger.exception(str(e))

        self.update_2Dscans()

    def update_2Dscans(self):
        try:
            if not self.h5module.isopen():
                self.h5module.open_file(self.h5module.file_path, 'a')
            scans = self.h5module.get_h5file_scans(self.h5module.root())
            # settings=[dict(scan_name=node._v_name,path=node._v_pathname, pixmap=nparray2Qpixmap(node.read()))),...]
            params = []
            children = [child.name() for child in self.settings.child(('scans')).children()]
            for scan in scans:
                if scan['scan_name'] not in children:
                    params.append({'name': scan['scan_name'], 'type': 'pixmap_check',
                                   'value': dict(data=scan['data'], checked=False, path=scan['path'],
                                                 info=scan['scan_name'])})
            self.settings.child(('scans')).addChildren(params)

            for child in self.settings.child(('scans')).children():
                if child.name() not in children:
                    val = child.value()
                    val.update(dict(checked=True))
                    child.setValue(val)
                    child.sigValueChanged.emit(child, child.value())

        except Exception as e:
            logger.exception(str(e))

    def update_h5file(self, h5file):
        if self.h5module is not None:
            self.h5module.h5file = h5file
        self.update_2Dscans()

    def update_status(self, txt, status_time=0, log_type=None):
        """
            Show the txt message in the status bar with a delay of status_time ms.

            =============== =========== =======================
            **Parameters**    **Type**    **Description**
            *txt*             string      The message to show
            *status_time*       int         the delay of showing
            *log_type*        string      the type of the log
            =============== =========== =======================
        """
        try:
            self.ui.statusbar.showMessage(txt, status_time)
            logger.info(txt)
        except Exception as e:
            logger.exception(str(e))


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    widg = QtWidgets.QWidget()
    prog = Navigator(widg)
    widg.show()
    sys.exit(app.exec_())
