__all__ = [
	"EBCOTCodec"
]

from copy import deepcopy
from multiprocessing import Pool
import numpy as np

from fpeg.base import Codec
from fpeg.config import read_config
from fpeg.funcs import parse_marker, cat_arrays_2d

config = read_config()

D = config.get("jpeg2000", "D")
G = config.get("jpeg2000", "G")
QCD = config.get("jpeg2000", "QCD")
mq_table = config.get("jpeg2000", "mq_table")

min_task_number = config.get("accelerate", "codec_min_task_number")
max_pool_size = config.get("accelerate", "codec_max_pool_size")

class EBCOTCodec(Codec):
	"""
	EBCOT codec
	"""

	def __init__(self,
							 name="EBCOT codec",
							 mode="encode",
							 D=D,
							 G=G,
							 QCD=QCD,
							 accelerated=False
							 ):
		"""
		Init and set attributes of a canonical EBCOT codec.

		Explicit Attributes
		-------------------
		name: str, optional
			Name of the codec.
		mode: str, optional
			Mode of the codec, must in ["encode", "decode"].
		G: integer, optional
			a parameter for calculate Kmax
		D: integer, optional
			Depth of graphic.
		epsilon_b:integer, must
			a parameter for calculate Kmax
		accelerated: bool, optional
			Whether the process would be accelerated by subprocess pool.

		"""
		super().__init__()

		self.name = name
		self.mode = mode
		# based on equation 10.22 in jpeg2000
		self.D = D
		self.G = G
		self.QCD = QCD
		self.accelerated = accelerated

		self.epsilon_b, _ = parse_marker(self.QCD)
		self.Kmax = max(0, self.G + self.epsilon_b - 1)
		self.min_task_number = min_task_number
		self.max_pool_size = max_pool_size

	def encode(self, X, **params):
		self.logs[-1] += self.formatter.message("Trying to encode received data.")
		try:
			self.epsilon_b = params["epsilon_b"]
		except KeyError:
			pass

		if self.accelerated:
			self.logs[-1] += self.formatter.message("Using multiprocess pool to accelerate EBCOT encoding.")
			inputs = [[x, self.D] for x in X]
			with Pool(min(self.task_number, self.max_pool_size)) as p:
				bitcodes = p.starmap(_EBCOT_encode, inputs)
		else:
			bitcodes = [_EBCOT_encode(x, self.D) for x in X]

		return bitcodes

	def decode(self, bitcodes, **params):
		self.logs[-1] += self.formatter.message("Trying to decode received data.")

		if self.accelerated:
			inputs = [[bitcode, self.D] for bitcode in bitcodes]
			with Pool(min(self.task_number, self.max_pool_size)) as p:
				X = p.starmap(_tile_decode, inputs)
		else:
			X = [_tile_decode(bitcode, self.D) for bitcode in bitcodes]

		# print(X[0])
		return X

def _EBCOT_encode(tile, D):
	"""
	EBCOT encode and decode part
	encode part:
	| _EBCOT_encode
		| _tile_encode
			| _band_encode
				| _embeddedBlockEncoder _MQencode
					 | 三个通道过程

	｜ _EBCOT_decode
		| _tile_decode
			| _band_decode
				| _block_decode
					| three decode passes and MQdecode
						|signdecode and runlengthdecode
	"""

	bitcode = list(_tile_encode(tile, D))

	# with open('test.bin', 'wb') as f:
	#   f.write(struct.pack(str(l)+'i', *bitcode))
	# streamonly = [int(i) for i in streamonly]
	# l = len(streamonly)
	# with open('streamonly.bin', 'wb') as f:
	#   f.write(struct.pack(str(l)+'i', *streamonly))

	return bitcode

def _MQencode(CX, D):
	"""
	MQ encode and decode part
	|MQencode
		|transferbyte
			|putbyte
	|MQdecode
		|fill_lsp

	"""
	PETTable, CXTable = deepcopy(mq_table)
	encoder = EBCOTparam()
	for i in range(len(D)):
		symbol = D[i][0]
		cxLabel = CX[i][0]
		expectedSymbol = CXTable[cxLabel][1]
		p = PETTable[CXTable[cxLabel][0]][3]  # PETTable [CXTable[cxLabel][0]]---[3]
		encoder.A = encoder.A - p
		if encoder.A < p:
			# Conditional exchange of MPS and LPS
			expectedSymbol = 1 - expectedSymbol
		if symbol == expectedSymbol:
			# assign MPS the upper sub-interval
			encoder.C = encoder.C + np.uint32(p)
		else:
			# assign LPS the lower sub-interval
			encoder.A = np.uint32(p)
		if encoder.A < 32768:
			if symbol == CXTable[cxLabel][1]:
				CXTable[cxLabel][0] = PETTable[CXTable[cxLabel][0]][0]
			else:
				CXTable[cxLabel][1] = CXTable[cxLabel][1] ^ PETTable[CXTable[cxLabel][0]][2]
				CXTable[cxLabel][0] = PETTable[CXTable[cxLabel][0]][1]
			while encoder.A < 32768:
				encoder.A = 2 * encoder.A
				encoder.C = 2 * encoder.C
				encoder.t = encoder.t - 1
				if encoder.t == 0:
					encoder = _transferbyte(encoder)
	encoder = _encode_end(encoder)
	return encoder

def _transferbyte(encoder):
	CPartialMask = np.uint32(133693440)  # 00000111111110000000000000000000
	CPartialCmp = np.uint32(4161273855)  # 11111000000001111111111111111111
	CMsbsMask = np.uint32(267386880)  # 27-20msbs标志位  00001111111100000000000000000000
	CMsbsCmp = np.uint32(4027580415)  # CMsbs的补码      11110000000011111111111111111111
	CCarryMask = np.uint32(2 ** 27)  # 取进位              00001000000000000000000000000000
	if encoder.T == 255:
		# 不能将任何进位传给T，需要位填充
		encoder = _putbyte(encoder)
		encoder.T = np.uint8((encoder.C & CMsbsMask) >> 20)  # 27-20位
		encoder.C = encoder.C & CMsbsCmp  #
		encoder.t = 7
	else:
		# 从C将任何进位传到T
		encoder.T = encoder.T + np.uint8((encoder.C & CCarryMask) >> 27)
		encoder.C = encoder.C ^ CCarryMask
		encoder = _putbyte(encoder)
		if encoder.T == 255:
			encoder.T = np.uint8((encoder.C & CMsbsMask) >> 20)
			encoder.C = encoder.C & CMsbsCmp
			encoder.t = 7
		else:
			encoder.T = np.uint8((encoder.C & CPartialMask) >> 19)
			encoder.C = encoder.C & CPartialCmp
			encoder.t = 8
	return encoder

def _putbyte(encoder):
	# 将T中的内容写入字节缓存
	if encoder.L >= 0:
		encoder.stream = np.append(encoder.stream, encoder.T)
	encoder.L = encoder.L + 1
	return encoder

def _MQ_decode(stream, CX):
	PETTable, CXTable = deepcopy(mq_table)
	# MQ decode initializtion
	encoder = EBCOTparam()
	encoder.A = np.uint16(0)
	encoder.C = np.uint32(0)
	encoder.t = np.uint8(0)
	encoder.T = np.uint8(0)
	encoder.L = np.int32(0)
	encoder.stream = stream
	encoder = _fill_lsb(encoder)
	encoder.C = encoder.C << encoder.t
	encoder = _fill_lsb(encoder)
	encoder.C = encoder.C << 7
	encoder.t = encoder.t - 7
	encoder.A = np.uint16(2 ** 15)

	# MQ decode procedure
	CActiveMask = np.uint32(16776960)  # 00000000111111111111111100000000
	CActiveCmp = np.uint32(4278190335)  # 11111111000000000000000011111111
	decodeD = []
	for i in range(len(CX)):
		cxLabel = CX[i][0]
		expectedSymbol = CXTable[cxLabel][1]
		p = PETTable[CXTable[cxLabel][0]][3]
		encoder.A = encoder.A - np.uint16(p)
		if encoder.A < np.uint16(p):
			expectedSymbol = 1 - expectedSymbol
		if ((encoder.C & CActiveMask) >> 8) < p:
			symbol = 1 - expectedSymbol
			encoder.A = np.uint16(p)
		else:
			symbol = expectedSymbol
			temp = ((encoder.C & CActiveMask) >> 8) - np.uint32(p)
			encoder.C = encoder.C & CActiveCmp
			encoder.C = encoder.C + np.uint32((np.uint32(temp << 8)) & CActiveMask)
		if encoder.A < 2 ** 15:
			if symbol == CXTable[cxLabel][1]:
				CXTable[cxLabel][0] = PETTable[CXTable[cxLabel][0]][0]
			else:
				CXTable[cxLabel][1] = CXTable[cxLabel][1] ^ PETTable[CXTable[cxLabel][0]][2]
				CXTable[cxLabel][0] = PETTable[CXTable[cxLabel][0]][1]
			while encoder.A < 2 ** 15:
				if encoder.t == 0:
					encoder = _fill_lsb(encoder)
				encoder.A = 2 * encoder.A
				encoder.C = 2 * encoder.C
				encoder.t = encoder.t - 1
		# print(i, symbol)
		decodeD.append([symbol])
	return decodeD

def _fill_lsb(encoder):
	encoder.t = 8
	if encoder.L == len(encoder.stream) or \
			(encoder.T == 255 and encoder.stream[encoder.L] > 143):
		encoder.C = encoder.C + 255
	else:
		if encoder.T == 255:
			encoder.t = 7
		encoder.T = encoder.stream[encoder.L]
		encoder.L = encoder.L + 1
		encoder.C = encoder.C + np.uint32((encoder.T) << (8 - encoder.t))
	return encoder

