# -*- coding: utf-8 -*-
#
# Copyright © keithleygui Project Contributors
# Licensed under the terms of the MIT License
# (see LICENSE.txt for details)

# system imports
from __future__ import division, print_function, absolute_import
import os.path as osp
import pkg_resources as pkgr
import visa
from PyQt5 import QtCore, QtWidgets, uic
from keithley2600 import FETResultTable
import numpy as np

# local imports
from keithleygui.pyqt_labutils import LedIndicator, SettingsWidget, ConnectionDialog
from keithleygui.pyqtplot_canvas import SweepDataPlot
from keithleygui.config.main import CONF

MAIN_UI_PATH = pkgr.resource_filename('keithleygui', 'main.ui')


class SMUSettingsWidget(SettingsWidget):

    SENSE_LOCAL = 0
    SENSE_REMOTE = 1

    def __init__(self, smu_name):
        SettingsWidget.__init__(self)

        self.smu_name = smu_name

        self.sense_type = self.addSelectionField('Sense type:', ['local (2-wire)', 'remote (4-wire)'])
        self.limit_i = self.addDoubleField('Current limit:', 0.1, 'A', limits=[0, 100])
        self.limit_v = self.addDoubleField('Voltage limit:', 200, 'V', limits=[0, 200])

        self.load_defaults()

    def load_defaults(self):

        if CONF.get(self.smu_name, 'sense') == 'SENSE_LOCAL':
            self.sense_type.setCurrentIndex(self.SENSE_LOCAL)
        elif CONF.get(self.smu_name, 'sense') == 'SENSE_REMOTE':
            self.sense_type.setCurrentIndex(self.SENSE_REMOTE)

        self.limit_i.setValue(CONF.get(self.smu_name, 'limiti'))
        self.limit_v.setValue(CONF.get(self.smu_name, 'limitv'))

    def save_defaults(self):

        if self.sense_type.currentIndex() == self.SENSE_LOCAL:
            CONF.set(self.smu_name, 'sense', 'SENSE_LOCAL')
        elif self.sense_type.currentIndex() == self.SENSE_REMOTE:
            CONF.set(self.smu_name, 'sense', 'SENSE_REMOTE')

        CONF.set(self.smu_name, 'limiti', self.limit_i.value())
        CONF.set(self.smu_name, 'limitv', self.limit_v.value())


# noinspection PyArgumentList
class SweepSettingsWidget(SettingsWidget):

    def __init__(self, keithley):
        SettingsWidget.__init__(self)

        self.keithley = keithley
        self.smu_list = list(self.keithley.SMU_LIST)

        while len(self.smu_list) < 2:
            self.smu_list.append('--')

        self.t_int = self.addDoubleField('Integration time:', 0.1, 's', [0.000016, 0.5])
        self.t_settling = self.addDoubleField('Settling time (auto = -1):', -1, 's', [-1, 100])
        self.sweep_type = self.addSelectionField('Sweep type:', ['Continuous', 'Pulsed'])
        self.smu_gate = self.addSelectionField('Gate SMU:', self.smu_list, 0)
        self.smu_drain = self.addSelectionField('Drain SMU:', self.smu_list, 1)

        self.load_defaults()

        self.smu_gate.currentIndexChanged.connect(self.on_smu_gate_changed)
        self.smu_drain.currentIndexChanged.connect(self.on_smu_drain_changed)

    def load_defaults(self):

        self.t_int.setValue(CONF.get('Sweep', 'tInt'))
        self.t_settling.setValue(CONF.get('Sweep', 'delay'))
        self.sweep_type.setCurrentIndex(int(CONF.get('Sweep', 'pulsed')))
        self.smu_gate.setCurrentText(CONF.get('Sweep', 'gate'))
        self.smu_drain.setCurrentText(CONF.get('Sweep', 'drain'))

    def save_defaults(self):
        CONF.set('Sweep', 'tInt', self.t_int.value())
        CONF.set('Sweep', 'delay', self.t_settling.value())
        CONF.set('Sweep', 'gate', self.smu_gate.currentText())
        CONF.set('Sweep', 'drain', self.smu_drain.currentText())

    @QtCore.pyqtSlot(int)
    def on_smu_gate_changed(self, int_smu):
        """Triggered when the user selects a different gate SMU. """

        if int_smu == 0 and len(self.smu_list) < 3:
            self.smu_drain.setCurrentIndex(1)
        elif int_smu == 1 and len(self.smu_list) < 3:
            self.smu_drain.setCurrentIndex(0)

    @QtCore.pyqtSlot(int)
    def on_smu_drain_changed(self, int_smu):
        """Triggered when the user selects a different drain SMU. """

        if int_smu == 0 and len(self.smu_list) < 3:
            self.smu_gate.setCurrentIndex(1)
        elif int_smu == 1 and len(self.smu_list) < 3:
            self.smu_gate.setCurrentIndex(0)


