"""Keysight thermocouple data acquisition system code
Notes
-----
The keysight thermocouple data acquisition code communicates opens a GUI that communicates with the Keysight DAQ970A
"""

###############
### IMPORTS ###
###############

import tkinter as tk
from tkinter import ttk
from tkinter import font
from tkinter import filedialog
from tkinter import messagebox

import matplotlib
matplotlib.use("TkAgg")
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

import numpy as np
import time
import threading
import pyvisa
import csv
import os

##############################################
########## GRAPHICAL USER INTERFACE ##########
##############################################

class ThermocoupleGUI(tk.Frame):

    ### Intialization of the GUI
    def __init__(self, master):
        tk.Frame.__init__(self, master)
        self.master = master
        self.master.title("THERMOCOUPLE MEASUREMENTS - DAQ970A")
        self.master.geometry("1100x600")

        self.create_widgets()
        self.pack()
        self.continueRunning = False # No reading from the start

        # Setting variables 
        self.numberScans = "INF"     # Number of scan sweeps to measure - no upper limit defined
        self.voltage = 2.00          # Voltage value to DAQ from -12 to 12
        self.Counter = 0

        # Prompt user to select the directory for saving files
        self.save_dir = self.saveSettingsFrame.saveEntryI.get()      # read path given in the GUI
        if not os.path.exists(self.save_dir):
            messagebox.showerror("Error", "Invalid directory path.")
            return

        # Connect to the instrument
        loc = "TCPIP0::169.254.173.160::inst0::INSTR"     # fixed adress givenb to Keysight DAQ970A
        try:
            self.rm = pyvisa.ResourceManager()            # try connect to instrument
            self.daq = self.rm.open_resource(loc)  
            print("Confirmed connection with: \n"+(self.daq.query("*IDN?")))
        except pyvisa.errors.VisaIOError:
            print("Device not connected. GUI will still open.")  # print in case no connection could be made - still opens the GUI

    # Creating widgets on the GUI
    def create_widgets(self):
        #The main frame is made up of 5 subframes
        self.channelSettingsFrame = channelSettings(self, title ="Channel Settings")
        self.channelSettingsFrame.grid(row=0, column=1, pady=(20,0), padx=(20,20), ipady=10, sticky='ew')

        self.saveSettingsFrame = saveSettings(self, title="Save Settings", browse_directory_func=self.browse_directory)
        self.saveSettingsFrame.grid(row=1, column=1, pady=(20,0), padx=(20,20), ipady=10, sticky='ew')
        # Store a reference to the saveEntryI widget in save_dir_entry
        self.save_dir_entry = self.saveSettingsFrame.saveEntryI

        self.thermocoupleSettingsFrame = thermocoupleSettings(self, title="Thermocouple Settings", update_func=self.updateList)
        self.thermocoupleSettingsFrame.grid(row=2, column=1, pady=(20,0), padx=(20,20), ipady=10, sticky='ew')

        self.recordSettingsFrame = recordSettings(self, title="Record Data", start_func=self.startTask, stop_func=self.stopTask)
        self.recordSettingsFrame.grid(row=3, column=1, pady=(20,0), padx=(20,20), ipady=10, sticky='ew')

        self.graphDataFrame = graphData(self)
        self.graphDataFrame.grid(row=0, rowspan=5, column=2, pady=(20,0), ipady=10, sticky="nsew")

        self.graphSettingsFrame = graphSettings(self, title="Visualization Settings", update_func=self.updateXlim)
        self.graphSettingsFrame.grid(row=4, column=1, pady=(20,0), padx=(20,20), ipady=10, sticky="ew")

    def startTask(self):
        # Start function to be executed when start button is pressed
        self.recordSettingsFrame.startButton["state"] = "disabled"   # disable start button when data is acquired, prevent user from starting second time

        # Make filenames for csv and npy files based on time stamp
        self.ts_0 = time.strftime("%Y%m%d%H%M%S")
        self.ts_1 = time.time()
        self.filename = os.path.join(self.save_dir, f"data_{self.ts_0}")

        # Shared flag to alert task if it should stop
        self.continueRunning = True

        # Setup channels configuration - read which channels should be read (user input) - default:101-120
        self.scanList = self.thermocoupleSettingsFrame.thermocoupleEntryI.get()

        # Extracting the channels to be read
        scanList2 = self.scanList.split(",")
        self.newList = []

        for elem in scanList2:
            if ":" in elem:
                for i in range(int(elem[:-4]), int(elem[4:])+1):
                    self.newList.append(i)
            else:
                self.newList.append(int(elem))

        self.displayList = self.newList    # all thermocouples are displayed by default
        self.displayIndex = [i+1 for i in range(len(self.newList))]

        # Update index from where to save data
        self.ind_save = self.Counter

        # Start communiction with DAQ970A
        self.string = "CONF:TEMP TC,K,(@{})".format(self.scanList)
        self.daq.write(self.string)

        # Setup scan list
        self.daq.write("ROUTE:SCAN " + "(@{})".format(self.scanList))     # Selects the channels to be included in the scan list

        # Reading format
        self.daq.write("FORMAT:READING:CHAN ON")       # Channel information is given

        # Channel delay
        self.daq.write("ROUT:CHAN:DELAY:AUTO ON")      # Automatic channel delay time

        # Setup when scanning starts and interval rate
        self.daq.write("TRIG:SOUR BUS")
        self.daq.write("INIT")

        # Start the scan and retrieve the scan time
        self.daq.write("SYST:TIME:SCAN?")
        print ("Measurements started at: "+self.daq.read())

        # default value for the xlim of the plot
        self.xlim = 0

        #spin off call to check 
        self.master.after(10, self.runTask)

    def runTask(self):
        # setup when scanning starts and interval rate
        # Redundancy in case of power failure
        try:
            self.daq.write("TRIG:SOUR BUS")
            self.daq.write("INIT")            # Intialize trigger sensitivity
            self.daq.write("*TRG")            # Trigger sweep
            fetch = self.daq.query("FETC?")   # Transfer all available readings to output buffer

            farray=np.fromstring(fetch,sep=',')
        except:
            farray=np.zeros(size=self.newList.size)

        # Initializing dataframe for new data
        if self.Counter == 0:
            # Intilialize dataframe
            self.df = np.empty(shape=(len(self.newList)+1,1))

            self.startTime = time.time()

            self.df[0,0] = 0
            self.df[1:,0] = farray[0::2]

            self.df_full = self.df

        # Adding new data to existing dataframe
        else:
            self.df_full = np.append(self.df_full, np.empty(shape=(len(self.newList)+1,1)), axis=1)
            self.df_full[0,-1] = time.time()-self.startTime
            self.df_full[1:,-1] = farray[0::2]

        # Save every 10 iterations
        if self.Counter % 10 == 0:
            self.master.after(10, self.save)
            print(self.newList)
            print(self.df_full[1:,-1])

        # Displaying only the desired channels (all by default but can be changed)
        df_disp = self.df_full[0]
        for i in self.displayIndex:
            df_disp = np.vstack((df_disp, self.df_full[i]))

        # Adjusting nuber of columns in the plot legend
        if np.floor(len(self.displayIndex) / 20) == 0: 
            ncol = 1
        elif np.floor(len(self.displayIndex) / 20) == 1:
            ncol = 2
        else:
            ncol = 3

        ### Update plot
        self.graphDataFrame.ax.cla()
        self.graphDataFrame.ax.set_title("Measured temperatures")
        self.graphDataFrame.ax.plot(df_disp[0,:]/60, df_disp[1:,:].T)
        self.graphDataFrame.ax.legend(self.displayList, fontsize=6, bbox_to_anchor=(1.02, 1), loc="upper left", ncol=ncol)

        self.graphDataFrame.ax.set_ylim([0,np.amax(df_disp[1:])+50])
        self.graphDataFrame.ax.set_xlim([self.xlim,max(df_disp[0,-1]/60,0.1)])
        self.graphDataFrame.ax.set_xlabel("Time [min]")
        self.graphDataFrame.ax.set_ylabel("Temperature [°C]")
        self.graphDataFrame.graph.draw()

        self.Counter +=1

        #check if the task should sleep or stop
        if self.continueRunning:
            self.master.after(10, self.runTask)
        else:
            self.recordSettingsFrame.startButton['state'] = 'active'

    def stopTask(self):
        # Call back for the "stop task" button
        self.continueRunning = False

        self.master.after(10, self.save)

    def save(self):
        # Function for saving the acquired data
        self.ts_new = time.time()

        if self.ts_new - self.ts_1 >= 120:
            self.ts_0 = time.strftime("%Y%m%d%H%M%S")
            self.filename = os.path.join(self.save_dir, f"data_{self.ts_0}")
            self.ts_1 = self.ts_new
            
            # Update index from where to save data
            self.ind_save = self.Counter

        # Save data to npy file
        np.save("{}.npy".format(self.filename), self.df_full[:,self.ind_save:])

        # Save data to CSV if self.continueRunning == False, to avoid too much delay
        # csv file only constructed when the acquisition is stopped
        if self.continueRunning == False:
            data = self.df_full.T
            header="Time"
            for i in self.newList:
                header=header+", TC "+str(i)
            np.savetxt("{}.csv".format(self.filename), data,  delimiter = ",",header=header)

        # New dataframe (np.array) to save data
        self.ind_save = self.Counter

    def updateList(self):
        # Function updating the thermocouples to be displayed
        self.dispList = self.thermocoupleSettingsFrame.thermocoupleEntryII.get()

        if self.dispList == "0":     # "0" as input leads to resetting the display list to all channels being scanned
            self.displayList = self.newList
            self.displayIndex = [i+1 for i in range(len(self.newList))]
        else:
            self.displayList = []
            self.displayIndex = []
            try:
                # additional safety to avoid crashing when input is incorrect
                for elem in self.dispList.split(","):
                    if ":" in elem:                                      # splitting in case a range of channels is defined
                        for i in range(int(elem[:-4]), int(elem[4:])+1):
                            ind = np.where(np.array(self.newList) == i)[0][0]
                            self.displayList.append(i)
                            self.displayIndex.append(ind+1)
                    else:
                        self.displayList.append(int(elem))
                        ind = np.where(np.array(self.newList) == int(elem))[0][0]
                        self.displayIndex.append(ind+1)
            except:
                # reset in case of incorrect input
                self.displayList = self.newList
                self.displayIndex = [i+1 for i in range(len(self.newList))]

    def updateXlim(self):
        # Function allowing to update the limit on the x-axis
        self.xlim = float(self.graphSettingsFrame.graphEntry.get())

    def browse_directory(self):
        # Function allowing to browse directories to change saving directory
        directory = filedialog.askdirectory()
        if directory:
            self.save_dir_entry.delete(0, tk.END)
            self.save_dir_entry.insert(0, directory)