def _tile_encode(tile, D, h=64, w=64):
	_depthOfDwt = D

	tile_cA = tile[0][:, :, 0]
	newBit, newStream = _band_encode(tile_cA, 'LL', h, w)
	bitcode = newBit
	streamOnly = newStream
	for i in range(1, _depthOfDwt+1):
		newBit, newStream = _band_encode(tile[i][0][:, :, 0], 'LH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][1][:, :, 0], 'HL', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][2][:, :, 0], 'HH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))

	tile_cA = tile[0][:, :, 1]
	newBit, newStream = _band_encode(tile_cA, 'LL', h, w)
	bitcode = np.hstack((bitcode, newBit))
	streamOnly = np.hstack((streamOnly, newStream))
	for i in range(1, _depthOfDwt+1):
		newBit, newStream = _band_encode(tile[i][0][:, :, 1], 'LH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][1][:, :, 1], 'HL', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][2][:, :, 1], 'HH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))

	tile_cA = tile[0][:, :, 2]
	newBit, newStream = _band_encode(tile_cA, 'LL', h, w)
	bitcode = np.hstack((bitcode, newBit))
	streamOnly = np.hstack((streamOnly, newStream))
	for i in range(1, _depthOfDwt+1):
		newBit, newStream = _band_encode(tile[i][0][:, :, 2], 'LH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][1][:, :, 2], 'HL', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
		newBit, newStream = _band_encode(tile[i][2][:, :, 2], 'HH', h, w)
		bitcode = np.hstack((bitcode, newBit))
		streamOnly = np.hstack((streamOnly, newStream))
	bitcode = np.hstack((bitcode, [2051]))
	return bitcode

def _band_encode(tile, bandMark, h=64, w=64, num=8):
	# 码流：[h, w, CX1, 2048, stream1, 2048, ..., CXn, streamn, 2048, 2049,CXn+1, streamn+1, 2048, ...,2050]
	h_cA, w_cA = np.shape(tile)
	h_left_over = h_cA % h
	w_left_over = w_cA % w
	cA_extend = np.pad(tile, ((0, h - h_left_over), (0, w - w_left_over)), 'constant')
	bitcode = [h_cA, w_cA]
	streamOnly = []
	for i in range(0, h_cA, h):
		for j in range(0, w_cA, w):
			codeBlock = cA_extend[i:i + h, j:j + w]
			CX, D, bitplanelength= _embeddedBlockEncoder(codeBlock, bandMark, h, w, num)
			encoder = _MQencode(CX, D)
			bitcode = np.hstack((bitcode, CX.flatten(), [2048], encoder.stream, [2048], bitplanelength,[2048]))
			streamOnly = np.hstack((streamOnly, encoder.stream))
		bitcode = np.hstack((bitcode, [2049]))
	bitcode = np.hstack((bitcode, [2050]))
	return (bitcode, streamOnly)

def _embeddedBlockEncoder(codeBlock, bandMark, h=64, w=64, num=8):
	S1 = np.zeros((h, w))
	S2 = np.zeros((h, w))
	S3 = np.zeros((h, w))
	MaxInCodeBlock = len(bin(int(np.max(abs(codeBlock)))))-2
	signs = (- np.sign(codeBlock) + 1) // 2  # positive: 0, negative: 1
	bitPlane = np.zeros((h,w,MaxInCodeBlock),dtype=np.uint8)
	for i in range(h):
		for j in range(w):
			# number = bin(np.abs(codeBlock[i][j], dtype=np.int64))[2:]
			number = bin(np.abs(codeBlock[i][j]))[2:]
			temp = [0] * (MaxInCodeBlock-len(number)) + [int(num) for num in number]
			bitPlane[i][j] = np.array(temp)
	bitPlane = np.transpose(bitPlane, (2, 0, 1))
	# For Test
	"""
	signs = np.zeros((8,8))
	bitPlane = np.zeros((2,8,8))
	bitPlane[0][1][1] = 1
	bitPlane[0][4][4] = 1
	bitPlane[1][0][2] = 1
	bitPlane[1][1] = np.array([0,1,0,0,1,1,0,0])
	bitPlane[1][2][2] = 1
	bitPlane[1][3][3] = 1
	bitPlane[1][4][5] = 1
	bitPlane[1][5] = np.array([0,0,0,0,1,1,0,1])
	bitPlane[1][6][6] = 1
	"""
	sizeofCXandD = h*w*MaxInCodeBlock *5
	CX = np.zeros((sizeofCXandD, 1), dtype=np.uint32)
	D = np.zeros((sizeofCXandD, 1), dtype=np.uint32)
	pointer = 0
	for i in range(MaxInCodeBlock ):
		######
		# three function need rename
		D, CX, S1, S3, pointer = _SignifiancePropagationPass(D, CX, S1, S3, pointer, bitPlane[i], bandMark, signs, w, h)
		D, CX, S2, pointer = _MagnitudeRefinementPass(D, CX, S1, S2, S3, pointer, bitPlane[i], w, h)
		D, CX, pointer, S1 = _CLeanUpPass(D, CX, S1, S3, pointer, bitPlane[i], bandMark, signs, w, h)
		S3 = np.zeros((h, w))
	CX_final = CX[0:pointer]
	D_final = D[0:pointer]
	return CX_final, D_final, MaxInCodeBlock

# three encode pass start here
# in the sequence of significancePass,magnitudepass,_cleanuppass.

def _SignifiancePropagationPass(D, CX, S1, S3, pointer, plane, bandMark, signs, w=64, h=64):
	# input S1:  list of significance, size 64*64
	# input CX: the list of context
	# plane: the value of bits at this plane
	# bandMark: LL, HL, HH, or LH
	# pointer: the pointer of the CX
	# S3: denote that the element has been coded
	# output: D, CX, S1, S3, pointer
	S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			for ii in range(4):
				row = 4 * i + ii
				if S1[row][col] != 0:
					continue  # is significant
				temp = S1extend[row][col] + S1extend[row + 1][col] + S1extend[row + 2][col] + S1extend[row][col + 1] + \
							 S1extend[row + 2][col + 1] + S1extend[row][col + 2] + S1extend[row + 1][col + 2] + S1extend[row + 2][
								 col + 2]
				if temp == 0:
					continue  # is insignificant
				tempCx = _ZeroCoding(S1extend[row:row + 3, col:col + 3], bandMark)
				D[pointer][0] = plane[row][col]
				CX[pointer][0] = tempCx
				pointer = pointer + 1
				S3[row][col] = 1  # mark that plane[row][col] has been coded
				if plane[row][col] == 1:  # _signcoding
					signComp, tempCx = _SignCoding(S1extend[row:row + 3, col:col + 3], signs[row][col])
					D[pointer][0] = signComp
					CX[pointer][0] = tempCx
					pointer = pointer + 1
					S1[row][col] = 1  # mark as significant
					S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
	return D, CX, S1, S3, pointer

def _MagnitudeRefinementPass(D, CX, S1, S2, S3, pointer, plane, w=64, h=64):
	S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			for ii in range(4):
				row = 4 * i + ii
				if S1[row][col] != 1 or S3[row][col] != 0:
					continue
				tempCx = _MagnitudeRefinementCoding(S1extend[row:row + 3, col:col + 3], S2[row][col])
				S2[row][col] = 1  # Mark that the element has been refined
				D[pointer][0] = plane[row][col]
				CX[pointer][0] = tempCx
				pointer = pointer + 1
	return D, CX, S2, pointer

def _CLeanUpPass(D, CX, S1, S3, pointer, plane, bandMark, signs, w=64, h=64):
	S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			ii = 0
			row = 4 * i
			tempSum = np.sum(S1extend[row:row + 6, col:col + 3]) + np.sum(S3[row:row + 4, col])
			# 整一列未被编码，都为非重要，且领域非重要
			if tempSum == 0:
				ii, tempD, tempCx = _RunLengthCoding(plane[row:row + 4, col])
				if len(tempD) == 1:
					D[pointer] = tempD
					CX[pointer] = tempCx
					pointer = pointer + 1
				else:
					D[pointer], D[pointer + 1], D[pointer + 2] = tempD[0], tempD[1], tempD[2]
					CX[pointer], CX[pointer + 1], CX[pointer + 2] = tempCx[0], tempCx[1], tempCx[2]
					pointer = pointer + 3
					# sign coding
					row = i * 4 + ii - 1
					signComp, tempCx = _SignCoding(S1extend[row:row + 3, col:col + 3], signs[row][col])
					D[pointer] = signComp
					CX[pointer] = tempCx
					pointer = pointer + 1
					S1[row][col] = 1
					S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
			while ii < 4:
				row = i * 4 + ii
				ii = ii + 1
				if S1[row][col] != 0 or S3[row][col] != 0:
					continue
				tempCx = _ZeroCoding(S1extend[row:row + 3, col:col + 3], bandMark)
				D[pointer] = plane[row][col]
				CX[pointer] = tempCx
				pointer = pointer + 1
				if plane[row][col] == 1:  # _signcoding
					signComp, tempCx = _SignCoding(S1extend[row:row + 3, col:col + 3], signs[row][col])
					D[pointer][0] = signComp
					CX[pointer][0] = tempCx
					pointer = pointer + 1
					S1[row][col] = 1  # mark as significant
					S1extend = np.pad(S1, ((1, 1), (1, 1)), 'constant')
	return D, CX, pointer, S1

# here is some function used by three passes

def _SignCoding(neighbourS1, sign):
	# input neighbourS1: size 3*3, matrix of significance
	# input sign
	# output: signComp,(equal: 0, not equal: 1) context
	if len(neighbourS1) == 3 and len(neighbourS1[0]) == 3:
		hstr = str(int(neighbourS1[1][0])) + str(int(neighbourS1[1][2]))
		vstr = str(int(neighbourS1[0][1])) + str(int(neighbourS1[2][1]))
		dict = {
			'00': 0, '1-1': 0, '-11': 0, '01': 1, '10': 1, '11': 1,
			'0-1': -1, '-10': -1, '-1-1': -1
		}
		h = dict[hstr]
		v = dict[vstr]
		hAndv = str(h) + str(v)
		hv2Sign = {
			'11': 0, '10': 0, '1-1': 0, '01': 0, '00': 0,
			'0-1': 1, '-11': 1, '-10': 1, '-1-1': 1
		}
		hv2Context = {
			'11': 13, '10': 12, '1-1': 11, '01': 10, '00': 9,
			'0-1': 10, '-11': 11, '-10': 12, '-1-1': 13
		}
		signPredict = hv2Sign[hAndv]
		context = hv2Context[hAndv]
		signComp = int(sign) ^ signPredict
	else:
		# self.logs[-1] += self.formatter.warning("_SignCoding: Size of neighbourS1 not valid")
		signComp = -1
		context = -1
		"""
		try:
			raise ValidationError('_SignCoding: Size of neighbourS1 not valid')
		except ValidationError as e:
			print(e.args)
			signComp = -1
			context = -1
		"""
	return signComp, context

def _ZeroCoding(neighbourS1, bandMark):
	# input neighbourS1: size 3*3, matrix of significance
	# input s2: whether it is the first time for Magnitude Refinement Coding
	# output: context
	if len(neighbourS1) == 3 and len(neighbourS1[0]) == 3:
		h = neighbourS1[1][0] + neighbourS1[1][2]
		v = neighbourS1[0][1] + neighbourS1[2][1]
		d = neighbourS1[0][0] + neighbourS1[0][2] + neighbourS1[2][0] + neighbourS1[2][2]
		if bandMark == 'LL' or bandMark == 'LH':
			if h == 2:
				cx = 8
			elif h == 1 and v >= 1:
				cx = 7
			elif h == 1 and v == 0 and d >= 1:
				cx = 6
			elif h == 1 and v == 0 and d == 0:
				cx = 5
			elif h == 0 and v == 2:
				cx = 4
			elif h == 0 and v == 1:
				cx = 3
			elif h == 0 and v == 0 and d >= 2:
				cx = 2
			elif h == 0 and v == 0 and d == 1:
				cx = 1
			else:
				cx = 0
		elif bandMark == 'HL':
			if v == 2:
				cx = 8
			elif v == 1 and h >= 1:
				cx = 7
			elif v == 1 and h == 0 and d >= 1:
				cx = 6
			elif v == 1 and h == 0 and d == 0:
				cx = 5
			elif v == 0 and h == 2:
				cx = 4
			elif v == 0 and h == 1:
				cx = 3
			elif v == 0 and h == 0 and d >= 2:
				cx = 2
			elif v == 0 and h == 0 and d == 1:
				cx = 1
			else:
				cx = 0
		elif bandMark == 'HH':
			hPlusv = h + v
			if d >= 3:
				cx = 8
			elif d == 2 and hPlusv >= 1:
				cx = 7
			elif d == 2 and hPlusv == 0:
				cx = 6
			elif d == 1 and hPlusv >= 2:
				cx = 5
			elif d == 1 and hPlusv == 1:
				cx = 4
			elif d == 1 and hPlusv == 0:
				cx = 3
			elif d == 0 and hPlusv >= 2:
				cx = 2
			elif d == 0 and hPlusv == 1:
				cx = 1
			else:
				cx = 0
		else:
			# self.logs[-1] += self.formatter.warning('_ZeroCoding: bandMark not valid')
			cx = -1
			"""
			try:
				raise ValidationError('_ZeroCoding: bandMark not valid')
			except ValidationError as e:
				print(e.args)
				cx = -1
			"""
	else:
		# self.logs[-1] += self.formatter.warning('_ZeroCoding: Size of neighbourS1 not valid')
		cx = -1
		"""
		try:
			raise ValidationError('_ZeroCoding: Size of neighbourS1 not valid')
		except ValidationError as e:
			print(e.args)
			cx = -1
		"""
	return cx

def _RunLengthCoding(listS1):
	# input listS1: size 1*4, list of significance
	# output n: number of elements encoded
	# output d: 0 means the _RunLengthCoding does not end.
	# [1, x, x] means the _RunLengthCoding ends and the position is indicated.
	# output cx: context
	if listS1.__len__() == 4:
		if listS1[0] == 0 and listS1[1] == 0 and listS1[2] == 0 and listS1[3] == 0:
			n = 4
			d = [0]
			cx = [17]
		elif listS1[0] == 1:
			n = 1
			d = [1, 0, 0]
			cx = [17, 18, 18]
		elif listS1[0] == 0 and listS1[1] == 1:
			n = 2
			d = [1, 0, 1]
			cx = [17, 18, 18]
		elif listS1[0] == 0 and listS1[1] == 0 and listS1[2] == 1:
			n = 3
			d = [1, 1, 0]
			cx = [17, 18, 18]
		elif listS1[0] == 0 and listS1[1] == 0 and listS1[2] == 0 and listS1[3] == 1:
			n = 4
			d = [1, 1, 1]
			cx = [17, 18, 18]
		else:
			# self.logs[-1] += self.formatter.warning('_RunLengthCoding: listS1 not valid')
			n, d, cx = 0, -1, -1
			"""
			try:
				raise ValidationError('_RunLengthCoding: listS1 not valid')
			except ValidationError as e:
				print(e.args)
				n, d, cx = 0, -1, -1
			"""
	else:
		# self.logs[-1] += self.formatter.warning('_RunLengthCoding: length of listS1 not valid')
		n, d, cx = 0, -1, -1
		"""
		try:
			raise ValidationError('_RunLengthCoding: length of listS1 not valid')
		except ValidationError as e:
			print(e.args)
			n, d, cx = 0, -1, -1
		"""
	return n, d, cx

def _MagnitudeRefinementCoding(neighbourS1, s2):
	# input neighbourS1: size 3*3, matrix of significance
	# input s2: whether it is the first time for Magnitude Refinement Coding
	# output: context
	if len(neighbourS1) == 3 and len(neighbourS1[0]) == 3:
		temp = np.sum(neighbourS1) - neighbourS1[1][1]
		if s2 == 1:
			cx = 16
		elif s2 == 0 and temp >= 1:
			cx = 15
		else:
			cx = 14
	else:
		# self.logs[-1] += self.formatter.warning('_MagnitudeRefinementCoding: Size of neighbourS1 not valid')
		cx = -1
		"""
		try:
			raise ValidationError('_MagnitudeRefinementCoding: Size of neighbourS1 not valid')
		except ValidationError as e:
			print(e.args)
			cx = -1
		"""
	return cx

def _encode_end(encoder):
	nbits = 27 - 15 - encoder.t
	encoder.C = encoder.C * np.uint32(2 ** encoder.t)
	while nbits > 0:
		encoder = _transferbyte(encoder)
		nbits = nbits - encoder.t
		encoder.C = encoder.C * np.uint32(2 ** encoder.t)
	encoder = _transferbyte(encoder)

	return encoder

def _tile_decode(codestream, D):
	_depthOfDWT = D
	temp = []
	for i in range(0, 9 * _depthOfDWT + 3):
		_index = codestream.index(2050)
		deStream = codestream[0:_index + 1]
		temp.append(_band_decode(deStream))
		codestream = codestream[_index + 1:]

	start1 = _depthOfDWT*3+1
	start2 = _depthOfDWT*6+2
	tile = [cat_arrays_2d([temp[0],
												temp[start1],
												temp[start2]])]

	for i in range(_depthOfDWT):

		tile.append((cat_arrays_2d([temp[3 * i + 1],
															 temp[3 * i + start1+1],
															 temp[3 * i + start2+1]]),
								 cat_arrays_2d([temp[3 * i + 2],
															 temp[3 * i + start1+2],
															 temp[3 * i + start2+2]]),
								 cat_arrays_2d([temp[3 * i + 3],
															 temp[3 * i + start1+3],
															 temp[3 * i + start2+3]])))

	return tile

def _band_decode(codestream, h=64, w=64, num=32):
	h_cA = codestream[0]
	w_cA = codestream[1]
	codestream = codestream[2:]
	h_num = h_cA // h + 1
	w_num = w_cA // w + 1
	band_extend = np.zeros((h_num * h, w_num * w))
	for i in range(0, h_num):
		for j in range(0, w_num):
			_index = codestream.index(2048)
			deCX = codestream[0:_index]
			deCX = np.resize(deCX, (_index + 1, 1))
			codestream = codestream[_index + 1:]
			_index = codestream.index(2048)
			deStream = codestream[0:_index]
			codestream = codestream[_index + 1:]
			num = codestream[0]
			codestream = codestream[2:]
			decodeD = _MQ_decode(deStream, deCX)
			band_extend[i * h:(i + 1) * h, j * w:(j + 1) * w] = _decode_block(decodeD, deCX, h, w, num)
		if codestream[0] != 2049:
			print("Error!")
		codestream = codestream[1:]
	if codestream[0] != 2050:
		print("Error!")
	print("finish one band")
	return band_extend[0:h_cA, 0:w_cA]

def _decode_block(D, CX, h=64, w=64, num=32):
	deS1 = np.uint32(np.zeros((h, w)))
	deS2 = np.uint32(np.zeros((h, w)))
	deS3 = np.uint32(np.zeros((h, w)))
	signs = np.uint32(np.zeros((h, w)))
	V = np.uint32(np.zeros((num, h, w)))
	deCode = np.zeros((h, w))
	pointer = 0
	for i in range(num):
		V[i, :, :], signs, deS1, deS3, pointer = _SignificancePassDecoding(V[i, :, :], D, CX, deS1, deS3, pointer, signs, w,
																																			 h)
		V[i, :, :], deS2, pointer = _MagnitudePassDecoding(V[i, :, :], D, deS1, deS2, deS3, pointer, w, h)
		V[i, :, :], deS1, deS3, signs, pointer = _CleanPassDecoding(V[i, :, :], D, CX, deS1, deS3, pointer, signs, w, h)
		deS3 = np.zeros((h, w))
	V = np.transpose(V, (1, 2, 0))
	tempV = np.zeros((h,w))
	for i in range(h):
		for j in range(w):
			tempV[i][j] = sum([V[i][j][k]*2**(num-1-k) for k in range(num)])

	for i in range(h):
		for j in range(w):
			deCode[i][j] = (1 - 2 * signs[i][j]) * tempV[i][j]
	return deCode

def _SignificancePassDecoding(V, D, CX, deS1, deS3, pointer, signs, w=64, h=64):
	S1extend = np.pad(deS1, ((1, 1), (1, 1)), 'constant')
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			for ii in range(4):
				row = 4 * i + ii
				temp = np.sum(S1extend[row:row + 3, col:col + 3]) - S1extend[row + 1][col + 1]
				if deS1[row][col] != 0 or temp == 0:
					continue
				###
				if pointer>=len(D):
					continue
				V[row][col] = D[pointer][0]
				pointer = pointer + 1
				deS3[row][col] = 1
				if V[row][col] == 1:
					signs[row][col] = _SignDecoding(D[pointer], CX[pointer], S1extend[row:row + 3, col:col + 3])
					pointer = pointer + 1
					deS1[row][col] = 1
					S1extend = np.pad(deS1, ((1, 1), (1, 1)), 'constant')
	return V, signs, deS1, deS3, pointer

def _MagnitudePassDecoding(V, D, deS1, deS2, deS3, pointer, w=64, h=64):
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			for ii in range(4):
				row = 4 * i + ii
				if deS1[row][col] != 1 or deS3[row][col] != 0:
					continue
				###
				if pointer>=len(D):
					continue
				V[row][col] = D[pointer][0]
				pointer = pointer + 1
				deS2[row][col] = 1
	return V, deS2, pointer

def _CleanPassDecoding(V, D, CX, deS1, deS3, pointer, signs, w=64, h=64):
	a = pointer
	S1extend = np.pad(deS1, ((1, 1), (1, 1)), 'constant')
	rounds = h // 4
	for i in range(rounds):
		for col in range(w):
			ii = 0
			row = 4 * i
			tempSum = np.sum(S1extend[row:row + 6, col:col + 3]) + np.sum(deS3[row:row + 4, col])
			# 整一列未被编码，都为非重要，且领域非重要
			if tempSum == 0:
				if CX.__len__() < pointer + 3:
					CXextend = np.pad(CX, (0, 2), 'constant')
					Dextend = np.pad(D, (0, 2), 'constant')
					tempCx = CXextend[pointer:pointer + 3]
					tempD = Dextend[pointer:pointer + 3]
				else:
					tempCx = CX[pointer:pointer + 3]
					tempD = D[pointer:pointer + 3]
				ii, tempV = _RunLengthDecoding(tempCx, tempD)
				if tempV == [0, 0, 0, 0]:
					V[row][col] = 0
					V[row + 1][col] = 0
					V[row + 2][col] = 0
					V[row + 3][col] = 0
					pointer = pointer + 1
				else:
					if tempV == [1]:
						V[row][col] = 1
						pointer = pointer + 3
					elif tempV == [0, 1]:
						V[row][col] = 0
						V[row + 1][col] = 1
						pointer = pointer + 3
					elif tempV == [0, 0, 1]:
						V[row][col] = 0
						V[row + 1][col] = 0
						V[row + 2][col] = 1
						pointer = pointer + 3
					elif tempV == [0, 0, 0, 1]:
						V[row][col] = 0
						V[row + 1][col] = 0
						V[row + 2][col] = 0
						V[row + 3][col] = 1
						pointer = pointer + 3
					# sign coding
					row = row + ii - 1
					###
					if pointer>=len(D):
						continue
					signs[row][col] = _SignDecoding(D[pointer], CX[pointer], S1extend[row:row + 3, col:col + 3])
					pointer = pointer + 1
					deS1[row][col] = 1
					S1extend = np.pad(deS1, ((1, 1), (1, 1)), 'constant')
			while ii < 4:
				row = i * 4 + ii
				ii = ii + 1
				if deS1[row][col] != 0 or deS3[row][col] != 0:
					continue
				###
				if pointer>=len(D):
					continue
				V[row][col] = D[pointer][0]
				pointer = pointer + 1
				deS3[row][col] = 1
				if V[row][col] == 1:
					signs[row][col] = _SignDecoding(D[pointer], CX[pointer], S1extend[row:row + 3, col:col + 3])
					pointer = pointer + 1
					deS1[row][col] = 1
					S1extend = np.pad(deS1, ((1, 1), (1, 1)), 'constant')
	return V, deS1, deS3, signs, pointer

def _RunLengthDecoding(CX, D):
	n = CX.__len__()
	wrong = 1
	if CX[0][0] == 17 and D[0][0] == 0 or CX[0][0] == 17 and CX[1][0] == 18 and CX[2][0] == 18 and D[0][0] == 1:
		wrong = 0
	if wrong == 0:
		if D[0][0] == 0:
			deLen = 4
			V = [0, 0, 0, 0]
		elif D[0][0] == 1 and D[1][0] == 0 and D[2][0] == 0:
			deLen = 1
			V = [1]
		elif D[0][0] == 1 and D[1][0] == 0 and D[2][0] == 1:
			deLen = 2
			V = [0, 1]
		elif D[0][0] == 1 and D[1][0] == 1 and D[2][0] == 0:
			deLen = 3
			V = [0, 0, 1]
		elif D[0][0] == 1 and D[1][0] == 1 and D[2][0] == 1:
			deLen = 4
			V = [0, 0, 0, 1]
		else:
			# self.logs[-1] += self.formatter.warning('_RunLengthDecoding: D not valid')
			deLen = -1
			V = [-1]
			"""
			try:
				raise ValidationError('_RunLengthDecoding: D not valid')
			except ValidationError as e:
				print(e.args)
				deLen = -1
				V = [-1]
			"""
	else:
		# self.logs[-1] += self.formatter.warning('_RunLengthDecoding: CX not valid')
		deLen = -1
		V = [-1]
		"""
		try:
			raise ValidationError('_RunLengthDecoding: CX not valid')
		except ValidationError as e:
			print(e.args)
			deLen = -1
			V = [-1]
		"""
	return deLen, V

def _SignDecoding(D, CX, neighbourS1):
	if neighbourS1.__len__() == 3 and neighbourS1[0].__len__() == 3:
		hstr = str(int(neighbourS1[1][0])) + str(int(neighbourS1[1][2]))
		vstr = str(int(neighbourS1[0][1])) + str(int(neighbourS1[2][1]))
		dict = {
			'00': 0, '1-1': 0, '-11': 0, '01': 1, '10': 1, '11': 1,
			'0-1': -1, '-10': -1, '-1-1': -1
		}
		h = dict[hstr]
		v = dict[vstr]
		hAndv = str(h) + str(v)
		hv2Sign = {
			'11': 0, '10': 0, '1-1': 0, '01': 0, '00': 0,
			'0-1': 1, '-11': 1, '-10': 1, '-1-1': 1
		}
		hv2Context = {
			'11': 13, '10': 12, '1-1': 11, '01': 10, '00': 9,
			'0-1': 10, '-11': 11, '-10': 12, '-1-1': 13
		}
		temp = hv2Sign[hAndv]
		deCX = hv2Context[hAndv]
		if deCX == CX:
			deSign = D[0] ^ temp
		else:
			# self.logs[-1] += self.formatter.warning('_SignDecoding: Context does not match. Error occurs.')
			deSign = -1
			"""
			try:
				raise ValidationError('_SignDecoding: Context does not match. Error occurs.')
			except ValidationError as e:
				print(e.args)
				deSign = -1
			"""
	else:
		# self.logs[-1] += self.formatter.warning('_SignDecoding: Size of neighbourS1 not valid')
		deSign = -1
		"""
		try:
			raise ValidationError('_SignDecoding: Size of neighbourS1 not valid')
		except ValidationError as e:
			print(e.args)
			deSign = -1
		"""
	return deSign

class EBCOTparam(object):
	"""
	EBCOT parameter is parameter used by MQ encode and decode processes

	initialized in MQ encoding
	interval length         A = 8000H
	Lower bound register        C = 0
	Current code byte number    L = -1
	Temporary byte buffer     T = 0
	Down-counter          t = 12

	"""

	def __init__(self):
		self.C = np.uint32(0)
		self.A = np.uint16(32768)
		self.t = np.uint8(12)
		self.T = np.uint8(0)
		self.L = np.int32(-1)
		self.stream = np.uint8([])

'''
改了decodeblock，_embeddedBlockEncoder，banddecode和bandencode改了num的预设值
'''

# if __name__ == "__main__":
# 	h, w = 64, 64
# 	testblock = np.zeros((h, w), dtype=np.uint8)
# 	for i in range(h):
# 		for j in range(w):
# 			testblock[i][j] = i*4
# 	#(bitcode, _) = _band_encode(testblock, "LL", h=64, w=64, num=8)
# 	#decodeblock = _band_decode(list(bitcode), h=64, w=64, num=32)
# 	CX, D, bitplanelength = _embeddedBlockEncoder(testblock, "LL", h, w, num=8)
# 	encoder = _MQencode(CX, D)
# 	decodeD = _MQ_decode(encoder.stream, CX)
# 	decodeblock = _decode_block(D, CX, h, w, num=8)

# 	a = 1


import torch
import torch.nn
import os,csv,datetime
import numpy as np
from cnn_structure import CNN
from torch.autograd import Variable
from sklearn.metrics import confusion_matrix

# MODEL_FOLDER = './model/'
MODEL_FOLDER = './'
DATA_FOLDER = './k1000_vec200/'
EMOTION = {'ne':0, 'ha':1, 'sa':2, 'an':3, 'di':4, 'su':5, 'or':6, 'nx':7}

file_list = os.listdir(MODEL_FOLDER)
# cnn = []
# for file in file_list:
# 	if file.endswith('.pt'):
# 		cnn.append(torch.load(MODEL_FOLDER+file))
input_data = []
input_tag = []
file_list = os.listdir(DATA_FOLDER)
begin_read = datetime.datetime.now()

for file in file_list:
	if file.endswith('.csv'):
		frames = []
		csvreader = csv.reader(open(DATA_FOLDER+file))
		for frame in csvreader:
			tmp = np.array(frame)
			tmp = [float(i) for i in tmp]
			frames.append(tmp)
		input_data.append(frames)
		input_tag.append(EMOTION[file[-16:-14]])
input_data = np.array(input_data)
input_tag = np.array(input_tag)
end_time = datetime.datetime.now()
print('time of read ',len(input_data),' files: ',datetime.datetime.now()-begin_read)

cnn = torch.load('./k_1000_vec200_featureCNN_8emo_fold0date1227.pt')

acc_count = 0
results = []
time = []
for i in range(len(input_data)):
	begin = datetime.datetime.now()
	data_tensor = Variable(torch.from_numpy(np.array([[input_data[i]]]))).cuda().float()
	# tag_tensor = Variable(torch.from_numpy([test_tag[i]])).cuda().float()
	test_output, lastlayer = cnn(data_tensor)
	tag_output = torch.max(test_output,1)[1].cuda().data.squeeze()
	# data_tensor = Variable(torch.from_numpy(np.array([[test_data[i]]]))).float()
	# test_output, lastlayer = cnn(data_tensor)
	# tag_output = torch.max(test_output,1)[1].data.squeeze()
	print(tag_output[0], ' time:',datetime.datetime.now()-begin)
	time.append(datetime.datetime.now()-begin)

	results.append(tag_output[0])
	if tag_output[0]==input_tag[i]:
		acc_count+=1
	acc = acc_count/(len(input_tag))
confusion = confusion_matrix(input_tag,results)
print(acc)
print(confusion)
# print('avg_time:',sum(time)/len(time))


import os
import re
import json
import glob
import tempfile
import argparse
import ast
import pandas as pd
import sys
import shipyard_utils as shipyard

from google.cloud import bigquery
from google.oauth2 import service_account
from google.api_core.exceptions import NotFound

EXIT_CODE_UNKNOWN_ERROR = 3
EXIT_CODE_INVALID_CREDENTIALS = 200
EXIT_CODE_INVALID_DATASET = 201
EXIT_CODE_INVALID_SCHEMA = 202
EXIT_CODE_SCHEMA_MISMATCH = 203
EXIT_CODE_FILE_NOT_FOUND = 204

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', dest='dataset', required=True)
    parser.add_argument('--table', dest='table', required=True)
    parser.add_argument(
        '--service-account',
        dest='service_account',
        required=True)
    parser.add_argument(
        '--upload-type',
        dest='upload_type',
        default='append',
        choices={
            'append',
            'overwrite'},
        required=False)
    parser.add_argument(
        '--source-file-name-match-type',
        dest='source_file_name_match_type',
        default='exact_match',
        choices={
            'exact_match',
            'regex_match'},
        required=False)
    parser.add_argument(
        '--source-file-name',
        dest='source_file_name',
        required=True)
    parser.add_argument(
        '--source-folder-name',
        dest='source_folder_name',
        default='',
        required=False)
    parser.add_argument(
        '--schema',
        dest='schema',
        default='',
        required=False
    )
    parser.add_argument(
        '--skip-header-rows',
        dest='skip_header_rows',
        default='',
        required=False
    )
    parser.add_argument('--quoted-newline', dest = 'quoted_newline', default = False, required = False)
    args = parser.parse_args()
    return args

def set_environment_variables(args):
    """
    Set GCP credentials as environment variables if they're provided via keyword
    arguments rather than seeded as environment variables. This will override
    system defaults.
    """
    credentials = args.service_account
    try:
        json_credentials = json.loads(credentials)
        fd, path = tempfile.mkstemp()
        print(f'Storing json credentials temporarily at {path}')
        with os.fdopen(fd, 'w') as tmp:
            tmp.write(credentials)
        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = path
        return path
    except Exception:
        print('Using specified json credentials file')
        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = credentials
        return

def string_to_boolean(value):
    if isinstance(value, bool):
        return value
    if value.lower() in ('true', 't', 'y'):
        return True
    elif value.lower() in ('false', 'f', 'n'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def find_all_local_file_names(source_folder_name):
    """
    Returns a list of all files that exist in the current working directory,
    filtered by source_folder_name if provided.
    """
    cwd = os.getcwd()
    cwd_extension = os.path.normpath(f'{cwd}/{source_folder_name}/**')
    file_names = glob.glob(cwd_extension, recursive=True)
    return [file_name for file_name in file_names if os.path.isfile(file_name)]

def find_all_file_matches(file_names, file_name_re):
    """
    Return a list of all file_names that matched the regular expression.
    """
    matching_file_names = []
    for _file in file_names:
        if re.search(file_name_re, _file):
            matching_file_names.append(_file)

    return matching_file_names

def combine_folder_and_file_name(folder_name, file_name):
    """
    Combine together the provided folder_name and file_name into one path
    variable.
    """
    combined_name = os.path.normpath(
        f'{folder_name}{"/" if folder_name else ""}{file_name}')
    combined_name = os.path.normpath(combined_name)

    return combined_name

def copy_from_csv(
        client,
        dataset,
        table,
        source_file_path,
        upload_type,
        schema=None,
        skip_header_rows=None,
        quoted_newline=False):
    """
    Copy CSV data into Bigquery table.
    """
    try:
        dataset_ref = client.dataset(dataset)
        table_ref = dataset_ref.table(table)
        job_config = bigquery.LoadJobConfig()
        if upload_type == 'overwrite':
            job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
        else:
            job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
        job_config.source_format = bigquery.SourceFormat.CSV
        if skip_header_rows:
            job_config.skip_leading_rows = skip_header_rows
        if schema:
            job_config.autodetect = False
            job_config.schema = format_schema(schema)
        else:
            job_config.autodetect = True
        if quoted_newline:
            job_config.allow_quoted_newlines = True
        with open(source_file_path, 'rb') as source_file:
            job = client.load_table_from_file(source_file, table_ref,
                                              job_config=job_config)
        job.result()
    except NotFound as nf_e:
        if 'Not found: Dataset' in str(nf_e):
            print(
                f'The dataset {dataset} could not be found. Please check for typos and try again')
            print(nf_e)
            sys.exit(EXIT_CODE_INVALID_DATASET)
    except Exception as e:
        if ('Invalid value for mode' or 'Invalid value for type') in str(e):
            print(
                'The provided schema was not valid. Please check to make sure that the provided schema matches the following format and accepted values. \n\n \
                Format: [["column","datatype"],["column","datatype","MODE"]] \
                Accepted Values: https://cloud.google.com/bigquery/docs/schemas')
            print(e)
            sys.exit(EXIT_CODE_INVALID_SCHEMA)
        if 'Provided Schema does not match' in str(e):
            print('The provided schema does not match the schema for the existing table. Please check your table and ensure that the column names and data types match up exactly.')
            print(e)
            sys.exit(EXIT_CODE_SCHEMA_MISMATCH)
        else:
            print(f'Failed to copy CSV {source_file_path} to BigQuery.')
            print(e)
            sys.exit(EXIT_CODE_UNKNOWN_ERROR)

    print(
        f'Successfully copied csv {source_file_path} to {dataset}.{table} on BigQuery')

def get_client(credentials):
    """
    Attempts to create the Google Bigquery Client with the associated
    environment variables
    """
    try:
        client = bigquery.Client()
        return client
    except Exception as e:
        print(f'Error accessing Google Bigquery with service account '
              f'{credentials}')
        print(e)
        sys.exit(EXIT_CODE_INVALID_CREDENTIALS)

def format_schema(schema):
    formatted_schema = []
    schema = ast.literal_eval(schema)
    for item in schema:
        schema_column = 'bigquery.SchemaField("'
        for value in item:
            schema_column += value + '","'
        schema_column += '")'
        formatted_schema.append(eval(schema_column))
    return formatted_schema

def main():
    args = get_args()
    tmp_file = set_environment_variables(args)
    dataset = args.dataset
    table = args.table
    upload_type = args.upload_type
    source_file_name = args.source_file_name
    source_folder_name = args.source_folder_name
    source_full_path = combine_folder_and_file_name(
        folder_name=f'{os.getcwd()}/{source_folder_name}',
        file_name=source_file_name)
    source_file_name_match_type = args.source_file_name_match_type
    schema = args.schema
    quoted_newline = shipyard.args.convert_to_boolean(args.quoted_newline)

    skip_header_rows = args.skip_header_rows

    if skip_header_rows:
        skip_header_rows = int(args.skip_header_rows)

    if tmp_file:
        client = get_client(tmp_file)
    else:
        client = get_client(args.service_account)

    if source_file_name_match_type == 'regex_match':
        file_names = find_all_local_file_names(source_folder_name)
        matching_file_names = find_all_file_matches(
            file_names, re.compile(source_file_name))
        print(f'{len(matching_file_names)} files found. Preparing to upload...')

        for index, file_name in enumerate(matching_file_names):
            print(f'Uploading file {index+1} of {len(matching_file_names)}')
            copy_from_csv(
                client=client,
                dataset=dataset,
                table=table,
                source_file_path=file_name,
                upload_type=upload_type,
                schema=schema,
                skip_header_rows=skip_header_rows,
                quoted_newline=quoted_newline)
    else:
        if not os.path.isfile(source_full_path):
            print(f'File {source_full_path} does not exist')
            sys.exit(EXIT_CODE_FILE_NOT_FOUND)

        copy_from_csv(
            client=client,
            dataset=dataset,
            table=table,
            source_file_path=source_full_path,
            upload_type=upload_type,
            schema=schema,
            skip_header_rows=skip_header_rows,
            quoted_newline=quoted_newline)

    if tmp_file:
        print(f'Removing temporary credentials file {tmp_file}')
        os.remove(tmp_file)

if __name__ == '__main__':
    main()


import pymongo

cliente = pymongo.MongoClient("mongodb://localhost:27017/")
database = cliente["bancoDados"]
pessoas = database["pessoas"]

pessoa1 = {"nome":"Gustavo","peso": 58}
insert = pessoas.insert_one(pessoa1)
print(insert.inserted_id)

listaDBs = cliente.list_database_names()
print(listaDBs)
listaCollections = database.list_collection_names()
print(listaCollections)


"""This module contains necessary function needed"""

# Import necessary modules
from imblearn.over_sampling import SMOTE, ADASYN
from imblearn.under_sampling import RandomUnderSampler, NearMiss

import numpy as np
import pandas as pd
import streamlit as st
import math
from sklearn.model_selection import cross_validate
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
import sqlite3
import joblib

@st.cache_data
def load_data(uploaded_file):
    """This function returns the preprocessed data"""
    df = pd.read_csv("Parkinsons-Detector/" + uploaded_file)

    # Rename the column names in the DataFrame.
    df.rename(columns = {"MDVP:Fo(Hz)": "AVFF",}, inplace = True)
    df.rename(columns = {"MDVP:Fhi(Hz)": "MAVFF",}, inplace = True)
    df.rename(columns = {"MDVP:Flo(Hz)": "MIVFF",}, inplace = True)

    # Perform feature and target split
    X  = df.drop(columns=['status','name'])
    y = df['status']

    return df, X, y

def train_model(algorithm, K, X, y):
    """This function trains the model and returns the model scores"""
    def cross_validation(algorithm, _X, _y, _cv):
        _scoring = ['accuracy', 'precision', 'recall', 'f1']
        results = cross_validate(estimator=algorithm,
                                X=_X,
                                y=_y,
                                cv=_cv,
                                scoring=_scoring,
                                return_train_score=True,
                                return_estimator=True
                                )

        return results

    result = cross_validation(algorithm,X,y,K)

    accuracy =  result['test_accuracy']
    precision = result['test_precision']
    recall = result['test_recall']
    f1 = result['test_f1']
    best_f1 = list(result['test_f1']).index(max(result['test_f1']))
    model=result["estimator"][best_f1]

    return accuracy, precision, recall, f1, model

def sampling_function(sampling, X, y):
    # Get model and model score
    if sampling == "Over samling":
        X, y = SMOTE().fit_resample(X, y)
    elif sampling == "Under samling":
        random_undersampler = RandomUnderSampler(sampling_strategy='auto')  # You can specify the sampling strategy
        near_miss = NearMiss(version=1)  # You can specify the version (1, 2, or 3)
        X, y = random_undersampler.fit_resample(X, y)  # Replace X and y with your data
    return X, y

def predict(features, algorithm):
    # Get model and model score
    conn = sqlite3.connect('Parkinsons-Detector/pakinson_admin.db')
    cursor = conn.cursor()
    exc = cursor.execute("SELECT model_path, accuracy FROM model WHERE m_default=? AND algorithm=?", ("1",algorithm))
    model = exc.fetchall()

    # Deserialize the model from the file
    model_path = "Parkinsons-Detector/models" + "/" + model[0][0]
    loaded_model = joblib.load(model_path)
    # Predict the value
    prediction = loaded_model.predict([features])

    return prediction, model[0][1]


from Crypto.PublicKey import RSA
from django.contrib.auth.models import User
from rest_framework import serializers

from app.models import *
import uuid

class UserRelationField(serializers.RelatedField):
    def to_representation(self, value):
        return '{}'.format(value.user.username)

class AllOthersRelationField(UserRelationField):
    def get_queryset(self):
        request = self.context.get('request', None)
        queryset = super(AllOthersRelationField, self).get_queryset()
        if not request or not queryset:
            return None
        return queryset.all().exclude(user=request.user)

class MyRelationField(UserRelationField):
    def get_queryset(self):
        request = self.context.get('request', None)
        queryset = super(MyRelationField, self).get_queryset()
        if not request or not queryset:
            return None
        return queryset.filter(user=request.user)

class UserFilteredPrimaryKeyRelatedField(serializers.PrimaryKeyRelatedField):
    def get_queryset(self):
        request = self.context.get('request', None)
        queryset = super(UserFilteredPrimaryKeyRelatedField, self).get_queryset()
        if not request or not queryset:
            return None
        return queryset.filter(owner=request.user)

class UserFilteredSlugRelatedField(serializers.SlugRelatedField):
    def get_queryset(self):
        request = self.context.get('request', None)
        queryset = super(UserFilteredSlugRelatedField, self).get_queryset()
        if not request or not queryset:
            return None
        return queryset.filter(owner=request.user)

class UserSerializer(serializers.HyperlinkedModelSerializer):
    class Meta:
        model = User
        fields = [
            'url',
            'username',
            'email',
        ]

class SearchSerializer(serializers.ModelSerializer):
    class Meta:
        model = Search
        fields = '__all__'

class UserKeysSerializer(serializers.ModelSerializer):
    user = UserSerializer(read_only=True)

    signing_key = UserFilteredSlugRelatedField(
        many=False,
        read_only=False,
        queryset=PrivateKey.objects,
        slug_field='secure_id',
        required=False
    )
    messaging_key = UserFilteredSlugRelatedField(
        many=False,
        read_only=False,
        queryset=PrivateKey.objects,
        slug_field='secure_id',
        required=False
    )

    class Meta:
        model = UserKeys
        fields = '__all__'

class HashSerializer(serializers.ModelSerializer):
    owner = serializers.PrimaryKeyRelatedField(read_only=True)
    class Meta:
        model = Hash
        fields = '__all__'

class SignatureSerializer(serializers.ModelSerializer):
    class Meta:
        model = Signature
        fields = '__all__'
        lookup_field = 'user'

class MessageSerializer(serializers.ModelSerializer):
    owner = serializers.PrimaryKeyRelatedField(read_only=True)
    file_to_decrypt = serializers.HiddenField(default='')
    recipient_public_key = AllOthersRelationField(
        queryset=UserKeys.objects
    )
    signing_key = MyRelationField(
        queryset=UserKeys.objects
    )

    class Meta:
        model = Message
        fields = '__all__'

class DecryptionSerializer(serializers.ModelSerializer):
    owner = serializers.PrimaryKeyRelatedField(read_only=True)
    recipient_public_key = MyRelationField(
        queryset=UserKeys.objects
    )
    signing_key = AllOthersRelationField(
        queryset=UserKeys.objects
    )
    content = serializers.HiddenField(default='')

    class Meta:
        model = Message
        fields = '__all__'

class PrivateKeySerializer(serializers.ModelSerializer):
    id = serializers.ReadOnlyField()
    secure_id = serializers.ReadOnlyField()
    key_from_bytes = serializers.SerializerMethodField()
    owner = UserSerializer(read_only=True)

    # attribute method
    def get_key_from_bytes(self, obj):
        return obj.content

    # generate a private key with RSA module
    def private_gen(self):
        key = RSA.generate(2048) # 2048 is secure enough for modern standards
        return key.export_key('PEM')

    # called on post to endpoint
    def create(self, validated_data):
        gen = self.private_gen()
        priv_key = PrivateKey(
            content=gen,
            owner=validated_data['owner'],
            secure_id=uuid.uuid4()
        )
        priv_key.save()
        return priv_key

    class Meta:
        model = PrivateKey
        fields = '__all__'
        depth = 1


import numpy as np
import theano
import theano.tensor as T
from collections import OrderedDict

def simulate_dynamics(initial_pos, initial_vel, stepsize, n_steps, energy_fn):
    """
    Return final (position, velocity) obtained after an `n_steps` leapfrog
    updates, using Hamiltonian dynamics.

    Parameters
    ----------
    initial_pos: shared theano matrix
        Initial position at which to start the simulation
    initial_vel: shared theano matrix
        Initial velocity of particles
    stepsize: shared theano scalar
        Scalar value controlling amount by which to move
    n_steps: shared theano scalar
        Scalar value controlling number of steps for which to run the integrator
    energy_fn: python function
        Python function, operating on symbolic theano variables, used to
        compute the potential energy at a given position.

    Returns
    -------
    rval1: theano matrix
        Final positions obtained after simulation
    rval2: theano matrix
        Final velocity obtained after simulation
    """

    def leapfrog(pos, vel, step):
        """
        Inside loop of Scan. Performs one step of leapfrog update, using
        Hamiltonian dynamics.

        Parameters
        ----------
        pos: theano matrix
            in leapfrog update equations, represents pos(t), position at time t
        vel: theano matrix
            in leapfrog update equations, represents vel(t - stepsize/2),
            velocity at time (t - stepsize/2)
        step: theano scalar
            scalar value controlling amount by which to move

        Returns
        -------
        rval1: [theano matrix, theano matrix]
            Symbolic theano matrices for new position pos(t + stepsize), and
            velocity vel(t + stepsize/2)
        rval2: dictionary
            Dictionary of updates for the Scan Op
        """
        # from pos(t) and vel(t-stepsize/2), compute vel(t+stepsize/2)
        dE_dpos = T.grad(energy_fn(pos).sum(), pos)
        new_vel = vel - step * dE_dpos
        # from vel(t+stepsize/2) compute pos(t+stepsize)
        new_pos = pos + step * new_vel
        return [new_pos, new_vel], {}

    # compute velocity at time-step: t + stepsize/2
    initial_energy = energy_fn(initial_pos)
    dE_dpos = T.grad(initial_energy.sum(), initial_pos)
    vel_half_step = initial_vel - 0.5 * stepsize * dE_dpos

    # compute position at time-step: t + stepsize
    pos_full_step = initial_pos + stepsize * vel_half_step

    # perform leapfrog updates: the scan op is used to repeatedly compute
    # vel(t + (m-1/2)*stepsize) and pos(t + m*stepsize) for m in [2,n_steps].
    (all_pos, all_vel), scan_updates = theano.scan(
        leapfrog,
        outputs_info=[
            dict(initial=pos_full_step),
            dict(initial=vel_half_step),
        ],
        non_sequences=[stepsize],
        n_steps=n_steps - 1)
    final_pos = all_pos[-1]
    final_vel = all_vel[-1]
    # NOTE: Scan always returns an updates dictionary, in case the
    # scanned function draws samples from a RandomStream. These
    # updates must then be used when compiling the Theano function, to
    # avoid drawing the same random numbers each time the function is
    # called. In this case however, we consciously ignore
    # "scan_updates" because we know it is empty.
    assert not scan_updates

    # The last velocity returned by scan is vel(t +
    # (n_steps - 1 / 2) * stepsize) We therefore perform one more half-step
    # to return vel(t + n_steps * stepsize)
    energy = energy_fn(final_pos)
    final_vel = final_vel - 0.5 * stepsize * T.grad(energy.sum(), final_pos)

    # return new proposal state
    return final_pos, final_vel

# start-snippet-1

def kinetic_energy(vel):
    """Returns the kinetic energy associated with the given velocity
    and mass of 1.

    Parameters
    ----------
    vel: theano matrix
        Symbolic matrix whose rows are velocity vectors.

    Returns
    -------
    return: theano vector
        Vector whose i-th entry is the kinetic entry associated with vel[i].

    """
    return 0.5 * (vel ** 2).sum(axis=0)

def hamiltonian(pos, vel, energy_fn):
    """
    Returns the Hamiltonian (sum of potential and kinetic energy) for the given
    velocity and position.

    Parameters
    ----------
    pos: theano matrix
        Symbolic matrix whose rows are position vectors.
    vel: theano matrix
        Symbolic matrix whose rows are velocity vectors.
    energy_fn: python function
        Python function, operating on symbolic theano variables, used tox
        compute the potential energy at a given position.

    Returns
    -------
    return: theano vector
        Vector whose i-th entry is the Hamiltonian at position pos[i] and
        velocity vel[i].
    """
    # assuming mass is 1
    return energy_fn(pos) + kinetic_energy(vel)

def metropolis_hastings_accept(energy_prev, energy_next, s_rng):
    """
    Performs a Metropolis-Hastings accept-reject move.

    Parameters
    ----------
    energy_prev: theano vector
        Symbolic theano tensor which contains the energy associated with the
        configuration at time-step t.
    energy_next: theano vector
        Symbolic theano tensor which contains the energy associated with the
        proposed configuration at time-step t+1.
    s_rng: theano.tensor.shared_randomstreams.RandomStreams
        Theano shared random stream object used to generate the random number
        used in proposal.

    Returns
    -------
    return: boolean
        True if move is accepted, False otherwise
    """
    ediff = energy_prev - energy_next
    return (T.exp(ediff) - s_rng.uniform(size=energy_prev.shape)) >= 0

def MJHMC_accept():

    return

def hmc_move(s_rng, positions, energy_fn, stepsize=0.1, n_steps=1):
    """
    This function performs one-step of Hybrid Monte-Carlo sampling. We start by
    sampling a random velocity from a univariate Gaussian distribution, perform
    `n_steps` leap-frog updates using Hamiltonian dynamics and accept-reject
    using Metropolis-Hastings.

    Parameters
    ----------
    s_rng: theano shared random stream
        Symbolic random number generator used to draw random velocity and
        perform accept-reject move.
    positions: shared theano matrix
        Symbolic matrix whose rows are position vectors.
    energy_fn: python function
        Python function, operating on symbolic theano variables, used to
        compute the potential energy at a given position.
    stepsize:  shared theano scalar
        Shared variable containing the stepsize to use for `n_steps` of HMC
        simulation steps.
    n_steps: integer
        Number of HMC steps to perform before proposing a new position.

    Returns
    -------
    rval1: boolean
        True if move is accepted, False otherwise
    rval2: theano matrix
        Matrix whose rows contain the proposed "new position"
    """
    # sample random velocity
    initial_vel = s_rng.normal(size=positions.shape)

    final_pos, final_vel = simulate_dynamics(
        initial_pos=positions,
        initial_vel=initial_vel,
        stepsize=stepsize,
        n_steps=n_steps,
        energy_fn=energy_fn
    )
    # accept/reject the proposed move based on the joint distribution
    accept = metropolis_hastings_accept(
        energy_prev=hamiltonian(positions, initial_vel, energy_fn),
        energy_next=hamiltonian(final_pos, final_vel, energy_fn),
        s_rng=s_rng
    )
    return accept, final_pos

def hmc_updates(positions,final_pos, accept):
    """def hmc_updates(positions, stepsize, avg_acceptance_rate, final_pos, accept,
                target_acceptance_rate, stepsize_inc, stepsize_dec,
                stepsize_min, stepsize_max, avg_acceptance_slowness):
    This function is executed after `n_steps` of HMC sampling
    (`hmc_move` function). It creates the updates dictionary used by
    the `simulate` function. It takes care of updating: the position
    (if the move is accepted), the stepsize (to track a given target
    acceptance rate) and the average acceptance rate (computed as a
    moving average).

    Parameters
    ----------
    positions: shared variable, theano matrix
        Shared theano matrix whose rows contain the old position
    stepsize: shared variable, theano scalar
        Shared theano scalar containing current step size
    avg_acceptance_rate: shared variable, theano scalar
        Shared theano scalar containing the current average acceptance rate
    final_pos: shared variable, theano matrix
        Shared theano matrix whose rows contain the new position
    accept: theano scalar
        Boolean-type variable representing whether or not the proposed HMC move
        should be accepted or not.
    target_acceptance_rate: float
        The stepsize is modified in order to track this target acceptance rate.
    stepsize_inc: float
        Amount by which to increment stepsize when acceptance rate is too high.
    stepsize_dec: float
        Amount by which to decrement stepsize when acceptance rate is too low.
    stepsize_min: float
        Lower-bound on `stepsize`.
    stepsize_min: float
        Upper-bound on `stepsize`.
    avg_acceptance_slowness: float
        Average acceptance rate is computed as an exponential moving average.
        (1-avg_acceptance_slowness) is the weight given to the newest
        observation.

    Returns
    -------
    rval1: dictionary-like
        A dictionary of updates to be used by the `HMC_Sampler.simulate`
        function.  The updates target the position, stepsize and average
        acceptance rate.

    """

    ## POSITION UPDATES ##
    # broadcast `accept` scalar to tensor with the same dimensions as
    # final_pos.
    #accept_matrix = accept.dimshuffle(0, *(('x',) * (final_pos.ndim - 1)))
    # if accept is True, update to `final_pos` else stay put
    #new_positions = T.switch(accept_matrix, final_pos, positions)
    #new_positions = T.switch(accept.ravel(), final_pos, positions).astype('float32')
    #new_positions = T.switch(accept[0].dimshuffle('x',0), final_pos, positions).astype('float32')
    new_positions = accept[0]*final_pos + (1-accept[0])*positions
    ## ACCEPT RATE UPDATES ##
    # perform exponential moving average
    '''
    mean_dtype = theano.scalar.upcast(accept.dtype, avg_acceptance_rate.dtype)
    new_acceptance_rate = T.add(
        avg_acceptance_slowness * avg_acceptance_rate,
        (1.0 - avg_acceptance_slowness) * accept.mean(dtype=mean_dtype))
    ## STEPSIZE UPDATES ##
    # if acceptance rate is too low, our sampler is too "noisy" and we reduce
    # the stepsize. If it is too high, our sampler is too conservative, we can
    # get away with a larger stepsize (resulting in better mixing).
    _new_stepsize = T.switch(avg_acceptance_rate > target_acceptance_rate,
                              stepsize * stepsize_inc, stepsize * stepsize_dec)
    # maintain stepsize in [stepsize_min, stepsize_max]
    new_stepsize = T.clip(_new_stepsize, stepsize_min, stepsize_max)
    return [(positions, new_positions),
            (stepsize, new_stepsize),
            (avg_acceptance_rate, new_acceptance_rate)]
    '''
    update = OrderedDict()
    update[positions] = new_positions.astype('float32')
    #return [(positions, new_positions)]
    return update

def wrapper_hmc(s_rng,energy_fn,dim=np.array([2,1]),L=10, beta = 0.1, epsilon = 0.1):

    """
    This should be the wrapper call that calls the various HMC definitions

    Parameters:
      Potential Energy -- function handle that captures the interest distrbution
      Number of Leap Frog Steps -- L (10)
      Momentum corruption parameter -- beta (0.1)
      Leapfrog Integrator step length -- epsilon (0.1)
      nsamples -- The number of samples to be generated (100)
    Returns:
      samples -- Samples generated
    """
    pos = np.random.randn(dim[0],dim[1]).astype('float32')
    vel = np.random.randn(dim[0],dim[1]).astype('float32')
    pos = theano.shared(pos,name='pos')
    vel = theano.shared(vel,name='vel')
    epsilon = theano.shared(epsilon,name='epsilon')
    L = theano.shared(L,'L')
    #pos, vel = simulate_dynamics(initial_pos=pos,initial_vel=vel,stepsize=epsilon,n_steps=L,energy_fn=energy_fn)
    accept, final_pos = hmc_move(s_rng=s_rng, positions=pos, energy_fn=energy_fn,stepsize=epsilon,n_steps=L)
    #Simulate updates
    simulate_updates = hmc_updates(positions=pos,final_pos=final_pos,accept=accept)
    simulate = theano.function([],[accept,simulate_updates[pos]],updates=simulate_updates)

    return simulate

def autocorrelation():
    X = T.tensor3().astype('float32')
    shape = X.shape
    #Assumes Length T, need to have a switch that also deals with (T/2)-1
    t_gap = T.arange(1,shape[2]-1)
    outputs_info = T.zeros((1,1,1),dtype='float32')

    #function def that computes the mean ac for each time lag
    def calc_ac(t_gap,output_t,X):
        return T.mean(X[:,:,:-t_gap]*X[:,:,t_gap:],dtype='float32',keepdims=True)

    #We will write a scan function that loops over the indices of the data tensor
    #and computes the autocorrelation
    result,updates = theano.scan(fn= calc_ac,
    #ac,updates = theano.scan(fn= lambda X,t_gap,ac: T.mean(X[:,:,:-t_gap]*X[:,:t_gap:]),
            outputs_info=[outputs_info],
            sequences=[t_gap],
            non_sequences=[X])
    #Append zero mean value of X to the front of the array and then return
    #Also, need to divide by the first element to scale the variances
    #For now though, let's do this in the main script
    theano_ac = theano.function(inputs=[X],outputs=[result],updates=updates)
    return theano_ac

def normed_autocorrelation(df):
    theano_ac = autocorrelation()

    Time = len(df)
    N, nbatch = df.loc[0]['X'].shape

    X = np.zeros((N,nbatch,Time))

    for tt in range(Time):
       X[:,:,tt] = df.loc[tt]['X']

    ac= theano_ac(X.astype('float32'))
    X_mean = np.mean(X**2,keepdims=True)[0][0]
    ac_squeeze = np.squeeze(ac[0])
    ac_squeeze = ac_squeeze/X_mean
    ac = np.vstack((1.,ac_squeeze.reshape(Time-2,1)))
    #This drops the last sample out of the data frame. Unclear, if this is the best way to do things but
    #it is the only way we can align the total number of samples from sample generation to
    #computing autocorrelation
    ac_df = df[:-1]
    ac_df.loc[:,'autocorrelation'] = ac
    return ac_df[['num energy', 'num grad', 'autocorrelation']]


"""
Tests for specific issues and pull requests
"""

import os
import tempfile
import difflib
from textwrap import dedent
import gffutils
from gffutils import feature
from gffutils import merge_criteria as mc
from nose.tools import assert_raises

def test_issue_79():
    gtf = gffutils.example_filename("keep-order-test.gtf")
    db = gffutils.create_db(
        gtf,
        "tmp.db",
        disable_infer_genes=False,
        disable_infer_transcripts=False,
        id_spec={"gene": "gene_id", "transcript": "transcript_id"},
        merge_strategy="create_unique",
        keep_order=True,
        force=True,
    )

    exp = open(gtf).read()
    obs = "\n".join([str(i) for i in db.all_features()])
    exp_1 = exp.splitlines(True)[0].strip()
    obs_1 = obs.splitlines(True)[0].strip()
    print("EXP")
    print(exp_1)
    print("OBS")
    print(obs_1)
    print("DIFF")
    print("".join(difflib.ndiff([exp_1], [obs_1])))
    assert obs_1 == exp_1

def test_issue_82():
    # key-val separator is inside an unquoted attribute value
    x = (
        "Spenn-ch12\tsgn_markers\tmatch\t2621812\t2622049\t.\t+\t.\t"
        "Alias=SGN-M1347;ID=T0028;Note=marker name(s): T0028 SGN-M1347 |identity=99.58|escore=2e-126"
    )
    y = feature.feature_from_line(x)
    assert y.attributes["Note"] == [
        "marker name(s): T0028 SGN-M1347 |identity=99.58|escore=2e-126"
    ]

    gffutils.create_db(gffutils.example_filename("keyval_sep_in_attrs.gff"), ":memory:")

def test_issue_85():
    # when start or stop was empty, #85 would fail Should now work with
    # blank fields
    f = feature.feature_from_line("\t".join([""] * 9))

    # or with "." placeholders
    f = feature.feature_from_line("\t".join(["."] * 9))

def test_issue_105():
    fn = gffutils.example_filename("FBgn0031208.gtf")
    home = os.path.expanduser("~")
    newfn = os.path.join(home, ".gffutils.test")
    with open(newfn, "w") as fout:
        fout.write(open(fn).read())
    f = gffutils.iterators.DataIterator(newfn)
    for i in f:
        pass
    os.unlink(newfn)

def test_issue_107():
    s = dedent(
        """
        chr1\t.\tgene\t10\t15\t.\t+\t.\tID=b;
        chr1\t.\tgene\t1\t5\t.\t-\t.\tID=a;
        chr2\t.\tgene\t25\t50\t.\t-\t.\tID=c;
        chr2\t.\tgene\t55\t60\t.\t-\t.\tID=d;
        """
    )
    tmp = tempfile.NamedTemporaryFile(delete=False).name
    with open(tmp, "w") as fout:
        fout.write(s + "\n")
    db = gffutils.create_db(tmp, ":memory:")
    interfeatures = list(
        db.interfeatures(db.features_of_type("gene", order_by=("seqid", "start")))
    )
    assert [str(i) for i in interfeatures] == [
        "chr1\tgffutils_derived\tinter_gene_gene\t6\t9\t.\t.\t.\tID=a,b;",
        "chr2\tgffutils_derived\tinter_gene_gene\t16\t54\t.\t-\t.\tID=c,d;",
    ]

def test_issue_119():
    # First file has these two exons with no ID:
    #
    #   chr2L FlyBase exon  8193  8589  .  +  .  Parent=FBtr0300690
    #   chr2L FlyBase exon  7529  8116  .  +  .  Name=CG11023:1;Parent=FBtr0300689,FBtr0300690
    #
    db0 = gffutils.create_db(gffutils.example_filename("FBgn0031208.gff"), ":memory:")

    # And this one, a bunch of reads with no IDs anywhere
    db1 = gffutils.create_db(
        gffutils.example_filename("F3-unique-3.v2.gff"), ":memory:"
    )

    # When db1 is updated by db0
    db2 = db1.update(db0)
    assert (
        db2._autoincrements == db1._autoincrements == {"exon": 2, "read": 112}
    ), db2._autoincrements

    assert len(list(db0.features_of_type("exon"))) == 6

    # Now we update that with db0 again
    db3 = db2.update(db0, merge_strategy="replace")

    # Using the "replace" strategy, we should have only gotten another 2 exons
    assert len(list(db3.features_of_type("exon"))) == 8

    # Make sure that the autoincrements for exons jumped by 2
    assert (
        db2._autoincrements == db3._autoincrements == {"exon": 4, "read": 112}
    ), db2._autoincrements

    # More isolated test, merging two databases each created from the same file
    # which itself contains only a single feature with no ID.
    tmp = tempfile.NamedTemporaryFile(delete=False).name
    with open(tmp, "w") as fout:
        fout.write("chr1\t.\tgene\t10\t15\t.\t+\t.\t\n")

    db4 = gffutils.create_db(tmp, tmp + ".db")
    db5 = gffutils.create_db(tmp, ":memory:")

    assert db4._autoincrements == {"gene": 1}
    assert db5._autoincrements == {"gene": 1}

    db6 = db4.update(db5)

    db7 = gffutils.FeatureDB(db4.dbfn)

    # both db4 and db6 should now have the same, updated autoincrements because
    # they both point to the same db.
    assert db6._autoincrements == db4._autoincrements == {"gene": 2}

    # But db5 was created independently and should have unchanged autoincrements
    assert db5._autoincrements == {"gene": 1}

    # db7 was created from the database pointed to by both db4 and db6. This
    # tests that when a FeatureDB is created it should have the
    # correctly-updated autoincrements read from the db
    assert db7._autoincrements == {"gene": 2}

def test_pr_131():
    db = gffutils.create_db(gffutils.example_filename("FBgn0031208.gff"), ":memory:")

    # previously would raise ValueError("No lines parsed -- was an empty
    # file provided?"); now just does nothing
    db2 = db.update([])

def test_pr_133():
    # Previously, merge_attributes would not deep-copy the values from the
    # second dict, and when the values are then modified, the second dict is
    # unintentionally modified.
    d1 = {"a": [1]}
    d2 = {"a": [2]}
    d1a = {"a": [1]}
    d2a = {"a": [2]}
    d3 = gffutils.helpers.merge_attributes(d1, d2)
    assert d1 == d1a, d1
    assert d2 == d2a, d2

def test_pr_139():
    db = gffutils.create_db(gffutils.example_filename("FBgn0031208.gff"), ":memory:")
    exons = list(db.features_of_type("exon"))
    inter = list(db.interfeatures(exons))

    # previously, the first exon's attributes would show up in subsequent merged features
    assert exons[0].attributes["Name"][0] not in inter[1].attributes["Name"]
    assert exons[0].attributes["Name"][0] not in inter[2].attributes["Name"]
    assert exons[0].attributes["Name"][0] not in inter[3].attributes["Name"]

def test_pr_144():
    # previously this would fail with:
    #   UnboundLocalError: local variable 'part' referenced before assignment
    f = gffutils.Feature(attributes={"a": [""]})

    # Make sure everything got converted correctly
    assert f.attributes["a"] == [""]
    assert str(f) == ".	.	.	.	.	.	.	.	a"
    g = gffutils.feature.feature_from_line(str(f))
    assert g == f

def test_pr_172():
    line = (
        "NC_049222.1\tGnomon\tgene\t209085\t282880\t.\t-\t.\t"
        'gene_id "ENPP1_3"; transcript_id ""; db_xref "GeneID:100856150";'
        'db_xref "VGNC:VGNC:40374"; gbkey "Gene"; gene "ENPP1"; '
        'gene_biotype "protein_coding";\n'
    )
    tmp = tempfile.NamedTemporaryFile(delete=False).name
    with open(tmp, "w") as fout:
        fout.write(line)
    db = gffutils.create_db(tmp, ":memory:")

def test_pr_171():
    q = gffutils.parser.Quoter()
    assert q.__missing__("\n") == "%0A"
    assert q.__missing__("a") == "a"

    assert q.__missing__("") == ""

def test_issue_129():

    # thanks @Brunox13 for the detailed notes on #129

    line = 'chr1\tdemo\tstart_codon\t69091\t69093\t.\t+\t.\tgene_id "demo";\n'
    tmp = tempfile.NamedTemporaryFile(delete=False).name
    with open(tmp, "w") as fout:
        fout.write(line)
    db = gffutils.create_db(tmp, ":memory:")

    # ASCII art to visualize each test (coords are along the top, from 69087 to
    # 69090). The tests slide a 4-bp region over the original 3-bp start codon.

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    # | | | |               Test feature
    res = list(db.region(region=("chr1", 69087, 69090), featuretype="start_codon"))
    assert len(res) == 0

    # NOTE: prior to #162, this did not return anything
    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #   | | | |             Test feature
    res = list(db.region(region=("chr1", 69088, 69091), featuretype="start_codon"))
    assert len(res) == 1

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #     | | | |           Test feature
    res = list(db.region(region=("chr1", 69089, 69092), featuretype="start_codon"))
    assert len(res) == 1

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #       | | | |         Test feature
    res = list(db.region(region=("chr1", 69090, 69093), featuretype="start_codon"))
    assert len(res) == 1

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #         | | | |       Test feature
    res = list(db.region(region=("chr1", 69091, 69094), featuretype="start_codon"))
    assert len(res) == 1

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #           | | | |     Test feature
    res = list(db.region(region=("chr1", 69092, 69095), featuretype="start_codon"))
    assert len(res) == 1

    # NOTE: priro to #162, this did not return anything
    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #             | | | |   Test feature
    res = list(db.region(region=("chr1", 69093, 69096), featuretype="start_codon"))
    assert len(res) == 1

    # 7 8 9 0 1 2 3 4 5 6 7
    #         | | |         Orig feature
    #               | | | | Test feature
    res = list(db.region(region=("chr1", 69094, 69097), featuretype="start_codon"))
    assert len(res) == 0

def test_issue_128():
    # In #128, some lines had separators of "; " and some with ";". The first
    # one in the file would win. Now the detection pays more attention to lines
    # with more attributes to make it work properly
    gff = gffutils.example_filename('gms2_example.gff3')
    db = gffutils.create_db(gff, ":memory:", force=True)
    expected = {
        'ID': ['1'],
        'Parent': ['gene_1'],
        'gene_type': ['native'],
        'partial': ['11'],
        'gc': ['33'],
        'length': ['363'],
    }
    assert dict(db['1'].attributes) == expected

def test_issue_157():
    # With the merge overhaul, children_bp incorrectly still used ignore_strand.
    db = gffutils.create_db(gffutils.example_filename('FBgn0031208.gff'), ":memory:")
    gene = next(db.features_of_type('gene'))
    children = list(db.children(gene, featuretype='exon'))

    # Modify the last one to have a different strand so we can test the
    # ignore_strand argument.
    children[-1].strand = '-'
    db.update(children[-1:], merge_strategy='replace')

    # and, since updating has been problematic in the past, double-check again
    # that the strand is changed in the db.
    assert list(db.children(gene, featuretype='exon'))[-1].strand == '-'
    cbp1 = db.children_bp(gene, child_featuretype='exon')

    # Previously this would give:
    #   TypeError: merge() got an unexpected keyword argument 'ignore_strand'
    #
    # Now changing to ValueError and suggesting a fix.
    assert_raises(ValueError, db.children_bp, gene, child_featuretype='exon', merge=True, ignore_strand=True)
    assert_raises(ValueError, db.children_bp, gene, ignore_strand=True, nonexistent=True)
    assert_raises(TypeError, db.children_bp, gene, nonexistent=True)

    # The way to do it now is the following (we can omit the mc.feature_type
    # since we're preselecting for exons anyway):
    db.children_bp(gene, child_featuretype='exon', merge=True, merge_criteria=(mc.overlap_end_inclusive))

def test_issue_159():
    db = gffutils.create_db(gffutils.example_filename('FBgn0031208.gff'), ":memory:")
    fasta = gffutils.example_filename('dm6-chr2L.fa')
    for transcript, seq in gffutils.helpers.canonical_transcripts(db, fasta):
        pass

def test_issue_164():
    # Something strange with the original annotation, but seems fine at least
    # after pasting in the offending genes from the GitHub comments.
    db = gffutils.create_db(
        gffutils.example_filename('sharr.gtf'),
        ':memory:',
        disable_infer_transcripts=True,
        disable_infer_genes=True,
        id_spec={'gene': 'gene_id', 'transcript': 'transcript_id'},
        merge_strategy='create_unique',
        keep_order=True)

def test_issue_166():
    # Added the new FeatureDB.seqids() method.
    db = gffutils.create_db(gffutils.example_filename('nonascii'), ':memory:')
    seqs = list(db.seqids())
    assert seqs == ['2L', '2R', '3L', '3R', 'X'], seqs

def test_issue_167():
    # Previously was causing sqlite3.InterfaceError
    db = gffutils.create_db(gffutils.example_filename('issue167.gff'), ':memory:')

def test_issue_174():
    db = gffutils.create_db(
        gffutils.example_filename('issue174.gtf'),
        ':memory:',
        merge_strategy='warning',
    )
    introns = [f for f in db.create_introns()]
    observed = [i.attributes['exon_number'] for i in introns]
    assert observed[7] == ['8', '9']
    assert observed[8] == ['10', '9']
    assert observed[9] == ['10', '11']

    # Now do the same thing, but with the new numeric_sort arg
    introns = [f for f in db.create_introns(numeric_sort=True)]
    observed = [i.attributes['exon_number'] for i in introns]
    assert observed[7] == ['8', '9']
    # This should be fixed:
    assert observed[8] == ['9', '10']
    assert observed[9] == ['10', '11']

def test_issue_181():
    db = gffutils.create_db(
        gffutils.example_filename('issue181.gff'),
        ':memory:')
    introns = db.create_introns()

    # This now warns that the provided ID key has multiple values.
    assert_raises(ValueError, db.update, introns)

    # The fix is to provide a custom intron ID converter.
    def intron_id(f):
        return ','.join(f['ID'])

    db.update(introns, id_spec={'intron': [intron_id]})


#!/usr/bin/env python
#
# Author: Greg Hellings - <ghelling@redhat.com> or <greg.hellings@gmail.com>
#
# Module to configure users in Jenkins authorized to use CLI
import xml.etree.ElementTree as ET
import os
from ansible.module_utils.basic import AnsibleModule

DOCUMENTATION = """
---
version_added: "2.1"
module: jenkins_cli_user
short_description: configure Jenkins CLI users with pub key
description:
  - This module configures admin users in Jenkins to utilize the specified
    SSH pubkey. Requires that role-based authentication be enabled and that
    a user be configured as an admin

options:
  jenkins_home:
    description:
     The root directory for the Jenkins install
    required: true
  jenkins_user:
    description:
      The name of the user to configure the SSH key for
  key_file:
    description:
      Path to the SSH keyfile to be listed as authorized
    required: true
  state:
    description:
      Currently limited to "present" - will create the user
    required: false

author: Gregory Hellings
"""

class UserList(object):
    def __init__(self, user_path):
        config = os.path.join(user_path, "users.xml")
        self.root = ET.parse(config).getroot()

    @property
    def users(self):
        if not hasattr(self, "_users"):
            mapping = self.root.find("idToDirectoryNameMap")
            if not mapping:
                return {}
            entries = mapping.getiterator("entry")
            if not entries:
                return {}
            users = {}
            for entry in entries:
                users[entry[0].text] = entry[1].text
            self._users = users
        return self._users

class User(object):
    def __init__(self, user_path, dirname):
        self.config = os.path.join(user_path, dirname, "config.xml")
        self.document = ET.parse(self.config)
        self.root = self.document.getroot()

    @property
    def keys(self):
        properties = self.root.find("properties")
        keys = properties.getiterator("authorizedKeys")
        return keys

    def add_key(self, pub_key):
        # If keys exist
        if self.keys:
            for key in self.keys:
                if pub_key not in str(key.text):
                    # No actual keys there, so create it, otherwise prepend
                    if key.text is None:
                        key.text = pub_key
                    else:
                        key.text = str(key.text) + "\n" + pub_key
                    return True
        # No keys exist
        else:
            properties = self.root.find("properties")
            xml = "org.jenkinsci.main.modules.cli.auth.ssh.UserPropertyImpl"
            ssh_auth = ET.SubElement(properties, xml)
            auth_key = ET.SubElement(ssh_auth, "authorizedKeys")
            auth_key.text = pub_key
            return True
        return False

    def save(self):
        self.document.write(self.config, encoding="UTF-8")

def main():
    module = AnsibleModule(
        argument_spec={
            "jenkins_home": {"required": True},
            "jenkins_user": {"required": True},
            "key": {"required": True},
            "state": {"choices": ["present"], "default": "present"},
        },
        supports_check_mode=False,
    )
    params = type("Params", (object,), module.params)
    user_path = os.path.join(params.jenkins_home, "users")
    changed = False

    # Check if user exists and find directory mapping
    users = UserList(user_path)
    if params.jenkins_user not in users.users.keys():
        module.fail_json(msg="No such user found." + users.users.keys())

    # Check user for configured key and add it
    user = User(user_path, users.users[params.jenkins_user])
    changed = user.add_key(params.key)
    if changed:
        user.save()
    module.exit_json(changed=changed)

main()


import wqpy.read
import aiohttp
import asyncio
import io

async def _basic_aquery(service_url, service_params):
  async with aiohttp.ClientSession() as session:
    async with session.get(service_url, params = service_params) as r:
      return(await r.text())

def multi_query(service_url, service_param_list, parse = True):
  query_args = [(service_url, params) for params in service_param_list]
  loop = asyncio.get_event_loop()
  outputs = loop.run_until_complete(
    asyncio.gather(*(_basic_aquery(*args) for args in query_args))
  )
  if parse:
    return([wqpy.read.read(io.StringIO(o)) for o in outputs])
  else:
    return([io.StringIO(o).read() for o in outputs])


from elegy.module import Module
import typing as tp

import haiku as hk
import jax.numpy as jnp
import numpy as np

from elegy import types

def _infer_shape(output_shape, dimensions):
    """
    Replaces the -1 wildcard in the output shape vector.

    This function infers the correct output shape given the input dimensions.

    Args:
        output_shape: Output shape.
        dimensions: List of input non-batch dimensions.

    Returns:
        Tuple of non-batch output dimensions.
    """
    # Size of input.
    n = np.prod(dimensions)
    # Size of output where defined.
    v = np.array(output_shape)
    m = abs(np.prod(v))
    # Replace wildcard.
    v[v == -1] = n // m
    return tuple(v)

class Reshape(Module):
    """
    Reshapes input Tensor, preserving the batch dimension.

    For example, given an input tensor with shape `[B, H, W, C, D]`:

    ```python
    B, H, W, C, D = range(1, 6)
    x = jnp.ones([B, H, W, C, D])
    ```

    The default behavior when `output_shape` is `(-1, D)` is to flatten
    all dimensions between `B` and `D`:

    ```python
    mod = elegy.nn.Reshape(output_shape=(-1, D))
    assert mod(x).shape == (B, H*W*C, D)
    ```

    You can change the number of preserved leading dimensions via
    `preserve_dims`:

    ```python
    mod = elegy.nn.Reshape(output_shape=(-1, D), preserve_dims=2)
    assert mod(x).shape == (B, H, W*C, D)
    mod = elegy.nn.Reshape(output_shape=(-1, D), preserve_dims=3)
    assert mod(x).shape == (B, H, W, C, D)
    mod = elegy.nn.Reshape(output_shape=(-1, D), preserve_dims=4)
    assert mod(x).shape == (B, H, W, C, 1, D)
    ```
    """

    def __init__(self, output_shape: types.Shape, preserve_dims: int = 1, **kwargs):
        """
        Constructs a `Reshape` module.

        Args:
            output_shape: Shape to reshape the input tensor to while preserving its
                first `preserve_dims` dimensions. When the special value -1 appears in
                `output_shape` the corresponding size is automatically inferred. Note
                that -1 can only appear once in `output_shape`.
                To flatten all non-batch dimensions use `Flatten`.
            preserve_dims: Number of leading dimensions that will not be reshaped.
            kwargs: Additional keyword arguments passed to Module.

        Raises:
            ValueError: If `preserve_dims` is not positive.
        """
        super().__init__(**kwargs)
        if preserve_dims <= 0:
            raise ValueError("Argument preserve_dims should be >= 1.")
        if output_shape.count(-1) > 1:
            raise ValueError("-1 can only occur once in `output_shape`.")

        self.output_shape = tuple(output_shape)
        self.preserve_dims = preserve_dims

    def call(self, inputs: np.ndarray) -> np.ndarray:
        """
        Arguments:
            inputs: the array to be reshaped.

        Returns:
            A reshaped array.
        """
        if inputs.ndim <= self.preserve_dims:
            return inputs

        if -1 in self.output_shape:
            reshaped_shape = _infer_shape(
                self.output_shape, inputs.shape[self.preserve_dims :]
            )
        else:
            reshaped_shape = self.output_shape
        shape = inputs.shape[: self.preserve_dims] + reshaped_shape
        return jnp.reshape(inputs, shape)

class Flatten(Reshape):
    """
    Flattens the input, preserving the batch dimension(s).

    By default, Flatten combines all dimensions except the first.
    Additional leading dimensions can be preserved by setting preserve_dims.

    ```python
    x = jnp.ones([3, 2, 4])
    flat = elegy.nn.Flatten()
    assert flat(x).shape == (3, 8)
    ```

    When the input to flatten has fewer than `preserve_dims` dimensions it is
    returned unchanged:

    ```python
    x = jnp.ones([3])
    assert flat(x).shape == (3,)
    ```
    """

    def __init__(
        self,
        preserve_dims: int = 1,
        name: tp.Optional[str] = None,
    ):
        super().__init__(output_shape=(-1,), preserve_dims=preserve_dims, name=name)


import numpy as _numpy

import cupy as _cupy
from cupy_backends.cuda.libs import cublas as _cublas
from cupy.cuda import device as _device

def gesv(a, b):
    """Solve a linear matrix equation using cusolverDn<t>getr[fs]().

    Computes the solution to a system of linear equation ``ax = b``.

    Args:
        a (cupy.ndarray): The matrix with dimension ``(M, M)``.
        b (cupy.ndarray): The matrix with dimension ``(M)`` or ``(M, K)``.

    Returns:
        cupy.ndarray:
            The matrix with dimension ``(M)`` or ``(M, K)``.

    Note: ``a`` and ``b`` will be overwritten.
    """
    from cupy_backends.cuda.libs import cusolver as _cusolver

    if a.ndim != 2:
        raise ValueError('a.ndim must be 2 (actual: {})'.format(a.ndim))
    if b.ndim not in (1, 2):
        raise ValueError('b.ndim must be 1 or 2 (actual: {})'.format(b.ndim))
    if a.shape[0] != a.shape[1]:
        raise ValueError('a must be a square matrix.')
    if a.shape[0] != b.shape[0]:
        raise ValueError('shape mismatch (a: {}, b: {}).'.
                         format(a.shape, b.shape))
    if a.dtype != b.dtype:
        raise TypeError('dtype mismatch (a: {}, b: {})'.
                        format(a.dtype, b.dtype))
    dtype = a.dtype
    if dtype == 'f':
        t = 's'
    elif dtype == 'd':
        t = 'd'
    elif dtype == 'F':
        t = 'c'
    elif dtype == 'D':
        t = 'z'
    else:
        raise TypeError('unsupported dtype (actual:{})'.format(a.dtype))
    helper = getattr(_cusolver, t + 'getrf_bufferSize')
    getrf = getattr(_cusolver, t + 'getrf')
    getrs = getattr(_cusolver, t + 'getrs')

    n = b.shape[0]
    nrhs = b.shape[1] if b.ndim == 2 else 1
    if a._f_contiguous:
        trans = _cublas.CUBLAS_OP_N
    elif a._c_contiguous:
        trans = _cublas.CUBLAS_OP_T
    else:
        raise ValueError('a must be F-contiguous or C-contiguous.')
    if not b._f_contiguous:
        raise ValueError('b must be F-contiguous.')

    handle = _device.get_cusolver_handle()
    dipiv = _cupy.empty(n, dtype=_numpy.int32)
    dinfo = _cupy.empty(1, dtype=_numpy.int32)
    lwork = helper(handle, n, n, a.data.ptr, n)
    dwork = _cupy.empty(lwork, dtype=a.dtype)
    # LU factrization (A = L * U)
    getrf(handle, n, n, a.data.ptr, n, dwork.data.ptr, dipiv.data.ptr,
          dinfo.data.ptr)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        getrf, dinfo)
    # Solves Ax = b
    getrs(handle, trans, n, nrhs, a.data.ptr, n,
          dipiv.data.ptr, b.data.ptr, n, dinfo.data.ptr)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        getrs, dinfo)

