import pyqtgraph as pg
from pyqtgraph.Qt import  QtGui
import numpy as np

from PyQt5.QtWidgets import *
from PyQt5.QtCore import *

from sortedcontainers import SortedSet


# ShotSelector class for choosing which the hdf5 files are import for plotting
class ShotSelector(pg.LayoutWidget):
    valueChanged = pyqtSignal()
    selectionChanged = pyqtSignal()

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

        self.nshots = 1

        self.setFixedHeight(100)

        self.current_idx_le = QLineEdit(self)

        self.current_idx_le.setMaximumWidth(30)
        self.current_idx_le.setValidator(QtGui.QIntValidator())

        self.current_idx_le.setText(str(-1))

        self.addWidget(self.current_idx_le)

        self.slider = QSlider(self)

        self.slider.setPageStep(1)
        self.slider.setOrientation(Qt.Horizontal)

        self.addWidget(self.slider, colspan=2)

        self.nextRow()

        self.addWidget(QLabel('index selector'))

        self.idx_select_le = QLineEdit(self)
        self.idx_select_le.setText(':')
        self.addWidget(self.idx_select_le)

        self.warning = QLabel()
        self.warning.setMargin(5)
        self.update_warning()
        self.addWidget(self.warning)

        self.update_nshots(self.nshots)

        self.idx_select_le.editingFinished.connect(self.update_selection)
        self.current_idx_le.editingFinished.connect(self.setSliderValue)
        self.slider.valueChanged.connect(self.setLabelValue)

    def update_nshots(self, nshots):
        self.nshots = nshots
        self.idx = np.arange(self.nshots)
        self.slider.setRange(0, self.nshots - 1)

        self.update_selection()
        self.setSliderValue()

    def update_warning(self, warning=''):
        if warning == '':
            self.warning.setStyleSheet("background-color: lightgreen")
            warning = 'all good'
        else:
            self.warning.setStyleSheet("background-color: red")

        self.warning.setText(warning)

    def update_selection(self):
        self.update_warning()

        slice_text = self.idx_select_le.text()
        slices = slice_text.split(',')

        self.idx_selected = SortedSet([])
        for s in slices:
            try:
                scope = locals()

                select = eval('self.idx[' + s + ']', scope)

                if isinstance(select, np.ndarray):
                    for idx in select:
                        self.idx_selected.add(idx)
                else:
                    self.idx_selected.add(select)

            except:
                self.update_warning('problem in selected indeces')
                return 0

        self.slider.setRange(0, len(self.idx_selected) - 1)

        if int(self.current_idx_le.text()) % self.nshots not in self.idx_selected:
            self.current_idx_le.setText(self.idx_selected[-1])

            self.update_warning('last index not in selection <br> -> setting last selected')

        self.selectionChanged.emit()

    def setLabelValue(self, value):
        newval = self.idx_selected[value]

        if newval != self.get_current_index():
            self.current_idx_le.setText(str(newval))

            self.valueChanged.emit()

    def setSliderValue(self):
        self.update_warning()

        value = int(self.current_idx_le.text())

        try:
            value_sl = self.idx_selected.index(value % len(self.idx))
            self.slider.setValue(value_sl)
        except ValueError:
            self.update_warning('set index not in selection <br> ignore')

    def get_current_index(self):
        return int(self.current_idx_le.text()) % self.nshots

    def get_selected_indices(self):
        return (np.array(self.idx_selected),)