__init__.py

# -*- coding: utf-8 -*-

#    Pyblab
#    Copyright (C) 2021  Marco Pizzocaro <m.pizzocaro@inrim.it>
#
#    This file is part of Pyblab.
#
#    Pyblab is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    Pyblab is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with Pyblab.  If not, see <https://www.gnu.org/licenses/>.

"""
Pyblab
~~~~~~~

A python module for management of laboratory experiments based on pyDAQmx.

Initially developed for INRIM IT-Yb1 optical clock.

*Changelog*
20210426
* added ability to add notes to outpput files (only before the experiment is run)

20210326
* added buffering=1 to the file open command -- hopefully it will help to keep the data flushed to disk

20210224
* moved the logic for closing scopes and controls to the cycle class.
This way cycles can close unexpected resources (e.g., other files, cameras, etc...)

20210209 -- Initial release risen from the ashes of old labedit source

20210722 -- exp.run return False if stopped.

Created on Tue Feb  9 11:46:58 2021

@author: Marco Pizzocaro
"""

__version__ = '20210722'


import configparser
import ctypes
from datetime import date, datetime
import io
import numpy
import numpy.lib.recfunctions as lrf
import os
import os.path
import socket
import string
import sys
import threading
import time


from PyDAQmx import *
import pyvisa as visa


import PySimpleGUI as sg



from .pattern import Pattern
from .acquisition import Acquisition
from .cycle import *
from .scope import *
from .synth import *



# default values for channels and aoms config file
default_exp = """
[DEFAULT]
invert	= False
delay	= 0
analog = False
factor = 1.
offset = 0.
minV = -10.
maxV = 10.
pass = 1.
sign = 1.
drift = 0.
"""

import __main__


class Experiment(object):
	""" High level interface for pyDAQmx and other experimental components,
	such as synthetizer, GUI, software calculations and data saving to file.
	
	See provided files for examples.
	
	"""
	def __init__(self, title= 'Pylab', name=None, version=None, save_file=True):
		
		
		
		self.acq = None
		self.config = configparser.ConfigParser(allow_no_value=True)
		self.controls = []
		self.cycles = [] # unique cycles
		self.cycles_with_repeates = [] # cycles may be called repeatedly
		self.is_running = False
		
		if name:
			self.name = name
		else:	
			self.name = os.path.basename(__main__.__file__) # name of the script rather of this file
		
		if version:
			self.version = version
		elif __main__.__version__:
			self.version = __main__.__version__
		else:
			self.version = None
			
		
			
		self.pattern = None
		self.seqs = []
		self.scopes = []
		self.synths = []
		self.thread = threading.Thread(target=self._run)
		self.title=title

		image = os.path.join(os.path.dirname(__file__), 'inrim-logo-small.png')

		self.layout = [
			[sg.Image(image)],
			[sg.Column([[sg.Button('Stop', size=(23,1))]], element_justification='c')]
			]
		self.window = None
		
		self.counterfile = "counter.npy"
		self.counter = 0
		
		self.file = None
		if save_file:
			self.open_file()
			
		self.notes = []


	def __enter__(self):
		return self

	def __exit__(self, type, value, traceback):
		self.clear()

	def _run(self):
		"""Threaded loop for cycles"""
		while self.is_running:
			for n, i in enumerate(self.cycles_with_repeates):
				i.pre()
				i.wait_for_data()
				i.post()

	def clear(self):
		self.is_running = False # should use an event
		if self.pattern:
			self.pattern.clear()
		if self.acq:
			self.acq.clear()
		
		# reset Ni board to default
		DAQmxResetDevice('dev1')
		
		if self.window:
			self.window.close()
		
		if self.file:
			self.file.write("# File closed on: " + str(datetime.now()) + "\n")	
			self.file.close()
		
		for c in self.cycles:
			c.close()
		
		# moved to cycle