def gels(a, b):
    """Solves over/well/under-determined linear systems.

    Computes least-square solution to equation ``ax = b` by QR factorization
    using cusolverDn<t>geqrf().

    Args:
        a (cupy.ndarray): The matrix with dimension ``(M, N)``.
        b (cupy.ndarray): The matrix with dimension ``(M)`` or ``(M, K)``.

    Returns:
        cupy.ndarray:
            The matrix with dimension ``(N)`` or ``(N, K)``.
    """
    from cupy_backends.cuda.libs import cusolver as _cusolver

    if a.ndim != 2:
        raise ValueError('a.ndim must be 2 (actual: {})'.format(a.ndim))
    if b.ndim == 1:
        nrhs = 1
    elif b.ndim == 2:
        nrhs = b.shape[1]
    else:
        raise ValueError('b.ndim must be 1 or 2 (actual: {})'.format(b.ndim))
    if a.shape[0] != b.shape[0]:
        raise ValueError('shape mismatch (a: {}, b: {}).'.
                         format(a.shape, b.shape))
    if a.dtype != b.dtype:
        raise ValueError('dtype mismatch (a: {}, b: {}).'.
                         format(a.dtype, b.dtype))

    dtype = a.dtype
    if dtype == 'f':
        t = 's'
    elif dtype == 'd':
        t = 'd'
    elif dtype == 'F':
        t = 'c'
    elif dtype == 'D':
        t = 'z'
    else:
        raise ValueError('unsupported dtype (actual: {})'.format(dtype))

    geqrf_helper = getattr(_cusolver, t + 'geqrf_bufferSize')
    geqrf = getattr(_cusolver, t + 'geqrf')
    trsm = getattr(_cublas, t + 'trsm')
    if t in 'sd':
        ormqr_helper = getattr(_cusolver, t + 'ormqr_bufferSize')
        ormqr = getattr(_cusolver, t + 'ormqr')
    else:
        ormqr_helper = getattr(_cusolver, t + 'unmqr_bufferSize')
        ormqr = getattr(_cusolver, t + 'unmqr')

    no_trans = _cublas.CUBLAS_OP_N
    if dtype.char in 'fd':
        trans = _cublas.CUBLAS_OP_T
    else:
        trans = _cublas.CUBLAS_OP_C

    m, n = a.shape
    mn_min = min(m, n)
    dev_info = _cupy.empty(1, dtype=_numpy.int32)
    tau = _cupy.empty(mn_min, dtype=dtype)
    cusolver_handle = _device.get_cusolver_handle()
    cublas_handle = _device.get_cublas_handle()
    one = _numpy.array(1.0, dtype=dtype)

    if m >= n:  # over/well-determined systems
        a = a.copy(order='F')
        b = b.copy(order='F')

        # geqrf (QR decomposition, A = Q * R)
        ws_size = geqrf_helper(cusolver_handle, m, n, a.data.ptr, m)
        workspace = _cupy.empty(ws_size, dtype=dtype)
        geqrf(cusolver_handle, m, n, a.data.ptr, m, tau.data.ptr,
              workspace.data.ptr, ws_size, dev_info.data.ptr)
        _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
            geqrf, dev_info)

        # ormqr (Computes Q^T * B)
        ws_size = ormqr_helper(
            cusolver_handle, _cublas.CUBLAS_SIDE_LEFT, trans, m, nrhs, mn_min,
            a.data.ptr, m, tau.data.ptr, b.data.ptr, m)
        workspace = _cupy.empty(ws_size, dtype=dtype)
        ormqr(cusolver_handle, _cublas.CUBLAS_SIDE_LEFT, trans, m, nrhs,
              mn_min, a.data.ptr, m, tau.data.ptr, b.data.ptr, m,
              workspace.data.ptr, ws_size, dev_info.data.ptr)
        _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
            ormqr, dev_info)

        # trsm (Solves R * X = (Q^T * B))
        trsm(cublas_handle, _cublas.CUBLAS_SIDE_LEFT,
             _cublas.CUBLAS_FILL_MODE_UPPER, no_trans,
             _cublas.CUBLAS_DIAG_NON_UNIT, mn_min, nrhs,
             one.ctypes.data, a.data.ptr, m, b.data.ptr, m)

        return b[:n]

    else:  # under-determined systems
        a = a.conj().T.copy(order='F')
        bb = b
        out_shape = (n,) if b.ndim == 1 else (n, nrhs)
        b = _cupy.zeros(out_shape, dtype=dtype, order='F')
        b[:m] = bb

        # geqrf (QR decomposition, A^T = Q * R)
        ws_size = geqrf_helper(cusolver_handle, n, m, a.data.ptr, n)
        workspace = _cupy.empty(ws_size, dtype=dtype)
        geqrf(cusolver_handle, n, m, a.data.ptr, n, tau.data.ptr,
              workspace.data.ptr, ws_size, dev_info.data.ptr)
        _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
            geqrf, dev_info)

        # trsm (Solves R^T * Z = B)
        trsm(cublas_handle, _cublas.CUBLAS_SIDE_LEFT,
             _cublas.CUBLAS_FILL_MODE_UPPER, trans,
             _cublas.CUBLAS_DIAG_NON_UNIT, m, nrhs,
             one.ctypes.data, a.data.ptr, n, b.data.ptr, n)

        # ormqr (Computes Q * Z)
        ws_size = ormqr_helper(
            cusolver_handle, _cublas.CUBLAS_SIDE_LEFT, no_trans, n, nrhs,
            mn_min, a.data.ptr, n, tau.data.ptr, b.data.ptr, n)
        workspace = _cupy.empty(ws_size, dtype=dtype)
        ormqr(cusolver_handle, _cublas.CUBLAS_SIDE_LEFT, no_trans, n, nrhs,
              mn_min, a.data.ptr, n, tau.data.ptr, b.data.ptr, n,
              workspace.data.ptr, ws_size, dev_info.data.ptr)
        _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
            ormqr, dev_info)

        return b

