"""
Classes and methods for working with Axona electrophysiological
recording data.
"""
import scipy, scipy.interpolate, scipy.stats, scipy.ndimage, scipy.signal
import numpy as np
import math
import os
import pickle
import fnmatch
from ephysiopy.common.utils import smooth
from ephysiopy.dacq2py.spikecalcs import SpikeCalcs

MAXSPEED = 4.0  # pos data speed filter in m/s
BOXCAR = 20  # this gives a 400ms smoothing window for pos averaging

empty_headers = {
	"tetrode" : os.path.join(os.path.dirname(__file__), "tetrode_header.pkl"),
	"pos" : os.path.join(os.path.dirname(__file__), "pos_header.pkl"),
	"set" : os.path.join(os.path.dirname(__file__), "set_header.pkl"),
	"eeg" : os.path.join(os.path.dirname(__file__), "eeg_header.pkl"),
	"egf" : os.path.join(os.path.dirname(__file__), "egf_header.pkl")
}

class IO(object):
	"""
	Axona data I/O. Also reads .clu files generated from KlustaKwik

	Parameters
	----------
	filename_root : str
		The fully-qualified filename
	"""

	tetrode_files = dict.fromkeys(["." + str(i) for i in range(1, 17)], [('ts', '>i'), ('waveform', '50b')])
	other_files = {'.pos': [('ts', '>i'), ('pos', '>8h')],
				   '.eeg': [('eeg', '=b')],
				   '.eeg2': [('eeg', '=b')],
				   '.egf': [('eeg', 'int16')],
				   '.egf2': [('eeg', 'int16')],
				   '.inp': [('ts', '>i4'), ('type', '>b'), ('value', '>2b')],
				   '.log': [('state', 'S3'), ('ts', '>i')],
				   '.stm': [('ts', '>i')]}

	axona_files = {**other_files, **tetrode_files} # this will only work in >= Python3.5

	def __init__(self, filename_root=''):
		self.filename_root = filename_root

	@staticmethod
	def getEmptyHeader(ftype: str)->dict:
		pname = empty_headers.get(ftype, '')
		if os.path.isfile(pname):
			with open(pname, 'rb') as f:
				return pickle.load(f)

	def getData(self, filename_root):
		"""
		Returns the data part of an Axona data file i.e. from "data_start" to
		"data_end"

		Parameters
		----------
		input :  str
			Fully qualified path name to the data file

		Returns
		-------
		output : ndarray
			The data part of whatever file was fed in
		"""
		n_samps = -1
		fType = os.path.splitext(filename_root)[1]
		if fType in self.axona_files:
			try:
				header = self.getHeader(filename_root)
				for key in header.keys():
					if len(fType) > 2:
						if fnmatch.fnmatch(key, 'num_*_samples'):
							n_samps = int(header[key])
					else:
						if key.startswith('num_spikes'):
							n_samps = int(header[key]) * 4
			except IOError:
				print('File type not recognised')
			f = open(filename_root, 'rb')
			data = f.read()
			st = data.find(b'data_start') + len('data_start')
			f.seek(st)
			if fType == '.log':
				f.seek(0)
			dt = np.dtype(self.axona_files[fType])
			a = np.fromfile(f, dtype=dt, count=n_samps)
			f.close()
		return a

	def getCluCut(self, tet):
		"""
		Load a clu file and return as an array of integers

		Parameters
		----------
		tet : int
			The tetrode the clu file relates to

		Returns
		-------
		out : ndarray
			Data read from the clu file
		"""
		filename_root = self.filename_root + '.' + 'clu.' + str(tet)
		dt = np.dtype([('data', '<i')])
		clu_data = np.loadtxt(filename_root, dtype=dt)
		return clu_data['data'][1::]  # first entry is number of clusters found

	def getCut(self, tet):
		"""
		Returns the cut file as a list of integers

		Parameters
		----------
		tet : int
			The tetrode the cut file relates to

		Returns
		-------
		out : ndarray
			The data read from the cut file
		"""
		a = []
		filename_root = self.filename_root + '_' + str(tet) + '.cut'
		if not os.path.exists(filename_root):
			cut = self.getCluCut(tet)
			return cut - 1
		with open(filename_root, 'r') as f:
			cut_data = f.read()
			f.close()
		tmp = cut_data.split('spikes: ')
		tmp1 = tmp[1].split('\n')
		cut = tmp1[1:]
		for line in cut:
			m = line.split()
			for i in m:
				a.append(int(i))
		return a

	def setHeader(self, filename_root: str, header: dict):
		"""
		Writes out the header to the specified file

		Parameters
		------------
		filename_root : str
			A fully qualified path to a file with the relevant suffix at
			the end (e.g. ".set", ".pos" or whatever)

		header : dict
			An empty version of which can be loaded using getEmptyHeader() above
		"""
		with open(filename_root, 'w') as f:
			for key, val in header.items():
				f.write(key)
				f.write(" ")
				if val is None:
					val = ""
				f.write(val)
				f.write('\r\n')
			f.write('data_start')
			f.write('\r\n')
			f.write('data_end')
			f.write('\r\n')

	def setData(self, filename_root: str, data: np.array):
		"""
		Writes Axona format data to the given filename

		Parameters
		----------
		filename_root : str
			The fully qualified filename including the suffix

		data : ndarray
			The data that will be saved
		"""
		fType = os.path.splitext(filename_root)[1]
		if fType in self.axona_files:
			f = open(filename_root, 'rb+')
			d = f.read()
			st = d.find(b'data_start') + len('data_start')
			f.seek(st)
			data.tofile(f)
			f.close()
			f = open(filename_root, 'a')
			f.write('\r\n')
			f.write('data_end')
			f.write('\r\n')
			f.close()

	def getHeader(self, filename_root):
		"""
		Reads and returns the header of a specified data file as a dictionary

		Parameters
		----------
		filename_root : str
			Fully qualified filename of Axona type

		Returns
		-------
		headerDict : dict
			key - value pairs of the header part of an Axona type file
		"""
		with open(filename_root, 'rb') as f:
			data = f.read()
			f.close()
		if os.path.splitext(filename_root)[1] != '.set':
			st = data.find(b'data_start') + len('data_start')
			header = data[0:st-len('data_start')-2]
		else:
			header = data
		headerDict = {}
		lines = header.splitlines()
		for line in lines:
			line = str(line.decode("ISO-8859-1")).rstrip()
			line = line.split(' ', 1)
			try:
				headerDict[line[0]] = line[1]
			except IndexError:
				headerDict[line[0]] = ''
		return headerDict

	def getHeaderVal(self, header, key):
		"""
		Get a value from the header as an int

		Parameters
		----------
		header : dict
			The header dictionary to read
		key : str
			The key to look up

		Returns
		-------
		value : int
			The value of `key` as an int
		"""
		tmp = header[key]
		val = tmp.split(' ')
		val = val[0].split('.')
		val = int(val[0])
		return val

class Pos(IO):
	"""
	Processs position data recorded with the Axona recording system

	Parameters
	----------
	filename_root : str
		The basename of the file i.e mytrial as opposed to mytrial.pos

	Notes
	-----
	Currently the only arg that does anything is 'cm' which will convert
	the xy data to cm, assuming that the pixels per metre value has been
	set correctly
	"""

	def __init__(self, filename_root, *args, **kwargs):
		self.filename_root = filename_root
		self.header = self.getHeader(filename_root + '.pos')
		self.setheader = None
		try:
			self.setheader = self.getHeader(filename_root + '.set')
		except Exception:
			pass
		self.posProcessed = False
		posData = self.getData(filename_root + '.pos')
		self.nLEDs = 1
		if self.setheader is not None:
			self.nLEDs = sum([self.getHeaderVal(self.setheader,'colactive_1'),
							  self.getHeaderVal(self.setheader,'colactive_2')])
		if self.nLEDs == 1:
			self.led_pos = np.ma.masked_values([posData['pos'][:,0],posData['pos'][:,1]],1023)
			self.led_pix = np.ma.masked_values(posData['pos'][:,4],1023)
		elif self.nLEDs == 2:
			self.led_pos = np.ma.masked_values([posData['pos'][:,0],posData['pos'][:,1],
											 posData['pos'][:,2],posData['pos'][:,3]],1023)
			self.led_pix = np.ma.masked_values([posData['pos'][:,4],posData['pos'][:,5]],1023)
		self.npos = len(self.led_pos[0])
		self.xy = np.ones([2,self.npos]) * np.nan
		self.dir = np.ones([self.npos]) * np.nan
		self.dir_disp = np.ones([self.npos]) * np.nan
		self.speed = np.ones([self.npos]) * np.nan
		self.pos_sample_rate = self.getHeaderVal(self.header, 'sample_rate')
		self._ppm = None
		if 'cm' in kwargs:
			self.cm = kwargs['cm']
		else:
			self.cm = False # if True return xy in cm, otherwise in pixels (see end of postprocesspos)

	@property
	def ppm(self):
		if self._ppm is None:
			try:
				self._ppm = self.getHeaderVal(self.header, 'pixels_per_metre')
			except IOError:
				self._ppm = None
		return self._ppm

	@ppm.setter
	def ppm(self, value):
		self._ppm = value
		self.posProcessed = False
		self.postprocesspos()

	def __getitem__(self, key):
		try:
			val = self.__dict__[key]
			return val
		except:
			pass

	def postprocesspos(self):
		"""
		Post processes position data

		Calls a few of the other methods in this class that deal with
		led swaps, impossibly fast points etc
		"""

		if self.posProcessed is True:
			return
		elif self.posProcessed is False:
			led_pos = self.led_pos
			led_pix = self.led_pix
			# as with AJ's implementation in mtint, calculate weights for a weighted mean
			# of the front and back leds for when trials have poorly tracked trials
			# NB could probably do this straight from the number of leds tracked which
			# is available in the raw data
			# need to mask all values of the array using logical or to replicate
			# mtints way of filling in missing values
			nLED_idx = self.nLEDs * 2
			led_pos[0:nLED_idx].__setmask__(led_pos[0:nLED_idx].mask.any(axis=0))
			weights = np.zeros(2)
			weights[0] = float(np.sum(np.nonzero(led_pos[0:2]), axis=1)[0]) / self.npos
			try:
				weights[1] = float(np.sum(np.nonzero(led_pos[2:4]), axis=1)[0]) / self.npos
			except IndexError:
				pass
			# need to deal with improperly tracked positions where the values are
			# plainly ridiculous
			# values less than 0 are masked
			led_pos[led_pos < 0] = np.ma.masked
			# deal with values outside the range of the tracked window
			led_pos[0, (led_pos[0] > int(self.header['max_x']))] = np.ma.masked
			led_pos[1, (led_pos[1] > int(self.header['max_y']))] = np.ma.masked
			# try and deal with other led if present
			try:
				led_pos[2, (led_pos[2] > int(self.header['max_x']))] = np.ma.masked
				led_pos[3, (led_pos[3] > int(self.header['max_y']))] = np.ma.masked
			except IndexError:
				pass
			if np.logical_and(np.any(np.nonzero(led_pix)), self.nLEDs==2):
				swap_list = self.ledswapFilter(led_pos, led_pix)
				tmp = led_pos[0:2, swap_list]
				led_pos[0:2, swap_list] = led_pos[2:4, swap_list]
				led_pos[2:4, swap_list] = tmp
				tmp = led_pix[0, swap_list]
				led_pix[0, swap_list] = led_pix[1, swap_list]
				led_pix[1, swap_list] = tmp
			ppm = self.ppm
			max_ppm_per_sample = MAXSPEED * ppm / self.pos_sample_rate
			led_pos = self.ledspeedFilter(led_pos,max_ppm_per_sample)
			led_pos = self.interpNans(led_pos)
			# get distances and angles of LEDs from rat
			pos1 = np.arange(0,self.npos)
			pos2 = np.arange(0,self.npos-1)
			if self.nLEDs == 1:
				self.xy[0:2,pos1] = led_pos[0:2,pos1]
				self.xy[0,:] = smooth(self.xy[0,:],BOXCAR,'flat')
				self.xy[1,:] = smooth(self.xy[1,:],BOXCAR,'flat')
				self.dir[pos2] = np.mod(((180/math.pi) * (np.arctan2(-self.xy[1,pos2+1] + self.xy[1,pos2],+self.xy[0,pos2+1]-self.xy[0,pos2]))), 360)
				self.dir[-1] = self.dir[-2]
				self.dir_disp = self.dir
			elif self.nLEDs == 2:
				lightBearings = np.zeros([2,1])
				lightBearings[0] = self.getHeaderVal(self.setheader,'lightBearing_1')
				lightBearings[1] = self.getHeaderVal(self.setheader,'lightBearing_2')
				front_back_xy_sm = np.zeros([4,self.npos])
				for i in range(len(front_back_xy_sm)):