class channelSettings(tk.LabelFrame):
    ### Creating widgets and default input for channel settings

    def __init__(self, parent, title):
        tk.LabelFrame.__init__(self, parent, text=title, labelanchor='n')
        self.parent = parent
        self.grid_columnconfigure(0, weight=1)
        self.xPadding = (30,30)
        self.create_widgets()

    def create_widgets(self):
        loc = "TCPIP0::169.254.173.160::inst0::INSTR"
        self.physicalChannelLabelI = ttk.Label(self, text="Physical Channel Input")
        self.physicalChannelLabelI.grid(row=0,  sticky='w', padx=self.xPadding, pady=(10,0))
        self.physicalChannelEntryI = ttk.Entry(self)
        self.physicalChannelEntryI.insert(0, loc)
        self.physicalChannelEntryI.grid(row=1,columnspan=2, sticky="ew", padx=self.xPadding)

class saveSettings(tk.LabelFrame):
    ### Creating widgets and default input for save settings

    def __init__(self, parent, title, browse_directory_func):
        tk.LabelFrame.__init__(self, parent, text=title, labelanchor='n')
        self.parent = parent
        self.browse_directory = browse_directory_func  # This line connects the browse_directory_func argument to the class attribute
        self.grid_columnconfigure(0, weight=1)
        self.xPadding = (30,30)
        self.path = os.getcwd()
        self.create_widgets()

    def create_widgets(self):
        self.saveLabelI = ttk.Label(self, text="Save directory")
        self.saveLabelI.grid(row=0,  sticky='w', padx=self.xPadding, pady=(10,0))

        self.saveEntryI = ttk.Entry(self, width=75)  # Adjust the width of the entry box
        self.saveEntryI.insert(0, self.path)
        self.saveEntryI.grid(row=1, column=0, columnspan=5, sticky="ew", padx=self.xPadding)

        self.saveButtonI = tk.Button(self, text="Browse", command=self.on_browse_button_click, width=10)  # Adjust the width of the button
        self.saveButtonI.grid(row=1, column=6, padx=10, sticky='ew')

    def on_browse_button_click(self):
        directory = filedialog.askdirectory()
        if directory:
            self.saveEntryI.delete(0, tk.END)  # Clear the entry box before inserting the new directory
            self.saveEntryI.insert(0, directory)  # Insert the selected directory into the entry box