class TransferSweepSettingsWidget(SettingsWidget):

    def __init__(self):
        SettingsWidget.__init__(self)

        self.vg_start = self.addDoubleField('Vg start:', 0, 'V')
        self.vg_stop = self.addDoubleField('Vg stop:', 0, 'V')
        self.vg_step = self.addDoubleField('Vg step:', 0, 'V')
        self.vd_list = self.addListField('Drain voltages:', [-5, -60])
        self.vd_list.setAcceptedStrings(['trailing'])

        self.load_defaults()

    def load_defaults(self):

        self.vg_start.setValue(CONF.get('Sweep', 'VgStart'))
        self.vg_stop.setValue(CONF.get('Sweep', 'VgStop'))
        self.vg_step.setValue(CONF.get('Sweep', 'VgStep'))
        self.vd_list.setValue(CONF.get('Sweep', 'VdList'))

    def save_defaults(self):
        CONF.set('Sweep', 'VgStart', self.vg_start.value())
        CONF.set('Sweep', 'VgStop', self.vg_stop.value())
        CONF.set('Sweep', 'VgStep', self.vg_step.value())
        CONF.set('Sweep', 'VdList', self.vd_list.value())


class OutputSweepSettingsWidget(SettingsWidget):

    def __init__(self):
        SettingsWidget.__init__(self)

        self.vd_start = self.addDoubleField('Vd start:', 0, 'V')
        self.vd_stop = self.addDoubleField('Vd stop:', 0, 'V')
        self.vd_step = self.addDoubleField('Vd step:', 0, 'V')
        self.vg_list = self.addListField('Gate voltages:', [0, -20, -40, -60])

        self.load_defaults()

    def load_defaults(self):

        self.vd_start.setValue(CONF.get('Sweep', 'VdStart'))
        self.vd_stop.setValue(CONF.get('Sweep', 'VdStop'))
        self.vd_step.setValue(CONF.get('Sweep', 'VdStep'))
        self.vg_list.setValue(CONF.get('Sweep', 'VgList'))

    def save_defaults(self):
        CONF.set('Sweep', 'VdStart', self.vd_start.value())
        CONF.set('Sweep', 'VdStop', self.vd_stop.value())
        CONF.set('Sweep', 'VdStep', self.vd_step.value())
        CONF.set('Sweep', 'VgList', self.vg_list.value())


class IVSweepSettingsWidget(SettingsWidget):

    def __init__(self, smu_list):
        SettingsWidget.__init__(self)

        self.v_start = self.addDoubleField('Vd start:', 0, 'V')
        self.v_stop = self.addDoubleField('Vd stop:', 0, 'V')
        self.v_step = self.addDoubleField('Vd step:', 0, 'V')
        self.smu_sweep = self.addSelectionField('Sweep SMU:', smu_list, 0)

        self.load_defaults()

    def load_defaults(self):

        self.v_start.setValue(CONF.get('Sweep', 'VStart'))
        self.v_stop.setValue(CONF.get('Sweep', 'VStop'))
        self.v_step.setValue(CONF.get('Sweep', 'VStep'))
        self.smu_sweep.setCurrentText(CONF.get('Sweep', 'smu_sweep'))

    def save_defaults(self):
        CONF.set('Sweep', 'VStart', self.v_start.value())
        CONF.set('Sweep', 'VStop', self.v_stop.value())
        CONF.set('Sweep', 'VStep', self.v_step.value())
        CONF.set('Sweep', 'smu_sweep', self.smu_sweep.currentText())