# 		for s in self.scopes:
# 			s.close()
# 			
# 		for t in self.controls:
# 			t.close()

	def open_file(self):
		today = date.today()
		now = datetime.now()
		
		# update daily counter 			
		
		try:	
			currentday, n = tuple(numpy.load(self.counterfile, allow_pickle=True)) # allow_pickle=True fix a bug with loading
			if today == currentday:
				n = n+1
			else:
				n = 0
				
			numpy.save(self.counterfile, [today, n])	
		except:
 			numpy.save(self.counterfile, [today, 0])
 			n = 0
			
		self.counter = n
		self.layout += [[sg.Text(str(self.counter), font=("Helvetica", 25), justification='center', size=(10,1))]]
				
		path = os.path.join("Data",today.strftime("%Y-%m-%d"))
		self.path = path
		
		if not os.path.exists(path):
			os.makedirs(path)			
		
		
					
		expname  = os.path.splitext(self.name)[0]
				
		hour = now.strftime('%Y%m%d_%H%M%S_')	
		basename = "{:03d}_".format(n) + hour + expname			
		filename = basename + '.dat'		
		filo = os.path.join(self.path,filename)
		
		self.started =  now.strftime('%H:%M:%S ')	
		
		
		
		self.file = open(filo, 'a', buffering=1)
		
		
		# writing file header
	
	def start_file(self):		
		self.file.write("# " + self.title + " data file.\n")
		self.file.write("# File generated on: " + str(datetime.now()) + "\n")		
		self.file.write("# With the script: " + self.name + "\n")
		if self.version:
			self.file.write("# Script version: " + self.version + "\n")
		self.file.write("# Pyblab version: " + __version__ + "\n")
		self.file.write("# Using cycles:\n")
		for i, x in enumerate(self.cycles):
			self.file.write("# " + str(i) + " " + x.name + " with columns as\n")
			if x.header:
				self.file.write(x.header + "\n")
			else:
				self.file.write("No columns specified.\n")
			
			
		for x in self.synths:
			self.file.write('# Synth ' +  x.name + ' with center freq ' + str(x.cf) + '\n')
		
		for x in self.notes:
			self.file.write('# ' + x + '\n')

	def start(self):
		"""Start the experiment (file, cycles, acqusition and pattern)."""
		if self.file:
			self.start_file()
		
		#pattern starts last
		if self.cycles and not self.is_running:
			self.is_running = True
			self.thread.start()

		if self.acq:
			self.acq.start()
		if self.pattern:
			self.pattern.start()
			
		print(str(self.counter) + ': ' + self.name)

	def run(self):
		"""Start the experiment (cycles, acqusition and pattern), wait for the program to be
		stopped by the user."""
		self.window = sg.Window(self.title + ' - ' + self.name, self.layout, return_keyboard_events=True, resizable=True, location=(25, 150), finalize=True, keep_on_top=True)

		self.start()
		clear = True
		
		while True:
			try:
				event, values = self.window.read(timeout=100.)
			except KeyboardInterrupt:
				event = 'KI'
				
			
			if event in (sg.WINDOW_CLOSED, 'Stop', 'F7:118', 'KI'):
				clear = False
				break
			
			if event == sg.TIMEOUT_KEY:
				if self.max_acq and self.acq.n > self.max_acq:
					break
			
			# check values on 'enter' keyboard press
			if event:
				if event == '\r':
					key = self.window.find_element_with_focus().Key
				else:
					key = event
					
				for t in self.controls:
					t.on_event(key, values)
					
					
		return clear


	def set_config(self, cfg):
		"""Read a configuration file of channels and AOMs.

		Parameters
		----------

		cfg : string
			name of the configuration file.

		Notes
		-----
		The configuration object (configparser.ConfigParser) is available as self.config.



		"""
		self.config.read_file(io.StringIO(default_exp))
		self.config.read(cfg)

	def set_pattern(self, seqstr=None, rate=10000, minV=-10., maxV=10., max_acq=None, ext_trig=None, retriggerable=False):
		"""Set the pattern to the National Instruments device using PyDAQmx.

		Parameters
		----------
		seqraw : string
			a string describing the required output analogous to the old
			labedit 'sequence matrix'. It will be read by numpy.genfromtxt
			into a structured array.
			The column 't' is expected to be the timing in milliseconds.
			Other column should reference channels as specified in the configuration file.


		rate : float
			the rate of the National Instrument clock in hertz.
			default = 10 kHz

		minV : float
			minimum voltage output for the analog task.
			default = -10.

		maxV : float
			maximum voltage output for the analog task.
			default = 10.

		"""

		self.max_acq = max_acq

		# transform the sequence string  in a numpy structured array
		self.pattern = Pattern()
		
		if self.seqs:
			aseqs = [numpy.genfromtxt(io.StringIO(x), dtype=None, names=True) for x in self.seqs]
			seq = lrf.stack_arrays(aseqs, usemask=False)
		else:	
			seq  = numpy.genfromtxt(io.StringIO(seqstr), dtype=None, names=True)


		# check which analog and digital channels are used
		sequence_channels = seq.dtype.names

		# intersection of configured channels and sequence channels, with the order of the config file
		self.channels = [x for x in self.config.sections() if x in sequence_channels]
		self.analog_channels = [x for x in self.channels if self.config.getboolean(x, 'analog')]
		self.digital_channels = [x for x in self.channels if x not in self.analog_channels]


		used_channels = self.analog_channels + self.digital_channels
		all_channels = used_channels + ['t']

		# check for unconfigured channels
		for x in sequence_channels:
			if x not in all_channels:
				msg = 'Channel ' + x + ' in the sequence file not found in the configuration files. Do you mispelled it?'
				print(msg)

		# daqmx lines
		# TODO check for unspecified lines
		# if a channel has no line specified this rise a ConfigParser.NoOptionError
		self.analog_lines = [self.config.get(x,'line') for x in self.analog_channels]
		self.digital_lines = [self.config.get(x,'line') for x in self.digital_channels]

		# extract delays
		delays = dict([(x, self.config.getfloat(x,'delay')) for x in used_channels])

		# apply transformations
		# analog lines
		for channel in self.analog_channels:
			factor = self.config.getfloat(channel,'factor')
			offset = self.config.getfloat(channel,'offset')
			maxV = self.config.getfloat(channel,'maxV')
			minV = self.config.getfloat(channel,'minV')


			# apply factor and offset
			seq[channel] += offset
			seq[channel] *= factor
			# apply max and min
			seq[channel] = seq[channel].clip(minV,maxV)
		# digital lines
		for channel in self.digital_channels:
			invert = self.config.getboolean(channel,'invert')

			seq[channel] ^= invert

		if ext_trig is not None:
			self.ext_trig_channel =  self.config.get(ext_trig,'line')
		else:
			self.ext_trig_channel = None


		self.sequence = seq

		# push data to pattern
		self.pattern.write(self.sequence, alines = self.analog_lines, anames = self.analog_channels,
			dlines = self.digital_lines, dnames = self.digital_channels,
			rate = rate, minV = minV, maxV = maxV, delays = delays, ext_trig=self.ext_trig_channel, retriggerable=retriggerable)
		
		



	def set_acquisition(self, anames, trigger, duration, rate=10000, minV=-10., maxV=10., edge='rising', terminal='single', num=1):
		"""Set up the acquisition.

		Parameters
		----------
		anames : string or list of string
			names of the channels to be acquired
			(as specified in the config file)

		trigger : string
			channel of the NI board to be used as trigger input
		
		duration : float
			duration in milliseconds of the acuisition

		rate : float
			the rate of the acquisition in hertz

		minV : float
			minimum voltage expected to be read
			default = -10.

		maxV : float
			maximum voltage expected to be read
			default = 10.


		edge : {'rising', 'falling'}
			trigger edge
			default = 'rising'

		terminal : {'diff', 'single''}
			type of aqcusition
			default = 'diff'

		num : int
			number of pulses before acquisition is complete
			default = 1

		"""

		self.acq = Acquisition()

		ed = (DAQmx_Val_Falling if edge == 'falling' else DAQmx_Val_Rising)
		ter = (DAQmx_Val_Diff if terminal == 'diff' else DAQmx_Val_RSE)

		# if a single string is given, replace it with a len 1 list
		if isinstance(anames, str):
			anames = [anames]

		alines = [self.config.get(x.strip(),'line') for x in anames]

		#self.acq_num = num
		#self.acquire = [self.config.get(x,'name') for x in acquire]

		# use config file to check trigger line
		trigger = self.config.get(trigger.strip(),'line')

		self.acq.rread(alines, anames, num, trigger, duration, terminal = ter, minV = minV, maxV = maxV, rate = rate, trigger_edge = ed)