class thermocoupleSettings(tk.LabelFrame):
    ### Creating widgets and default input for thermocouple settings

    def __init__(self, parent, title, update_func):
        tk.LabelFrame.__init__(self, parent, text=title, labelanchor='n')
        self.parent = parent
        self.update = update_func
        self.grid_columnconfigure(0, weight=1)
        self.xPadding = (30,30)
        self.create_widgets()

    def create_widgets(self):
        self.thermocoupleLabelI = ttk.Label(self, text="Thermocouples list (e.g., 101:103,105)")
        self.thermocoupleLabelI.grid(row=0, column=0,  sticky='w', padx=self.xPadding, pady=(10,0))

        self.thermocoupleEntryI = ttk.Entry(self)
        self.thermocoupleEntryI.insert(0, "101:120")
        self.thermocoupleEntryI.grid(row=1, column=0, sticky="ew", padx=self.xPadding)

        self.thermocoupleLabelII = ttk.Label(self, text="Display thermocouples")
        self.thermocoupleLabelII.grid(row=0, column=1,  sticky='w', padx=self.xPadding, pady=(10,0))

        self.thermocoupleEntryII = ttk.Entry(self)
        self.thermocoupleEntryII.insert(0, "0")
        self.thermocoupleEntryII.grid(row=1, column=1, sticky="ew", padx=self.xPadding)

        self.thermocoupleButton = tk.Button(self, text="Update", command=self.update, width=15)
        self.thermocoupleButton.grid(row=1, column=2, padx=10, sticky="nsew")