# noinspection PyArgumentList
class KeithleyGuiApp(QtWidgets.QMainWindow):
    """ Provides a GUI for transfer and output sweeps on the Keithley 2600."""

    QUIT_ON_CLOSE = True

    def __init__(self, keithley):
        super(self.__class__, self).__init__()
        # load user interface layout from .ui file
        uic.loadUi(MAIN_UI_PATH, self)

        self.keithley = keithley
        self.smu_list = list(self.keithley.SMU_LIST)
        self.sweep_data = None

        # create sweep settings panes
        self.transfer_sweep_settings = TransferSweepSettingsWidget()
        self.output_sweep_settings = OutputSweepSettingsWidget()
        self.iv_sweep_settings = IVSweepSettingsWidget(self.smu_list)
        self.general_sweep_settings = SweepSettingsWidget(self.keithley)

        self.tabWidgetSweeps.widget(0).layout().addWidget(self.transfer_sweep_settings)
        self.tabWidgetSweeps.widget(1).layout().addWidget(self.output_sweep_settings)
        self.tabWidgetSweeps.widget(2).layout().addWidget(self.iv_sweep_settings)
        self.groupBoxSweepSettings.layout().addWidget(self.general_sweep_settings)

        # create tabs for smu settings
        self.smu_tabs = []
        for smu_name in self.smu_list:
            tab = SMUSettingsWidget(smu_name)
            self.tabWidgetSettings.addTab(tab, smu_name)
            self.smu_tabs.append(tab)

        # create plot widget
        self.canvas = SweepDataPlot()
        self.gridLayout2.addWidget(self.canvas)

        # create LED indicator
        self.led = LedIndicator(self)
        self.statusBar.addPermanentWidget(self.led)
        self.led.setChecked(False)

        # create connection dialog
        self.connectionDialog = ConnectionDialog(self, self.keithley, CONF)

        # restore last position and size
        self.restore_geometry()

        # update GUI status and connect callbacks
        self.actionSaveSweepData.setEnabled(False)
        self.connect_ui_callbacks()
        self.on_load_default()
        self.update_gui_connection()

        # connection update timer: check periodically if keithley is connected
        self.connection_status_update = QtCore.QTimer()
        self.connection_status_update.timeout.connect(self.update_gui_connection)
        self.connection_status_update.start(10000)  # 10 sec

    @staticmethod
    def _string_to_vd(string):
        try:
            return float(string)
        except ValueError:
            if 'trailing' in string:
                return 'trailing'
            else:
                raise ValueError('Invalid drain voltage.')

    def closeEvent(self, event):
        if self.QUIT_ON_CLOSE:
            self.exit_()
        else:
            self.hide()

# =============================================================================
# GUI setup
# =============================================================================

    def restore_geometry(self):
        x = CONF.get('Window', 'x')
        y = CONF.get('Window', 'y')
        w = CONF.get('Window', 'width')
        h = CONF.get('Window', 'height')

        self.setGeometry(x, y, w, h)

    def save_geometry(self):
        geo = self.geometry()
        CONF.set('Window', 'height', geo.height())
        CONF.set('Window', 'width', geo.width())
        CONF.set('Window', 'x', geo.x())
        CONF.set('Window', 'y', geo.y())

    def connect_ui_callbacks(self):
        """Connect buttons and menus to callbacks."""
        self.pushButtonRun.clicked.connect(self.on_sweep_clicked)
        self.pushButtonAbort.clicked.connect(self.on_abort_clicked)

        self.actionSettings.triggered.connect(self.connectionDialog.open)
        self.actionConnect.triggered.connect(self.on_connect_clicked)
        self.actionDisconnect.triggered.connect(self.on_disconnect_clicked)
        self.action_Exit.triggered.connect(self.exit_)
        self.actionSaveSweepData.triggered.connect(self.on_save_clicked)
        self.actionLoad_data_from_file.triggered.connect(self.on_load_clicked)
        self.actionSaveDefaults.triggered.connect(self.on_save_default)
        self.actionLoadDefaults.triggered.connect(self.on_load_default)

