from PyQt6.QtGui import *
from PyQt6.QtWidgets import *
from PyQt6.QtCore import *

import multiprocessing
import multiprocessing.managers

import time
from datetime import datetime
import traceback,sys,os
import numpy as np
import pyqtgraph as pg
import import_txt

# Get the current script's directory
current_dir = os.path.dirname(os.path.abspath(__file__))
# Get the parent directory by going one level up
parent_dir = os.path.dirname(current_dir)
# Add the parent directory to sys.path
sys.path.append(parent_dir)
from drivers import Lakeshore336

from design_files.LS336_design import Ui_MainWindow


class WorkerSignals(QObject):
    '''
    Defines the signals available from a running worker thread.
    Supported signals are:
    finished:        No data
    error:        tuple (exctype, value, traceback.format_exc() )
    result:        object data returned from processing, anything
    progress:        int indicating % progress
    '''
    finished = pyqtSignal()
    error = pyqtSignal(tuple)
    result = pyqtSignal(object)
    progress = pyqtSignal(list)


class Worker(QRunnable):
    '''
    Worker thread
    Inherits from QRunnable to handler worker thread setup, signals and wrap-up.
    :param callback: The function callback to run on this worker thread. Supplied args and
                     kwargs will be passed through to the runner.
    :type callback: function
    :param args: Arguments to pass to the callback function
    :param kwargs: Keywords to pass to the callback function
    '''

    def __init__(self, fn, *args, **kwargs):
        super(Worker, self).__init__()

        # Store constructor arguments (re-used for processing)
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.signals = WorkerSignals()

        # Add the callback to our kwargs
        self.kwargs['progress_callback'] = self.signals.progress

    @pyqtSlot()
    def run(self):
        '''
        Initialise the runner function with passed args, kwargs.
        '''

        # Retrieve args/kwargs here; and fire processing using them
        try:
            result = self.fn(*self.args, **self.kwargs)
        except:
            traceback.print_exc()
            exctype, value = sys.exc_info()[:2]
            self.signals.error.emit((exctype, value, traceback.format_exc()))
        else:
            self.signals.result.emit(result)  # Return the result of the processing
        finally:
            self.signals.finished.emit()  # Done

def get_float(Qline,default = 0): #gets value from QLineEdit and converts it to float. If text is empty or cannot be converted, it returns "default" which is 0, if not specified
    try:
        out = float(Qline.text())
    except:
        out = default
    return(out)