#                    front_back_xy_sm[i,pos1] = scipy.signal.convolve(led_pos[i, pos1], np.ones(BOXCAR) / BOXCAR, mode='same')
					front_back_xy_sm[i,pos1] = smooth(led_pos[i,pos1],BOXCAR,'flat')
				correction = lightBearings[0]
				self.dir[pos1] = np.mod((180/math.pi) * (np.arctan2(-front_back_xy_sm[1,pos1]+front_back_xy_sm[3,pos1],
								  +front_back_xy_sm[0,pos1]-front_back_xy_sm[2,pos1])-correction),360)
				# get xy from smoothed individual lights weighting for reliability
				self.xy[0,pos1] = (weights[0]*front_back_xy_sm[0,pos1] + weights[1]*front_back_xy_sm[2,pos1]) / np.sum(weights)
				self.xy[1,pos1] = (weights[0]*front_back_xy_sm[1,pos1] + weights[1]*front_back_xy_sm[3,pos1]) / np.sum(weights)
				self.dir_disp[pos2] = np.mod(((180/math.pi) * (np.arctan2(-self.xy[1,pos2+1] + self.xy[1,pos2],+self.xy[0,pos2+1]-self.xy[0,pos2]))) ,360)
				self.dir_disp[-1] = self.dir_disp[-2]

			if self.cm:
				self.xy = self.xy / ppm * 100 # xy now in cm
			# calculate speed based on distance
			self.speed[pos2] = np.sqrt(np.sum(np.power(np.diff(self.xy),2),0))
			self.speed[self.npos-1] = self.speed[-1]
			self.speed = self.speed * (100 * self.pos_sample_rate / ppm) # *100 to get into cm/s
			if np.isnan(self.speed[-1]):
				self.speed[-1] = 0

			self.posProcessed = True

	def ledspeedFilter(self,led_pos,max_ppm_per_sample):
		"""
		Filters for impossibly fast tracked points

		Parameters
		----------
		led_pos : np.ma.MaskedArray
			Masked led_pos array [x1,y1,x2,y2]
		max_ppm_per_sample : int
			The maximum distance (in pixels) that points are allowed to jump

		Returns
		-------
		led_pos : np.ma.MaskedArray
			The filtered data
		"""
		max_ppms_sqd = max_ppm_per_sample ** 2
		for i in range(0,len(led_pos),2):
			ok_pos = led_pos[i,:]
			prev_pos = ok_pos[0:-1]
			cur_pos = ok_pos[1:]
			pix_per_sample_sqd = (np.power((np.subtract(led_pos[i,cur_pos], led_pos[i,prev_pos])),2) + np.power((np.subtract(led_pos[i+1,cur_pos], led_pos[i+1,prev_pos])),2)) / np.power(np.subtract(cur_pos,prev_pos),2)
			pix_per_sample_sqd = np.insert(pix_per_sample_sqd, -1, 0)
			led_pos[i:i+2,pix_per_sample_sqd > max_ppms_sqd] = np.ma.masked
		return led_pos

	def ledswapFilter(self,led_pos,led_pix):
		"""
		Checks for led swapping in 2-spot mode

		Parameters
		----------
		led_pos : np.ma.MaskedArray
			Masked array of dims [4 x nPosSamples]
		led_pix : np.ma.MaskedArray
			The number of pixels tracked for each LED.
			Has dimensions [2 x nPosSamples]

		Returns
		-------
		swap_list : ndarray
			Array of swapped positions
		"""
		thresh = 5
		mean_npix = led_pix.mean(axis=1).data
		std_npix = led_pix.std(axis=1).data
		pos = np.arange(1,led_pix.shape[1])
		#calculate distances
		dist12 = np.sqrt(np.nansum(((np.squeeze(led_pos[0:2,pos])-np.squeeze(led_pos[2:4,pos-1]))**2),axis=0))
		dist11 = np.sqrt(np.nansum(((np.squeeze(led_pos[0:2,pos])-np.squeeze(led_pos[0:2,pos-1]))**2),axis=0))
		dist21 = np.sqrt(np.nansum(((np.squeeze(led_pos[2:4,pos])-np.squeeze(led_pos[0:2,pos-1]))**2),axis=0))
		dist22 = np.sqrt(np.nansum(((np.squeeze(led_pos[2:4,pos])-np.squeeze(led_pos[2:4,pos-1]))**2),axis=0))
		switched = np.logical_or(np.logical_and((dist12 < dist11 - thresh).data,led_pos[2,pos].mask),(dist21 < dist22-thresh).data)
		z11 = (mean_npix[0] - led_pix[0,pos]) / std_npix[0]
		z12 = (led_pix[0,pos] - mean_npix[1]) / std_npix[1]
		shrunk = z11 > z12
		swap_list = np.nonzero(np.logical_and(switched, shrunk.data))[0] + 1
		return swap_list

	def interpNans(self,led_pos):
		"""
		Interpolates over missing values

		Parameters
		----------
		led_pos : np.ma.MaskedArray
			The LED positions

		Returns
		-------
		output : np.ma.MaskedArray
			The smoothed, unmasked array
		"""
		for i in range(0,len(led_pos),2):
			missing = led_pos[i:i+2].mask.any(axis=0)
			ok = np.logical_not(missing)
			ok_idx = ok.ravel().nonzero()[0]#gets the indices of ok poses
			missing_idx = missing.ravel().nonzero()[0]#get the indices of missing poses
			good_data = led_pos.data[i,ok_idx]
			good_data1 = led_pos.data[i+1,ok_idx]
			led_pos.data[i,missing_idx] = np.interp(missing_idx,ok_idx,good_data)#,left=np.min(good_data),right=np.max(good_data)
			led_pos.data[i+1,missing_idx] = np.interp(missing_idx,ok_idx,good_data1) #,left=np.min(good_data1),right=np.max(good_data1) y coord
		#unmask the array
		led_pos.mask = 0
		return led_pos
	def filterPos(self, filterDict):
		"""
		Filters position data depending on the specified filter

		Parameters
		----------
		filterDict : dict
			Contains the type(s) of filter to be used and the
			range of values to filter for. Values are pairs specifying the range
			of values to filter for NB can take multiple filters and iteratively apply them
			legal values are:
			* 'dir' - the directional range to filter for NB this can contain 'w','e','s' or 'n'
			* 'speed' - min and max speed to filter for
			* 'xrange' - min and max values to filter x pos values
			* 'yrange' - same as xrange but for y pos
			* 'time' - the times to keep / remove specified in ms

		Returns
		--------
		pos_index_to_keep : ndarray
			The position indices that should be kept
		"""
		if filterDict is None:
			return
		nSamples = int(self.header['num_pos_samples'])
		bool_arr = np.ones(shape=(len(filterDict), nSamples), dtype=np.bool)
		for idx, key in enumerate(filterDict):
			if isinstance(filterDict[key], str):
				if len(filterDict[key]) == 1 and 'dir' in key:
					if 'w' in filterDict[key]:
						filterDict[key] = (135, 225)
					elif 'e' in filterDict[key]:
						filterDict[key] = (315, 45)
					elif 's' in filterDict[key]:
						filterDict[key] = (225, 315)
					elif 'n' in filterDict[key]:
						filterDict[key] = (45, 135)
				else:
					raise ValueError("filter must contain a key / value pair")
			if 'speed' in key:
				if filterDict[key][0] > filterDict[key][1]:
					raise ValueError("First value must be less than the second one")
				else:
					bool_arr[idx,:] = np.logical_and(self.speed > filterDict[key][0],
									 self.speed < filterDict[key][1])
			elif 'dir' in key:
				if filterDict[key][0] < filterDict[key][1]:
					bool_arr[idx,:] = np.logical_and(self.dir > filterDict[key][0],
										 self.dir < filterDict[key][1])
				else:
					bool_arr[idx,:] = np.logical_or(self.dir > filterDict[key][0],
										self.dir < filterDict[key][1])
			elif 'xrange' in key:
				bool_arr[idx, :] = np.logical_and(self.xy[0, :] > filterDict[key][0],
										self.xy[0, :] < filterDict[key][1])
			elif 'yrange' in key:
				bool_arr[idx, :] = np.logical_and(self.xy[1, :] > filterDict[key][0],
										self.xy[1, :] < filterDict[key][1])
			elif 'time' in key:
				# takes the form of 'from' - 'to' times in SECONDS such that only pos's between these ranges are KEPT
				filterDict[key] = filterDict[key]  * self.pos_sample_rate
				if filterDict[key].ndim == 1:
					bool_arr[idx, filterDict[key][0]:filterDict[key][1]] = False
				else:
					for i in filterDict[key]:
						bool_arr[idx, i[0]:i[1]] = False
				bool_arr = ~bool_arr
			else:
				print("Unrecognised key in dict")
				pass
		return np.expand_dims(np.any(~bool_arr, axis=0), 0)

class Tetrode(IO, SpikeCalcs):
	"""
	Processes tetrode files recorded with the Axona recording system

	Mostly this class deals with interpolating tetrode and position timestamps
	and getting indices for particular clusters.

	Parameters
	---------
	filename_root : str
		The fully qualified name of the file without it's suffix
	tetrode : int
		The number of the tetrode
	volts : bool, optional
		Whether to convert the data values volts. Default True
	"""

	def __init__(self, filename_root, tetrode, volts=True):
		self.filename_root = filename_root
		self.tetrode = tetrode
		self.volts = volts
		self.header = self.getHeader(self.filename_root + '.' + str(tetrode))
		data = self.getData(filename_root + '.' + str(tetrode))
		self.spk_ts = data['ts'][::4]
		self.nChans = self.getHeaderVal(self.header, 'num_chans')
		self.samples = self.getHeaderVal(self.header, 'samples_per_spike')
		self.nSpikes = self.getHeaderVal(self.header, 'num_spikes')
		self.posSampleRate = self.getHeaderVal(self.getHeader(self.filename_root + '.' + 'pos'), 'sample_rate')
		self.waveforms = data['waveform'].reshape(self.nSpikes, self.nChans, self.samples)
		del data
		if volts:
			set_header = self.getHeader(self.filename_root + '.set')
			gains = np.zeros(4)
			st = (tetrode - 1) * 4
			for i, g in enumerate(np.arange(st, st+4)):
				gains[i] = int(set_header['gain_ch_' + str(g)])
			ADC_mv = int(set_header['ADC_fullscale_mv'])
			scaling = (ADC_mv/1000.) / gains
			self.scaling = scaling
			self.gains = gains
			self.waveforms = (self.waveforms / 128.) * scaling[:,np.newaxis]# waveforms now in volts
		self.timebase = self.getHeaderVal(self.header, 'timebase')
		try:
			cut = np.array(self.getCut(self.tetrode), dtype=int)
			self.cut = cut
			self.clusters = np.unique(self.cut)
		except IOError:
			try:
				cut = self.getCluCut(self.tetrode)
				cut = np.array(cut) - 1
				self.cut = cut
				self.clusters = np.unique(self.cut)
			except IOError:
				self.cut = None
		self.pos_samples = None

	def getSpkTS(self):
		"""
		Return all the timestamps for all the spikes on the tetrode
		"""
		return np.ma.compressed(self.spk_ts)

	def getClustTS(self, cluster=None):
		"""
		Returns the timestamps for a cluster on the tetrode

		Parameters
		----------
		cluster : int
			The cluster whose timestamps we want

		Returns
		-------
		clustTS : ndarray
			The timestamps

		Notes
		-----
		If None is supplied as input then all timestamps for all clusters
		is returned i.e. getSpkTS() is called
		"""
		if cluster is None:
			clustTS = self.getSpkTS()
		else:
			if self.cut is None:
				try:
					cut = np.array(self.getCut(self.tetrode),dtype=int)
				except IOError:
					cut = self.getCluCut(self.tetrode)
					cut = np.array(cut) - 1
				self.cut = cut
			self.getSpkTS()
			clustTS = np.ma.compressed(self.spk_ts[self.cut==cluster])
		return clustTS

	def getPosSamples(self):
		"""
		Returns the pos samples at which the spikes were captured
		"""
		self.pos_samples = np.floor(self.getSpkTS() / float(self.timebase) * self.posSampleRate).astype(int)
		return np.ma.compressed(self.pos_samples)

	def getClustSpks(self, cluster):
		"""
		Returns the waveforms of `cluster`

		Parameters
		----------
		cluster : int
			The cluster whose waveforms we want

		Returns
		-------
		waveforms : ndarray
			The waveforms on all 4 electrodes of the tgtrode so the shape of
			the returned array is [nClusterSpikes, 4, 50]
		"""
		if self.cut is None:
			self.getClustTS(cluster)
		return self.waveforms[self.cut==cluster, :, :]#taking the mean of this along axis=0 gives mean waveform on each channel for a cluster

	def getClustIdx(self, cluster):
		"""
		Get the indices of the position samples corresponding to the cluster

		Parameters
		----------
		cluster : int
			The cluster whose position indices we want

		Returns
		-------
		pos_samples : ndarray
			The indices of the position samples, dtype is int
		"""
		if self.cut is None:
			try:
				cut = np.array(self.getCut(self.tetrode), dtype=int)
			except IOError:
				cut = self.getCluCut(self.tetrode)
				cut = np.array(cut) - 1
			self.cut = cut
		if self.pos_samples is None:
			self.getPosSamples()
		return self.pos_samples[self.cut == cluster].astype(int)

	def getUniqueClusters(self):
		"""
		Returns the unique clusters
		"""
		if self.cut is None:
			try:
				cut = np.array(self.getCut(self.tetrode), dtype=int)
			except IOError:
				cut = self.getCluCut(self.tetrode)
				cut = np.array(cut) - 1
			self.cut = cut
		else:
			cut = self.cut
		return np.unique(cut)