# =============================================================================
# Measurement callbacks
# =============================================================================

    def apply_smu_settings(self):
        """
        Applies SMU settings to Keithley before a measurement.
        Warning: self.keithley.reset() will reset those settings.
        """
        for tab in self.smu_tabs:

            smu = getattr(self.keithley, tab.smu_name)

            if tab.sense_type.currentIndex() == tab.SENSE_LOCAL:
                smu.sense = smu.SENSE_LOCAL
            elif tab.sense_type.currentIndex() == tab.SENSE_REMOTE:
                smu.sense = smu.SENSE_REMOTE

            lim_i = tab.limit_i.value()
            smu.source.limiti = lim_i
            smu.trigger.source.limiti = lim_i

            lim_v = tab.limit_v.value()
            smu.source.limitv = lim_v
            smu.trigger.source.limitv = lim_v

    @QtCore.pyqtSlot()
    def on_sweep_clicked(self):
        """ Start a transfer measurement with current settings."""

        if self.keithley.busy:
            msg = ('Keithley is currently used by another program. ' +
                   'Please try again later.')
            QtWidgets.QMessageBox.information(self, str('error'), msg)

            return

        self.apply_smu_settings()

        params = dict()

        if self.tabWidgetSweeps.currentIndex() == 0:
            self.statusBar.showMessage('    Recording transfer curve.')
            # get sweep settings
            params['sweep_type'] = 'transfer'
            params['VgStart'] = self.transfer_sweep_settings.vg_start.value()
            params['VgStop'] = self.transfer_sweep_settings.vg_stop.value()
            params['VgStep'] = self.transfer_sweep_settings.vg_step.value()
            params['VdList'] = self.transfer_sweep_settings.vd_list.value()

        elif self.tabWidgetSweeps.currentIndex() == 1:
            self.statusBar.showMessage('    Recording output curve.')
            # get sweep settings
            params['sweep_type'] = 'output'
            params['VdStart'] = self.output_sweep_settings.vd_start.value()
            params['VdStop'] = self.output_sweep_settings.vd_stop.value()
            params['VdStep'] = self.output_sweep_settings.vd_step.value()
            params['VgList'] = self.output_sweep_settings.vg_list.value()

        elif self.tabWidgetSweeps.currentIndex() == 2:
            self.statusBar.showMessage('    Recording IV curve.')
            # get sweep settings
            params['sweep_type'] = 'iv'
            params['VStart'] = self.iv_sweep_settings.v_start.value()
            params['VStop'] = self.iv_sweep_settings.v_stop.value()
            params['VStep'] = self.iv_sweep_settings.v_step.value()
            smusweep = self.iv_sweep_settings.smu_sweep.currentText()
            params['smu_sweep'] = getattr(self.keithley, smusweep)

        else:
            return

        # get general sweep settings
        smu_gate = self.general_sweep_settings.smu_gate.currentText()
        smu_drain = self.general_sweep_settings.smu_drain.currentText()
        params['tInt'] = self.general_sweep_settings.t_int.value()
        params['delay'] = self.general_sweep_settings.t_settling.value()
        params['smu_gate'] = getattr(self.keithley, smu_gate)
        params['smu_drain'] = getattr(self.keithley, smu_drain)
        params['pulsed'] = bool(self.general_sweep_settings.sweep_type.currentIndex())

        # check if integration time is valid, return otherwise
        freq = self.keithley.localnode.linefreq

        if not 0.001/freq < params['tInt'] < 25.0/freq:
            msg = ('Integration time must be between 0.001 and 25 ' +
                   'power line cycles of 1/(%s Hz).' % freq)
            QtWidgets.QMessageBox.information(self, str('error'), msg)

            return

        # create measurement thread with params dictionary
        self.measureThread = MeasureThread(self.keithley, params)
        self.measureThread.finished_sig.connect(self.on_measure_done)

        # run measurement
        self._gui_state_busy()
        self.measureThread.start()

    def on_measure_done(self, sd):
        self.statusBar.showMessage('    Ready.')
        self._gui_state_idle()
        self.actionSaveSweepData.setEnabled(True)

        self.sweep_data = sd
        self.canvas.plot(self.sweep_data)
        if not self.keithley.abort_event.is_set():
            self.on_save_clicked()

    @QtCore.pyqtSlot()
    def on_abort_clicked(self):
        """
        Aborts current measurement.
        """
        self.keithley.abort_event.set()

