from itertools import cycle

from numpy import array, nan, isnan
from pyqtgraph import mkPen, mkBrush

from nmreval.dsc.hodge import tau_hodge
from nmreval.lib.colors import Tab10
from ..Qt import QtWidgets, QtCore
from .._py.tnmh_dialog import Ui_DSCEvalDialog
from ..lib.pg_objects import PlotItem, RegionItem

from nmreval.data import DSC, Points


class TgCalculator(QtWidgets.QWizard, Ui_DSCEvalDialog):
    newData = QtCore.pyqtSignal(dict, str)

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

        self.setupUi(self)

        self._management = management
        self._colors = cycle(Tab10)

        self._dsc = {}
        self._plots = {}
        self._tg_value = {}
        self._fit = {}
        self._hodge = {
            'onset': (
                PlotItem(x=[], y=[], pen=None, symbol='o', symbolBrush=Tab10.TabBlue.rgb(), name='Onset'),
                None,
                (PlotItem(x=[], y=[], pen=mkPen({'color': Tab10.TabBlue.rgb()})),
                 PlotItem(x=[], y=[], pen=None, symbol='o', symbolBrush=Tab10.TabBlue.rgb())),
                None,
            ),
            'midpoint': (
                PlotItem(x=[], y=[], pen=None, symbol='s', symbolBrush=Tab10.TabOrange.rgb(), name='Midpoint'),
                None,
                (PlotItem(x=[], y=[], pen=mkPen({'color': Tab10.TabOrange.rgb()})),
                 PlotItem(x=[], y=[], pen=None, symbol='s', symbolBrush=Tab10.TabOrange.rgb())),
                None,
            ),
            'end': (
                PlotItem(x=[], y=[], pen=None, symbol='t', symbolBrush=Tab10.TabGreen.rgb(), name='End'),
                None,
                (PlotItem(x=[], y=[], pen=mkPen({'color': Tab10.TabGreen.rgb()})),
                 PlotItem(x=[], y=[], pen=None, symbol='t', symbolBrush=Tab10.TabGreen.rgb())),
                None,
            ),
            'inflection': (
                PlotItem(x=[], y=[], pen=None, symbol='d', symbolBrush=Tab10.TabRed.rgb(), name='Inflection'),
                None,
                (PlotItem(x=[], y=[], pen=mkPen({'color': Tab10.TabRed.rgb()})),
                 PlotItem(x=[], y=[], pen=None, symbol='d', symbolBrush=Tab10.TabRed.rgb())),
                None,
            ),
            'fictive': (
                PlotItem(x=[], y=[], pen=None, symbol='t1', symbolBrush=Tab10.TabPurple.rgb(), name='Fictive'),
                None,
                (PlotItem(x=[], y=[], pen=mkPen({'color': Tab10.TabPurple.rgb()})),
                 PlotItem(x=[], y=[], pen=None, symbol='t1', symbolBrush=Tab10.TabPurple.rgb())),
                None,
            ),
        }
        self._lines = {}
        self.tau_plot.getPlotItem().addLegend()
        for plt, _, fitplt, _ in self._hodge.values():
            self.tau_plot.addItem(plt)
            self.tghodge_graph.addItem(fitplt[0])
            self.tghodge_graph.addItem(fitplt[1])
        self.tau_plot.setLogMode(y=True)
        self.tghodge_graph.setLogMode(y=True)

        self.limits = RegionItem(), RegionItem()
        for lim in self.limits:
            self.dsc_plot.addItem(lim)
        self._limitless = True

        self.add_sets()

        self.tnmh_graph_check.stateChanged.connect(lambda state: self.tnmh_graph_combo.setEnabled(not bool(state)))
        self.hodge_graph_check.stateChanged.connect(lambda state: self.hodge_graph_combo.setEnabled(not bool(state)))

        self.next_button.clicked.connect(lambda: self.stackedWidget.setCurrentIndex((self.stackedWidget.currentIndex() + 1) % 3))
        self.back_button.clicked.connect(lambda: self.stackedWidget.setCurrentIndex((self.stackedWidget.currentIndex() + 2) % 3))

        self.listWidget.itemChanged.connect(self.change_visibility)

    def __call__(self):
        self.clear()
        self._colors = cycle(Tab10)
        self.add_sets()

        return self

    def clear(self):
        self.listWidget.clear()
        self.tg_tree.clear()
        self.tnmh_tree.clear()

        for plots in self._plots.values():
            for val in plots:
                self.dsc_plot.removeItem(val)
                self.tnmh_graphics.removeItem(val)

        for key, plt in self._hodge.items():
            plt[0].setData(x=[], y=[])
            plt[2][0].setData(x=[], y=[])
            plt[2][1].setData(x=[], y=[])
            self._hodge[key] = (plt[0], None, plt[2], None)

        self._dsc = {}
        self._plots = {}
        self._tg_value = {}
        self._lines = {}
        self._fit = {}

        self.stackedWidget.setCurrentIndex(0)

    def add_sets(self):
        for w in (self.tnmh_graph_combo, self.hodge_graph_combo):
            w.clear()
            for graphs in self._management.graphs.list():
                w.addItem(graphs[1], userData=graphs[0])

        min_x = 10_000_000
        max_x = -10_000_000
        for (key, name), c in zip(self._management.active_sets, self._colors):
            data = self._management[key].data
            if not isinstance(data, DSC):
                continue

            min_x = min(min_x, data.x.min())
            max_x = max(max_x, data.x.max())

            item = QtWidgets.QListWidgetItem(name)
            item.setCheckState(QtCore.Qt.CheckState.Checked)
            item.setData(QtCore.Qt.ItemDataRole.UserRole, key)
            item.setForeground(mkBrush(c.rgb()))
            self.listWidget.addItem(item)

            self._dsc[key] = (data, None)

            data_plot = PlotItem(x=data.x, y=data.y, pen=mkPen(c.rgb()))
            self.dsc_plot.addItem(data_plot)

            glass = PlotItem()
            glass.set_line(style=2, color=c)
            self.dsc_plot.addItem(glass)

            liquid = PlotItem()
            liquid.set_line(style=2, color=c)
            self.dsc_plot.addItem(liquid)

            tangent = PlotItem()
            tangent.set_line(style=2, color=c)
            self.dsc_plot.addItem(tangent)

            tg_plot = PlotItem(pen=None, symbolBrush=c.rgb(), symbol='o')
            self.dsc_plot.addItem(tg_plot)

            fictive_cp = PlotItem(pen=mkPen(c.rgb()))
            self.tnmh_graphics.addItem(fictive_cp)

            tnmh_fit = PlotItem()
            tnmh_fit.set_line(style=2, color=c)
            self.tnmh_graphics.addItem(tnmh_fit)

            self._plots[key] = (data_plot, tg_plot, glass, liquid, tangent, fictive_cp, tnmh_fit)
            self._tg_value[key] = {
                'onset': (nan, nan),
                'midpoint': (nan, nan),
                'end': (nan, nan),
                'inflection': (nan, nan),
                # 'fictive': (nan, nan),
            }

        if self._limitless and max_x != -10000000 and min_x != 10000000 :
            dist = max_x - min_x
            self.limits[0].setRegion((min_x, min_x+min(0.1*dist, 5)))
            self.limits[1].setRegion((max_x-min(5, 0.1*dist), max_x))
            self._limitless = False

    @QtCore.pyqtSlot(name='on_calctg_button_clicked')
    def calc_tg(self):
        baselines = tuple(lim.getRegion() for lim in self.limits)
        if baselines[0][0] > baselines[1][0]:
            baselines = baselines[1], baselines[0]

        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)
            if item.checkState() == QtCore.Qt.CheckState.Unchecked:
                continue

            key = item.data(QtCore.Qt.ItemDataRole.UserRole)
            plot = self._plots[key]
            data, _ = self._dsc[key]

            tg_results, glass, liquid, tangent = data.glass_transition(*baselines)
            self._lines[key] = (glass, liquid, tangent)

            for i, line in enumerate((glass, liquid, tangent)):
                plot[i+2].setData(x=line.x, y=line.y)

            self._tg_value[key].update(tg_results)

        self._update_tg_plots()

    def _update_tg_plots(self):
        self.tg_tree.clear()
        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)

            tree_item = QtWidgets.QTreeWidgetItem([item.text()])
            values = self._tg_value.get(item.data(QtCore.Qt.ItemDataRole.UserRole))

            if values is not None:
                for name, pos in values.items():
                    child_item = QtWidgets.QTreeWidgetItem([f'{name.capitalize()}: {pos[0]:.2f} K'])
                    tree_item.addChild(child_item)

                self.tg_tree.addTopLevelItem(tree_item)

            key = item.data(QtCore.Qt.ItemDataRole.UserRole)
            plot = self._plots[key]
            data, _ = self._dsc[key]

            plot[1].setData(array(list(self._tg_value[key].values())))

    @QtCore.pyqtSlot(name='on_tg_export_button_clicked')
    def export_tg(self):
        ret_dic = {}
        for key, tg in self._tg_value.items():
            tgx = [x for x, y in tg.values()]
            tgy = [y for x, y in tg.values()]
            if self.tg_export_check.isChecked():
                tg_pts = Points(x=tgx, y=tgy, name=self._management[key].name + ' (Tg)', value=self._management[key].value)
            else:
                tg_pts = None

            line = []
            if self.tglines_export_check.isChecked():
                line = self._lines.get(key, [])

            ret_dic[key] = (tg_pts, line)

        self.newData.emit(ret_dic, 'tg')

    @QtCore.pyqtSlot(QtWidgets.QListWidgetItem)
    def change_visibility(self, item: QtWidgets.QListWidgetItem):
        is_checked = bool(item.checkState())
        plot = self._plots[item.data(QtCore.Qt.ItemDataRole.UserRole)]
        for val in plot:
            val.setVisible(is_checked)

    def get_fictive(self, key, baselines):
        plot = self._plots[key]
        data, _ = self._dsc[key]

        cp, tg = data.get_fictive_cp(*baselines)

        plot[5].setData(cp.x, cp.y)
        self._dsc[key] = (data, cp)

        return cp

    @QtCore.pyqtSlot(name='on_tnhm_fitbutton_clicked')
    def make_tnmh(self):
        baselines = tuple(lim.getRegion() for lim in self.limits)
        if baselines[0][0] > baselines[1][0]:
            baselines = baselines[1], baselines[0]

        self.tnmh_tree.clear()
        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)
            if item.checkState() == QtCore.Qt.CheckState.Unchecked:
                continue

            key = item.data(QtCore.Qt.ItemDataRole.UserRole)

            data = self.get_fictive(key, baselines)

            res = data.calculate_tnmh([60, 0.5, 1, 2e5], *baselines, return_fictive=False)
            self._fit[key] = res

            plot = self._plots[key]
            plot[-1].setData(res.x, res.y)

        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)

            tree_item = QtWidgets.QTreeWidgetItem([item.text()])
            values = self._fit.get(item.data(QtCore.Qt.ItemDataRole.UserRole))

            if values is not None:
                child_item = QtWidgets.QTreeWidgetItem([values.parameter_string()])
                tree_item.addChild(child_item)

                self.tnmh_tree.addTopLevelItem(tree_item)

    @QtCore.pyqtSlot(name='on_tnmh_export_button_clicked')
    def export_tnmh(self):
        ret_dic = {}
        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)
            if item.checkState() == QtCore.Qt.CheckState.Unchecked:
                continue

            key = item.data(QtCore.Qt.ItemDataRole.UserRole)

            cp = None
            if self.fictive_export_check.isChecked():
                _, cp = self._dsc[key]

            line = None

            if self.tnmhfit_export_check.isChecked():
                line = self._fit.get(key)

            ret_dic[key] = (cp, line)

        ret_dic['graph'] = '' if self.tnmh_graph_check.isChecked() else self.tnmh_graph_combo.currentData()

        self.newData.emit(ret_dic, 'tnmh')

    @QtCore.pyqtSlot(name='on_hodge_button_clicked')
    def hodge(self):
        for tg_type, (plot, data, fitplots, fit) in self._hodge.items():

            m = []
            for idx in range(self.listWidget.count()):
                item = self.listWidget.item(idx)
                if item.checkState() == QtCore.Qt.CheckState.Unchecked:
                    continue

                key = item.data(QtCore.Qt.ItemDataRole.UserRole)
                data, _ = self._dsc[key]
                try:
                    tg_value = self._tg_value[key][tg_type][0]
                    if isnan(tg_value):
                        continue
                except KeyError:
                    continue
                    
                m.append([tg_value, data.value])

            if len(m) > 1:
                data, fit = tau_hodge(*array(m).T)
                data.name = f'{data.name} ({tg_type.capitalize()})'
                plot.setData(data.x, data.y)
                fitplots[0].setData(fit.x, fit.y)
                fitplots[1].setData(fit.x_data, fit.y_data)

                self._hodge[tg_type] = (plot, data, fitplots, fit)

    @QtCore.pyqtSlot(name='on_export_hodge_button_clicked')
    def export_hodge(self):
        ret_dic = {}

        for cb in (self.onset_check, self.mid_check, self.end_check, self.inflection_check, self.fictive_check):
            if cb.isChecked():
                item = cb.text().lower()
                data = self._hodge.get(item)
                if data[1] is not None:
                    ret_dic[item] = data[1]

        ret_dic['graph'] = '' if self.hodge_graph_check.isChecked() else self.hodge_graph_combo.currentData()

        self.newData.emit(ret_dic, 'hodge')

    def close(self) -> bool:
        self.clear()
        return super().close()