class EEG(IO):
	"""
	Processes eeg data collected with the Axona recording system

	Parameters
	---------
	filename_root : str
		The fully qualified filename without the suffix
	eeg_file, egf : int
		Whether to read the 'eeg' file or the 'egf' file. 0 is False, 1 is True
	"""

	def __init__(self, filename_root, eeg_file=1, egf=0):
		self.showfigs = 0
		self.filename_root = filename_root
		if egf == 0:
			denom = 128.0 #used below to normalise data
			if eeg_file == 1:
				eeg_suffix = '.eeg'
			else:
				eeg_suffix = '.eeg' + str(eeg_file)
		elif egf == 1:
			denom = 128.0 # used below to normalise data
			if eeg_file == 1:
				eeg_suffix = '.egf'
			else:
				eeg_suffix = '.egf' + str(eeg_file)
		self.header = self.getHeader(self.filename_root + eeg_suffix)
		self.eeg = self.getData(filename_root + eeg_suffix)['eeg']
		# sometimes the eeg record is longer than reported in the 'num_EEG_samples'
		# value of the header so eeg record should be truncated to match 'num_EEG_samples'
		# TODO: this could be taken care of in the IO base class
		if egf:
			self.eeg = self.eeg[0:int(self.header['num_EGF_samples'])]
		else:
			self.eeg = self.eeg[0:int(self.header['num_EEG_samples'])]
		if egf == 1:
			# self.eeg = self.eeg.view(np.int8).reshape(self.eeg.shape+(2,)) # fast!
			self.eeg = self.eeg[:, 1]
		self.sample_rate = int(self.getHeaderVal(self.header, 'sample_rate'))
		set_header = self.getHeader(self.filename_root + '.set')
		eeg_ch = int(set_header['EEG_ch_1']) - 1
		if eeg_ch < 0:
			eeg_ch = 0
		eeg_gain = int(set_header['gain_ch_' + str(eeg_ch)])
		# EEG polarity is determined by the "mode_ch_n" key in the setfile
		# where n is the channel # for the eeg. The possibles values to these
		# keys are as follows:
		#	0 = Signal
		#	1 = Ref
		#	2 = -Signal
		#	3 = -Ref
		#	4 = Sig-Ref
		#	5 = Ref-Sig
		#	6 = grounded
		# So if the EEG has been recorded with -Signal (2) then the recorded polarity
		# is inverted with respect to that in the brain
		eeg_mode = int(set_header['mode_ch_' + set_header['EEG_ch_1']])
		polarity = 1 # ensure it always has a value
		if eeg_mode == 2:
			polarity = -1
		ADC_mv = float(set_header['ADC_fullscale_mv'])
		scaling = (ADC_mv/1000.) * eeg_gain
		self.scaling = scaling
		self.gain = eeg_gain
		self.polarity = polarity
		self.eeg = (self.eeg / denom) * scaling * polarity# eeg now in microvolts
		self.EEGphase = None
		# x1 / x2 are the lower and upper limits of the eeg filter
		self.x1 = 6
		self.x2 = 12

	def eegfilter(self, E=None):
		"""
		Filter the LFP signal

		Filters the eeg using a 251-tap bandpass (6-12Hz) blackman filter
		between the values given in x1 and x2: defaults to filtering between theta
		frequency (6-12Hz)

		Parameters
		----------
		E : ndarray
			The LFP signal

		Returns
		-------
		filtEEG : ndarray
			The filtered LFP signal
		"""
		if E is None:
			E = self.eeg
		nyquist = self.sample_rate / 2.
		eegfilter = scipy.signal.firwin(int(self.sample_rate) + 1, [self.x1/nyquist, self.x2/nyquist], window='black', pass_zero=False)
		filtEEG = scipy.signal.filtfilt(eegfilter, [1], E.ravel(), padtype='odd')
		if np.ma.is_masked(self.eeg):
			mask = np.ma.getmask(self.eeg)
			return np.ma.masked_where(mask, filtEEG)
		else:
			return filtEEG

	def thetaAmpPhase(self, fx=None):
		"""
		Extracts the amplitude and instantaneous frequency of the LFP

		Parameters
		----------
		fs : ndarray
			The filtered LFP signal, typically filtered using eegfilter()
		"""
		if fx is None:
			fx = self.eegfilter(E=self.eeg)
		analytic = scipy.signal.hilbert(fx)
		self.EEGphase = np.angle(analytic)
		self.UWphase = np.unwrap(self.EEGphase)
		tmp = np.append(self.UWphase, np.nan)
		self.EEGinstfreq = np.diff(tmp) * (self.sample_rate / (2*np.pi))
		self.thAmp = np.abs(analytic)

	def nextpow2(self, val):
		"""
		Calculates the next power of 2 that will hold val

		Parameters
		----------
		val : int
			The value to find the next power of 2 of

		Returns
		-------
		out : ndarray
			The next power of 2 that contains val
		"""
		val = val - 1
		val = (val >> 1) | val
		val = (val >> 2) | val
		val = (val >> 4) | val
		val = (val >> 8) | val
		val = (val >> 16) | val
		val = (val >> 32) | val
		return np.log2(val + 1)

class Stim(dict, IO):
	"""
	Processes the stimulation data recorded using Axona

	Parameters
	----------
	filename_root : str
		The fully qualified filename without the suffix
	"""
	def __init__(self, filename_root, *args, **kwargs):
		self.update(*args, **kwargs)
		self.filename_root = filename_root
		stmData = self.getData(filename_root + '.stm')
		self.__setitem__('on', stmData['ts'])
		stmHdr = self.getHeader(filename_root + '.stm')
		for k,v in stmHdr.items():
			self.__setitem__(k, v)
		tb = int(self['timebase'].split(' ')[0])
		self.timebase = tb

	def update(self, *args, **kwargs):
		d = dict(*args, **kwargs)
		for k in d.keys():
			self[k] = d[k]

	def __getitem__(self, key):
		try:
			val = dict.__getitem__(self, key)
			return val
		except KeyError:
			print('KeyError')

	def __setitem__(self, key, val):
		dict.__setitem__(self, key, val)

	def getTS(self):
		"""
		Gets the timestamps of the on events
		"""
		return self['on'] / int(self.timebase / 1000)# in ms

	def getPosIdx(self):
		"""
		Gets the position indices of the on events
		"""
		scale = self.timebase / float(self['posSampRate'])
		return self['on'] / scale

	def getEEGIdx(self):
		"""
		Gets the EEG indices of the on events
		"""
		scale = self.timebase / float(self['eegSampRate'])
		return (self['on'] / scale).astype(int)

	def getEGFIdx(self):
		"""
		Gets the EGF indices of the on events
		"""
		scale = self.timebase / float(self['egfSampRate'])
		return (self['on'] / scale).astype(int)


import os
import cv2
import shutil
import numpy as np

font = cv2.FONT_HERSHEY_SIMPLEX
PREDEFINE_LEN = 6