# =============================================================================
# Interface callbacks
# =============================================================================

    @QtCore.pyqtSlot()
    def on_connect_clicked(self):
        self.keithley.connect()
        self.update_gui_connection()
        if not self.keithley.connected:
            msg = ('Keithley cannot be reached at %s. ' % self.keithley.visa_address
                   + 'Please check if address is correct and Keithley is ' +
                   'turned on.')
            QtWidgets.QMessageBox.information(self, str('error'), msg)

    @QtCore.pyqtSlot()
    def on_disconnect_clicked(self):
        self.keithley.disconnect()
        self.update_gui_connection()
        self.statusBar.showMessage('    No Keithley connected.')

    @QtCore.pyqtSlot()
    def on_save_clicked(self):
        """Show GUI to save current sweep data as text file."""
        prompt = 'Save as .txt file.'
        filename = 'untitled.txt'
        formats = 'Text file (*.txt)'
        filepath, _ = QtWidgets.QFileDialog.getSaveFileName(self, prompt, filename,
                                                            formats)
        if len(filepath) < 4:
            return
        self.sweep_data.save(filepath)

    @QtCore.pyqtSlot()
    def on_load_clicked(self):
        """Show GUI to load sweep data from file."""
        prompt = 'Please select a data file.'
        filepath, _ = QtWidgets.QFileDialog.getOpenFileName(self, prompt)
        if not osp.isfile(filepath):
            return

        self.sweep_data = FETResultTable()
        self.sweep_data.load(filepath)

        self.canvas.plot(self.sweep_data)
        self.actionSaveSweepData.setEnabled(True)

    @QtCore.pyqtSlot()
    def on_save_default(self):
        """Saves current settings from GUI as defaults."""

        # save sweep settings
        self.transfer_sweep_settings.save_defaults()
        self.output_sweep_settings.save_defaults()
        self.iv_sweep_settings.save_defaults()
        self.general_sweep_settings.save_defaults()

        # save smu specific settings
        for tab in self.smu_tabs:
            tab.save_defaults()

    @QtCore.pyqtSlot()
    def on_load_default(self):
        """Load default settings to interface."""

        # load sweep settings
        self.transfer_sweep_settings.load_defaults()
        self.output_sweep_settings.load_defaults()
        self.iv_sweep_settings.load_defaults()
        self.general_sweep_settings.load_defaults()

        # smu settings
        for tab in self.smu_tabs:
            tab.load_defaults()

    @QtCore.pyqtSlot()
    def exit_(self):
        self.keithley.disconnect()
        self.connection_status_update.stop()
        self.save_geometry()
        self.deleteLater()

