import numpy as np
from pyqtgraph import mkBrush, mkPen

from nmreval.data.points import Points
from nmreval.nmr.relaxation import RelaxationEvaluation
from nmreval.nmr.coupling import *
from nmreval.distributions import *

from ..lib.pg_objects import PlotItem
from ..Qt import QtCore, QtWidgets, QtGui
from .._py.t1dialog import Ui_t1dialog
from ..lib.forms import FormWidget, SelectionWidget
from ..lib.utils import busy_cursor


class QT1Widget(QtWidgets.QDialog, Ui_t1dialog):
    magnitude = {0: 1e6, 1: 1e3, 2: 1}
    temp_conversion = {0: lambda xx: 1000 / xx, 1: lambda xx: xx, 2: lambda xx: 1 / xx}
    time_conversion = {0: lambda yy: yy, 1: lambda yy: 1/yy}

    newData = QtCore.pyqtSignal(list, str)

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

        self.setupUi(self)

        self.connected_figure = ''
        self.name = ''

        self.t1calculator = RelaxationEvaluation()

        self.sd_parameter = []
        self.sdmodels = [Debye, ColeCole, ColeDavidson, KWW, HavriliakNegami, LogGaussian]
        for i in self.sdmodels:
            self.specdens_combobox.addItem(i.name)
        self.specdens_combobox.currentIndexChanged.connect(self.update_specdens)

        self.cp_parameter = []
        self.coupling = [Quadrupolar, Czjzek, HomoDipolar, Constant]
        for i in self.coupling:
            self.coupling_combobox.addItem(i.name)
        self.coupling_combobox.currentIndexChanged.connect(self.update_coupling)

        self.temp_combobox.currentIndexChanged.connect(self.update_dimensions)
        self.t1_combobox.currentIndexChanged.connect(self.update_dimensions)
        self.lineEdit_2.textChanged.connect(lambda: self.determine_minimum(self.interpol_combobox.currentIndex()))
        self.lineEdit_3.textChanged.connect(lambda: self.determine_minimum(self.interpol_combobox.currentIndex()))

        self.conv_x = QT1Widget.temp_conversion[self.temp_combobox.currentIndex()]
        self.conv_y = QT1Widget.time_conversion[self.t1_combobox.currentIndex()]

        self.minimum = (1, np.inf)
        self.min_pos = PlotItem(x=np.array([]), y=np.array([]),
                                symbol='+', symbolBrush=mkBrush(color='r'), symbolPen=mkPen(color='r'), symbolSize=14)
        self.parabola = PlotItem(x=np.array([]), y=np.array([]))

        self.lineEdit_2.setValidator(QtGui.QDoubleValidator())
        self.lineEdit_3.setValidator(QtGui.QDoubleValidator())

        self.freq_combox.currentIndexChanged.connect(lambda x: self.update_model())
        self.freq_spinbox.valueChanged.connect(lambda x: self.update_model())

        self.update_specdens(0)
        self.update_coupling(0)

    def set_graphs(self, graphs: list):
        self.graph_combobox.clear()
        for g in graphs:
            self.graph_combobox.addItem(g[1], userData=g[0])

    def set_data(self, x, y, name=''):
        x = self.conv_x(x)
        y = self.conv_y(y)

        sortidx = np.argsort(x)
        x = x[sortidx]
        y = y[sortidx]

        left_b = max(np.argmin(y)-2, 0)
        right_b = min(np.argmin(y)+3, len(x)-1)

        self.lineEdit_2.blockSignals(True)
        self.lineEdit_2.setText('{:.2f}'.format(x[left_b]))
        self.lineEdit_2.blockSignals(False)
        self.lineEdit_3.blockSignals(True)
        self.lineEdit_3.setText('{:.2f}'.format(x[right_b]))
        self.lineEdit_3.blockSignals(False)

        self.t1calculator.set_data(x, y)

        self.determine_minimum(self.interpol_combobox.currentIndex())
        self.name = name

    @property
    def frequency(self):
        return self.freq_spinbox.value() * QT1Widget.magnitude[self.freq_combox.currentIndex()]

    @QtCore.pyqtSlot(int)
    def update_specdens(self, idx: int):
        self.sd_parameter = []
        while self.verticalLayout_3.count():
            item = self.verticalLayout_3.takeAt(0)
            try:
                item.widget().deleteLater()
            except AttributeError:
                pass

        if self.sdmodels[idx].parameter is not None:
            for name in self.sdmodels[idx].parameter:
                _temp = FormWidget(parent=self, name=name, fixable=True)
                _temp.value = 1
                _temp.setChecked(True)
                _temp.valueChanged.connect(self.update_specdens_parameter)
                _temp.stateChanged.connect(self.update_specdens_parameter)
                self.verticalLayout_3.addWidget(_temp)
                self.sd_parameter.append(_temp)

        self.t1calculator.set_distribution(self.sdmodels[idx])
        self.update_model()

    @QtCore.pyqtSlot()
    def update_specdens_parameter(self):
        new_p = []
        for p in self.sd_parameter:
            new_p.append(p.value)

        self.update_model()

    def update_sddisplay(self, values):
        try:
            for i, v, in enumerate(values):
                self.sd_parameter[i].blockSignals(True)
                self.sd_parameter[i].value = '{:.3g}'.format(round(v, 3))
                self.sd_parameter[i].blockSignals(False)
        except IndexError:
            pass

    @QtCore.pyqtSlot(int)
    def update_coupling(self, idx: int):
        self.cp_parameter = []
        while self.verticalLayout_4.count():
            item = self.verticalLayout_4.takeAt(0)
            try:
                item.widget().deleteLater()
            except AttributeError:
                pass

        curr_coupling = self.coupling[idx]
        if curr_coupling.parameter is not None:
            for i, name in enumerate(curr_coupling.parameter):
                if curr_coupling.unit is not None:
                    u = curr_coupling.unit[i]
                    name = name+'/'+u if u else name
                _temp = FormWidget(parent=self, name=name, fixable=True)
                self.verticalLayout_4.addWidget(_temp)
                _temp.value = 1
                _temp.setChecked(True)
                _temp.valueChanged.connect(self.update_coupling_parameter)
                _temp.stateChanged.connect(self.update_coupling_parameter)
                self.cp_parameter.append(_temp)

        if self.coupling[idx].choice is not None:
            for (name, kw_name, opts) in self.coupling[idx].choice:
                widgt = SelectionWidget(name, kw_name, opts)
                widgt.comboBox.currentIndexChanged.connect(self.update_coupling_parameter)
                self.verticalLayout_4.addWidget(widgt)
                self.cp_parameter.append(widgt)

        self.update_coupling_parameter()

    @QtCore.pyqtSlot()
    def update_coupling_parameter(self):
        new_p = []
        new_kw = {}
        for pp in self.cp_parameter:
            if isinstance(pp, FormWidget):
                new_p.append(pp.value)
            else:
                new_kw.update(pp.value)
        new_coupling = self.coupling[self.coupling_combobox.currentIndex()]
        self.t1calculator.set_coupling(new_coupling, parameter=new_p, keywords=new_kw)

        self.update_model()

    @QtCore.pyqtSlot(int, name='on_temp_combobox_currentIndexChanged')
    @QtCore.pyqtSlot(int, name='on_t1_combobox_currentIndexChanged')
    def update_dimensions(self, idx: int):
        _x = self.t1calculator.x
        _y = self.t1calculator.y

        if self.sender() == self.temp_combobox:
            _x = self.conv_x(self.t1calculator.x)
            self.conv_x = self.temp_conversion[idx]

        else:
            _y = self.conv_y(self.t1calculator.y)
            self.conv_y = self.time_conversion[idx]

        self.set_data(_x, _y, name=self.name)

    @QtCore.pyqtSlot(tuple)
    def t1min_picked(self, pos: tuple):
        self.min_pos.setData(x=[pos[0]], y=[pos[1]])
        self.update_min(self.conv_x(pos[0]), pos[1])

    def update_min(self, x, y):
        self.label_13.setText(f'{x:.4g} K')
        self.label_12.setText(f' {y:.3e} s')
        if y > np.min(self.t1calculator.y):
            self.label_12.setStyleSheet('QLabel {color: red; font-weight: bold}')
        else:
            self.label_12.setStyleSheet('QLabel {}')
        self.min_pos.setData(x=[self.conv_x(x)], y=[self.conv_y(y)])

        self.minimum = x, y
        self.update_model()

    @QtCore.pyqtSlot(int, name='on_interpol_combobox_currentIndexChanged')
    def determine_minimum(self, idx):
        if idx == 0:
            self.checkBox_interpol.setChecked(False)
            self.checkBox_interpol.hide()
            self.frame.hide()
            m, i_func = self.t1calculator.calculate_t1_min(interpolate=None)
        else:
            self.checkBox_interpol.show()
            self.frame.show()
            try:
                m, i_func = self.t1calculator.calculate_t1_min(interpolate=idx,
                                                               trange=(float(self.lineEdit_2.text()),
                                                                       float(self.lineEdit_3.text())))
            except ValueError:
                m, i_func = self.t1calculator.calculate_t1_min(interpolate=None)

        self.update_min(*m)

        if i_func is not None:
            self.parabola.setData(x=self.conv_x(i_func[0]), y=self.conv_y(i_func[1]))
        else:
            self.parabola.setData(x=np.array([]), y=np.array([]))

        return m

    @QtCore.pyqtSlot(name='on_t1min_toolButton_clicked')
    def set_min(self):
        m = self.determine_minimum(self.interpol_combobox.currentIndex())
        self.update_min(x=m[0], y=m[1])

    def update_model(self):
        sd_args, sd_fix = self.get_sd_values()
        cp_args, cp_kwargs, cp_fix = self.get_cp_values()

        if (len(sd_fix)-sum(sd_fix)) + (len(cp_fix)-sum(cp_fix)) > 1:
            QtWidgets.QMessageBox.information(self, 'Too many free variables',
                                              'More than one parameter is variable.\n'
                                              'Only one can be determined from a minimum.')
            return

        notfix = None
        var_idx = 0
        if not all(sd_fix):
            notfix = 'distribution'
            var_idx = sd_fix.index(False)

        if not all(cp_fix):
            notfix = 'coupling'
            var_idx = cp_fix.index(False)

        if not np.isfinite(self.t1calculator.t1min[1]):
            return

        with busy_cursor():
            calc_stretching, mini = self.t1calculator.get_increase(height=self.minimum[1],
                                                                   idx=var_idx, mode=notfix,
                                                                   omega=2*np.pi*self.frequency,
                                                                   dist_parameter=sd_args, prefactor=cp_args,
                                                                   coupling_kwargs=cp_kwargs)

        self.label_t1min.setText(f'{mini:.4g} s')

        if notfix:
            forms = self.sd_parameter if notfix == 'distribution' else self.cp_parameter
            forms[var_idx].blockSignals(True)
            forms[var_idx].value = f'{calc_stretching:.4g}'
            forms[var_idx].blockSignals(False)

    @QtCore.pyqtSlot(name='on_calc_pushButton_clicked')
    def calculate_correlations(self):
        sd_args, _ = self.get_sd_values()
        cp_args, cp_kwargs, _ = self.get_cp_values()
        tau_mode = ['fit', 'peak', 'mean', 'logmean'][self.tau_combox.currentIndex()]
        corr, opts = self.t1calculator.correlation_from_t1(omega=2*np.pi*self.frequency, dist_parameter=sd_args,
                                                           coupling_param=cp_args, coupling_kwargs=cp_kwargs,
                                                           mode=tau_mode, interpolate=self.checkBox_interpol.isChecked())

        name = self.name + '-' + str(self.t1calculator) + '('
        name += ','.join([f'{a:.3g}' for a in sd_args])
        name += ')-' + tau_mode + ' tau'
        new_data = Points(x=self.conv_x(corr[:, 0]), y= corr[:, 1], name=name)
        new_data.update(opts)
        if self.graph_checkbox.isChecked():
            gid = ''
        else:
            gid = self.graph_combobox.currentData()
        self.newData.emit([new_data], gid)

    def get_sd_values(self):
        sd_args = []
        sd_fix = []
        for i, p in enumerate(self.sd_parameter):
            sd_args.append(p.value)
            sd_fix.append(p.isChecked())

        return sd_args, sd_fix

    def get_cp_values(self):
        cp_args = []
        cp_kwargs = {}
        cp_fix = []
        for i, p in enumerate(self.cp_parameter):
            try:
                cp_fix.append(p.isChecked())
                cp_args.append(p.value)
            except AttributeError:
                # SelectionWidget has no isChecked()
                cp_kwargs.update(p.value)

        return cp_args, cp_kwargs, cp_fix

    @QtCore.pyqtSlot(int, name='on_graph_checkbox_stateChanged')
    def changed_state(self, checked):
        self.graph_combobox.setEnabled(checked != QtCore.Qt.Checked)