class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, *args, **kwargs):
        # Get the current script's directory
        self.current_dir = os.path.dirname(os.path.abspath(__file__))
        # Get the parent directory by going one level up
        self.parent_dir = os.path.dirname(current_dir)

        #establish connection to global variables
        try: #try to connect to global variables
            manager = multiprocessing.managers.BaseManager(address=('localhost',5001), authkey=b'')
            manager.connect()
            manager.register('sync_LS_336')
            self.sync_LS_336 = manager.sync_LS_336()
        except: #open global variables, if no connection can be made (i.e. it is not running). Then connect to it
            # subprocess.call(['D:\\Python instrument drivers\\env\\Scripts\\python.exe', 'D:\\Python instrument drivers\\StandAlones\\global_variables.py'])
            self.global_vars = QProcess()
            self.global_vars.start(self.current_dir+"\\env\\Scripts\\python.exe", [self.current_dir+'\\global_variables.py'])
            manager.connect()
            manager.register('sync_LS_336')
            self.sync_LS_336 = manager.sync_LS_336()
            print('!!!\nI opened global variables myself. If you close me, global variables will shut down too. Consider starting global variables in own instance for more security\n!!!')
        
        #fill in variables, if they are not defined in global variables
        self.sync_LS_336.update({'setT':0, 'ramprate':'0','Contr_Ch':0, 'T':[0,0,0,0,0], 'Range':0,'PID':[0,0,0]})        

        #import Gui from QT designer file
        super(MainWindow, self).__init__(*args, **kwargs)
        self.setupUi(self)

        #setup plot
        self.graphWidget.setBackground('w')
        self.graphWidget.setTitle("Temperature")
        self.graphWidget.setLabel('left', 'Temperature [K]')
        self.graphWidget.setLabel('bottom', 'Time (H)')
        axis = pg.DateAxisItem()
        self.graphWidget.setAxisItems({'bottom':axis})

        temp = [time.time(),time.time()-1]
        pen1 = pg.mkPen(color=(255, 0, 0), width=2)
        pen2 = pg.mkPen(color=(0, 0, 255), width=2)
        pen3 = pg.mkPen(color=(0, 255, 0), width=2)
        pen4 = pg.mkPen(color=(255, 255, 0), width=2)
        pen5 = pg.mkPen(color=(0, 0, 0), width=2)
        self.plot_A = self.graphWidget.plot(temp,[1,0],pen = pen1, name = 'Ch: A')
        self.plot_B = self.graphWidget.plot(temp,[1,0],pen = pen2, name = 'Ch: B')
        self.plot_C = self.graphWidget.plot(temp,[1,0],pen = pen3, name = 'Ch: C')
        self.plot_D = self.graphWidget.plot(temp,[1,0],pen = pen4, name = 'Ch: D')
        self.plot_P = self.graphWidget.plot(temp,[1,0],pen = pen5, name = 'Power')
        self.plot_P.hide()
        self.graphWidget.addLegend()
        
        #set up pyQT threadpool
        self.threadpool = QThreadPool()

        #define and standard threads.
        worker_save = Worker(self.save)
        self.threadpool.start(worker_save)

        #define signals and slots
        self.actionSet_default.triggered.connect(self.set_default)
        self.actionReset_default.triggered.connect(self.read_default)
        self.button_setPID.clicked.connect(self.set_PID)
        self.button_connect.clicked.connect(self.start_meas)
        self.line_Nplot.editingFinished.connect(self.set_Npoints)
        self.line_saveInterval.editingFinished.connect(self.change_timing)
        self.checkBox_A.stateChanged.connect(self.plot_hide)
        self.checkBox_B.stateChanged.connect(self.plot_hide)
        self.checkBox_C.stateChanged.connect(self.plot_hide)
        self.checkBox_D.stateChanged.connect(self.plot_hide)
        self.checkBox_pwr.stateChanged.connect(self.plot_hide)
        self.checkBox_disableplots.stateChanged.connect(self.set_displot)
        self.line_setT.editingFinished.connect(self.set_T)
        self.line_Ramprate.editingFinished.connect(self.set_rate)
        self.comboBox_Channel.currentIndexChanged.connect(self.set_channel)
        self.comboBox_range.currentIndexChanged.connect(self.set_range)
        

        #define constants
        self.Temperature = np.zeros((1,5)) #store temperature and power data
        self.t = [time.time()] #store timestamps
        self.t1 = [datetime.now()] #store timestamps with higher precision
        self.last_save = self.t1[-1] #timestamp of last write-to-file event
        self.Npoints = 200 #number of point to plot
        self.running = True #true while app is running
        self.disable_plot = False #constant to disable plot to improve performance. Is changed by checkbox checkBox_disableplots
        self.timing_save = 5 #save intervall [s]
        self.set_old = [0,0,0] #variable to save the 'old' set values to compare them to the global variables. Since the length is only 3, it differs from set_new in the first iteration. This ensures that new parameters are send to the device
        self.set_new = [0,0,0,0,0] #variable to save the new set values to compare them to the old ones
        self.lines_config_float = [self.line_setP,self.line_setI,self.line_setD,self.line_setT,self.line_Ramprate,self.line_Nplot]#is used for config file
        self.lines_config_strings = [self.line_devAdr,self.line_filePath,self.line_saveInterval]#is used for config file
        self.checkboxes_config = [self.checkBox_A, self.checkBox_B, self.checkBox_C, self.checkBox_D, self.checkBox_pwr,self.checkBox_disableplots,self.checkBox_save]#is used for config file

        #read default values from fonfig and set them in gui
        self.read_default()
        #write values from gui to global variables. 
        self.set_T()
        self.set_rate()
        self.set_channel()
        self.set_range()
        self.set_PID()
        #update save intervall to the gui value
        self.change_timing()


    def start_meas(self):
        #Connect to device
        address = self.line_devAdr.text()
        self.LS = Lakeshore336.LakeShore336(address)
        #start thread for communication with device
        self.worker = Worker(self.update_T)  
        self.worker.signals.progress.connect(self.update_gui)
        self.threadpool.start(self.worker)

    def update_T(self, progress_callback):      
    #get values from device and write them to global variables. Checks if global variables changed from last iteration. Also pass it to upddate_gui with emit(T)
         while self.running == True:            
            for i,n in enumerate(['setT','ramprate','Contr_Ch','Range','PID']): #get new set values from global variables and compare to old ones. 
                self.set_new[i] = self.sync_LS_336.get(n)
            
            if self.set_new != self.set_old:             #if a button is clicked or global variables are changed self.changed is set to true and new parameters are send to device
                self.LS.conf_outp(out=1, mode=1, inp=self.set_new[2]+1, powup=0) #configures output channel 1
                self.LS.set_Ramp(out=1, ON=1, ramp=self.set_new[1])            #configues ramp settings, if Rate=0 the ramp is turned off
                self.LS.set_T(out = 1, T = self.set_new[0])                   #sets temperature setpoint
                self.LS.turn_on_outp(out=1, range = self.set_new[3])          #turn heater on or off, according to Range
                self.LS.conf_pid(1,self.set_new[4][0],self.set_new[4][1],self.set_new[4][2]) #set PID values
                self.update_setValues(self.set_new)

            T = self.LS.read(['A','B','C','D'])    #read Temperature data from all 4 channels
            T.append(0)
            self.sync_LS_336.update({'T':T})
            progress_callback.emit(T)
            self.set_old = self.set_new[:]      #List needs to be sliced so that only values are taken and not just a pointer is created  
            # time.sleep(0.1)

         del(self.LS)    #disconnect device when self.running is set to False

    def update_gui(self,T):
        #set numbers
        self.T_A.setText(str(T[0]))
        self.T_B.setText(str(T[1]))
        self.T_C.setText(str(T[2]))
        self.T_D.setText(str(T[3]))

        #Create database for plotting
        self.Temperature = np.vstack([self.Temperature, np.array(T)])
        x = range(len(self.Temperature))
        self.t.append(time.time())
        self.t1.append(datetime.now())

        #plot
        if self.disable_plot == False:
            self.plot_A.setData(self.t[-self.Npoints:],self.Temperature[-self.Npoints:,0])
            self.plot_B.setData(self.t[-self.Npoints:],self.Temperature[-self.Npoints:,1])
            self.plot_C.setData(self.t[-self.Npoints:],self.Temperature[-self.Npoints:,2])
            self.plot_D.setData(self.t[-self.Npoints:],self.Temperature[-self.Npoints:,3])
            self.plot_P.setData(self.t[-self.Npoints:],self.Temperature[-self.Npoints:,4])

    def update_setValues(self,setV):             
        #sets setvalues obtained from update_T in gui ['setT','ramprate','Contr_Ch','Range','PID']
        self.line_setT.setText(f"{setV[0]}")
        self.line_Ramprate.setText(f"{setV[1]}")
        self.comboBox_Channel.setCurrentIndex(setV[2])
        self.comboBox_range.setCurrentIndex(setV[3])
        self.line_setP.setText(f"{setV[4][0]}")
        self.line_setI.setText(f"{setV[4][1]}")
        self.line_setD.setText(f"{setV[4][2]}")

    def set_T(self):
        #updates the set temperature in global variables. The change will be detected by update T and it will be passed to the device
        self.sync_LS_336.update({'setT':get_float(self.line_setT)})
    
    def set_rate(self):
        #updates the ramprate in global variables. The change will be detected by update T and it will be passed to the device
        self.sync_LS_336.update({'ramprate':get_float(self.line_Ramprate)})

    def set_channel(self):
        #updates the control channel in global variables. The change will be detected by update T and it will be passed to the device
        self.sync_LS_336.update({'Contr_Ch':self.comboBox_Channel.currentIndex()})

    def set_range(self):
        #updates the reange in global variables. The change will be detected by update T and it will be passed to the device
        self.sync_LS_336.update({'Range':self.comboBox_range.currentIndex()})

    def set_PID(self):
        #updates the PID values in global variables. The change will be detected by update T and it will be passed to the device
        self.sync_LS_336.update({'PID':[get_float(self.line_setP),get_float(self.line_setI),get_float(self.line_setD)]})


    def set_Npoints(self):
        #sets the number of points to plot
        self.Npoints = int(self.line_Nplot.text())

    def set_displot(self):
        #sets variable to disable plot so checkbox state does not need be read out every iteration
        self.disable_plot = self.checkBox_disableplots.isChecked()
     
    def plot_hide(self):
        #shows or hides plots according to the checkboxes next to the plot area
        boxes = [self.checkBox_A, self.checkBox_B, self.checkBox_C, self.checkBox_D, self.checkBox_pwr]
        plots = [self.plot_A, self.plot_B, self.plot_C, self.plot_D, self.plot_P]

        for b,p in zip(boxes,plots):
            if b.isChecked() == True:
                p.show()
            else:
                p.hide()
    
    def change_timing(self):
    #updates the timing which is used for "save". If no value it given, it is set to 1 s
        self.timing_save = get_float(self.line_saveInterval,1)

    def save(self, progress_callback):
        #if save checkbox is checked it writes measurement values to file specified in line.filePath. There the full path including file extension must be given.
        while self.running == True:
            time.sleep(self.timing_save) #wait is at beginning so first point is not corrupted when app just started.
            if self.checkBox_save.isChecked() == True and self.t1[-1] > self.last_save:
                #write only, if there is a new timestamp
                path = self.line_filePath.text()
                if os.path.isfile(path) == False:
                    with open(path,'a') as file:
                        file.write('date\tCh:A[K]\tCh:B[K]\tCh:C[K]\tCh:D[K]\tPower[%]\n')
                file = open(path,'a')
                #file.write(time.strftime("%Y-%m-%d_%H-%M-%S",time.localtime(self.t[-1]))+'\t') #original timestamp
                file.write(self.t1[-1].strftime("%Y-%m-%d_%H-%M-%S.%f")+'\t')
                for d in self.Temperature[-1]:
                    file.write(f"{d}\t")
                file.write('\n')
                self.last_save = self.t1[-1]
                file.close

    def set_default(self):
        #saves current set values to txt file in subdirectory configs. All entries that are saved are defined in self.lines_config
        #Additionally control channel of LS336 is saved. Overwrites old values in config file.
        current_dir = os.path.dirname(os.path.abspath(__file__))
        path = current_dir+'\\configs\\LS336_config.txt' #To make shure the config file is at the right place, independent from where the program is started the location of the file is retrieved
        file = open(path,'w')
        for l in self.lines_config_float:
            temp = f"{get_float(l)}"
            file.write(temp+'\t')
        for l in self.lines_config_strings:
            file.write(l.text()+'\t')
        for c in self.checkboxes_config:
            file.write(str(c.isChecked())+'\t')
        file.write(str(self.comboBox_Channel.currentIndex()))
        file.write('\n')
        file.close

    def read_default(self):
        #reads default values from config file in subdirectory config and sets the values in gui. Then self.change is set to true so values are send
        #to device. (If no config file exists, it does nothing.)
        current_dir = os.path.dirname(os.path.abspath(__file__))
        path = current_dir+'\\configs\\LS336_config.txt' #To make shure the config file is read from the right place, independent from where the program is started the location of the file is retrieved
        try: #exit function if config file does not exist
            vals = import_txt.read_raw(path)
        except:
            print('no config file found on')
            print(path)
            return 
        formats = ['.2f', '.2f', '.2f','.2f','.2f','.0f']

        for l,v,f in zip(self.lines_config_float,vals[0],formats):
            v = float(v)        #convert string in txt to float, so number can be formatted according to "formats" when it's set
            l.setText(format(v,f))
        
        for l,v in zip(self.lines_config_strings,vals[0][len(self.lines_config_float):]):
            l.setText(v)

        for c,v in zip(self.checkboxes_config,vals[0][len(self.lines_config_float)+len(self.lines_config_strings):]):
            c.setChecked(v == 'True')

        self.comboBox_Channel.setCurrentIndex(int(vals[0][-1]))  

        self.change = True
        
    def closeEvent(self,event): #when window is closed self.running is set to False, so all threads stop
        self.running = False
        time.sleep(1)
        event.accept()



app = QApplication(sys.argv)

window = MainWindow()
window.show()
app.exec()