def _batched_posv(a, b):
    from cupy_backends.cuda.libs import cusolver as _cusolver
    import cupyx.cusolver

    if not cupyx.cusolver.check_availability('potrsBatched'):
        raise RuntimeError('potrsBatched is not available')

    dtype = _numpy.promote_types(a.dtype, b.dtype)
    dtype = _numpy.promote_types(dtype, 'f')

    if dtype == 'f':
        potrfBatched = _cusolver.spotrfBatched
        potrsBatched = _cusolver.spotrsBatched
    elif dtype == 'd':
        potrfBatched = _cusolver.dpotrfBatched
        potrsBatched = _cusolver.dpotrsBatched
    elif dtype == 'F':
        potrfBatched = _cusolver.cpotrfBatched
        potrsBatched = _cusolver.cpotrsBatched
    elif dtype == 'D':
        potrfBatched = _cusolver.zpotrfBatched
        potrsBatched = _cusolver.zpotrsBatched
    else:
        msg = ('dtype must be float32, float64, complex64 or complex128'
               ' (actual: {})'.format(a.dtype))
        raise ValueError(msg)

    a = a.astype(dtype, order='C', copy=True)
    ap = _cupy._core._mat_ptrs(a)
    lda, n = a.shape[-2:]
    batch_size = int(_numpy.prod(a.shape[:-2]))

    handle = _device.get_cusolver_handle()
    uplo = _cublas.CUBLAS_FILL_MODE_LOWER
    dev_info = _cupy.empty(batch_size, dtype=_numpy.int32)

    # Cholesky factorization
    potrfBatched(handle, uplo, n, ap.data.ptr, lda, dev_info.data.ptr,
                 batch_size)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        potrfBatched, dev_info)

    b_shape = b.shape
    b = b.conj().reshape(batch_size, n, -1).astype(dtype, order='C', copy=True)
    bp = _cupy._core._mat_ptrs(b)
    ldb, nrhs = b.shape[-2:]
    dev_info = _cupy.empty(1, dtype=_numpy.int32)

    # NOTE: potrsBatched does not currently support nrhs > 1 (CUDA v10.2)
    # Solve: A[i] * X[i] = B[i]
    potrsBatched(handle, uplo, n, nrhs, ap.data.ptr, lda, bp.data.ptr, ldb,
                 dev_info.data.ptr, batch_size)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        potrsBatched, dev_info)

    # TODO: check if conj() is necessary when nrhs > 1
    return b.conj().reshape(b_shape)