class boxing:
    def __init__(self, path, label_path):
        self.Path = path  #####视频路径
        self.draw = False  #####画框的标志
        self.pt0 = None  #####矩形框左上角坐标
        self.pt1 = None  #####矩形框右下角坐标
        self._boxes = np.zeros([0,4],dtype='int')  #####用来存储每一个框的坐标信息
        self._all_boxes = np.zeros([0,5],dtype='int')
        self._all_cls = np.zeros([0],dtype='int')
        self.base_name = os.path.basename(self.Path)
        self.base_name = os.path.splitext(self.base_name)[0]
        self.label_path = os.path.join(label_path, self.base_name + '.txt')#bbox 文件
        self.label_path_cls = os.path.join(label_path, self.base_name + '_cls' + '.txt')#分类文件
        self.cur_idx = 0
        self.cur_len = 0
        self.ratio = 1
        if os.path.exists(self.label_path):
            self._all_boxes = np.loadtxt(self.label_path,dtype='int')
            self._all_boxes = self._all_boxes.reshape([-1,5])
        if os.path.exists(self.label_path_cls):
            self._all_cls = np.loadtxt(self.label_path_cls,dtype='int')
            self._all_cls = self._all_cls.reshape([-1,1])

    def _on_mouse(self, event, x, y, flags, para):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.draw = True
            self.pt0 = np.array([x,y])
        elif event == cv2.EVENT_LBUTTONUP:
            self.draw = False
            self.pt1 = np.array([x,y])
            xx = np.concatenate((self.pt0,self.pt1),0).reshape([1,4])
            # self._boxes.append(np.concatenate((self.pt0,self.pt1),0))
            self._boxes = np.concatenate((self._boxes,xx))
        elif event == cv2.EVENT_MOUSEMOVE:  #########鼠标移动的时候把画框轨迹显示出来
            self.pt1 = np.array([x,y])

    def _drawing_box(self, img):
        canvas = cv2.copyMakeBorder(img, 0, 16, 0, 0, cv2.BORDER_CONSTANT, value=(192, 192, 192))
        anno_progress = str(self.cur_idx) + '/' + str(self.cur_len - PREDEFINE_LEN)
        cv2.putText(canvas,anno_progress,(0,50),font,2,(0,255,0),2)

        if self._cls == 0:
            cv2.putText(canvas,'NEG',(img.shape[1]//2,img.shape[0]//2),font,2,(0,255,0),2)
        elif self._cls == 1:
            cv2.putText(canvas,'POS',(img.shape[1]//2,img.shape[0]//2),font,2,(0,0,255),2)
        else:
            cv2.putText(canvas,'IGNORE',(img.shape[1]//2,img.shape[0]//2),font,2,(255,0,0),2)

        for box in self._boxes:
            bpt0 = (int(box[0]),int(box[1]))
            bpt1 = (int(box[2]),int(box[3]))
            cv2.rectangle(canvas, bpt0, bpt1, (255, 0, 0), 2)

        ###画每一个框显示轨迹
        if self.draw:
            cv2.rectangle(canvas, tuple(self.pt0), tuple(self.pt1), (255, 0, 0), thickness=2)
        return canvas

    def start(self):
        #加载所有的视频
        video_cap = cv2.VideoCapture(self.Path)
        frame_list = []
        ratio = 1
        standard_height = 0
        standard_width = 1
        while True:
            is_read,img = video_cap.read()
            if is_read:
                ratio = 600.0 / max(img.shape[1], img.shape[0])
                standard_height = img.shape[0]
                standard_width = img.shape[1]
                img = cv2.resize(img, (int(float(img.shape[1]) * ratio), int(float(img.shape[0]) * ratio)))
                frame_list.append(img)
            else:
                break
        self.ratio = ratio
        self.cur_len = len(frame_list)
        continue_anno = True
        win_name = "action_window"
        cv2.namedWindow(win_name)
        cv2.setMouseCallback(win_name, self._on_mouse)
        while continue_anno:
            img_stack = []
            for i in range(PREDEFINE_LEN):
                img_stack.append(frame_list[self.cur_idx+i])

            loop = True

            ##重新读入label
            self._boxes = self._all_boxes[self._all_boxes[:, 0] == self.cur_idx][:, 1:]
            self._boxes = self._boxes.astype('float32') * self.ratio

            if self._all_cls.shape[0] == 0:
                self._all_cls = np.ones([len(frame_list)],dtype='int')*-1

            self._cls = self._all_cls[self.cur_idx]

            while loop:
                for s_img in img_stack:
                    cp_img = s_img.copy()
                    cp_img = self._drawing_box(cp_img)
                    cv2.imshow(win_name, cp_img)
                    k = cv2.waitKey(200)
                    if k == ord('s'):
                        t_boxes = np.array(self._boxes)
                        t_boxes = np.floor(t_boxes/ratio)
                        t_boxes = t_boxes.astype('int')
                        t_boxes[t_boxes<0] = 0
                        tmp_boxes_w = t_boxes[:,[0,2]]
                        tmp_boxes_h = t_boxes[:,[1,3]]
                        # 调换位置
                        for child_box in tmp_boxes_w:
                            if child_box[0] > child_box[1]:
                                child_box[0], child_box[1] = child_box[1], child_box[0]
                        for child_box in tmp_boxes_h:
                            if child_box[0] > child_box[1]:
                                child_box[0], child_box[1] = child_box[1], child_box[0]
                        tmp_boxes_w[tmp_boxes_w>=standard_width] = standard_width-1
                        tmp_boxes_h[tmp_boxes_h>=standard_height] = standard_height-1
                        t_boxes[:,[0,2]] = tmp_boxes_w
                        t_boxes[:,[1,3]] = tmp_boxes_h
                        add_idx = np.ones([t_boxes.shape[0],1],dtype='int')*self.cur_idx
                        t_boxes = np.concatenate((add_idx,t_boxes),1)

                        if t_boxes.shape[0] != 0:
                            self._all_boxes = self._all_boxes[~(self._all_boxes[:, 0] == self.cur_idx)]
                            self._all_boxes = np.concatenate((self._all_boxes,t_boxes),0)
                            np.savetxt(self.label_path,self._all_boxes,fmt='%d')

                            self._all_cls[self.cur_idx] = 1
                            np.savetxt(self.label_path_cls, self._all_cls, fmt='%d')

                        # self.cur_idx += 1
                        # self.cur_idx = min(self.cur_idx,len(frame_list)-1)
                        loop = False
                        break
                    if k == ord('c'):
                        self._boxes = np.zeros([0, 4], dtype='int')
                        break
                    if k == ord('0'):#neg
                        loop = False
                        if not os.path.exists(self.label_path):
                            np.savetxt(self.label_path, np.ones([0,5],dtype=np.int32), fmt='%d')
                        else:
                            self._all_boxes = self._all_boxes[~(self._all_boxes[:, 0] == self.cur_idx)]
                            np.savetxt(self.label_path, self._all_boxes, fmt='%d')
                        self._all_cls[self.cur_idx] = 0
                        np.savetxt(self.label_path_cls, self._all_cls, fmt='%d')
                        break

                    if k == ord('r'):#neg
                        loop = False
                        if not os.path.exists(self.label_path):
                            np.savetxt(self.label_path, np.ones([0,5],dtype=np.int32), fmt='%d')
                        else:
                            self._all_boxes = self._all_boxes[~(self._all_boxes[:, 0] == self.cur_idx)]
                            np.savetxt(self.label_path, self._all_boxes, fmt='%d')
                        self._all_cls[self.cur_idx] = -1
                        np.savetxt(self.label_path_cls, self._all_cls, fmt='%d')
                        break

                    if k == ord('a'):
                        self.cur_idx -= 1
                        self.cur_idx = max(0,self.cur_idx)
                        loop = False
                        break
                    if k == ord('d'):
                        self.cur_idx += 1
                        self.cur_idx = min(self.cur_idx,len(frame_list)-6)
                        loop = False
                        break
                    if k == ord('q'):
                        loop = False
                        continue_anno = False
                        break
            # cv2.destroyAllWindows()

if __name__ == '__main__':
    tar_path = r'/media/hzh/work/workspace/data/fighting/positive'
    label_path = tar_path+'_label'

    os.makedirs(label_path,exist_ok=True)
    sub_video_list = os.listdir(tar_path)
    suffix = '.mp4'
    for i, sub_video in enumerate(sub_video_list):
        if sub_video.endswith(suffix):
            print(i + 1, "--->%d" % len(sub_video_list))
            full_sub_video = os.path.join(tar_path, sub_video)
            instance_ = boxing(full_sub_video,label_path)
            instance_.start()


"""Shared utilities for ot3 hardware control."""
from typing import Dict, Iterable, List, Tuple, TypeVar
from opentrons.config.types import OT3MotionSettings, OT3CurrentSettings, GantryLoad
from opentrons.hardware_control.types import (
    OT3Axis,
    OT3AxisKind,
    OT3AxisMap,
    CurrentConfig,
    OT3SubSystem,
    OT3Mount,
)
import numpy as np

from opentrons_hardware.firmware_bindings.constants import NodeId
from opentrons_hardware.hardware_control.motion_planning import (
    AxisConstraints,
    SystemConstraints,
    Coordinates,
    Move,
    CoordinateValue,
)
from opentrons_hardware.hardware_control.tool_sensors import ProbeTarget
from opentrons_hardware.hardware_control.motion_planning.move_utils import (
    unit_vector_multiplication,
)
from opentrons_hardware.hardware_control.motion import (
    create_step,
    NodeIdMotionValues,
    create_home_step,
    MoveGroup,
    MoveStopCondition,
)

# TODO: These methods exist to defer uses of NodeId to inside
# method bodies, which won't be evaluated until called. This is needed
# because the robot server doesn't have opentrons_ot3_firmware as a dep
# which is where they're defined, and therefore you can't have references
# to NodeId that are interpreted at import time because then the robot
# server tests fail when importing hardware controller. This is obviously
# terrible and needs to be fixed.

def axis_nodes() -> List["NodeId"]:
    return [
        NodeId.gantry_x,
        NodeId.gantry_y,
        NodeId.head_l,
        NodeId.head_r,
        NodeId.pipette_left,
        NodeId.pipette_right,
        NodeId.gripper_z,
    ]

def node_axes() -> List[OT3Axis]:
    return [
        OT3Axis.X,
        OT3Axis.Y,
        OT3Axis.Z_L,
        OT3Axis.Z_R,
        OT3Axis.P_L,
        OT3Axis.P_R,
        OT3Axis.Z_G,
    ]

def home_axes() -> List[OT3Axis]:
    return [
        OT3Axis.P_L,
        OT3Axis.P_R,
        OT3Axis.Z_L,
        OT3Axis.Z_R,
        OT3Axis.Z_G,
        OT3Axis.X,
        OT3Axis.Y,
    ]

def axis_to_node(axis: OT3Axis) -> "NodeId":
    anm = {
        OT3Axis.X: NodeId.gantry_x,
        OT3Axis.Y: NodeId.gantry_y,
        OT3Axis.Z_L: NodeId.head_l,
        OT3Axis.Z_R: NodeId.head_r,
        OT3Axis.P_L: NodeId.pipette_left,
        OT3Axis.P_R: NodeId.pipette_right,
        OT3Axis.Z_G: NodeId.gripper_z,
        OT3Axis.G: NodeId.gripper_g,
    }
    return anm[axis]

def node_to_axis(node: "NodeId") -> OT3Axis:
    nam = {
        NodeId.gantry_x: OT3Axis.X,
        NodeId.gantry_y: OT3Axis.Y,
        NodeId.head_l: OT3Axis.Z_L,
        NodeId.head_r: OT3Axis.Z_R,
        NodeId.pipette_left: OT3Axis.P_L,
        NodeId.pipette_right: OT3Axis.P_R,
        NodeId.gripper_z: OT3Axis.Z_G,
        NodeId.gripper_g: OT3Axis.G,
    }
    return nam[node]

def node_is_axis(node: "NodeId") -> bool:
    try:
        node_to_axis(node)
        return True
    except KeyError:
        return False

def axis_is_node(axis: OT3Axis) -> bool:
    try:
        axis_to_node(axis)
        return True
    except KeyError:
        return False

def sub_system_to_node_id(sub_sys: OT3SubSystem) -> "NodeId":
    """Convert a sub system to a NodeId."""
    nam = {
        OT3SubSystem.gantry_x: NodeId.gantry_x,
        OT3SubSystem.gantry_y: NodeId.gantry_y,
        OT3SubSystem.head: NodeId.head,
        OT3SubSystem.pipette_left: NodeId.pipette_left,
        OT3SubSystem.pipette_right: NodeId.pipette_right,
        OT3SubSystem.gripper: NodeId.gripper,
    }
    return nam[sub_sys]

def get_current_settings(
    config: OT3CurrentSettings,
    gantry_load: GantryLoad,
) -> OT3AxisMap[CurrentConfig]:
    conf_by_pip = config.by_gantry_load(gantry_load)
    currents = {}
    for axis_kind in [OT3AxisKind.P, OT3AxisKind.X, OT3AxisKind.Y, OT3AxisKind.Z]:
        for axis in OT3Axis.of_kind(axis_kind):
            currents[axis] = CurrentConfig(
                conf_by_pip["hold_current"][axis_kind],
                conf_by_pip["run_current"][axis_kind],
            )
    return currents

def get_system_constraints(
    config: OT3MotionSettings,
    gantry_load: GantryLoad,
) -> "SystemConstraints[OT3Axis]":
    conf_by_pip = config.by_gantry_load(gantry_load)
    constraints = {}
    for axis_kind in [OT3AxisKind.P, OT3AxisKind.X, OT3AxisKind.Y, OT3AxisKind.Z]:
        for axis in OT3Axis.of_kind(axis_kind):
            constraints[axis] = AxisConstraints.build(
                conf_by_pip["acceleration"][axis_kind],
                conf_by_pip["max_speed_discontinuity"][axis_kind],
                conf_by_pip["direction_change_speed_discontinuity"][axis_kind],
            )
    return constraints

def _convert_to_node_id_dict(
    axis_pos: Coordinates[OT3Axis, CoordinateValue],
) -> NodeIdMotionValues:
    target: NodeIdMotionValues = {}
    for axis, pos in axis_pos.items():
        if axis_is_node(axis):
            target[axis_to_node(axis)] = np.float64(pos)
    return target

def create_move_group(
    origin: Coordinates[OT3Axis, CoordinateValue],
    moves: List[Move[OT3Axis]],
    present_nodes: Iterable[NodeId],
    stop_condition: MoveStopCondition = MoveStopCondition.none,
) -> Tuple[MoveGroup, Dict[NodeId, float]]:
    pos = _convert_to_node_id_dict(origin)
    move_group: MoveGroup = []
    for move in moves:
        unit_vector = move.unit_vector
        for block in move.blocks:
            distances = unit_vector_multiplication(unit_vector, block.distance)
            node_id_distances = _convert_to_node_id_dict(distances)
            velocities = unit_vector_multiplication(unit_vector, block.initial_speed)
            accelerations = unit_vector_multiplication(unit_vector, block.acceleration)
            step = create_step(
                distance=node_id_distances,
                velocity=_convert_to_node_id_dict(velocities),
                acceleration=_convert_to_node_id_dict(accelerations),
                duration=block.time,
                present_nodes=present_nodes,
                stop_condition=stop_condition,
            )
            for ax in pos.keys():
                pos[ax] += node_id_distances.get(ax, 0)
            move_group.append(step)
    return move_group, {k: float(v) for k, v in pos.items()}

def create_home_group(
    distance: Dict[OT3Axis, float], velocity: Dict[OT3Axis, float]
) -> MoveGroup:
    node_id_distances = _convert_to_node_id_dict(distance)
    node_id_velocities = _convert_to_node_id_dict(velocity)
    step = create_home_step(
        distance=node_id_distances,
        velocity=node_id_velocities,
    )
    move_group: MoveGroup = [step]
    return move_group

AxisMapPayload = TypeVar("AxisMapPayload")

def axis_convert(
    axis_map: Dict["NodeId", AxisMapPayload], default_value: AxisMapPayload
) -> OT3AxisMap[AxisMapPayload]:
    ret: OT3AxisMap[AxisMapPayload] = {k: default_value for k in node_axes()}
    for node, value in axis_map.items():
        if node_is_axis(node):
            ret[node_to_axis(node)] = value
    return ret

_sensor_node_lookup: Dict[OT3Mount, ProbeTarget] = {
    OT3Mount.LEFT: NodeId.pipette_left,
    OT3Mount.RIGHT: NodeId.pipette_right,
    OT3Mount.GRIPPER: NodeId.gripper,
}

def sensor_node_for_mount(mount: OT3Mount) -> ProbeTarget:
    return _sensor_node_lookup[mount]


import numpy as np
from scipy.ndimage.filters import gaussian_filter1d as gf
import inflect
from ..initial_param.kinect_para import Kinect_para

class Clasp_spread(object):
    """
    """
    def __init__(self):
        self.cnt         = 0
        self.handsdown   = True   # True : hands down, False : hands up
        self.clasp_cnt   = 0
        self.spread_cnt  = 0
        self.spread_time = 0
        self.hold        = 5
        self.first       = True
        self.cnvt        = inflect.engine()  # converting numerals into ordinals
        self.kpm         = Kinect_para()
        self.elbowstus   = {}
        self.elbowstus['clasp'] = False
        self.elbowstus['spread'] = False
        self.mode = 'clasp'
        # save in log
        self.holdtime = []
        # self.claspsuc = 0
        # default parameters
        self.cnt     = 0
        self.do      = False
        self.err     = []
        self.errsum  = []
        self.evalstr = ''
        self.eval    = ''

    def state_update(self, joints, kpm):
        """ update accoding to each frame data
        """
        if joints[kpm.LElbow_y] < joints[kpm.SpineShld_y]-20 and joints[kpm.RElbow_y] < joints[kpm.SpineShld_y]-20:
            self.clasp_cnt += 1
            self.spread_cnt = 0
        elif joints[kpm.LElbow_y] > joints[kpm.SpineShld_y]-20 and joints[kpm.RElbow_y] > joints[kpm.SpineShld_y]-20:
            self.spread_cnt += 1
            self.clasp_cnt = 0
        if self.clasp_cnt >= self.hold and not self.handsdown:
            self.handsdown = True
            self.elbowstus['clasp'] == False
        if self.spread_cnt >= self.hold and self.handsdown:
            self.handsdown = False
            self.first = False
            self.do    = True
            self.spread_time = 0
            # self.elbowstus['spread'] == False

    def bodystraight(self, joints, kpm, th=20):
        """ check whether body is straight or not
        """
        torso_z = np.mean([joints[kpm.SpineBase_z], joints[kpm.SpineMid_z]])
        if torso_z-joints[kpm.Neck_z] > th and torso_z-joints[kpm.Head_z] > th:
            self.evalstr = 'please stand straight.'
            self.eval = 'please stand straight.'
            return False
        return True

    def clasp(self, joints, kpm, spread_th=60, elbow_th=75):
        """ arms clasp state
        """
        if self.clasp_cnt == self.hold:
            self.elbowstus['clasp'] = False
            if not self.first:
                self.holdtime.append(self.spread_time/30.)
                self.cnt += 1
                print ('count: %s' %self.cnt)
                if not self.elbowstus['spread']:
                    if self.spread_time < spread_th:
                        self.evalstr = 'Elbows should be behind your head far enough\n'
                        self.eval += 'Elbows should be behind your head far enough\n'
                        self.err.append('At the '+self.cnvt.ordinal(self.cnt)+ \
                                        ' time spread is not good. Elbows should be behind your head far enough')
                        self.errsum.append('Elbows should be behind your head far enough')
                    self.elbowstus['spread'] = True

                if self.eval == '':
                    self.evalstr = 'Repitition done: Well done'
                else:
                    self.evalstr = 'Repitition done.\n'+self.eval
                    self.eval = ''
        if np.abs(joints[kpm.LElbow_x]-joints[kpm.RElbow_x]) < 75:
            self.elbowstus['clasp'] = True

    def spread(self, joints, kpm, spread_th=30):
        """ arms spread state
        """
        if self.spread_cnt == self.hold:
            self.elbowstus['spread'] = False
            if not self.elbowstus['clasp']:
                self.evalstr = 'When you raise your arms, elbows should be close to each other.\n'
                self.eval = 'When you raise your arms, elbows should be close to each other.\n'
                self.err.append('At the '+self.cnvt.ordinal(self.cnt)+ ' time, clasp is not good')
                self.errsum.append('Elbows should be close to each other.')
                self.elbowstus['clasp'] = True

            if self.eval == '':
                self.evalstr = 'Subsequence done: Clasp well done'
            else:
                self.evalstr = 'Subsequence done.\n'+self.eval

        if (joints[kpm.LElbow_z]+joints[kpm.RElbow_z])/2 > joints[kpm.Head_z]:
            self.spread_time += 1
        if self.spread_time >= spread_th:
            self.elbowstus['spread'] = True

    def run(self, joints):
        """ Joint : denoised unified joint
        """
        if self.handsdown:
            # print 'hands down'
            self.mode = 'clasp'
            self.clasp(joints, self.kpm)
        else:
            # print 'hands up'
            self.mode = 'spread'
            self.spread(joints, self.kpm)
            self.bodystraight(joints, self.kpm)
        self.state_update(joints, self.kpm)


from datetime import datetime
from math import ceil

import numpy as np
import tensorflow as tf
from sklearn.base import BaseEstimator, ClassifierMixin

from common.util.log_helper import LogHelper

# he_init = tf.contrib.layers.variance_scaling_initializer()
dim_fasttext = 300
n_birnn_layers = 2

class ESIM(BaseEstimator, ClassifierMixin):
    """
    https://arxiv.org/abs/1609.06038
    """

    def __init__(self, name=None, ckpt_path=None, word_trainable=False, lstm_layers=1, num_neurons=[128, 128, 32],
                 pos_weight=None, optimizer='adam', learning_rate=0.001, batch_size=128,
                 activation='relu', initializer='he', num_epoch=100, dropout_rate=None,
                 max_check_without_progress=10, show_progress=1, tensorboard_logdir=None, random_state=None,
                 vocab_size=None, n_outputs=3, device=None, stance_trainable=False, domain_trainable=True,
                 suffix_trainable=True, protocol_trainable=True, max_gpu_memory=0.5):

        self.ckpt_path = ckpt_path
        self.word_trainable = word_trainable
        self.lstm_layers = lstm_layers
        self.num_neurons = num_neurons
        self.optimizer = optimizer
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.activation = activation
        self.num_epoch = num_epoch
        self.initializer = initializer
        self.dropout_rate = dropout_rate
        self.max_check_without_progress = max_check_without_progress
        self.show_progress = show_progress
        self.random_state = random_state
        self.tensorboard_logdir = tensorboard_logdir
        self.vocab_size = vocab_size
        self.n_outputs = n_outputs
        self.pos_weight = pos_weight
        self.name = name
        self.device = device
        self.stance_trainable = stance_trainable
        self.domain_trainable = domain_trainable
        self.suffix_trainable = suffix_trainable
        self.protocol_trainable = protocol_trainable
        self.max_gpu_memory = max_gpu_memory
        self.word_embedding = None
        self.domain_embedding = None
        self.suffix_embedding = None
        self.protocol_embedding = None
        self.stance_embedding = None
        self.stance_embedding_size = None
        self.suffix_embedding_size = None
        self.protocol_embedding_size = None
        self.domain_embedding_size = None
        self._graph = None
        self._classes = None
        self._session = None
        self._initializer = None
        self._optimizer = None
        self._activation = None
        self.logger = LogHelper.get_logger(self.name)

    def __reduce__(self):
        """
        For pickle's usage
        :return:
        """
        return (ESIM, (
            self.name, self.ckpt_path, self.word_trainable, self.lstm_layers, self.num_neurons, self.pos_weight,
            self.optimizer, self.learning_rate, self.batch_size, self.activation, self.initializer, self.num_epoch,
            self.dropout_rate, self.max_check_without_progress, self.show_progress, self.tensorboard_logdir,
            self.random_state, self.vocab_size, self.n_outputs, self.device, self.stance_trainable,
            self.domain_trainable, self.suffix_trainable, self.protocol_trainable, self.max_gpu_memory))

    def lstm_cell(self, hidden_size):
        lstm = tf.nn.rnn_cell.BasicLSTMCell(hidden_size)
        if self.dropout_rate:
            lstm = tf.nn.rnn_cell.DropoutWrapper(lstm, input_keep_prob=self.keep_prob, output_keep_prob=self.keep_prob)
        return lstm

    def gru_cell(self, num_neuron):

        gru = tf.contrib.rnn.GRUCell(num_neuron)
        if self.dropout_rate:
            gru = tf.contrib.rnn.DropoutWrapper(gru, input_keep_prob=self.keep_prob, output_keep_prob=self.keep_prob)
        return gru

    def _bidirectional_rnn(self, inputs, inputs_length, num_units, scope=None):

        with tf.variable_scope(scope or 'birnn'):

            if self.lstm_layers == 1:
                rnn_cells_fw = self.lstm_cell(num_units)
                rnn_cells_bw = self.lstm_cell(num_units)
            else:
                rnn_cells_fw = tf.nn.rnn_cell.MultiRNNCell([self.lstm_cell(n) for n in num_units])
                rnn_cells_bw = tf.nn.rnn_cell.MultiRNNCell([self.lstm_cell(n) for n in num_units])

            ((fw_outputs, bw_outputs), (fw_states, bw_states)) = tf.nn.bidirectional_dynamic_rnn(rnn_cells_fw,
                                                                                rnn_cells_bw, inputs,
                                                                                inputs_length,
                                                                                dtype=tf.float32)
            outputs = tf.concat([fw_outputs, bw_outputs], axis=2)

            if self.lstm_layers > 1:
                fw_states = fw_states[self.lstm_layers - 1]
                bw_states = bw_states[self.lstm_layers - 1]
        return outputs, fw_states, bw_states

    def _esim(self, head_output_steps, body_output_steps, head_sent_sizes, body_sent_sizes):
        """
        Perform ESIM for each sentence pair
        :param head_output_steps: Output time steps of BiLSTM for one claim sentence but with batch size. batch_size * sents * words * embeddings
        :param body_output_steps: Output time steps of BiLSTM for one evidence sentence but with batch size. batch_size * sents * words * embeddings
        :param head_sent_sizes: Sentence sizes of the claim sentence but with batch size. batch_size * sents
        :param body_sent_sizes: Sentence sizes of the evidence sentence but with batch size. batch_size * sents
        :return:
        """
        with tf.variable_scope("ESIM_single_sentence_pair", reuse=True):
            # batch_size * sents * head_words * body_words
            attention_matrix = tf.matmul(head_output_steps, body_output_steps, transpose_b=True)
            # batch_size * sents * head_words * body_words
            soft_attention_matrix_for_head = tf.nn.softmax(attention_matrix, 3)
            # batch_size * sents * body_words * head_words
            soft_attention_matrix_for_body = tf.nn.softmax(tf.transpose(attention_matrix, [0, 1, 3, 2]), 3)

            # batch_size * sents * head_words * embed
            head_attended = tf.matmul(soft_attention_matrix_for_head, body_output_steps, name="attend_head")

            # batch_size * sents * body_words * embed
            body_attended = tf.matmul(soft_attention_matrix_for_body, head_output_steps, name="attend_body")

            batch_size, num_sents, max_head_words, embed_size = tf.unstack(tf.shape(head_attended))
            flat_head_sent_sizes = tf.reshape(head_sent_sizes, [batch_size * num_sents])
            flat_head_seq_masks = tf.sequence_mask(flat_head_sent_sizes, max_head_words)
            head_seq_masks = tf.reshape(flat_head_seq_masks, [batch_size, num_sents, max_head_words])
            # batch_size * sents * head_words
            head_masks_per_word = tf.where(head_seq_masks, tf.ones_like(head_seq_masks, dtype=tf.float32),
                                           tf.zeros_like(head_seq_masks, dtype=tf.float32))
            encode_size = 2 * self.num_neurons[0]
            # now batch_size * sents * head_words * embed
            head_masks = tf.tile(tf.expand_dims(head_masks_per_word, 3), [1, 1, 1, encode_size])
            masked_head_attended = head_attended * head_masks

            _, _, max_body_words, _ = tf.unstack(tf.shape(body_attended))
            flat_body_sent_sizes = tf.reshape(body_sent_sizes, [batch_size * num_sents])
            flat_body_seq_masks = tf.sequence_mask(flat_body_sent_sizes, max_body_words)
            body_seq_masks = tf.reshape(flat_body_seq_masks, [batch_size, num_sents, max_body_words])
            # batch_size * sents * body_words
            body_masks_per_word = tf.where(body_seq_masks, tf.ones_like(body_seq_masks, dtype=tf.float32),
                                           tf.zeros_like(body_seq_masks, dtype=tf.float32))
            # now batch_size * sents * body_words * embed
            body_masks = tf.tile(tf.expand_dims(body_masks_per_word, 3), [1, 1, 1, encode_size])
            masked_body_attended = body_attended * body_masks

            # batch_size * sents * head_words * (4 * embed)
            head_concat = tf.concat(
                [head_output_steps, masked_head_attended, tf.abs(tf.subtract(masked_head_attended, head_output_steps)),
                 tf.multiply(head_output_steps, masked_head_attended)], 3, name="concat_head_attended")
            # batch_size * sents * body_words * (4 * embed)
            body_concat = tf.concat(
                [body_output_steps, masked_body_attended, tf.abs(tf.subtract(masked_body_attended, body_output_steps)),
                 tf.multiply(body_output_steps, masked_body_attended)], 3, name="concat_body_attended")
        return head_concat, body_concat, head_masks_per_word, body_masks_per_word

    def _trainable_alignment(self, sum_heads_encoded, sum_bodies_encoded, encode_size, projection_size=100):
        """
        Generate the alignments over the evidences by calculating the cosine similarity of the sum of encoded words of each claim-evidence pair
        :param sum_heads_encoded: batch_size * encoded_size
        :param sum_bodies_encoded: batch_size * sents * encoded_size
        :param encode_size: encode size, last rank of shape
        :return: batch_size * sents * 1
        """
        with tf.variable_scope("trainable_alignment") as scope:
            batch_size, sents_size, _ = tf.unstack(tf.shape(sum_bodies_encoded))
            flat_sum_bodies_encoded = tf.reshape(sum_bodies_encoded, [batch_size * sents_size, encode_size])

            def _projection(inputs, num_units, reuse=None):
                return tf.contrib.layers.fully_connected(inputs, num_units, reuse=reuse, scope=scope)

            heads_projection = _projection(sum_heads_encoded, projection_size)
            flat_bodies_projection = _projection(flat_sum_bodies_encoded, projection_size, True)
            heads_projection = tf.expand_dims(heads_projection, 1)
            bodies_projection = tf.reshape(flat_bodies_projection, [batch_size, sents_size, projection_size])
            vector_attn = tf.reduce_sum(
                tf.multiply(tf.nn.l2_normalize(bodies_projection, 2), tf.nn.l2_normalize(heads_projection, 2)), axis=2,
                keepdims=True)
        return vector_attn

    def _align_without_softmax(self, sum_heads_encoded, sum_bodies_encoded):
        """
        Generate the alignments over the evidences by calculating the cosine similarity of the sum of encoded words of each claim-evidence pair
        :param sum_heads_encoded: batch_size * encoded_size
        :param sum_bodies_encoded: batch_size * sents * encoded_size
        :return: batch_size * sents * 1
        """
        _tmp_heads_encoded = tf.expand_dims(sum_heads_encoded, 1)
        vector_attn = tf.reduce_sum(
            tf.multiply(tf.nn.l2_normalize(sum_bodies_encoded, 2), tf.nn.l2_normalize(_tmp_heads_encoded, 2)), axis=2,
            keepdims=True)
        return vector_attn

    def _align(self, sum_heads_encoded, sum_bodies_encoded):
        """
        Generate the alignments over the evidences by calculating the cosine similarity of the sum of encoded words of each claim-evidence pair
        :param sum_heads_encoded: batch_size * encoded_size
        :param sum_bodies_encoded: batch_size * sents * encoded_size
        :return: batch_size * sents * 1
        """
        _tmp_heads_encoded = tf.expand_dims(sum_heads_encoded, 1)
        vector_attn = tf.reduce_sum(
            tf.multiply(tf.nn.l2_normalize(sum_bodies_encoded, 2), tf.nn.l2_normalize(_tmp_heads_encoded, 2)), axis=2,
            keepdims=True)
        return tf.nn.softmax(vector_attn, axis=1)

    def _multi_rnn(self, heads_embeddings, bodies_embeddings, h_sent_sizes, b_sent_sizes, b_sizes, stances,
                   domain_names, suffixes, protocols, scope=None, epsilon=1e-12):
        """
        :param heads_embeddings: batch_size * 1 * h_words * embed
        :param bodies_embeddings: batch_size * sents * b_words * embed
        :param h_sent_sizes: batch_size * 1
        :param b_sent_sizes: batch_size * sents
        :param scope:
        :param epsilon: avoid dividing zero
        :return:
        """
        with tf.variable_scope(scope or "ESIM_rnn"):
            with tf.variable_scope("local_inference_modelling") as scope:
                (batch_size, body_sent_size, body_word_size, embed_size) = tf.unstack(tf.shape(bodies_embeddings))
                flat_bodies_embeddings = tf.reshape(bodies_embeddings,
                                                    [batch_size * body_sent_size, body_word_size, self.embedding_size])
                flat_b_sent_sizes = tf.reshape(b_sent_sizes, [batch_size * body_sent_size], name="flat_b_sent_sizes")
                flat_bodies_encoded, _, _ = self._bidirectional_rnn(flat_bodies_embeddings, flat_b_sent_sizes,
                                                                    self.num_neurons[0])
                encode_size = 2 * self.num_neurons[0]
                bodies_encoded = tf.reshape(flat_bodies_encoded,
                                            (batch_size, body_sent_size, body_word_size, encode_size),
                                            name="bodies_encoded")
                # now batch_size * h_words * embed
                flat_heads_embeddings = tf.squeeze(heads_embeddings, 1, name="flat_heads_embeddings")
                # now batch_size
                flat_h_sent_sizes = tf.squeeze(h_sent_sizes, 1, name="flat_h_sent_sizes")
                scope.reuse_variables()
                # batch_size * h_words * encode_size
                heads_encoded, _, _ = self._bidirectional_rnn(flat_heads_embeddings, flat_h_sent_sizes,
                                                              self.num_neurons[0])
                # batch_size * sents * h_words * encode_size
                tiled_heads_encoded = tf.tile(tf.expand_dims(heads_encoded, 1), [1, body_sent_size, 1, 1])
                # batch_size * sents
                tiled_h_sent_sizes = tf.tile(h_sent_sizes, [1, body_sent_size])

                heads_attended, bodies_attended, head_masks_per_word, body_masks_per_word = self._esim(
                    tiled_heads_encoded, bodies_encoded, tiled_h_sent_sizes, b_sent_sizes)

            with tf.variable_scope("inference_composition") as scope:
                attended_size = 4 * encode_size
                (batch_size, b_sent_size, h_word_size, _) = tf.unstack(tf.shape(heads_attended))
                (_, _, b_word_size, _) = tf.unstack(tf.shape(bodies_attended))
                reshape_expand_h_sent_sizes = tf.reshape(tf.tile(h_sent_sizes, [1, b_sent_size]),
                                                         [batch_size * b_sent_size], name="reshape_expand_h_sent_sizes")
                reshape_b_sent_sizes = tf.reshape(b_sent_sizes, [batch_size * b_sent_size], name="reshape_b_sent_sizes")
                flat_heads_attended = tf.reshape(heads_attended,
                                                 [batch_size * body_sent_size, h_word_size, attended_size],
                                                 name="flat_heads_attended")
                flat_bodies_attended = tf.reshape(bodies_attended,
                                                  [batch_size * body_sent_size, b_word_size, attended_size],
                                                  name="flat_bodies_attended")
                # (batch_size * sents) * h_words * output_dim
                flat_heads_outputs, flat_heads_fw, flat_heads_bw = self._bidirectional_rnn(flat_heads_attended,
                                                                                reshape_expand_h_sent_sizes,
                                                                                self.num_neurons[1])
                scope.reuse_variables()
                # (batch_size * sents) * b_words * output_dim
                flat_bodies_outputs, flat_bodies_fw, flat_bodies_bw = self._bidirectional_rnn(flat_bodies_attended,
                                                                                reshape_b_sent_sizes,
                                                                                self.num_neurons[1])
                # (batch_size * sents) * 1
                epsilon_matrix = tf.fill([batch_size * b_sent_size, 1], epsilon)
                # all (batch_size * sents) * output_dim
                flat_heads_mean = tf.reduce_sum(flat_heads_outputs, 1) / tf.maximum(
                    tf.reshape(tf.cast(tiled_h_sent_sizes, tf.float32), [batch_size * b_sent_size, 1]), epsilon_matrix)
                flat_heads_max = tf.reduce_max(flat_heads_outputs, 1)
                flat_bodies_mean = tf.reduce_sum(flat_bodies_outputs, 1) / tf.maximum(
                    tf.reshape(tf.cast(b_sent_sizes, tf.float32), [batch_size * b_sent_size, 1]), epsilon_matrix)
                flat_bodies_max = tf.reduce_max(flat_bodies_outputs, 1)

                # (batch_size * sents) * (4 * output_dim)
                flat_output_concat = tf.concat([flat_heads_mean, flat_heads_max, flat_bodies_mean, flat_bodies_max], 1)
            # MLP
            with tf.variable_scope("MLP", reuse=tf.AUTO_REUSE):
                output = flat_output_concat
                for layer in range(self.mlp_layers):
                    if self.dropout_rate:
                        output = tf.layers.dropout(output, rate=self.dropout_rate, training=self._training)
                    output = tf.layers.dense(output, self.num_neurons[n_birnn_layers + layer],
                                             activation=self._activation,
                                             kernel_initializer=self._initializer,
                                             name="hidden{}".format(layer + 1))
                output = tf.reshape(output, [batch_size, b_sent_size, self.num_neurons[-1]])
                stance_embeddings = self._add_embedding(stances, "stance_embedding_lookup", "stance_embedding",
                                                        self.stance_embedding, self.stance_trainable)
                suffix_embeddings = self._add_embedding(suffixes, "suffix_embedding_lookup", "suffix_embedding",
                                                        self.suffix_embedding, self.suffix_trainable)
                domain_embeddings = self._add_embedding(domain_names, "domain_embedding_lookup", "domain_embedding",
                                                        self.domain_embedding, self.domain_trainable)
                protocol_embeddings = self._add_embedding(protocols, "protocol_embedding_lookup", "protocol_embedding",
                                                          self.protocol_embedding, self.protocol_trainable)
                output_with_features = tf.concat(
                    [output, stance_embeddings, domain_embeddings, suffix_embeddings, protocol_embeddings], 2)
                # batch_size * sents * 3
                logits_per_sent = tf.layers.dense(output_with_features, self.n_outputs,
                                                  kernel_initializer=self._initializer, name="logits_per_sent")
                mask_condition = tf.sequence_mask(b_sizes, maxlen=b_sent_size)
                mask_condition = tf.tile(tf.expand_dims(mask_condition, axis=2), [1, 1, self.n_outputs])
                logits_per_sent = tf.where(mask_condition, logits_per_sent,
                                           tf.zeros([batch_size, b_sent_size, self.n_outputs]))
                sents_count_with_epsilon = tf.maximum(
                    tf.cast(b_sizes, tf.float32), tf.zeros([batch_size], dtype=tf.float32) + epsilon)
                tiled_sents_count_with_epsilon = tf.tile(tf.expand_dims(sents_count_with_epsilon, axis=1),
                                                         [1, self.n_outputs])
                logits = tf.reduce_sum(logits_per_sent, axis=1, name="logits") / tiled_sents_count_with_epsilon
                probabilities = tf.nn.softmax(logits, name="probabilities")
        return logits, probabilities

    def _add_embedding(self, inputs, scope, name, embeddings, trainable):

        with tf.variable_scope(scope):
            embedding = tf.get_variable(initializer=embeddings, dtype=tf.float32, trainable=trainable, name=name)
            inputs_embedded = tf.nn.embedding_lookup(embedding, inputs)
            return inputs_embedded

    def _ann(self, head_inputs, body_inputs, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes, stances, domain_names,
             suffixes, protocols):

        with tf.variable_scope("embedding_lookup") as scope:
            word_embedding_var_name = "word_embedding"
            heads_embeddings = self._add_embedding(head_inputs, scope, word_embedding_var_name, self.word_embedding,
                                                   self.word_trainable)
            scope.reuse_variables()
            body_embeddings = self._add_embedding(body_inputs, scope, word_embedding_var_name, self.word_embedding,
                                                  self.word_trainable)
        logits, probabilities = self._multi_rnn(heads_embeddings, body_embeddings, h_sent_sizes, b_sent_sizes, b_sizes,
                                                stances, domain_names, suffixes, protocols)
        return logits, probabilities

    def _construct_graph(self):

        if self.random_state:
            tf.set_random_seed(self.random_state)
            np.random.seed(self.random_state)

        if self._initializer is None:
            if self.initializer == 'he':
                self._initializer = tf.contrib.layers.variance_scaling_initializer()

        if self._activation is None:
            if self.activation == 'relu':
                self._activation = tf.nn.relu

        if self._optimizer is None:
            if self.optimizer == 'adam':
                self._optimizer = tf.train.AdamOptimizer

        self.embed_size = len(self.word_embedding[0]) if self.word_embedding is not None else 0
        self.mlp_layers = len(self.num_neurons) - 2
        self.embedding_size = self.embed_size

        X_heads = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_heads")
        X_bodies = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_bodies")
        X_head_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="head_sizes")
        X_body_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="body_sizes")
        X_head_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="head_sent_sizes")
        X_body_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="body_sent_sizes")
        X_stances = tf.placeholder(dtype=tf.int32, shape=[None, None], name="stances")
        X_domain_names = tf.placeholder(dtype=tf.int32, shape=[None, None], name="domain_names")
        X_suffixes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="suffixes")
        X_protocols = tf.placeholder(dtype=tf.int32, shape=[None, None], name="protocols")

        y_ = tf.placeholder(tf.int32, shape=[None], name="y")
        y_one_hot = tf.one_hot(y_, self.n_outputs, on_value=1.0, off_value=0.0, axis=-1, dtype=tf.float32)

        if self.dropout_rate:
            self._training = tf.placeholder_with_default(False, shape=[], name="training")
            self.keep_prob = tf.cond(self._training, lambda: tf.constant(1 - self.dropout_rate),
                                     lambda: tf.constant(1.0))
        else:
            self._training = None

        logits, probabilities = self._ann(X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes,
                                          X_body_sent_sizes, X_stances, X_domain_names, X_suffixes, X_protocols)

        if self.pos_weight is None:
            xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_, logits=logits)
        else:
            xentropy = tf.nn.weighted_cross_entropy_with_logits(y_one_hot, logits, self.pos_weight)
        loss = tf.reduce_mean(xentropy, name="loss")
        variables = tf.trainable_variables()
        for v in variables:
            self.logger.debug(self.name + ": " + v.name)

        optimizer = self._optimizer(learning_rate=self.learning_rate)
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            training_op = optimizer.minimize(loss)

        correct = tf.nn.in_top_k(logits, y_, 1)
        accuracy = tf.reduce_mean(tf.cast(correct, tf.float32), name="accuracy")
        _, predicts = tf.nn.top_k(logits, k=1, sorted=False)
        confusion_matrix = tf.confusion_matrix(y_, predicts, num_classes=self.n_outputs, name="confusion_matrix")

        init = tf.global_variables_initializer()
        saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES))

        if self.tensorboard_logdir:
            now = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
            tb_logdir = self.tensorboard_logdir + "/run{}".format(now)
            cost_summary = tf.summary.scalar("validation_loss", loss)
            acc_summary = tf.summary.scalar("validation_accuracy", accuracy)
            merged_summary = tf.summary.merge_all()
            file_writer = tf.summary.FileWriter(
                tb_logdir, tf.get_default_graph())

            self._merged_summary = merged_summary
            self._file_writer = file_writer

        self._X_head, self._X_body, self._X_h_sizes, self._X_b_sizes, self._X_h_sent_sizes, self._X_b_sent_sizes, self.y = X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes, X_body_sent_sizes, y_
        self._X_domain_names = X_domain_names
        self._X_suffixes = X_suffixes
        self._X_protocols = X_protocols
        self._X_stances = X_stances
        self._logits = logits
        self._probabilities = probabilities
        self._loss = loss
        self._training_op = training_op
        self._accuracy = accuracy
        self._confusion_matrix = confusion_matrix
        self._init, self._saver = init, saver

    def get_batch_attention(self, h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes, domain_names, suffixes,
                            protocols, stances, y=None):

        num_batches = ceil(h_np.shape[0] / self.batch_size)

        for i in range(num_batches):
            if (i + 1) * self.batch_size > h_np.shape[0]:
                batch_h_np = h_np[i * self.batch_size:h_np.shape[0], :]
                batch_b_np = b_np[i * self.batch_size:h_np.shape[0], :]
                batch_h_sizes = h_sizes[i * self.batch_size:h_np.shape[0]]
                batch_b_sizes = b_sizes[i * self.batch_size:h_np.shape[0]]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:h_np.shape[0], :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:h_np.shape[0], :]
                batch_domain_names = domain_names[i * self.batch_size:h_np.shape[0], :]
                batch_suffixes = suffixes[i * self.batch_size:h_np.shape[0], :]
                batch_protocols = protocols[i * self.batch_size:h_np.shape[0], :]
                batch_stances = stances[i * self.batch_size:h_np.shape[0], :]
                if y is not None:
                    batch_y = y[i * self.batch_size:h_np.shape[0]]
            else:
                batch_h_np = h_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_np = b_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_h_sizes = h_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_b_sizes = b_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_domain_names = domain_names[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_suffixes = suffixes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_protocols = protocols[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_stances = stances[i * self.batch_size:(i + 1) * self.batch_size, :]
                if y is not None:
                    batch_y = y[i * self.batch_size:(i + 1) * self.batch_size]
            if y is not None:
                yield (batch_h_np, batch_b_np, batch_h_sizes, batch_b_sizes, batch_h_sent_sizes, batch_b_sent_sizes,
                       batch_domain_names, batch_suffixes, batch_protocols, batch_stances, batch_y)
            else:
                yield (batch_h_np, batch_b_np, batch_h_sizes, batch_b_sizes, batch_h_sent_sizes, batch_b_sent_sizes,
                       batch_domain_names, batch_suffixes, batch_protocols, batch_stances)

    def cal_f1_macro(self, confusion_matrix):
        """
        calculate f1 macro
        :param confusion_matrix:
        :return: f1 macro score
        """

        self.logger.info("\n" + str(confusion_matrix))
        diag = np.diag(confusion_matrix).astype(np.float32)
        num_golds = np.sum(confusion_matrix, axis=0).astype(np.float32)
        num_predicts = np.sum(confusion_matrix, axis=1).astype(np.float32)
        precisions = np.divide(diag, num_golds, out=np.zeros_like(diag), where=num_golds != 0)
        recalls = np.divide(diag, num_predicts, out=np.zeros_like(diag), where=num_predicts != 0)

        average_precision = np.mean(precisions)
        average_recall = np.mean(recalls)
        f1_macro = 2 * average_precision * average_recall / (average_recall + average_precision)
        self.logger.info("average_precision: {:.4f}, average_recall: {:.4f} f1_macro: {:.4f}".format(average_precision,
                                                                                average_recall,
                                                                                f1_macro))

        return f1_macro

    def fit(self, X_dict, y):
        self.logger.debug("training...")
        X = X_dict['X_train']
        valid_X = X_dict['X_valid']
        y_valid = X_dict['y_valid']
        self.word_embedding = X_dict['word_embedding']
        self.stance_embedding = X_dict['stance_embedding']
        self.domain_embedding = X_dict['domain_embedding']
        self.protocol_embedding = X_dict['protocol_embedding']
        self.suffix_embedding = X_dict['suffix_embedding']
        self.stance_embedding_size = self.stance_embedding.shape[-1]
        self.suffix_embedding_size = self.suffix_embedding.shape[-1]
        self.protocol_embedding_size = self.protocol_embedding.shape[-1]
        self.domain_embedding_size = self.domain_embedding.shape[-1]
        self.close_session()
        y = np.array(y)

        if len(y.shape) == 2:
            y = np.argmax(y, axis=1)

        self._classes, _ = np.unique(y, return_counts=True)

        self._graph = tf.Graph()
        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']
        stances = X['stance']
        domain_names = X['domain_name']
        suffixes = X['suffix']
        protocols = X['protocol']
        with self._graph.as_default():
            self._construct_graph()

        checks_without_progress = 0
        # best_acc = np.inf
        best_f1_macro = 0
        best_parameters = None
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=self.max_gpu_memory)
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self._session = tf.Session(graph=self._graph,
                                   config=tf.ConfigProto(gpu_options=gpu_options)
                                   # config=config
                                   )

        with self._session.as_default() as sess:
            self._init.run()
            num_instances = h_np.shape[0]
            for epoch in range(self.num_epoch):
                losses = []
                accs = []
                rnd_idx = np.random.permutation(num_instances)
                # batch_num = 0
                for rnd_indices in np.array_split(rnd_idx, num_instances // self.batch_size):

                    X_head_batch, X_body_batch, X_h_sizes_batch, X_b_sizes_batch, X_h_sent_sizes_batch, X_b_sent_sizes_batch, y_batch = \
                        h_np[rnd_indices], b_np[rnd_indices], \
                        h_sizes[rnd_indices], b_sizes[rnd_indices], \
                        h_sent_sizes[rnd_indices], b_sent_sizes[rnd_indices], \
                        y[rnd_indices]
                    X_stance_batch = stances[rnd_indices]
                    X_domain_name_batch = domain_names[rnd_indices]
                    X_suffix_batch = suffixes[rnd_indices]
                    X_protocol_batch = protocols[rnd_indices]
                    y_batch = np.asarray(y_batch)
                    feed_dict = {self._X_head: X_head_batch, self._X_body: X_body_batch,
                                 self._X_h_sizes: X_h_sizes_batch, self._X_b_sizes: X_b_sizes_batch,
                                 self._X_h_sent_sizes: X_h_sent_sizes_batch, self._X_b_sent_sizes: X_b_sent_sizes_batch,
                                 self._X_domain_names: X_domain_name_batch, self._X_suffixes: X_suffix_batch,
                                 self._X_protocols: X_protocol_batch, self._X_stances: X_stance_batch,
                                 self.y: y_batch}
                    if self._training is not None:
                        feed_dict[self._training] = True
                    train_acc, _, loss = sess.run([self._accuracy, self._training_op, self._loss], feed_dict=feed_dict)
                    losses.append(loss)
                    accs.append(train_acc)
                average_loss = sum(losses) / len(losses)
                average_acc = sum(accs) / len(accs)

                if valid_X is not None and y_valid is not None:

                    self.logger.debug("validation phase")
                    valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes, valid_h_sent_sizes, valid_b_sent_sizes = \
                        valid_X['h_np'], valid_X['b_np'], valid_X[
                            'h_sizes'], valid_X['b_sizes'], valid_X['h_sent_sizes'], valid_X['b_sent_sizes']

                    valid_stances = valid_X['stance']
                    valid_domain_names = valid_X['domain_name']
                    valid_suffixes = valid_X['suffix']
                    valid_protocols = valid_X['protocol']
                    batches = self.get_batch_attention(valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes,
                                                       valid_h_sent_sizes, valid_b_sent_sizes, valid_domain_names,
                                                       valid_suffixes, valid_protocols, valid_stances, y_valid)

                    batch_losses = []
                    batch_accuracies = []
                    valid_cm = np.zeros(shape=(self.n_outputs, self.n_outputs), dtype=np.int32)
                    for (
                            valid_h_batch, valid_b_batch, valid_h_sizes_batch, valid_b_sizes_batch,
                            valid_h_sent_sizes_batch, valid_b_sent_sizes_batch, valid_domain_name_batch,
                            valid_suffix_batch, valid_protocol_batch, valid_stance_batch, valid_y_batch) in batches:

                        feed_dict_valid = {self._X_head: valid_h_batch, self._X_body: valid_b_batch,
                                           self._X_h_sizes: valid_h_sizes_batch, self._X_b_sizes: valid_b_sizes_batch,
                                           self._X_h_sent_sizes: valid_h_sent_sizes_batch,
                                           self._X_b_sent_sizes: valid_b_sent_sizes_batch,
                                           self._X_domain_names: valid_domain_name_batch,
                                           self._X_suffixes: valid_suffix_batch,
                                           self._X_protocols: valid_protocol_batch, self._X_stances: valid_stance_batch,
                                           self.y: valid_y_batch}
                        if self.tensorboard_logdir:
                            val_acc_batch, val_loss_batch, cm, summary = sess.run(
                                [self._accuracy, self._loss, self._confusion_matrix, self._merged_summary],
                                feed_dict=feed_dict_valid)
                            self._file_writer.add_summary(summary, epoch)
                        else:
                            val_acc_batch, val_loss_batch, cm = sess.run(
                                [self._accuracy, self._loss, self._confusion_matrix], feed_dict=feed_dict_valid)

                        batch_losses.append(val_loss_batch)
                        batch_accuracies.append(val_acc_batch)
                        valid_cm = np.add(valid_cm, cm)

                    val_f1_macro = self.cal_f1_macro(valid_cm)
                    val_loss = sum(batch_losses) / len(batch_losses)
                    val_acc = sum(batch_accuracies) / len(batch_accuracies)

                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            self.logger.info(
                                "Epoch: {} Current training accuracy: {:.4f} ,Current training loss: {:.6f} Validation Accuracy: {:.4f} Validation Loss{:.6f}".format(
                                    epoch + 1, average_acc, average_loss, val_acc, val_loss))
                    if val_f1_macro > best_f1_macro:
                        best_f1_macro = val_f1_macro
                        checks_without_progress = 0
                        self.logger.info("f1_macro has been improved!")
                        best_parameters = self._get_model_parameters()
                        self.save(self.ckpt_path)
                    else:
                        checks_without_progress += 1
                    if checks_without_progress > self.max_check_without_progress:
                        self.logger.info("Stopping Early! Loss has not improved in {} epoches".format(
                            self.max_check_without_progress))
                        break
                else:
                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            self.logger.info("Epoch: {} Current training accuracy: {:.4f}".format(
                                epoch + 1, train_acc))

            if best_parameters:
                self._restore_model_parameters(best_parameters)
            self.save(self.ckpt_path)
            return self

    def predict_proba(self, X_dict, restore_param_required=True):
        import pickle
        self.logger.debug("testing...")
        X = X_dict['X_test']
        self.word_embedding = X_dict['word_embedding']
        self.stance_embedding = X_dict['stance_embedding']
        self.domain_embedding = X_dict['domain_embedding']
        self.protocol_embedding = X_dict['protocol_embedding']
        self.suffix_embedding = X_dict['suffix_embedding']
        if restore_param_required:
            self.restore_model(self.ckpt_path)

        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']
        stances = X['stance']
        domain_names = X['domain_name']
        suffixes = X['suffix']
        protocols = X['protocol']
        batches = self.get_batch_attention(h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes, domain_names,
                                           suffixes, protocols, stances)
        probabilities = []
        with self._session.as_default():
            for (pred_h_batch, pred_b_batch, pred_h_sizes_batch, pred_b_sizes_batch, pred_h_sent_sizes_batch,
                 pred_b_sent_sizes_batch, pred_domain_batch, pred_suffix_batch, pred_protocol_batch,
                 pred_stance_batch) in batches:
                predictions_batch = self._probabilities.eval(
                    feed_dict={self._X_head: pred_h_batch, self._X_body: pred_b_batch,
                               self._X_h_sizes: pred_h_sizes_batch, self._X_b_sizes: pred_b_sizes_batch,
                               self._X_h_sent_sizes: pred_h_sent_sizes_batch,
                               self._X_b_sent_sizes: pred_b_sent_sizes_batch,
                               self._X_domain_names: pred_domain_batch, self._X_suffixes: pred_suffix_batch,
                               self._X_protocols: pred_protocol_batch, self._X_stances: pred_stance_batch})
                for prediction in predictions_batch:
                    probabilities.append(prediction)
        np_probas = np.asarray(probabilities)
        # save predictions
        _prediction_save_file = self.ckpt_path + "_predictions.p"
        with open(_prediction_save_file, 'wb') as f:
            pickle.dump(np_probas, f, protocol=pickle.HIGHEST_PROTOCOL)
            self.logger.info("predictions saved in " + _prediction_save_file)
        return np_probas

    def predict(self, X_dict, restore_param_required=True):

        predictions = np.argmax(self.predict_proba(X_dict, restore_param_required), axis=1)
        return np.reshape(predictions, (-1,))

    def _get_model_parameters(self):

        with self._graph.as_default():
            gvars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        return {gvar.op.name: value for gvar, value in zip(gvars, self._session.run(gvars))}

    def _restore_model_parameters(self, model_parameters):

        gvar_names = list(model_parameters.keys())

        assign_ops = {gvar_name: self._graph.get_operation_by_name(gvar_name + "/Assign") for gvar_name in gvar_names}

        init_values = {gvar_name: assign_op.inputs[1] for gvar_name, assign_op in assign_ops.items()}
        feed_dict = {init_values[gvar_name]: model_parameters[gvar_name] for gvar_name in gvar_names}
        self._session.run(assign_ops, feed_dict=feed_dict)

    def close_session(self):
        if self._session:
            self._session.close()

    def save(self, path):
        self._saver.save(self._session, path)

    def restore_model(self, path):
        self._graph = tf.Graph()
        with self._graph.as_default():
            self._construct_graph()
            gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=self.max_gpu_memory)
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True
            self._session = tf.Session(graph=self._graph,
                                       config=tf.ConfigProto(gpu_options=gpu_options)
                                       # config=config
                                       )
            # sess = tf.Session(graph=self._graph)
            with self._session.as_default() as sess:
                self._init.run()
                sess.run(tf.tables_initializer())
                self._saver.restore(sess, path)
            # self._session = sess
        return self


import numpy as np
import pytest

from jina import Document, DocumentArray
from jina.drivers.multimodal import MultiModalDriver
from jina.executors.encoders.multimodal import BaseMultiModalEncoder
from jina.types.document.multimodal import MultimodalDocument

@pytest.fixture(scope='function')
def embedding():
    class EmbeddingFactory(object):
        def create(self, dim, seed):
            np.random.seed(seed)
            return np.random.random([dim])

    return EmbeddingFactory()

@pytest.fixture(scope='function')
def embeddings(embedding):
    return [
        embedding.create(dim=16, seed=1),
        embedding.create(dim=16, seed=2),
        embedding.create(dim=24, seed=3),
    ]

@pytest.fixture(scope='function')
def doc_with_multimodal_chunks(embeddings):
    doc = MultimodalDocument()
    chunk1 = Document()
    chunk2 = Document()
    chunk3 = Document()
    chunk1.modality = 'visual1'
    chunk2.modality = 'visual2'
    chunk3.modality = 'textual'
    chunk1.embedding = embeddings[0]
    chunk2.embedding = embeddings[1]
    chunk3.embedding = embeddings[2]
    doc.chunks.append(chunk1)
    doc.chunks.append(chunk2)
    doc.chunks.append(chunk3)
    return doc

class MockMultiModalEncoder(BaseMultiModalEncoder):
    def __init__(self, positional_modality, *args, **kwargs):
        super().__init__(positional_modality=positional_modality, *args, **kwargs)

    def encode(self, *data: 'np.ndarray', **kwargs) -> 'np.ndarray':
        visual1 = data[(self.positional_modality.index('visual1'))]
        visual2 = data[(self.positional_modality.index('visual2'))]
        textual = data[(self.positional_modality.index('textual'))]
        return np.concatenate((visual1, visual2, textual), axis=1)

@pytest.fixture
def mock_multimodal_encoder():
    return MockMultiModalEncoder(positional_modality=['visual1', 'visual2', 'textual'])

class SimpleMultiModalDriver(MultiModalDriver):
    def __init__(self, *args, **kwargs):
        import logging

        super().__init__(*args, **kwargs)
        self.test_logger = logging.getLogger('test multimodal driver')

    @property
    def logger(self):
        return self.test_logger

    @property
    def exec_fn(self):
        return self._exec_fn

@pytest.fixture
def simple_multimodal_driver():
    return SimpleMultiModalDriver()

def test_multimodal_driver(
    simple_multimodal_driver, mock_multimodal_encoder, doc_with_multimodal_chunks
):
    simple_multimodal_driver.attach(executor=mock_multimodal_encoder, runtime=None)
    simple_multimodal_driver._apply_all(DocumentArray([doc_with_multimodal_chunks]))
    doc = doc_with_multimodal_chunks
    assert len(doc.chunks) == 3
    visual1 = doc.chunks[0]
    visual2 = doc.chunks[1]
    textual = doc.chunks[2]
    assert (
        doc.embedding.shape[0]
        == visual1.embedding.shape[0]
        + visual2.embedding.shape[0]
        + textual.embedding.shape[0]
    )

@pytest.fixture(scope='function')
def doc_with_multimodal_chunks_wrong(embeddings):
    doc = MultimodalDocument()
    chunk1 = Document()
    chunk2 = Document()
    chunk3 = Document()
    chunk1.modality = 'visual'
    chunk2.modality = 'visual'
    chunk3.modality = 'textual'
    chunk1.embedding = embeddings[0]
    chunk2.embedding = embeddings[1]
    chunk3.embedding = embeddings[2]
    doc.chunks.append(chunk1)
    doc.chunks.append(chunk2)
    doc.chunks.append(chunk3)
    return doc

def test_multimodal_driver_assert_one_chunk_per_modality(
    simple_multimodal_driver, mock_multimodal_encoder, doc_with_multimodal_chunks_wrong
):
    simple_multimodal_driver.attach(executor=mock_multimodal_encoder, runtime=None)
    assert not doc_with_multimodal_chunks_wrong.is_valid

@pytest.fixture
def mock_multimodal_encoder_shuffled():
    return MockMultiModalEncoder(positional_modality=['visual2', 'textual', 'visual1'])

def test_multimodal_driver_with_shuffled_order(
    simple_multimodal_driver,
    mock_multimodal_encoder_shuffled,
    doc_with_multimodal_chunks,
):
    simple_multimodal_driver.attach(
        executor=mock_multimodal_encoder_shuffled, runtime=None
    )
    simple_multimodal_driver._apply_all(DocumentArray([doc_with_multimodal_chunks]))
    doc = doc_with_multimodal_chunks
    assert len(doc.chunks) == 3
    visual1 = doc.chunks[2]
    visual2 = doc.chunks[0]
    textual = doc.chunks[1]
    control = np.concatenate([visual2.embedding, textual.embedding, visual1.embedding])
    test = doc.embedding
    np.testing.assert_array_equal(control, test)