from __future__ import print_function
from typing import List, Dict
import json
import os
import pathlib

from qtpy import QtGui, uic, QtCore, QtWidgets
from pyqtgraph.widgets.SpinBox import SpinBox

from . data_source import RectangularDataSelection

try:
    from chisurf import logging
except:
    import logging
    logging.basicConfig()


class SurfacePlotWidget(QtWidgets.QWidget):

    _selections = list()  # type: List[RectangularDataSelection]
    axis_settings = dict()  # type: Dict[str, Dict[str, float]]

    @property
    def scale_x(self):
        if bool(self.checkBoxLogX.isChecked()):
            return "log"
        else:
            return "lin"

    @scale_x.setter
    def scale_x(self, v):
        if v == "log":
            self.checkBoxLogX.setChecked(True)
        else:
            self.checkBoxLogX.setChecked(False)

    @property
    def scale_y(self):
        if bool(self.checkBoxLogY.isChecked()):
            return "log"
        else:
            return "lin"

    @scale_y.setter
    def scale_y(self, v):
        if v == "log":
            self.checkBoxLogY.setChecked(True)
        else:
            self.checkBoxLogY.setChecked(False)

    @property
    def scale_z(self):
        if self.checkBoxLogZ.isChecked():
            return "log"
        else:
            return "lin"

    @scale_z.setter
    def scale_z(self, v):
        if v == "log":
            self.checkBoxLogZ.setChecked(True)
        else:
            self.checkBoxLogZ.setChecked(False)

    @property
    def normed_hist_x(self):
        return bool(self.checkBoxNormX.isChecked())

    @property
    def normed_hist_y(self):
        return bool(self.checkBoxNormY.isChecked())

    @property
    def normed_hist_z(self):
        return bool(self.checkBoxNormZ.isChecked())

    @property
    def p1(self):
        idx = self.comboBoxSelX.currentIndex()
        name = self.comboBoxSelX.currentText()
        return idx, str(name)

    @property
    def p2(self):
        idx = self.comboBoxSelY.currentIndex()
        name = self.comboBoxSelY.currentText()
        return idx, str(name)

    @property
    def p3(self):
        idx = self.comboBoxSelZ.currentIndex()
        name = self.comboBoxSelZ.currentText()
        return idx, str(name)

    @property
    def binsX(self):
        return int(self.spinBoxBin1DX.value())

    @property
    def binsY(self):
        return int(self.spinBoxBin1DY.value())

    @property
    def binsZ(self):
        return int(self.spinBoxBin1DZ.value())

    @property
    def bins2X(self):
        return int(self.spinBoxBin2DX.value())

    @property
    def bins2Y(self):
        return int(self.spinBoxBin2DY.value())

    @property
    def n_xhist_1d(self):
        return int(self.spinBoxBin1DX.value())

    @n_xhist_1d.setter
    def n_xhist_1d(self, v):
        self.spinBoxBin1DX.setValue(int(v))

    @property
    def n_yhist_1d(self):
        return int(self.spinBoxBin1DY.value())

    @n_yhist_1d.setter
    def n_yhist_1d(self, v):
        self.spinBoxBin1DY.setValue(int(v))

    @property
    def n_zhist_1d(self):
        return int(self.spinBoxBin1DZ.value())

    @n_zhist_1d.setter
    def n_zhist_1d(self, v):
        self.spinBoxBin1DZ.setValue(int(v))

    @property
    def n_xhist_2d(self):
        return int(self.spinBoxBin2DX.value())

    @n_xhist_2d.setter
    def n_xhist_2d(self, v):
        self.spinBoxBin2DX.setValue(int(v))

    @property
    def n_yhist_2d(self):
        return int(self.spinBoxBin2DY.value())

    @n_yhist_2d.setter
    def n_yhist_2d(self, v):
        self.spinBoxBin2DY.setValue(int(v))

    @property
    def x_range(self):
        return float(self.spinBoxXmin.value()), \
               float(self.spinBoxXmax.value())

    @property
    def xmin(self):
        return self.x_range[0]

    @xmin.setter
    def xmin(self, v):
        self.spinBoxXmin.setValue(v)

    @property
    def xmax(self):
        return self.x_range[1]

    @xmax.setter
    def xmax(self, v):
        self.spinBoxXmax.setValue(v)

    @property
    def y_range(self):
        return float(self.spinBoxYmin.value()), \
               float(self.spinBoxYmax.value())

    @property
    def ymin(self):
        return float(self.spinBoxYmin.value())

    @ymin.setter
    def ymin(self, v):
        self.spinBoxYmin.setValue(v)

    @property
    def ymax(self):
        return float(self.spinBoxYmax.value())

    @ymax.setter
    def ymax(self, v):
        self.spinBoxYmax.setValue(v)

    @property
    def z_range(self):
        return float(self.spinBoxZmin.value()), \
               float(self.spinBoxZmax.value())

    @property
    def zmin(self):
        return float(self.spinBoxZmin.value())

    @zmin.setter
    def zmin(self, v):
        self.spinBoxZmin.setValue(v)

    @property
    def zmax(self):
        return float(self.spinBoxZmax.value())

    @zmax.setter
    def zmax(self, v):
        self.spinBoxZmax.setValue(v)

    def __init__(self, parent=None):
        super(SurfacePlotWidget, self).__init__()
        self.parent = parent
        logging.log(0, "Initializing SurfacePlotWidget")
        #########################
        # GUI
        #########################
        ui_file = pathlib.Path(__file__).parent / 'plot_control.ui'
        self.spinBoxXmin = SpinBox()
        self.spinBoxXmax = SpinBox()
        self.spinBoxYmin = SpinBox()
        self.spinBoxYmax = SpinBox()
        self.spinBoxZmin = SpinBox()
        self.spinBoxZmax = SpinBox()

        uic.loadUi(str(ui_file.as_posix()), self)
        self.horizontalLayout.addWidget(self.spinBoxXmin)
        self.horizontalLayout.addWidget(self.spinBoxXmax)
        self.horizontalLayout_2.addWidget(self.spinBoxYmin)
        self.horizontalLayout_2.addWidget(self.spinBoxYmax)
        self.horizontalLayout_3.addWidget(self.spinBoxZmin)
        self.horizontalLayout_3.addWidget(self.spinBoxZmax)

        # Auto complete for selectors
        self.comboBoxSelX.completer().setCompletionMode(QtWidgets.QCompleter.PopupCompletion)
        self.comboBoxSelX.setInsertPolicy(QtWidgets.QComboBox.NoInsert)
        self.comboBoxSelY.completer().setCompletionMode(QtWidgets.QCompleter.PopupCompletion)
        self.comboBoxSelY.setInsertPolicy(QtWidgets.QComboBox.NoInsert)
        self.comboBoxSelZ.completer().setCompletionMode(QtWidgets.QCompleter.PopupCompletion)
        self.comboBoxSelZ.setInsertPolicy(QtWidgets.QComboBox.NoInsert)
        #########################
        # Actions
        #########################
        # Generic action
        self.actionUpdatePlots.triggered.connect(self.parent.update_plots)

        # Auto range
        self.actionAuto_range_x.triggered.connect(self.onAutoRangeX)
        self.actionAuto_range_x.triggered.connect(self.onUpdate_axis_scales)
        self.actionAuto_range_y.triggered.connect(self.onAutoRangeY)
        self.actionAuto_range_y.triggered.connect(self.onUpdate_axis_scales)
        self.actionAuto_range_z.triggered.connect(self.onAutoRangeZ)
        self.actionAuto_range_z.triggered.connect(self.onUpdate_axis_scales)
        self.actionAuto_range_z.triggered.connect(self.auto_selection_range)
        self.actionUpdate_axis_scales.triggered.connect(self.onUpdate_axis_scales)
        self.actionAuto_range_x.triggered.connect(self.onUpdate_axis_scales)

        # Selection table
        self.actionSelectionTableClicked.triggered.connect(self.onSelectionTableClicked)
        self.actionSave_selection.triggered.connect(self.onSave_selection)
        self.actionLoad_selection.triggered.connect(self.onLoad_selection)
        self.actionClear_Selection.triggered.connect(self.onClearSelection)
        self.actionAdd_Selection.triggered.connect(self.onAddSelection)
        self.actionSave_Burst_IDs.triggered.connect(self.parent.onSaveBurstIDs)

        # Change axis range
        self.spinBoxXmin.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)
        self.spinBoxXmax.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)
        self.spinBoxYmin.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)
        self.spinBoxYmax.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)
        self.spinBoxZmin.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)
        self.spinBoxZmax.sigValueChanged.connect(self.actionUpdate_axis_scales.trigger)

        # Change parameter plotted on axis
        self.actionX_axis_changed.triggered.connect(self.onX_axis_changed)
        self.actionY_axis_changed.triggered.connect(self.onY_axis_changed)
        self.actionZ_axis_changed.triggered.connect(self.onZ_axis_changed)

        # Update axis settings
        self.actionUpdate_x_axis_settings.triggered.connect(self.onUpdate_x_axis_settings)
        self.actionUpdate_y_axis_settings.triggered.connect(self.onUpdate_y_axis_settings)
        self.actionUpdate_z_axis_settings.triggered.connect(self.onUpdate_z_axis_settings)

    def set_axis_settings(self, name, amin, amax, scale, bins_1d, bins_2d):
        self.axis_settings[str(name)] = {
            "n_bins_1d": float(bins_1d),
            "min": float(amin),
            "max": float(amax),
            "scale": str(scale)
        }
        if bins_2d is not None:
            self.axis_settings[str(name)].update(
                {
                    "n_bins_2d": int(bins_2d)
                }
            )
        logging.log(0, f"Axis settings updated for {name}: {self.axis_settings[str(name)]}")

    def onUpdate_x_axis_settings(self):
        logging.log(0, "onUpdate_x_axis_settings")
        self.set_axis_settings(
            self.p1[1],
            self.xmin, self.xmax,
            self.scale_x,
            self.n_xhist_1d,
            self.n_xhist_2d
        )

    def onUpdate_y_axis_settings(self):
        logging.log(0, "onUpdate_y_axis_settings")
        self.set_axis_settings(
            self.p2[1],
            self.ymin, self.ymax,
            self.scale_y,
            self.n_yhist_1d,
            self.n_yhist_2d
        )

    def onUpdate_z_axis_settings(self):
        logging.log(0, "onUpdate_z_axis_settings")
        self.set_axis_settings(
            self.p3[1],
            self.zmin, self.zmax,
            self.scale_z,
            self.n_zhist_1d,
            None
        )

    def onX_axis_changed(self):
        _, name = self.p1
        if name in self.axis_settings:
            d = self.axis_settings[name]
            self.n_xhist_1d = d.get('n_bins_1d', 50)
            self.n_xhist_2d = d.get('n_bins_2d', 50)
            self.xmin = d.get('min', self.parent.xmin)
            self.xmax = d.get('max', self.parent.xmax)
            self.scale_x = d.get('scale', "lin")
            logging.log(0, f"X axis changed to settings: {d}")
        else:
            logging.log(0, f"X axis settings for {name} not found. Using auto range.")
            self.onAutoRangeX()

        self.parent.update_plots()

    def onY_axis_changed(self):
        _, name = self.p2
        if name in self.axis_settings:
            d = self.axis_settings[name]
            self.n_yhist_1d = d.get('n_bins_1d', 50)
            self.n_yhist_2d = d.get('n_bins_2d', 50)
            self.ymin = d.get('min', self.parent.ymin)
            self.ymax = d.get('max', self.parent.ymax)
            self.scale_y = d.get('scale', "lin")
            logging.log(0, f"Y axis changed to settings: {d}")
        else:
            logging.log(0, f"Y axis settings for {name} not found. Using auto range.")
            self.onAutoRangeY()

        self.parent.update_plots()

    def onZ_axis_changed(self):
        _, name = self.p3
        if name in self.axis_settings:
            d = self.axis_settings[name]
            self.n_zhist_1d = d.get('n_bins_1d', 50)
            self.zmin = d.get('min', self.parent.zmin)
            self.zmax = d.get('max', self.parent.zmax)
            self.scale_z = d.get('scale', "lin")
            logging.log(0, f"Z axis changed to settings: {d}")
        else:
            logging.log(0, f"Z axis settings for {name} not found. Using auto range.")
            self.onAutoRangeZ()

        self.parent.update_plots()

    def onUpdate_axis_scales(self):
        self.parent.g_yplot.set_axis_scale("left", self.scale_y)
        self.parent.g_xplot.set_axis_scale("bottom", self.scale_x)
        self.parent.g_zplot.set_axis_scale("bottom", self.scale_z)
        # self.parent.g_xyplot.set_axis_scale("bottom", self.scale_x)
        # self.parent.g_xyplot.set_axis_scale("left", self.scale_y)
        self.parent.update_plots()
        logging.log(0, "Axis scales updated")

    def auto_selection_range(self):
        z = self.parent.z_values
        m = z.mean()
        sd = z.std()
        self.parent.selection_z.set_range(m - 2 * sd, m + 2 * sd)
        logging.log(0, f"Auto selection range set to: {(m - 2 * sd, m + 2 * sd)}")

    def update(self):
        super(SurfacePlotWidget, self).update()
        self.actionUpdatePlots.blockSignals(True)
        self.actionUpdate_axis_scales.blockSignals(True)
        pn = self.parent.data_source.parameter_names
        self.comboBoxSelX.clear()
        self.comboBoxSelY.clear()
        self.comboBoxSelZ.clear()
        self.comboBoxSelX.addItems(pn)
        self.comboBoxSelY.addItems(pn)
        self.comboBoxSelZ.addItems(pn)
        self.actionUpdatePlots.blockSignals(False)
        self.actionUpdate_axis_scales.blockSignals(False)
        self.actionUpdatePlots.trigger()
        logging.log(0, "Updated parameter selectors and triggered plot update")

    def onClearSelection(self):
        logging.log(0, "onClearSelection")
        self.tableWidget.setRowCount(0)
        self.parent.update_plots()

    def onSave_selection(self):
        logging.log(0, "onSave_selection")
        l = [s.__dict__ for s in self.get_selections()]
        fn = QtWidgets.QFileDialog.getSaveFileName(
            None,
            "Selection JSON",
            self.parent.working_path,
            'All files (*.selection.json)'
        )[0]
        with open(fn, "w") as fp:
            json.dump(l, fp=fp, indent=4)
        logging.log(0, f"Selection saved to file: {fn}")

    def onLoad_selection(self):
        fn = QtWidgets.QFileDialog.getOpenFileName(
            None,
            "Selection JSON",
            self.parent.working_path,
            'All files (*.selection.json)'
        )[0]
        with open(fn, "r") as fp:
            d = json.load(fp)
            for selection in d:
                self.addSelection(
                    selection['parameter_idx'],
                    selection['lower'],
                    selection['upper'],
                    selection['invert'],
                    selection['enabled'],
                    selection['name']
                )
        logging.log(0, f"Selections loaded from file: {fn}")

    def onAutoRangeX(self):
        logging.log(0, "onAutoRangeX")
        self.spinBoxXmin.blockSignals(True)
        self.spinBoxXmax.blockSignals(True)
        self.spinBoxXmin.setValue(self.parent.xmin)
        self.spinBoxXmax.setValue(self.parent.xmax)
        self.spinBoxXmin.blockSignals(False)
        self.spinBoxXmax.blockSignals(False)

    def onAutoRangeY(self):
        logging.log(0, "onAutoRangeY")
        self.spinBoxYmin.blockSignals(True)
        self.spinBoxYmax.blockSignals(True)
        self.spinBoxYmin.setValue(self.parent.ymin)
        self.spinBoxYmax.setValue(self.parent.ymax)
        self.spinBoxYmin.blockSignals(False)
        self.spinBoxYmax.blockSignals(False)

    def onAutoRangeZ(self):
        logging.log(0, "onAutoRangeZ")
        self.spinBoxZmin.blockSignals(True)
        self.spinBoxZmax.blockSignals(True)
        self.spinBoxZmin.setValue(self.parent.zmin)
        self.spinBoxZmax.setValue(self.parent.zmax)
        self.spinBoxZmin.blockSignals(False)
        self.spinBoxZmax.blockSignals(False)

    def onSelectionTableClicked(self):
        logging.log(0, "onSelectionTableClicked")
        row = self.tableWidget.currentRow()
        self.tableWidget.removeRow(row)
        self.parent.update_plots()

    def addSelection(self, idx, xmin, xmax, invert=False, enabled=True, name=""):

        table = self.tableWidget
        row = table.rowCount()
        table.setRowCount(row + 1)

        tmp = QtWidgets.QTableWidgetItem("%s" % name)
        tmp.setFlags(QtCore.Qt.ItemIsEnabled)
        tmp.setData(1, idx)
        table.setItem(row, 0, tmp)

        tmp = QtWidgets.QTableWidgetItem()
        tmp.setData(0, xmin)
        tmp.setFlags(QtCore.Qt.ItemIsEnabled)
        table.setItem(row, 1, tmp)

        tmp = QtWidgets.QTableWidgetItem()
        tmp.setFlags(QtCore.Qt.ItemIsEnabled)
        tmp.setData(0, xmax)
        table.setItem(row, 2, tmp)

        cb_invert_x = QtWidgets.QCheckBox(table)
        table.setCellWidget(row, 3, cb_invert_x)
        cb_invert_x.setChecked(invert)

        cb_enable_x = QtWidgets.QCheckBox(table)
        table.setCellWidget(row, 4, cb_enable_x)
        cb_enable_x.setChecked(enabled)
        self.parent.update_plots()

        # Actions for selection checkbox
        cb_enable_x.stateChanged.connect(self.actionUpdatePlots.trigger)
        cb_invert_x.stateChanged.connect(self.actionUpdatePlots.trigger)
        logging.log(0, f"Added selection for parameter index {idx} with range ({xmin}, {xmax}), invert={invert}, enabled={enabled}")

    def onAddSelection(self):
        idx, name = self.p3
        xsel = self.parent.selection_z.get_range()
        xmin = float(min(xsel))
        xmax = float(max(xsel))
        self.addSelection(idx, xmin, xmax, False, True, name)
        logging.log(0, f"onAddSelection: Added selection for {name} with range ({xmin}, {xmax})")

    def get_selections(self):
        selections = list()
        table = self.tableWidget
        n_rows = int(table.rowCount())
        for r in range(n_rows):
            idx = int(table.item(r, 0).data(1))
            name = str(table.item(r, 0).data(0))
            lower = float(table.item(r, 1).data(0))
            upper = float(table.item(r, 2).data(0))
            invert = bool(table.cellWidget(r, 3).checkState())
            enabled = bool(table.cellWidget(r, 4).checkState())
            selections.append(
                RectangularDataSelection(
                    parameter_idx=idx,
                    lower=lower,
                    upper=upper,
                    invert=invert,
                    enabled=enabled,
                    name=name
                )
            )
        logging.log(0, f"get_selections: Retrieved {len(selections)} selections")
        return selections