# =============================================================================
# Interface states
# =============================================================================

    def update_gui_connection(self):
        """Check if Keithley is connected and update GUI."""
        if self.keithley.connected and not self.keithley.busy:
            try:
                test = self.keithley.localnode.model
                self._gui_state_idle()
            except (visa.VisaIOError, visa.InvalidSession, OSError):
                self.keithley.disconnect()
                self._gui_state_disconnected()

        elif self.keithley.connected and self.keithley.busy:
            self._gui_state_busy()

        elif not self.keithley.connected:
            self._gui_state_disconnected()

    def _gui_state_busy(self):
        """Set GUI to state for running measurement."""

        self.pushButtonRun.setEnabled(False)
        self.pushButtonAbort.setEnabled(True)

        self.actionConnect.setEnabled(False)
        self.actionDisconnect.setEnabled(False)

        self.statusBar.showMessage('    Measuring.')
        self.led.setChecked(True)

    def _gui_state_idle(self):
        """Set GUI to state for IDLE Keithley."""

        self.pushButtonRun.setEnabled(True)
        self.pushButtonAbort.setEnabled(False)

        self.actionConnect.setEnabled(False)
        self.actionDisconnect.setEnabled(True)
        self.statusBar.showMessage('    Ready.')
        self.led.setChecked(True)

    def _gui_state_disconnected(self):
        """Set GUI to state for disconnected Keithley."""

        self.pushButtonRun.setEnabled(False)
        self.pushButtonAbort.setEnabled(False)

        self.actionConnect.setEnabled(True)
        self.actionDisconnect.setEnabled(False)
        self.statusBar.showMessage('    No Keithley connected.')
        self.led.setChecked(False)


# noinspection PyUnresolvedReferences
class MeasureThread(QtCore.QThread):

    started_sig = QtCore.Signal()
    finished_sig = QtCore.Signal(object)

    def __init__(self, keithley, params):
        QtCore.QThread.__init__(self)
        self.keithley = keithley
        self.params = params

    def __del__(self):
        self.wait()

    def run(self):
        self.started_sig.emit()
        sweep_data = None

        if self.params['sweep_type'] == 'transfer':
            sweep_data = self.keithley.transferMeasurement(
                    self.params['smu_gate'], self.params['smu_drain'],
                    self.params['VgStart'], self.params['VgStop'], self.params['VgStep'],
                    self.params['VdList'], self.params['tInt'], self.params['delay'],
                    self.params['pulsed']
                    )
        elif self.params['sweep_type'] == 'output':
            sweep_data = self.keithley.outputMeasurement(
                    self.params['smu_gate'], self.params['smu_drain'],
                    self.params['VdStart'], self.params['VdStop'], self.params['VdStep'],
                    self.params['VgList'], self.params['tInt'], self.params['delay'],
                    self.params['pulsed']
                    )

        elif self.params['sweep_type'] == 'iv':
            direction = np.sign(self.params['VStop'] - self.params['VStart'])
            stp = direction * abs(self.params['VStep'])
            sweeplist = np.arange(self.params['VStart'], self.params['VStop'] + stp, stp)
            v, i = self.keithley.voltageSweepSingleSMU(
                    self.params['smu_sweep'], sweeplist, self.params['tInt'],
                    self.params['delay'], self.params['pulsed']
                    )

            self.keithley.beeper.beep(0.3, 2400)
            self.keithley.reset()

            params = {'sweep_type': 'iv', 't_int': self.params['tInt'],
                      'delay': self.params['delay'], 'pulsed': self.params['pulsed']}

            sweep_data = FETResultTable(['Voltage', 'Current'], ['V', 'A'],
                                        np.array([v, i]).transpose(), params)

        self.finished_sig.emit(sweep_data)


def run():

    import sys
    import argparse
    from keithley2600 import Keithley2600

    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', help='increase output verbosity',
                        action='store_true')
    args = parser.parse_args()
    if args.verbose:
        import keithley2600
        keithley2600.log_to_screen()

    keithley_address = CONF.get('Connection', 'VISA_ADDRESS')
    visa_library = CONF.get('Connection', 'VISA_LIBRARY')
    keithley = Keithley2600(keithley_address, visa_library)

    app = QtWidgets.QApplication(sys.argv)

    keithley_gui = KeithleyGuiApp(keithley)
    keithley_gui.show()
    app.exec_()


if __name__ == '__main__':

    run()