# 	def set_server(self, address='localhost', port=10003):
# 		# udp socket
# 		# no need to open the socket here.
# 		# self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
# 		self.server_address = (address, port)

	def get_synth(self, name):
		"""Get a Synth object from the configuration file.

		Parameters
		----------

		name : string
			name of the section of the config file to read as a synth.
			specifying address, channel, freq and pass.

		"""

		address = self.config.get(name, 'address')
		channel = self.config.getint(name, 'channel')
		freq = self.config.getfloat(name, 'freq')
		pas  = self.config.getint(name, 'pass')

		return Synth(address, channel, freq, pas, name=name)


	def add_cycle(self, what, seq=None, rep=1, synth=None, name=None, aux_synth=None, **kwargs):
		
		if not name:
			name = str(len(self.cycles)+1)
		
		
		
		
		if synth:
			s = self.get_synth(synth)
			self.synths += [s]
			#a = what(self.acq, synth=s, filo=self.file, name=name, **kwargs)
		else:
			s = None
		
		#print('aux', aux_synth)
		if aux_synth:
			auxs = self.get_synth(aux_synth)
			#print(auxs)
		else:
			auxs=None
		
		a = what(self.acq, synth=s, filo=self.file, name=name, aux_synth=auxs, **kwargs)

		
#		if synth:
#			s = self.get_synth(synth)
#			self.synths += [s]
#			a = what(self.acq, synth=s, filo=self.file, name=name, **kwargs)
#			
#		else:
#			a = what(self.acq, filo=self.file, name=name, **kwargs)

		self.cycles += [a]
		self.cycles_with_repeates += [a]*rep
		
		if seq:
			self.seqs += [seq]*rep
		
		return a
	
	def add_scopes(self, **kwargs):
		for i, c in enumerate(self.cycles):
			
			title = c.name + " - {:03d}".format(self.counter) + " " + self.started + " - " + self.name 
			
			
			
			s = c.set_scope(title=title, **kwargs)

			# arrange plots on screen
			# experimental!
			s.fig.canvas.manager.window.move(i*(s.size[0]+20) + 400, 25)

			
		
			t = c.set_control(title=title, **kwargs)
			
			
			
			# IMPORTANT! add the control layout to the main window layout
			self.layout += t.layout

			self.scopes +=[s]
			self.controls += [t]
			
			
	
	def add_note(self, x):
		self.notes += [x]
		
		
	def shuffle_cycles(self, shuffle = [0,2,1,3,4,6,5,7]):
		"""
		Reorder the cycle and sequence list.
		This is inspired by PTB-style interleaving.
		"""
		
		if len(self.cycles_with_repeates) != len(self.seqs):
			raise Exception("Cannot shuffle experiment with different numbers of sequences and cycles")
		
		if len(self.cycles_with_repeates) != len(shuffle):
			raise Exception("Shuffle sequence has different number of elements than original sequence")
			
		
		new_cycles = [self.cycles_with_repeates[x] for x in shuffle]
		new_seqs = [self.seqs[x] for x in shuffle]
		
		
		self.cycles_with_repeates = new_cycles
		self.seqs = new_seqs