def posv(a, b):
    """Solve the linear equations A x = b via Cholesky factorization of A,
    where A is a real symmetric or complex Hermitian positive-definite matrix.

    If matrix ``A`` is not positive definite, Cholesky factorization fails
    and it raises an error.

    Note: For batch input, NRHS > 1 is not currently supported.

    Args:
        a (cupy.ndarray): Array of real symmetric or complex hermitian
            matrices with dimension (..., N, N).
        b (cupy.ndarray): right-hand side (..., N) or (..., N, NRHS).
    Returns:
        x (cupy.ndarray): The solution (shape matches b).
    """
    from cupy_backends.cuda.libs import cusolver as _cusolver

    _util = _cupy.linalg._util
    _util._assert_cupy_array(a, b)
    _util._assert_stacked_2d(a)
    _util._assert_stacked_square(a)

    if a.ndim > 2:
        return _batched_posv(a, b)

    dtype = _numpy.promote_types(a.dtype, b.dtype)
    dtype = _numpy.promote_types(dtype, 'f')

    if dtype == 'f':
        potrf = _cusolver.spotrf
        potrf_bufferSize = _cusolver.spotrf_bufferSize
        potrs = _cusolver.spotrs
    elif dtype == 'd':
        potrf = _cusolver.dpotrf
        potrf_bufferSize = _cusolver.dpotrf_bufferSize
        potrs = _cusolver.dpotrs
    elif dtype == 'F':
        potrf = _cusolver.cpotrf
        potrf_bufferSize = _cusolver.cpotrf_bufferSize
        potrs = _cusolver.cpotrs
    elif dtype == 'D':
        potrf = _cusolver.zpotrf
        potrf_bufferSize = _cusolver.zpotrf_bufferSize
        potrs = _cusolver.zpotrs
    else:
        msg = ('dtype must be float32, float64, complex64 or complex128'
               ' (actual: {})'.format(a.dtype))
        raise ValueError(msg)

    a = a.astype(dtype, order='F', copy=True)
    lda, n = a.shape

    handle = _device.get_cusolver_handle()
    uplo = _cublas.CUBLAS_FILL_MODE_LOWER
    dev_info = _cupy.empty(1, dtype=_numpy.int32)

    worksize = potrf_bufferSize(handle, uplo, n, a.data.ptr, lda)
    workspace = _cupy.empty(worksize, dtype=dtype)

    # Cholesky factorization
    potrf(handle, uplo, n, a.data.ptr, lda, workspace.data.ptr,
          worksize, dev_info.data.ptr)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        potrf, dev_info)

    b_shape = b.shape
    b = b.reshape(n, -1).astype(dtype, order='F', copy=True)
    ldb, nrhs = b.shape

    # Solve: A * X = B
    potrs(handle, uplo, n, nrhs, a.data.ptr, lda, b.data.ptr, ldb,
          dev_info.data.ptr)
    _cupy.linalg._util._check_cusolver_dev_info_if_synchronization_allowed(
        potrs, dev_info)

    return _cupy.ascontiguousarray(b.reshape(b_shape))