class recordSettings(tk.LabelFrame):
    ### Creating widgets and default input for record settings

    def __init__(self, parent, title, start_func, stop_func):
        tk.LabelFrame.__init__(self, parent, text=title, labelanchor='n')
        self.parent = parent
        self.start = start_func
        self.stop = stop_func
        self.grid_columnconfigure(0, weight=1)
        self.xPadding = (30,30)
        self.create_widgets()

    def create_widgets(self):
        self.startLabel = ttk.Label(self, text="Start recording")
        self.startLabel.grid(row=0, column=0, padx=self.xPadding, pady=(10,0))

        self.startButton = tk.Button(self, text="Start", command=self.start, width=15)
        self.startButton.grid(row=1, column=0, padx=10, pady=(5, 0), sticky='nsew')

        self.stopLabel = ttk.Label(self, text="Stop recording")
        self.stopLabel.grid(row=0, column=1, padx=self.xPadding, pady=(10,0))

        self.stopButton = tk.Button(self, text="Stop", command=self.stop, width=15)
        self.stopButton.grid(row=1, column=1, padx=10, pady=(5, 0), sticky='nsew')

class graphData(tk.Frame):
    ### Creating plot area

    def __init__(self, parent):
        tk.Frame.__init__(self, parent)
        self.create_widgets()

    def create_widgets(self):
        self.graphTitle = ttk.Label(self, text="Voltage Input")
        self.fig = Figure(figsize=(7,5), dpi=100)
        self.ax = self.fig.add_subplot(1,1,1)
        self.ax.set_xlabel("Time [s]")
        self.ax.set_ylabel("Temperature [°C]")

        box = self.ax.get_position()
        self.ax.set_position([box.x0, box.y0, box.width * 0.9, box.height * 1])

        self.graph = FigureCanvasTkAgg(self.fig, self)
        self.graph.draw()
        self.graph.get_tk_widget().pack()

class graphSettings(tk.Frame):
    ### Creating widgets for updating the graph settings

    def __init__(self, parent, title, update_func):
        tk.LabelFrame.__init__(self, parent, text=title, labelanchor='n')
        self.update = update_func
        self.grid_columnconfigure(0, weight=1)
        self.xPadding = (30,30)
        self.create_widgets()

    def create_widgets(self):
        self.xlimLabel = ttk.Label(self, text="Xlim")
        self.xlimLabel.grid(row=0, column=1,  sticky='w', padx=self.xPadding, pady=(10,0))

        self.graphEntry = ttk.Entry(self)
        self.graphEntry.insert(0, "0")
        self.graphEntry.grid(row=0, column=2, sticky="ew", padx=self.xPadding, pady=(10,0))

        self.xlimButton = tk.Button(self, text="Update", command=self.update, width=15)
        self.xlimButton.grid(row=0, column=3, padx=10, pady=(5, 0), sticky='nsew')


if __name__ == "__main__":
    root = tk.Tk()
    app = ThermocoupleGUI(root)
    app.mainloop()