#===================================================================#
#       Tool Name: Zaius FRU Tool                                   #
#         Version: 0.3                                              #
#         Edit by: Chris Liu 2017/08/05                             #
#===================================================================#
import sys
import datetime
import subprocess
import os
import time
import ctypes
import serial
import argparse

global VER
VER = "0.3"

global DEBUG_MODE
DEBUG_MODE = True

global FAIL_CONTINUE
FAIL_CONTINUE = True

global ROOT_DIR
ROOT_DIR = os.getcwd()

global LOG_DIR
LOG_DIR = os.path.join(ROOT_DIR, "Log")

global FRU_DIR
FRU_DIR = os.path.join(ROOT_DIR, "FRU")

global COM_PORT
COM_PORT = "COM8"

global FRU_BIN
FRU_BIN = bytearray(128)

global PN
PN = "42001726-04"

global SN
SN = "ZMBFTJ163500004"

global MAC
MAC = "18:66:DA:E4:E1:B7"

FONT_NONE   = 0
FONT_WHITE  = 1
FONT_RED    = 2
FONT_GREEN  = 3
FONT_YELLOW = 4
#===============================================================================
def Banner(msg):
	line_0 = "#" + "="*78 + "#"
	line_1 = "#" + " "*78 + "#"
	tmp_str = "#" + msg.center(78) + "#"

	if(sys.platform == "win32"):
		ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x01 | 0x02 | 0x08)
		print("")
		print(line_0)
		print(line_1)
		print(tmp_str)
		print(line_1)
		print(line_0)
		print("")
		ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x01 | 0x02 | 0x04)
	else:
		print("")
		print("\033[35;1m%s\033[0m"%(line_0))
		print("\033[35;1m%s\033[0m"%(line_1))
		print("\033[35;1m%s\033[0m"%(tmp_str))
		print("\033[35;1m%s\033[0m"%(line_1))
		print("\033[35;1m%s\033[0m"%(line_0))
		print("")
#===============================================================================
def Log(msg, color = FONT_WHITE):
	tmp = "[%s] %s\n"%(datetime.datetime.strftime(datetime.datetime.now(), "%y/%m/%d %H:%M:%S"), msg)

	try:
		f = open(os.path.join(LOG_DIR, LOG_FILE), "a")
		f.write(tmp)
		f.close()
	except:
		print("Logging Error!!")
		return

	tmp = tmp[:-1]

	if(color == FONT_RED):
		if(sys.platform == "win32"):
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x04 | 0x08)
			print(tmp)
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x01 | 0x02 | 0x04)
		else:
			print("\033[31;1m%s\033[0m"%(tmp))
	elif(color == FONT_GREEN):
		if(sys.platform == "win32"):
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x02 | 0x08)
			print(tmp)
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x01 | 0x02 | 0x04)
		else:
			print("\033[32;1m%s\033[0m"%(tmp))
	elif(color == FONT_YELLOW):
		if(sys.platform == "win32"):
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x02 | 0x04 | 0x08)
			print(tmp)
			ctypes.windll.kernel32.SetConsoleTextAttribute(ctypes.windll.kernel32.GetStdHandle(-11), 0x01 | 0x02 | 0x04)
		else:
			print("\033[33;1m%s\033[0m"%(tmp))
	elif(color == FONT_NONE):
		pass
	else:
		try:
			print(tmp)
		except:
			print("Logging Error!!")
#===============================================================================
def Find_Offset():
	'''Find FRU Offset'''

	global DEBUG_MODE
	global FRU_BIN
	global Internal_Use_Area_Begin
	global Internal_Use_Area_End
	global Internal_Use_Area_Checksum_Index
	global Board_Info_Area_Begin
	global Board_Info_Area_End
	global Mfg_Date_Time_Begin
	global Mfg_Date_Time_End
	global Board_Manufacturer_Type
	global Board_Manufacturer_Length
	global Board_Manufacturer_Begin
	global Board_Manufacturer_End
	global Board_Product_Name_Type
	global Board_Product_Name_Length
	global Board_Product_Name_Begin
	global Board_Product_Name_End
	global Board_Serial_Number_Type
	global Board_Serial_Number_Length
	global Board_Serial_Number_Begin
	global Board_Serial_Number_End
	global Board_Part_Number_Type
	global Board_Part_Number_Length
	global Board_Part_Number_Begin
	global Board_Part_Number_End
	global Board_Area_Checksum_Index

	Internal_Use_Area_Begin = FRU_BIN[0x01] * 8
	Internal_Use_Area_End = Internal_Use_Area_Begin + FRU_BIN[Internal_Use_Area_Begin + 0x02] - 1
	Internal_Use_Area_Checksum_Index = Internal_Use_Area_End

	Board_Info_Area_Begin = FRU_BIN[0x03] * 8
	Board_Info_Area_End = Board_Info_Area_Begin + FRU_BIN[Board_Info_Area_Begin + 1] * 8 - 1

	Mfg_Date_Time_Begin = Board_Info_Area_Begin + 3
	Mfg_Date_Time_End = Board_Info_Area_Begin + 5

	Board_Manufacturer_Type = Decode_Type(FRU_BIN[Mfg_Date_Time_End + 1])
	Board_Manufacturer_Length = FRU_BIN[Mfg_Date_Time_End + 1] & 0b00111111
	Board_Manufacturer_Begin = Mfg_Date_Time_End + 2
	Board_Manufacturer_End = Board_Manufacturer_Begin + Board_Manufacturer_Length - 1

	Board_Product_Name_Type = Decode_Type(FRU_BIN[Board_Manufacturer_End + 1])
	Board_Product_Name_Length = FRU_BIN[Board_Manufacturer_End + 1] & 0b00111111
	Board_Product_Name_Begin = Board_Manufacturer_End + 2
	Board_Product_Name_End = Board_Product_Name_Begin + Board_Product_Name_Length - 1

	Board_Serial_Number_Type = Decode_Type(FRU_BIN[Board_Product_Name_End + 1])
	Board_Serial_Number_Length = FRU_BIN[Board_Product_Name_End + 1] & 0b00111111
	Board_Serial_Number_Begin = Board_Product_Name_End + 2
	Board_Serial_Number_End = Board_Serial_Number_Begin + Board_Serial_Number_Length - 1

	Board_Part_Number_Type = Decode_Type(FRU_BIN[Board_Serial_Number_End + 1])
	Board_Part_Number_Length = FRU_BIN[Board_Serial_Number_End + 1] & 0b00111111
	Board_Part_Number_Begin = Board_Serial_Number_End + 2
	Board_Part_Number_End = Board_Part_Number_Begin + Board_Part_Number_Length - 1

	Board_Area_Checksum_Index = Board_Info_Area_End

	if(DEBUG_MODE):
		Log("Checksum ==>", FONT_YELLOW)
		Log("Board_Area_Checksum(0x%02X) = 0x%02X (0x%02X ~ 0x%02X)"%(Board_Area_Checksum_Index, FRU_BIN[Board_Area_Checksum_Index], Board_Info_Area_Begin, Board_Info_Area_End), FONT_WHITE)
		Log("Internal_Use_Area_Checksum(0x%02X) = 0x%02X (0x%02X ~ 0x%02X)"%(Internal_Use_Area_Checksum_Index, FRU_BIN[Internal_Use_Area_Checksum_Index], Internal_Use_Area_Begin, Internal_Use_Area_End), FONT_WHITE)

		Log("Internal Use Area ==>", FONT_YELLOW)
		Log("Internal_Use_Area_Begin:0x%02X"%Internal_Use_Area_Begin, FONT_WHITE)
		Log("Internal_Use_Area_End:0x%02X"%Internal_Use_Area_End, FONT_WHITE)

		Log("Board Info Area ==>", FONT_YELLOW)
		Log("Board_Info_Area_Begin:0x%02X"%Board_Info_Area_Begin, FONT_WHITE)
		Log("Board_Info_Area_End:0x%02X"%Board_Info_Area_End, FONT_WHITE)
		Log("Mfg_Date_Time_Begin:0x%02X"%Mfg_Date_Time_Begin, FONT_WHITE)
		Log("Mfg_Date_Time_End:0x%02X"%Mfg_Date_Time_End, FONT_WHITE)
		Log("Board_Manufacturer_Type:%s"%Board_Manufacturer_Type, FONT_WHITE)
		Log("Board_Manufacturer_Length:%d"%Board_Manufacturer_Length, FONT_WHITE)
		Log("Board_Manufacturer_Begin:0x%02X"%Board_Manufacturer_Begin, FONT_WHITE)
		Log("Board_Manufacturer_End:0x%02X"%Board_Manufacturer_End, FONT_WHITE)
		Log("Board_Product_Name_Type:%s"%Board_Product_Name_Type, FONT_WHITE)
		Log("Board_Product_Name_Length:%d"%Board_Product_Name_Length, FONT_WHITE)
		Log("Board_Product_Name_Begin:0x%02X"%Board_Product_Name_Begin, FONT_WHITE)
		Log("Board_Product_Name_End:0x%02X"%Board_Product_Name_End, FONT_WHITE)
		Log("Board_Serial_Number_Type:%s"%Board_Serial_Number_Type, FONT_WHITE)
		Log("Board_Serial_Number_Length:%d"%Board_Serial_Number_Length, FONT_WHITE)
		Log("Board_Serial_Number_Begin:0x%02X"%Board_Serial_Number_Begin, FONT_WHITE)
		Log("Board_Serial_Number_End:0x%02X"%Board_Serial_Number_End, FONT_WHITE)
		Log("Board_Part_Number_Type:%s"%Board_Part_Number_Type, FONT_WHITE)
		Log("Board_Part_Number_Length:%d"%Board_Part_Number_Length, FONT_WHITE)
		Log("Board_Part_Number_Begin:0x%02X"%Board_Part_Number_Begin, FONT_WHITE)
		Log("Board_Part_Number_End:0x%02X"%Board_Part_Number_End, FONT_WHITE)

		Manufacture_Name = FRU_BIN[Board_Manufacturer_Begin:Board_Manufacturer_End + 1].decode(Board_Manufacturer_Type)
		Product_Name = FRU_BIN[Board_Product_Name_Begin:Board_Product_Name_End + 1].decode(Board_Product_Name_Type)
		Serial_Number = FRU_BIN[Board_Serial_Number_Begin:Board_Serial_Number_End + 1].decode(Board_Serial_Number_Type)
		Part_Number = FRU_BIN[Board_Part_Number_Begin:Board_Part_Number_End + 1].decode(Board_Part_Number_Type)

		Log("Board Information ==>", FONT_YELLOW)
		Log("Manufacture Name: %s"%(Manufacture_Name), FONT_WHITE)
		Log("Product Name: %s"%(Product_Name), FONT_WHITE)
		Log("Serial Number: %s"%(Serial_Number), FONT_WHITE)
		Log("Part Number: %s"%(Part_Number), FONT_WHITE)

	return True
#===============================================================================
def Decode_Type(x):
	tmp = (x & 0b11000000) >> 6
	if(tmp == 0b00):
		return "Binary"
	elif(tmp == 0b01):
		return "BCD Plus"
	elif(tmp == 0b10):
		return "6BIT_ASCII"
	elif(tmp == 0b11):
		return "ASCII"
#===============================================================================
def Program_MFG_Time():
	'''Program Manufacture Date/Time'''

	global FRU_BIN
	global Mfg_Date_Time_Begin

	Mfg_Date_Time_B0 = Mfg_Date_Time_Begin
	Mfg_Date_Time_B1 = Mfg_Date_Time_Begin + 1
	Mfg_Date_Time_B2 = Mfg_Date_Time_Begin + 2

	Base_Time = datetime.datetime(1996, 1, 1, 0, 0)
	System_Time = datetime.datetime.now()
	Delta_Minutes = int((System_Time - Base_Time).total_seconds()/60)

	FRU_BIN[Mfg_Date_Time_B0] = (Delta_Minutes & 0x0000FF)
	FRU_BIN[Mfg_Date_Time_B1] = (Delta_Minutes & 0x00FF00) >> 8
	FRU_BIN[Mfg_Date_Time_B2] = (Delta_Minutes & 0xFF0000) >> 16

	if(DEBUG_MODE):
		Log("Mfg_Date_Time_B0 (0x%02X) = 0x%02X"%(Mfg_Date_Time_B0, FRU_BIN[Mfg_Date_Time_B0]), FONT_YELLOW)
		Log("Mfg_Date_Time_B1 (0x%02X) = 0x%02X"%(Mfg_Date_Time_B1, FRU_BIN[Mfg_Date_Time_B1]), FONT_YELLOW)
		Log("Mfg_Date_Time_B2 (0x%02X) = 0x%02X"%(Mfg_Date_Time_B2, FRU_BIN[Mfg_Date_Time_B2]), FONT_YELLOW)

	return True
#===============================================================================
def Program_PN():
	'''Program Board Part Number'''

	global PN
	global Board_Part_Number_Type
	global Board_Part_Number_Length
	global Board_Part_Number_Begin
	global Board_Part_Number_End

	FRU_BIN[Board_Part_Number_Begin:Board_Part_Number_End + 1] = bytes(PN.ljust(Board_Part_Number_Length, "\0"), Board_Part_Number_Type)

	if(DEBUG_MODE):
		Log("Board_Part_Number (0x%02X ~ 0x%02X): %s"%(Board_Part_Number_Begin, Board_Part_Number_End, FRU_BIN[Board_Part_Number_Begin:Board_Part_Number_End + 1].decode(Board_Part_Number_Type)), FONT_YELLOW)

	return True
#===============================================================================
def Program_SN():
	'''Program Board Serial Number'''

	global SN
	global Board_Serial_Number_Type
	global Board_Serial_Number_Length
	global Board_Serial_Number_Begin
	global Board_Serial_Number_End

	FRU_BIN[Board_Serial_Number_Begin:Board_Serial_Number_End + 1] = bytes(SN.ljust(Board_Serial_Number_Length, "\0"), Board_Serial_Number_Type)

	if(DEBUG_MODE):
		Log("Board_Serial_Number (0x%02X ~ 0x%02X): %s"%(Board_Serial_Number_Begin, Board_Serial_Number_End, FRU_BIN[Board_Serial_Number_Begin:Board_Serial_Number_End + 1].decode(Board_Serial_Number_Type)), FONT_YELLOW)

	return True
#===============================================================================
def Program_MAC():
	'''Program MAC Address'''

	global MAC
	global Internal_Use_Area_Begin

	mac_list = MAC.split(":")
	for i in range(len(mac_list)):
		FRU_BIN[Internal_Use_Area_Begin + 3 + i] = int(mac_list[i], 16)

	if(DEBUG_MODE):
		Log("MAC Address (0x%02X ~ 0x%02X):"%(Internal_Use_Area_Begin + 3, Internal_Use_Area_Begin + 8), FONT_YELLOW)
		Log("Byte1: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 3]), FONT_YELLOW)
		Log("Byte2: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 4]), FONT_YELLOW)
		Log("Byte3: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 5]), FONT_YELLOW)
		Log("Byte4: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 6]), FONT_YELLOW)
		Log("Byte5: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 7]), FONT_YELLOW)
		Log("Byte6: 0x%02X"%(FRU_BIN[Internal_Use_Area_Begin + 8]), FONT_YELLOW)

	return True
#===============================================================================
def Program_Checksum():
	'''Program Board Area Checksum & Internal Use Area Checksum'''

	global FRU_BIN
	global Internal_Use_Area_Checksum_Index
	global Board_Area_Checksum_Index
	global Internal_Use_Area_Begin
	global Internal_Use_Area_End
	global Board_Info_Area_Begin
	global Board_Info_Area_End

	Internal_Use_Area_Checksum = (~sum([FRU_BIN[i] for i in range(Internal_Use_Area_Begin, Internal_Use_Area_End + 1) if(i != Internal_Use_Area_Checksum_Index)]) + 1)%256
	FRU_BIN[Internal_Use_Area_Checksum_Index] = Internal_Use_Area_Checksum

	Board_Area_Checksum = (~sum([FRU_BIN[i] for i in range(Board_Info_Area_Begin, Board_Info_Area_End + 1) if(i != Board_Area_Checksum_Index)]) + 1)%256
	FRU_BIN[Board_Area_Checksum_Index] = Board_Area_Checksum

	if(DEBUG_MODE):
		Log("Internal_Use_Area_Checksum (0x%02X) = 0x%02X"%(Internal_Use_Area_Checksum_Index, Internal_Use_Area_Checksum), FONT_YELLOW)
		Log("Board_Area_Checksum (0x%02X) = 0x%02X"%(Board_Area_Checksum_Index, Board_Area_Checksum), FONT_YELLOW)

	return True
#===============================================================================
def Input_CMD_BMC(cmd, delay = 1):
	Log("Input BMC Command: %s"%(cmd), FONT_WHITE)

	try:
		s = serial.Serial(port = COM_PORT, baudrate = 115200, timeout = delay)
	except:
		Log("Open BMC COM Port (%s) Fail!!"%(COM_PORT), FONT_RED)
		sys.exit(-1)

	error_list = [
		"ucd9000 0-0064",
		"ftgmac100 1e660000.ethernet",
		"ucd9000:",
		"random: nonblocking pool is initialized",
	]

	msg = []
	cmd = "%s\r"%(cmd)

	for retry in range(5):
		flag_retry = True

		s.write(cmd.encode(errors = "ignore"))
		ret = s.readlines()
		for i in range(len(ret)):
			ret[i] = ret[i].decode(errors = "ignore").strip()

			if("ttyS4: 1 input overrun(s)" in ret[i]):
				flag_retry = False

		if(flag_retry):
			break

	#Remove error message...
	for i in range(len(ret)):
		flag = True
		for error in error_list:
			if(error in ret[i]):
				flag = False
				break

		if(flag):
			msg.append(ret[i])
			if(DEBUG_MODE):
				Log("ret[%02d] %s"%(len(msg)-1, ret[i]), FONT_WHITE)

	s.close()

	return msg
#===============================================================================
def Enter_BMC_Linux():
	'''Enter BMC Linux System'''

	print("Waiting for Enter BMC Linux System!!")
	timeout_5s = 0
	for i in range(timeout_5s):
		print(("%s secs...")%(5*(timeout_5s - i)))
		time.sleep(5)

	for i in range(10):
		ret = Input_CMD_BMC("pwd")
		if("/home/root" in ret[1]):
			Log("Enter_BMC_Linux", FONT_YELLOW)
			break

		if(i <= 9):
			time.sleep(10)
			print(("%s secs...")%(10*(10 - i)))
		else:
			Log("Enter_BMC_Linux Fail (TIMEOUT!!)", FONT_RED)
			return False

	Log("Enter_BMC_Linux Pass", FONT_GREEN)
	return True
#===============================================================================
def Save_FRU():
	'''Save FRU Data to BIN File'''

	global FRU_BIN
	global SN

	flag_1 = False
	flag_2 = False

	#Save FRU File
	filename = "FRU_%s_%s.bin"%(SN, datetime.datetime.strftime(datetime.datetime.now(), "%y%m%d-%H%M%S"))
	f = open(os.path.join(FRU_DIR, filename), "wb")
	f.write(FRU_BIN)
	f.close()
	Log("Save FRU to %s"%(filename), FONT_YELLOW)

	#Transfer FRU from Host PC to BMC via TFTP
	Input_CMD_BMC("ifconfig eth1 192.168.0.101;tftp -g -r /FRU/%s 192.168.0.120"%(filename))

	#Write FRU via dd
	ret = Input_CMD_BMC("dd if=%s of=/sys/bus/i2c/devices/0-0050/eeprom bs=1 count=128 skip=0"%(filename))
	for i in ret:
		if("128+0 records out" in i):
			flag_1 = True

	#Read FRU via hexdump
	ret = Input_CMD_BMC("hexdump -C -v -n 128 /sys/bus/i2c/devices/0-0050/eeprom")
	eeprom_bin = Hexdump_to_BIN(ret[1:])

	if(eeprom_bin == FRU_BIN):
		flag_2 = True

	if(flag_1 and flag_2):
		Log("Save_FRU Pass", FONT_GREEN)
		return True
	else:
		if(flag_1 == False):
			Log("Save_FRU Fail (01. Write FRU)", FONT_RED)
		if(flag_2 == False):
			Log("Save_FRU Fail (02. Compare FRU)", FONT_RED)
		return False
#===============================================================================
def Hexdump_to_BIN(msg):
	'''Hexdump Message to BIN'''

	bin_array = bytearray()

	for line in range(len(msg)):
		try:
			msg_list = msg[line].split("|", maxsplit = 1)[0].split()[1:]
			for index in range(len(msg_list)):
				bin_array.append(int(msg_list[index], 16))
		except:
			return 0

	return bin_array
#===============================================================================
def main():
	global LOG_FILE
	global FRU_BIN
	global PN
	global SN
	global MAC
	global COM_PORT

	Banner("Zaius FRU Tool, By Foxconn CESBG-EPDI-TE, Version: %s"%(VER))

	opts = argparse.ArgumentParser(description = "Zaius FRU Tool, By Foxconn CESBG-EPDI-TE, Version: %s"%(VER))
	opts.add_argument("-v", "--version", action = "version", version = VER, help = "Show Tool Version")
	opts.add_argument("-p", "--pn", required = False, default = "42001726-04", help = "Part Number")
	opts.add_argument("-s", "--sn", required = False, default = "ZMBFTJ163500004", help = "Serial Number")
	opts.add_argument("-m", "--mac", required = False, default = "00:1F:E2:FD:EB:4A", help = "MAC Address")
	opts.add_argument("-c", "--com", required = False, default = "COM8", help = "COM Port")

	args = opts.parse_args()

	SN = args.sn
	PN = args.pn
	MAC = args.mac
	COM_PORT = args.com

	LOG_FILE = "Log_%s_%s_%s.txt"%(os.path.basename(__file__).split(".")[0], SN, datetime.datetime.strftime(datetime.datetime.now(), "%y%m%d-%H%M%S"))

	if(DEBUG_MODE):
		Log("DEBUG_MODE", FONT_YELLOW)
	if(FAIL_CONTINUE):
		Log("FAIL_CONTINUE", FONT_YELLOW)

	Log("Part Number = %s"%(PN), FONT_YELLOW)
	Log("Serial Number = %s"%(SN), FONT_YELLOW)
	Log("MAC Address = %s"%(MAC), FONT_YELLOW)
	Log("COM Port = %s"%(COM_PORT), FONT_YELLOW)

	if(PN not in ["42001726-04", "1025129-01", "1025130-01"]):
		Log("Wrong Part Number!! (%s)"%(PN), FONT_RED)
		sys.exit(-1)

	if(len(SN) != 15 or SN[0:3] != "ZMB" or SN[3:6] != "FTJ"):
		Log("Wrong Serial Number!! (%s)"%(SN), FONT_RED)
		sys.exit(-1)

	if(len(MAC) != 17 or MAC[:8] != "00:1F:E2"):
		Log("Wrong MAC Address!! (%s)"%(MAC), FONT_RED)
		sys.exit(-1)

	if(SN[0:3] == "ZMB"):
		Log("Board: Zaius MB", FONT_YELLOW)
		fru_file = "Zaius_MB_FRU.BIN"

	try:
		f = open(fru_file, "rb")
		FRU_BIN = bytearray(f.read()[:128])
		f.close()
	except:
		Log("Open FRU File Fail (%s)"%(fru_file), FONT_RED)
		sys.exit(-1)

	test_sequence = [
		Enter_BMC_Linux,
		Find_Offset,
		Program_MFG_Time,
		Program_SN,
		Program_PN,
		Program_MAC,
		Program_Checksum,
		Save_FRU,
	]

	test_result = True
	result_msg = []

	test_start = datetime.datetime.now()
	Log("Test Start...", FONT_YELLOW)
	for test_item in test_sequence:
		Log("="*58, FONT_NONE)
		Banner(test_item.__doc__)
		Log("Test Item: %s (%s)"%(test_item.__doc__, test_item.__name__), FONT_YELLOW)
		#time.sleep(1)
		if(test_item() == False):
			test_result = False
			result_msg.append((test_item.__name__, False))
			if(FAIL_CONTINUE):
				input("%s Fail!! Press ENTER to Continue..."%(test_item.__doc__))
			else:
				break
		else:
			result_msg.append((test_item.__name__, True))
		#time.sleep(1)
		Log("="*58, FONT_NONE)
	Log("Test End...", FONT_YELLOW)
	test_end = datetime.datetime.now()

	print("")
	Log("Test Start: %s"%(str(test_start)), FONT_YELLOW)
	Log("Test End:   %s"%(str(test_end)), FONT_YELLOW)
	Log("Test Time:  %s"%(str(test_end - test_start)), FONT_YELLOW)
	print("")
	for (item_name, result) in result_msg:
		if(result):
			msg = item_name.ljust(52, "-") + "[PASS]"
			Log(msg, FONT_GREEN)
		else:
			msg = item_name.ljust(52, "-") + "[FAIL]"
			Log(msg, FONT_RED)
	print("")

	if(test_result):
		Log("Test Pass!!", FONT_GREEN)
		sys.exit(0)
	else:
		Log("Test Fail!!", FONT_RED)
		sys.exit(-1)
#===============================================================================
if(__name__ == "__main__"):
	try:
		main()
	except Exception as e:
		print("ERROR: %s"%(str(e)))
		sys.exit(-1)
	sys.exit(0)


from dotenv import load_dotenv
from datetime import datetime
import json
import sys
import os
load_dotenv()
sys.path.append(os.path.abspath(os.getenv("system_path")))
from lib import rabbit_mq, mongo

try:
    connection = rabbit_mq.create_connection()
    channel = connection.channel()

    channel.queue_declare(queue='mongoB_jobs_save', durable=True)

    client = mongo.create_connection()

    def callback(ch, method, properties, body):
        db = mongo.set_db(client)
        coll = mongo.set_collection(db, "QuoraInteriorDesigner")
        # print(" [x] Received %r" % json.loads(body))
        payload = json.loads(body)
        print("\n\n",payload)
        # print(type(payload))

        #coll.insert_many([payload])
        coll.find_one_and_update({'job_link': payload['job_link']}, {'$set': payload, "$setOnInsert": { "date": datetime.utcnow()}}, upsert= True)

    # channel.basic_consume(jobs, 'mongo_jobs_save',  no_ack=False)
    channel.basic_consume(
       queue='mongoB_jobs_save', on_message_callback=callback, auto_ack=False)

    print(' [*] Waiting for messages. To exit press CTRL+C')
    channel.start_consuming()

except Exception as e:
    error = {
        "status": "Error occured while saving jobs.....",
        "errorMsg": e
    }
    print("Error: ",error)


'''

'''
import os, subprocess, time
from sched import scheduler
import RunVasp
from comMethods import finishCheck, convergeCheck, getNSW, joinLists, flat

class JobManager:
    """ This class is responsible for scheduling and watching the VASP calculations.  It starts
        the low-precision relaxation, waits for all structures to either timeout or complete,
        reports how many structures converged and how many did not, then starts the
        normal-precision relaxation for all structures that converged.  It repeats this process
        for the normal-precision relaxation and the DOS runs.  An instance of the RunVasp class
        is used to prepare all the directories, retrieve the needed files, and submit the jobs to
        the supercomputer. """

    def __init__(self, user, atomSets,siteConcList,ediffLowScale,jobstime):
        """ CONSTRUCTOR """
        self.jobstime = jobstime
        self.atomSets = atomSets
        self.vaspRunner = RunVasp.RunVasp(user, self.atomSets, siteConcList,ediffLowScale,jobstime)

    def reportFinished(self, jobIds):
        """ Returns true if all the jobs with IDs given in 'jobIds' have finished VASP
            calculations, false otherwise. """
        devnull = open(os.devnull, 'w')
        for jobid in jobIds:
            proc = subprocess.Popen(['squeue', '--job', jobid], stdout=subprocess.PIPE, stderr=devnull)
            output = proc.communicate()[0].split()
            if len(output) != 8 and len(output) != 0:   # It will list either all the headers or
                return False                            # sometimes an error and outputs nothing.
                                                        # The error in this case is an "invalid
                                                        # job id" error because the job is no
        return True                                     # longer on the supercomputer.

    def reportLowStats(self, vstructsToStart):
        """ Displays the percentage of structures that converged during the low-precision VASP
            calculations.  Also displays the number of structures that converged and the number
            of structures that didn't. """
        for i in xrange(len(self.atomSets)):
            subprocess.call(['echo','\nFor atom ' + self.atomSets[i] + ':'])
            total = 0
            converged = 0
            atomSetDir = os.getcwd() + '/' + self.atomSets[i]
            if os.path.isdir(atomSetDir):
                for structure in vstructsToStart[i]:
                    structureDir = atomSetDir + '/' + structure
                    if os.path.isdir(structureDir):
                        if finishCheck(structureDir) and convergeCheck(structureDir, getNSW(structureDir)):
                            total += 1
                            converged += 1
                        else:
                            total += 1

            if total == 0:
                subprocess.call(['echo','\tNo structures were submitted for ' + self.atomSets[i]])
            else:
                percent = float(float(converged) / float(total)) * 100.0
                notConverged = total - converged
                subprocess.call(['echo','\t%5.2f %% of the structures converged during low-precision relaxation.' % percent])
                subprocess.call(['echo','\t%d structures converged.' % converged])
                subprocess.call(['echo','\t%d structures did not converge.' % notConverged])

    def reportDOSStats(self, vstructsToStart):
        """ Reports the percentage of structures that converged during the Density of States VASP
            calculations.  Also reports the number of structures that converged and the number
            that didn't. """
        for i in xrange(len(self.atomSets)):
            subprocess.call(['echo','\nFor atom ' + self.atomSets[i] + ':'])
            total = 0
            converged = 0
            atomSetDir = os.getcwd() + '/' + self.atomSets[i]
            if os.path.isdir(atomSetDir):
                for struct in vstructsToStart[i]:
                    structDir = atomSetDir + '/' + struct
                    if os.path.isdir(structDir):
                        dosDir = structDir + '/DOS'
                        if os.path.isdir(dosDir):
			    #if os.path.isdir(dosDir + '/vasp.x'):
                            #    subprocess.call(['rm',dosDir + '/vasp.x'])
                            if finishCheck(dosDir) and convergeCheck(dosDir, 2):
                                total += 1
                                converged += 1
                            else:
                                total += 1
            if total == 0:
                subprocess.call(['echo','\tNo structures were submitted for ' + self.atomSets[i]])
            else:
                percent = float(float(converged) / float(total)) * 100.0
                notConverged = total - converged
                subprocess.call(['echo','\t%5.2f %% of the structures converged during the DOS run.' % percent])
                subprocess.call(['echo','\t%d structures converged.' % converged])
                subprocess.call(['echo','\t%d structures did not converge.' % notConverged])

    def reportNormalStats(self, vstructsToStart):
        """ Reports the percentage of structures that converged during normal-precision VASP
            calculations.  Also reports the number of structures that converged and the number
            that didn't. """
        for i in xrange(len(self.atomSets)):
            subprocess.call(['echo','\nFor atom ' + self.atomSets[i] + ':'])
            total = 0
            converged = 0
            atomSetDir = os.getcwd() + '/' + self.atomSets[i]
            if os.path.isdir(atomSetDir):
                for struct in vstructsToStart[i]:
                    structDir = atomSetDir + '/' + struct
                    if os.path.isdir(structDir):
                        normalDir = structDir + '/normal'
                        if os.path.isdir(normalDir):
                            if finishCheck(normalDir) and convergeCheck(normalDir, getNSW(normalDir)):
                                total += 1
                                converged += 1
                            else:
                                total += 1
            if total == 0:
                subprocess.call(['echo','\tNo structures were submitted for ' + self.atomSets[i]])
            else:
                percent = float(float(converged) / float(total)) * 100.0
                notConverged = total - converged
                subprocess.call(['echo','\t%5.2f %% of the structures converged during normal-precision relaxation.' % percent])
                subprocess.call(['echo','\t%d structures converged.' % converged])
                subprocess.call(['echo','\t%d structures did not converge.' % notConverged])

    def runDOSJobs(self, vstructsToStart, vstructsRestart):
        """ Starts the Density of States VASP calculations for all of the structures in
            'vstructsToStart' and waits for all of the jobs to finish. It checks on the jobs every ten
            minutes. """
        subprocess.call(['echo','\nStarting DOS run. . .\n'])
        vstructsToRun = joinLists([vstructsRestart,vstructsToStart])
        self.vaspRunner.runJobs('DOS', vstructsToStart, vstructsToRun)
        s = scheduler(time.time, time.sleep)
        finished = False
        start_time = time.time()
        event_time = start_time
        while not finished:
            event_time += 60 #seconds
            s.enterabs(event_time, 1, self.reportFinished, ([self.vaspRunner.getCurrentJobIds()]))
            s.run()
            finished = self.reportFinished(self.vaspRunner.getCurrentJobIds())
        self.reportDOSStats(vstructsToRun)

    def runHexMono(self):
        subprocess.call(['echo','\nPreparing and running hexagonal metal monolayers directories\n']) #bch
        jobs = self.vaspRunner.makeRunHexMono()
        return jobs

    def runLowJobs(self, vstructsToStart,vstructsRestart):
        """ Starts the low-precision VASP calculations for all of the structures in 'vstructsToStart'
            and waits for all of the jobs to finish. It checks on the jobs every ten minutes. """
        subprocess.call(['echo','\nPreparing directories for VASP. . .\n'])
        vstructsToRun = joinLists([vstructsRestart,vstructsToStart])
        self.vaspRunner.prepareForVasp(vstructsToStart)
        self.vaspRunner.prepareRestarts(vstructsRestart) #contcar->poscar, incar and job file

        s = scheduler(time.time, time.sleep)

        if len(flat(vstructsToRun)) != 0:
            subprocess.call(['echo','\nStarting low-precision ionic relaxation. . .\n'])
            self.vaspRunner.runJobs('low',vstructsToStart,vstructsToRun)

            finished = False
            start_time = time.time()
            event_time = start_time
            while not finished:
                event_time += 5
                s.enterabs(event_time, 1, self.reportFinished, ([self.vaspRunner.getCurrentJobIds()]))
                s.run()
                finished = self.reportFinished(self.vaspRunner.getCurrentJobIds())

        #self.reportLowStats(vstructsToRun)

    def runNormalJobs(self, vstructsToStart,vstructsRestart):
        """ Starts the normal-precision VASP calculations for all of the structures in 'vstructsToStart'
            and waits for all of the jobs to finish. It checks on the jobs every ten minutes. """
        subprocess.call(['echo','\nStarting normal-precision ionic relaxation. . .\n'])
        vstructsToRun = joinLists([vstructsRestart,vstructsToStart])
        self.vaspRunner.runJobs('accurate', vstructsToStart,vstructsToRun)

        s = scheduler(time.time, time.sleep)

        finished = False
        start_time = time.time()
        event_time = start_time
        while not finished or len(self.vaspRunner.getCurrentJobIds())==0:
            event_time += 60
            s.enterabs(event_time, 1,self.reportFinished(self.vaspRunner.getCurrentJobIds())) #bch: was self.reportFinished, ([self.vaspRunner.getCurrentJobIds()]))
            s.run()
            finished = self.reportFinished(self.vaspRunner.getCurrentJobIds())

        self.reportNormalStats(vstructsToRun)

    def runSingleAtoms(self):
        subprocess.call(['echo','\nPreparing and running single atom directories\n']) #bch
        jobs = self.vaspRunner.makeRunSingleDirectories()
        return jobs


import unittest
import collections
import wiktextract
from wiktextract.wiktionary import WiktionaryConfig, parse_wiktionary
from wikitextprocessor import Wtp

class LongTests(unittest.TestCase):

    def test_long(self):
        # Just parse through the data and make sure that we find some words
        # This takes about 0.5 minutes.

        langs = collections.defaultdict(int)
        words = collections.defaultdict(int)
        poses = collections.defaultdict(int)
        num_transl = 0
        num_redirects = 0

        def word_cb(data):
            nonlocal num_transl
            nonlocal num_redirects
            if "redirect" in data:
                assert isinstance(data["redirect"], str)
                word = data["title"]
                words[word] += 1
                num_redirects += 1
                return
            word = data["word"]
            assert word
            words[word] += 1
            lang = data["lang"]
            pos = data["pos"]
            assert word and lang and pos
            langs[lang] += 1
            poses[pos] += 1
            if data.get("translations"):
                num_transl += 1
            for sense in data.get("senses", ()):
                if sense.get("translations"):
                    num_transl += 1

        path = "tests/test-pages-articles.xml.bz2"
        print("Parsing test data")
        ctx = Wtp()
        config = WiktionaryConfig(capture_languages=["English", "Finnish",
                                                     "Spanish", "German",
                                                     "Chinese", "Japanese",
                                                     "Italian", "Portuguese",
                                                     "Translingual"],
                                  capture_translations=True,
                                  capture_pronunciation=True,
                                  capture_linkages=True,
                                  capture_compounds=True,
                                  capture_redirects=True)
        parse_wiktionary(ctx, path, config, word_cb, None)
        print("Test data parsing complete")
        assert num_redirects > 0
        assert len(words) > 100
        assert all(x < 50 for x in words.values())
        assert langs["English"] > 0
        assert langs["Finnish"] > 0
        assert langs["Translingual"] > 0
        assert len(langs.keys()) == 9
        assert len(poses.keys()) <= len(wiktextract.PARTS_OF_SPEECH)
        assert sum(poses.values()) == sum(langs.values())
        assert sum(words.values()) == sum(poses.values()) + num_redirects
        assert num_transl > 0


import unittest
from pathlib import Path
from modules.parsers.spycloud import SpyCloudParser
from modules.collectors.spycloud.collector import SpyCloudCollector

class SpyCloudParserTest(unittest.TestCase):
    def test_parse(self):
        path = 'tests/fixtures/data_anonymized_spycloud.csv'
        tc = SpyCloudCollector()
        statuscode, df = tc.collect(Path(path))
        assert statuscode == "OK"
        tp = SpyCloudParser()
        idf = tp.parse(df)
        assert idf
        # print([ i for i in idf ])
        for i in idf:
            if "error_msg" in i.dict() and i.error_msg:
                print("error_msg: %s" % i.error_msg)
                print("orig_line: %s" % i.original_line)


#!/usr/bin/env python
# coding=utf-8
"""
Created on April 24, 2017

@author: yytang

1. get the first page url
2. get the pictures of the current page
3. get the next page url
"""

import scrapy
from scrapy.selector import Selector

from comicsCrawler.items import ComicsCrawlerItem
from libs.misc import *

class CartoonmadSpider(scrapy.Spider):
    """
    classdocs

    example: https://www.cartoonmad.com/comic/119400011126001.html
    """
    dom = 'www.cartoonmad.com'
    name = get_spider_name_from_domain(dom)
    allowed_domains = [dom]

    # custom_settings = {
    #     'DOWNLOAD_DELAY': 0.3,
    # }

    def __init__(self, *args, **kwargs):
        super(CartoonmadSpider, self).__init__(*args, **kwargs)
        self.root_path = kwargs['root_path']
        urls = kwargs['start_urls']
        self.start_urls = [self.polish_url(url) for url in urls]
        print(self.start_urls)

    def polish_url(self, url):
        url = url.strip('\n').strip()
        # pattern = 'https://e-hentai.org/[\w]/421552/4a24a76b83'
        # url = re.search(pattern, url).group(0)
        return url

    def get_title(self, response):
        sel = Selector(response)
        title = sel.xpath('//title/text()').extract()[0]
        title = title.split('-')[0].strip()
        return title

    def get_episodes(self, response):
        sel = Selector(response)
        episodes = []
        aa = sel.xpath('//tr/td/a')

        def episode_filter(s):
            p = 'href=\"/comic/\d{6,}.html'
            return True if re.search(p, s) else False

        aa = [a for a in aa if episode_filter(a.extract())]
        for a in aa:
            episode_url = a.xpath('@href').extract()[0]
            episode_url = response.urljoin(episode_url)
            episode_title = a.xpath('text()').extract()[0]
            episodes.append((episode_url, episode_title))
        return episodes

    def parse(self, response):
        title = self.get_title(response)
        episodes = self.get_episodes(response)

        for episode_url, episode_title in episodes:
            request = scrapy.Request(episode_url, callback=self.parse_image)
            request.meta['title'] = '{0}/{1}'.format(title, episode_title)
            request.meta['page_id'] = 1
            yield request

    def parse_image(self, response):
        title = response.meta['title']
        page_id = response.meta['page_id']
        image_name = '{0}/{1:03d}.jpg'.format(title, page_id)
        image_name = os.path.join('cartoonmad', image_name)
        image_path = os.path.join(self.root_path, image_name)
        html = str(response.body)
        if not os.path.isfile(image_path):
            image_url = re.search('img src="(http://[\w.]*?cartoonmad\.com.+?)"', html).group(1)
            item = ComicsCrawlerItem()
            item['image_name'] = image_name
            item['Referer'] = response.url
            item['image_urls'] = [image_url]
            yield item

        match = re.search('a href="(\d+[^"]*)', html)
        if match:
            next_page_url = response.urljoin(match.group(1))
            request = scrapy.Request(next_page_url, callback=self.parse_image)
            request.meta['title'] = title
            request.meta['page_id'] = page_id + 1
            yield request


# Firebase Web API
firebase_config = {
    "apiKey": "",
    "authDomain": "",
    "databaseURL": "",
    "projectId": "",
    "storageBucket": "",
    "messagingSenderId": "",
    "appId": ""
}

# Path to Firebase Admin Conifg JSON file
path_to_json = ""

# Twilio Authentication
twilio_auth = ("Twilio Auth", "Twilio Contact")

# Mobile Contact Numbers

twilio_contact = "Twilio Contact Number for WhatsApp"

twilio_sms_contact = "Twilio Contact Number for SMS"

mobile_contact = "Mobile User Phone Number"

web_contact = "Web App User Phone Number"


#!/usr/bin/python

from neopixel import *
import random
import datetime

MIN_SIZE = 2
MAX_SIZE = 10

class Stick:

    def __init__(self, startPos, size, ledCount):
        self.size = size
        self.ledCount = ledCount
        self.startPos = startPos
        self.endPos = self.startPos + size
        self.velocity = random.randint(1, 1000)
        self.direction = bool(random.getrandbits(1))
        self.lastMove = datetime.datetime.now()
        self.setRandomColor()

    def edgeDetection(self):
        if ((self.startPos <= 0 and not self.direction) or (self.endPos >= self.ledCount and self.direction)):
            self.direction = not self.direction

    def setRandomColor(self):
        self.color = Color(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    def step(self):
        timeDelta = datetime.datetime.now() - self.lastMove
        millSecs = int(timeDelta.total_seconds() * 1000)
        if (millSecs > self.velocity):
            self.edgeDetection()

            if (self.direction):
                self.startPos += 1
                self.endPos += 1
            else:
                self.startPos -= 1
                self.endPos -= 1
            self.lastMove = datetime.datetime.now()

class Colliders:

    def __init__(self, settings):
        self.settings = settings
        self.sticks = []
        #Temp
        for i in range(15):
            self.spawnSingle()

    def isSpaceAvailable(self, startPos, endPos):
        for stick in self.sticks:
            if ((startPos >= stick.startPos or endPos >= stick.startPos)
                and (startPos <= stick.endPos or endPos <= stick.endPos)):
                return False
        return True

    def spawnSingle(self):
        retries = 10
        spaceIsFree = False

        while (not spaceIsFree and retries > 0):
            size = random.randint(MIN_SIZE, MAX_SIZE)
            startPos = random.randint(0, self.settings.ledCount)
            spaceIsFree = self.isSpaceAvailable(startPos, startPos + size)
            if (spaceIsFree):
                stick = Stick(startPos, size, self.settings.ledCount)
                self.sticks.append(stick)
                break
            retries -= 1

    def deleteSingle(self):
        self.sticks.pop()

    def drawAll(self):
        self.clear()
        for stick in self.sticks:
            for i in range(stick.startPos, stick.endPos):
                self.settings.strip.setPixelColor(i, stick.color)

    def clear(self):
        for i in range(0, self.settings.ledCount):
            self.settings.strip.setPixelColorRGB(i, 0, 0, 0)

    def collisionDetection(self, cStick, currentIndex):
        for i, stick in enumerate(self.sticks):
            if ((i != currentIndex)
                and (((cStick.endPos >= stick.startPos and cStick.startPos < stick.startPos)
                    and ((cStick.direction and (not stick.direction
                             or (stick.direction and cStick.velocity < stick.velocity)))
                        or (not cStick.direction and not stick.direction and cStick.velocity > stick.velocity)))
                    or ((cStick.startPos <= stick.endPos and cStick.endPos > stick.endPos)
                    and ((not cStick.direction and (stick.direction
                            or (not stick.direction and cStick.velocity < stick.velocity)))
                        or (cStick.direction and stick.direction and cStick.velocity > stick.velocity))))):
                if (cStick.direction != stick.direction):
                    currentDirection = cStick.direction
                    cStick.direction = not currentDirection
                    stick.direction = currentDirection
                currentLastMov = cStick.lastMove
                currentVelocity = cStick.velocity
                cStick.velocity = stick.velocity
                cStick.lastMove = stick.lastMove
                stick.velocity = currentVelocity
                stick.lastMove = currentLastMov
                stick.setRandomColor()
                cStick.setRandomColor()

    def step(self):
        for i, stick in enumerate(self.sticks):
            self.collisionDetection(stick, i)
            stick.step()
        self.drawAll()
        self.settings.strip.show()

    def show(self, settings = None):
        if (settings != None):
            self.settings = settings
        self.step()


# -----------------------------------------------------------------------------
# Copyright (c) 2014--, The Qiita Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------
from __future__ import division

from collections import defaultdict

from tornado.web import authenticated, HTTPError
from tornado.gen import coroutine, Task

from qiita_core.util import execute_as_transaction
from qiita_db.artifact import Artifact
from qiita_pet.handlers.base_handlers import BaseHandler
from qiita_pet.handlers.util import check_access

html_error_message = "An error occurred %s %s
%s"

def _to_int(value):
    """Transforms `value` to an integer

    Parameters
    ----------
    value : str or int
        The value to transform

    Returns
    -------
    int
        `value` as an integer

    Raises
    ------
    HTTPError
        If `value` cannot be transformed to an integer
    """
    try:
        res = int(value)
    except ValueError:
        raise HTTPError(500, "%s cannot be converted to an integer" % value)
    return res

class PreprocessingSummaryHandler(BaseHandler):
    @execute_as_transaction
    def _get_template_variables(self, preprocessed_data_id, callback):
        """Generates all the variables needed to render the template

        Parameters
        ----------
        preprocessed_data_id : int
            The preprocessed data identifier
        callback : function
            The callback function to call with the results once the processing
            is done

        Raises
        ------
        HTTPError
            If the preprocessed data does not have a log file
        """
        # Get the objects and check user privileges
        ppd = Artifact(preprocessed_data_id)
        study = ppd.study
        check_access(self.current_user, study, raise_error=True)

        # Get the return address
        back_button_path = self.get_argument(
            'back_button_path',
            '/study/description/%d?top_tab=preprocessed_data_tab&sub_tab=%s'
            % (study.id, preprocessed_data_id))

        # Get all the filepaths attached to the preprocessed data
        files_tuples = ppd.filepaths

        # Group the files by filepath type
        files = defaultdict(list)
        for _, fp, fpt in files_tuples:
            files[fpt].append(fp)

        try:
            log_path = files['log'][0]
        except KeyError:
            raise HTTPError(500, "Log file not found in preprocessed data %s"
                                 % preprocessed_data_id)

        with open(log_path, 'U') as f:
            contents = f.read()
            contents = contents.replace('\n', '
')
            contents = contents.replace('\t', '    ')

        title = 'Preprocessed Data: %d' % preprocessed_data_id

        callback((title, contents, back_button_path))

    @authenticated
    @coroutine
    @execute_as_transaction
    def get(self, preprocessed_data_id):
        ppd_id = _to_int(preprocessed_data_id)

        title, contents, back_button_path = yield Task(
            self._get_template_variables, ppd_id)

        self.render('text_file.html', title=title, contents=contents,
                    back_button_path=back_button_path)


# -*- coding: utf-8 -*-
"""
-------------------------------------------------
# @Project  :爬虫.py
# @File     :爬虫可用
# @Date     :2021/4/14 11:37
# @Author   :dyh
# @Email    :dingyihangheu@163.com
# @Software :PyCharm
# Description :

-------------------------------------------------
"""
import requests
import re
import os
import time

# 获取图片url连接
def get_parse_page(pn,name):

    for i in range(int(pn)):
        # 1.获取网页
        print('正在获取第{}页'.format(i+1))

        # 百度图片首页的url
        # name是你要搜索的关键词
        # pn是你想下载的页数

        url = 'https://image.baidu.com/search/flip?tn=baiduimage&ie=utf-8&word=%s&pn=%d' %(name,i*20)

        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/53.0.2785.104 Safari/537.36 Core/1.53.4843.400 QQBrowser/9.7.13021.400'}

        # 发送请求，获取相应
        response = requests.get(url, headers=headers)
        html = response.content.decode()
        # print(html)

        # 2.正则表达式解析网页
        # "objURL":"http://n.sinaimg.cn/sports/transform/20170406/dHEk-fycxmks5842687.jpg"
        results = re.findall('"objURL":"(.*?)",', html) # 返回一个列表

        # 根据获取到的图片链接，把图片保存到本地
        save_to_txt(results, name, i)

# 保存图片到本地
def save_to_txt(results, name, i):

    j = 0
    # 在当目录下创建文件夹
    if not os.path.exists('./' + name):
        os.makedirs('./' + name)

    # 下载图片
    for result in results:
        print('正在保存第{}个'.format(j))
        try:
            pic = requests.get(result, timeout=10)
            # time.sleep(1)
        except:
            print('当前图片无法下载')
            j += 1
            continue

        # 可忽略，这段代码有bug
        # file_name = result.split('/')
        # file_name = file_name[len(file_name) - 1]
        # print(file_name)
        #
        # end = re.search('(.png|.jpg|.jpeg|.gif)$', file_name)
        # if end == None:
        #     file_name = file_name + '.jpg'

        # 把图片保存到文件夹
        file_full_name = './' + name + '/' + str(i) + '-' + str(j) + '.jpg'
        with open(file_full_name, 'wb') as f:
            f.write(pic.content)

        j += 1

# 主函数
if __name__ == '__main__':

    name = input('请输入你要下载的关键词：')
    pn = input('你想下载前几页（1页有60张）:')
    get_parse_page(pn, name)

BS_预处理_爬取图片.py

BS_预处理_爬取图片.py


# -*- coding: utf-8 -*-
"""
Created on Sun Jul 29 19:29:57 2018

@author: s144299
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

#unet model
class double_conv(nn.Module):
    '''(conv => BN => ReLU) * 2'''
    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size = (3,3), padding = (1,1)),
            nn.BatchNorm2d(out_ch, eps=1e-05, momentum=0.1, affine=True),
            #nn.Dropout(0.2),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_ch, out_ch, kernel_size = (3,3), padding = (1,1)),
            nn.BatchNorm2d(out_ch, eps=1e-05, momentum=0.1, affine=True),
            #nn.Dropout(0.2),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x

class up_conv(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(up_conv, self).__init__()

        if bilinear: #either use bilinear
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else: #else ise transposed conv
            self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)

        self.double_conv1 = double_conv(in_ch, out_ch)

    def forward(self, x1, x2):
        """ x1 is the current input and x2 is the copied from downsampling """
        x1 = self.up(x1)
        #calculate difference for padding
        diffX = x1.size()[2] - x2.size()[2]
        diffY = x1.size()[3] - x2.size()[3]
        #pad input from downsampling
        x2 = F.pad(input = x2, pad = (diffX // 2, int(diffX / 2),
                                      diffY // 2, int(diffY / 2)),
                   mode = 'constant',
                   value = 0
            )
        #concatenate along channel axis
        x1 = torch.cat([x2, x1], dim=1)
        #double convolution
        x1 = self.double_conv1(x1)
        return x1

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes):
        super(UNet, self).__init__()

        self.in_conv = double_conv(n_channels, 32)

        self.double_conv1 = double_conv(32, 64)
        self.double_conv2 = double_conv(64, 128)
        self.double_conv3 = double_conv(128, 256)

        self.up_conv1 = up_conv(384, 128)
        self.up_conv2 = up_conv(192, 64)
        self.up_conv3 = up_conv(96, 32)

        self.out_conv =  nn.Conv2d(32, n_classes-1, kernel_size = (1,1)) #ignore background

        self.max_pool = nn.MaxPool2d(kernel_size = (2,2), stride = (2,2))

        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x1 = self.in_conv(x)

        x1 = self.max_pool(x1)
        x1 = self.dropout(x1)
        x2 = self.double_conv1(x1)

        x2 = self.max_pool(x2)
        x2 = self.dropout(x2)
        x3 = self.double_conv2(x2)

        x3 = self.max_pool(x3)
        x3 = self.dropout(x3)
        x4 = self.double_conv3(x3)

        #upsampling
        x = self.up_conv1(x4, x3)
        x = self.dropout(x)
        x = self.up_conv2(x, x2)
        x = self.dropout(x)
        x = self.up_conv3(x, x1)
        x = self.dropout(x)

        #out
        x = self.out_conv(x)

        return self.sigmoid(x)

def get_model(n_channels, n_classes):
    model = UNet(n_channels, n_classes)
    if torch.cuda.is_available(): model.cuda()
    return model


from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import accuracy_score
from sklearn.neighbors import KNeighborsClassifier
from nltk.corpus import stopwords
from sklearn.grid_search import GridSearchCV
from sklearn.ensemble import VotingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier

from sklearn.ensemble import RandomForestClassifier
from sklearn.naive_bayes import MultinomialNB

import time

import matplotlib.pyplot as plt; plt.rcdefaults()
import numpy as np
import matplotlib.pyplot as plt
from sklearn import tree

start_time = time.time()
log=open('classification_demo.txt','w') # test file

def loadData(fname):

    labels=[]
    reviews=[]
    f=open(fname)
    for line in f:
        rating,vote,review=line.strip().split('\t')
        reviews.append(review.lower())
        labels.append(rating.lower())
    f.close()
    return reviews,labels

question_train,labels_train=loadData('train.txt')
question_test,labels_test=loadData('test.txt')

#Build a counter based on the training dataset
counter = CountVectorizer(stop_words=stopwords.words('english'))
counter.fit(question_train)

#count the number of times each term appears in a document and transform each doc into a count vector
counts_train = counter.transform(question_train)#transform the training data
counts_test = counter.transform(question_test)#transform the testing data

#    Classifiers Used:
#    1. KNN
#    2.Logistic Regression
#    3.DecisionTreeClassifier
#    4.RandomForestClassifier
#    5. Bayesian

KNN_classifier=KNeighborsClassifier()
LREG_classifier=LogisticRegression()
DT_classifier = DecisionTreeClassifier()
RF_classifier=RandomForestClassifier()
NB_classifier=MultinomialNB()

predictors=[('knn',KNN_classifier),('lreg',LREG_classifier),('dt',DT_classifier),('rf', RF_classifier),('nb',NB_classifier)]
VT=VotingClassifier(predictors)

#=======================================================================================
#build the parameter grid
KNN_grid = [{'n_neighbors': [1,3,5,7,9,11,13,15,17], 'weights':['uniform','distance']}]

#build a grid search to find the best parameters
gridsearchKNN = GridSearchCV(KNN_classifier, KNN_grid, cv=5)

#run the grid search
gridsearchKNN.fit(counts_train,labels_train)

#=======================================================================================

#build the parameter grid
DT_grid = [{'max_depth': [3,4,5,6,7,8,9,10,11,12],'criterion':['gini','entropy']}]

#build a grid search to find the best parameters
gridsearchDT  = GridSearchCV(DT_classifier, DT_grid, cv=5)

#run the grid search
gridsearchDT.fit(counts_train,labels_train)

DT = tree.DecisionTreeClassifier()
treeClf = DT.fit(counts_train,labels_train)
tree.export_graphviz(treeClf, out_file='DTree.dot')
#=======================================================================================

#build the parameter grid
LREG_grid = [ {'C':[0.5,1,1.5,2],'penalty':['l1','l2']}]

#build a grid search to find the best parameters
gridsearchLREG  = GridSearchCV(LREG_classifier, LREG_grid, cv=5)

#run the grid search
gridsearchLREG.fit(counts_train,labels_train)

#=======================================================================================
#build the parameter grid
RF_grid = [{'n_estimators':[500,1000,1500,2000,2500,3000], 'criterion':['gini','entropy']}]

#build a grid search to find the best parameters
gridsearchRF = GridSearchCV(RF_classifier, RF_grid, cv=5)

#run the grid search
gridsearchRF.fit(counts_train,labels_train)

#=======================================================================================
#build the parameter grid
NB_grid = [{'alpha':[0,1.0,5.0,10.0,20.0,30.0,40.0], 'fit_prior':[True,False]}]

#build a grid search to find the best parameters
gridsearchNB = GridSearchCV(NB_classifier, NB_grid, cv=5)

#run the grid search
gridsearchNB.fit(counts_train,labels_train)

#=======================================================================================

VT.fit(counts_train,labels_train)

#use the VT classifier to predict
predicted=VT.predict(counts_test)

#print the accuracy
print (accuracy_score(predicted,labels_test))
log.write(str(accuracy_score(predicted,labels_test))+"\n");

algos=[];
acc = [];
m = 0;

for params, mean_score, scores in gridsearchKNN.grid_scores_:
    print (params, mean_score)
    log.write(str(params)+str(mean_score)+"\n")
    if mean_score>m:
        m = mean_score

acc.append(m)
m=0

for params, mean_score, scores in gridsearchDT.grid_scores_:
    print (params, mean_score)
    log.write(str(params)+str(mean_score)+"\n")
    if mean_score>m:
        m = mean_score

acc.append(m)
m=0

for params, mean_score, scores in gridsearchLREG.grid_scores_:
    print (params, mean_score)
    log.write(str(params)+str(mean_score)+"\n")
    if mean_score>m:
        m = mean_score

acc.append(m)
m=0

for params, mean_score, scores in gridsearchRF.grid_scores_:
    print (params, mean_score)
    log.write(str(params)+str(mean_score)+"\n")
    if mean_score>m:
        m = mean_score

acc.append(m)
m=0

for params, mean_score, scores in gridsearchNB.grid_scores_:
    print (params, mean_score)
    log.write(str(params)+str(mean_score)+"\n")
    if mean_score>m:
        m = mean_score

acc.append(m)
m=0

print(str(algos))
print(str(acc))

print("--- %s seconds ---" % round(time.time() - start_time, 2))

algos = ['KNN', 'DT', 'LREG', 'RF', 'NB']
y_pos = np.arange(len(algos))

plt.bar(y_pos, acc, align='center', alpha=0.5)
plt.xticks(y_pos, algos)
plt.ylabel('Accuracy')
plt.title('Algorithms v Accuracy')

plt.show()


#Input: Files loaded by loadtxt: input .raw file with directional data (created by getmoments.py) and 'directions_sph.inp'. Make sure these correspond.
#       Maximum n value of spherical harmonic Y_{n}^{m}
#Output: formatted coefficients of directional data in dir-velmoments.raw up to maximun n value specified using Clenshaw-Curtis quadrature for polar angles
#Syntax: python3 sph_coeffs.py <.raw file>  <maxn>

import numpy as np
import scipy as scp
from scipy import special
from scipy import integrate
from scipy import fft
import sys
import warnings

scp.version.full_version

data=np.loadtxt(sys.argv[1]) #only reads third coulmn.
#data=np.loadtxt('dir-velmoments.raw')
angles=np.loadtxt("directions_sph.inp") #assumes that these are the directions that correspond to the directional values in
                                        #'dir-velmoments.raw'
maxn=int(sys.argv[2])
#maxn=4
#converting to spherical polar coordinates (NOTE: NOT a geographic system like in create_pm3dfile.py)
# r=np.sqrt(np.power(data[:,0],2) + np.power(data[:,1],2) + np.power(data[:,2],2))
# phipoints=np.arccos(np.divide(data[:,2],r)) #range of arccos = [-pi/2,pi/2]

#This method assumes that the directions are created by equally dividing theta and phi into an equal number of points
value=data[:,3]
value=value/9.274E-41 #converting to atomic units
thetapoints=angles[:,0] #note that the range here is usually [-180,180]
phipoints=angles[:,1] #note that the range here is usually [-90,90]
phipoints=90-phipoints #converting from geographic to regular spherical coordinates. Now range is [90,0]
thetapoints=thetapoints*np.pi/180.0 #converting to radians
phipoints=phipoints*np.pi/180.0

if(len(thetapoints) != len(phipoints)):
    sys.exit("Error: The number of theta and phi angles not equal!")

if(len(value) != len(thetapoints)):
    sys.exit("Error: There are more or less data points than angles! Do not use directions from directions_sph.inp")

len_theta=int(np.sqrt(len(thetapoints))) #program only works if these two are equal
len_phi=int(np.sqrt(len(phipoints)))

if(len_phi != len_theta):
    sys.exit('Number of grid points of theta and phi not equal.')

#function for getting real spherical harmonics
def realharm(m,n,theta,phi):
    if(m==0):
        return np.real(scp.special.sph_harm(m,n,theta,phi)) #the returned datatype is float and not complex float

    #General formula of real spherical harmonic
    output=(1/np.sqrt(2))*np.array(scp.special.sph_harm(-np.abs(m),n,theta,phi) + np.sign(m)*(-1)**(np.abs(m))*scp.special.sph_harm(np.abs(m),n,theta,phi))
    if(m < 0):
        output=complex(0,1)*output

    return np.real(output) #the returned datatype is float and not complex float

#function that integrates the product of two functions f1 and f2 using the Clenshaw-Curtis quadrature for polar angles phi and trapezoidal rule for azimulthal theta
#Requires that the input array f1 and f2 have dimension not greater than two
def integrate(f1,f2,theta,phi):
    if((np.ndim(f1)>2)or(np.ndim(f2)>2)):
        sys.exit('Dimension of array(s) to be integrated larger than two')

    len_theta=int(np.sqrt(len(thetapoints)))
    len_phi=int(np.sqrt(len(phipoints)))

   #All the subsequent if statements are to properly subset the numpy arrays when for values of dimension = 1,2 of f1 and f2 (4 cases)
    if((np.ndim(f1)==2)and(np.ndim(f2)==2)): #will return array with one row and len_phi columns
        #print('a')
        temp=np.zeros(shape=(len(f1),len_phi))
        for i in range(len_phi):
            temp[:,i]=scp.integrate.trapz(np.multiply(f1[:,i::len_phi],f2[:,i::len_phi]),theta[::len_phi]) #integrating over azimuthal theta

        tempdct=scp.fft.dct(temp,type=1) #flipping the array important otherwise the fourier coefficients will be all wrong

        integral=tempdct[:,0]
        kmax=(len_phi-2)//2
        k=1
        while(k <= kmax):
            integral=integral + 2.0*tempdct[:,2*k]/(1-(2*k)**2)
            k=k+1

        integral=integral/(len_phi-1)

    if((np.ndim(f1)==1)and(np.ndim(f2)==1)): #will return scalar
        #print('b')
        temp=np.zeros(shape=(len_phi))
        for i in range(len_phi):
            temp[i]=scp.integrate.trapz(np.multiply(f1[i::len_phi],f2[i::len_phi]),theta[::len_phi]) #integrating over azimuthal theta

        tempdct=scp.fft.dct(temp,type=1) #flipping the array important otherwise the fourier coefficients will be all wrong

        integral=tempdct[0]
        kmax=(len_phi-2)//2
        k=1
        while(k <= kmax):
            integral=integral + 2.0*tempdct[2*k]/(1-(2*k)**2)
            k=k+1

        integral=integral/(len_phi-1)

    if((np.ndim(f1)==1)and(np.ndim(f2)==2)): #will return array with one row and len_phi columns
        #print('c')
        temp=np.zeros(shape=(len(f1),len_phi))
        for i in range(len_phi):
            temp[:,i]=scp.integrate.trapz(np.multiply(f1[i::len_phi],f2[:,i::len_phi]),theta[::len_phi]) #integrating over azimuthal theta

        tempdct=scp.fft.dct(temp,type=1) #flipping the array important otherwise the fourier coefficients will be all wrong

        integral=tempdct[:,0]
        kmax=(len_phi-2)//2
        k=1
        while(k <= kmax):
            integral=integral + 2.0*tempdct[:,2*k]/(1-(2*k)**2)
            k=k+1

        integral=integral/(len_phi-1)

    if((np.ndim(f1)==2)and(np.ndim(f2)==1)): #will return array with one row and len_phi columns
        #print('d')
        temp=np.zeros(shape=(len(f1),len_phi))
        for i in range(len_phi):
            temp[:,i]=scp.integrate.trapz(np.multiply(f1[:,i::len_phi],f2[i::len_phi]),theta[::len_phi]) #integrating over azimuthal theta

        tempdct=scp.fft.dct(temp,type=1) #flipping the array important otherwise the fourier coefficients will be all wrong

        integral=tempdct[:,0]
        kmax=(len_phi-2)//2
        k=1
        while(k <= kmax):
            integral=integral + 2.0*tempdct[:,2*k]/(1-(2*k)**2)
            k=k+1

        integral=integral/(len_phi-1)

    return integral

#-------------------------------------------------------------------------------------------------------------------------------------------------

#Integration for normalization coefficient of input data
norm=integrate(value,value,thetapoints,phipoints)
norm

#Vector of spherical harmonics. Order = [(0,0),(-1,1),(0,1),(1,1),(-2,2),(-1,2),(0,2)] where (m,n)and so on
#Real spherical harmonics defined here: https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
sph_harm=np.zeros(shape=(((maxn+1)**2),len(thetapoints)))
counter=0
for i in range(maxn+1):
    for j in range(-i,i+1):
        sph_harm[counter]=np.array(realharm(j,i,thetapoints,phipoints))
        counter=counter+1

#Vector of norms. norm[i] 0 norm of sph_harm[i].
#Norm for testting. do not need to divide by norm as they should be exactly equal to one
sph_norm=integrate(sph_harm,sph_harm,thetapoints,phipoints)

#Vector of coeffs. coeff[i] is the coeff corresponding to spherical harmonic in sph_harm[i]
coeff=integrate(sph_harm,value,thetapoints,phipoints)
coeff

#Dividing all coeffs by the  respective spherical harmonic norms
#coeff=np.divide(coeff,sph_norm)

#Normalized (sum of squares = 1) coeffs of input data
coeff_normalized=coeff/np.sqrt(norm)

#List of orbital names
sph_harm_names=list() #empty list
if(maxn>=0):
    sph_harm_names=['s']
if(maxn>=1):
    sph_harm_names.extend(['px','py','pz'])
if(maxn>=2):
    sph_harm_names.extend(['dxy','dyz','dz2', 'dxz', 'dx2-y2'])
if(maxn>=3):
    sph_harm_names.extend(['fy(3x2-y2)', 'fxyz', 'fyz2', 'fz3', 'fxz2', 'fz(x2-y2)', 'fx(x2-3y2)'])
if(maxn>=3):
    sph_harm_names.extend(['gxy(x2-y2)', 'gzy3', 'gz2xy', 'gz3y', 'gz4', 'gz3x', 'gz2(x2-y2)','gzx3','gx4+y4'])
if(maxn>=4):
    templist=['higher harmonic']*((maxn+1)**2-25)
    sph_harm_names.extend(templist)

#printing to stdout
print("{0:19s}  {1:14s}   {2:18s}".format('Spherical Harmonic','Coeff','Coeff (normalized)'))
for i in range(len(coeff)):
    #print(sph_harm_names[i]+'     '+str(coeff[i])+'   '+str(coeff_normalized[i]))
    print("{0:19s} {1: 12.8e}  {2: 12.10e}".format(sph_harm_names[i], coeff[i], coeff_normalized[i]))

#Integrations without defining the function integrate:
#Spherical norms:
# tempsph_norm=np.zeros(shape=(len(sph_harm),len_phi))
# for i in range(len_phi):
    # tempsph_norm[:,i]=scp.integrate.trapz(np.multiply(sph_harm[:,i::len_phi],sph_harm[:,i::len_phi]),thetapoints[::len_phi]) #integrating over azimuthal theta
#
# tempdct_sph_norm=scp.fft.dct(np.flip(tempsph_norm),type=1) #flipping the arraw important otherwise the fourier coefficients will be all wrong
#
# sph_norm=tempdct_sph_norm[:,0]
# kmax=(len_phi-2)//2
# k=1
# while(k <= kmax):
    # sph_norm=sph_norm + 2.0*tempdct_sph_norm[:,2*k]/(1-(2*k)**2)
    # k=k+1
#
# sph_norm=sph_norm/(len_phi-1)

#Coefficients:
# tempcoeff=np.zeros(shape=(len(sph_harm),len_phi))
# for i in range(len_phi):
    # tempcoeff[:,i]=scp.integrate.trapz(np.multiply(sph_harm[:,i::len_phi],value[i::len_phi]),thetapoints[::len_phi]) #integrating over azimuthal theta
#
# tempdct_coeff=scp.fft.dct(np.flip(tempcoeff),type=1) #flipping the arraw important otherwise the fourier coefficients will be all wrong
#
# tempdct_coeff
# coeff=tempdct_coeffs[:,0]
# kmax=(len_phi-2)//2
# k=1
# while(k <= kmax):
    # coeff=coeff + 2.0*tempdct_coeff[:,2*k]/(1-(2*k)**2)
    # k=k+1
#
# coeff=coeff/(len_phi-1)
# coeff


#Alexander_Máni_Einarsson
#19/09/19
#aukaverkefni_(while)

"""
#liður 1
tala = 18
while tala < 40 and tala > 17:
    tala = tala + 2
    print(tala, end =" ")

#liður 2
tala = 15
summa = 0
while tala < 346 and tala > 14:
    summa = summa + tala
    tala = tala + 2
print(summa)

#liður 3
tala = int(input("Sláðu inn tala sem er minni en hunndrað"))
tala = tala - 1
while tala < 100:
    tala = tala + 1
    print(tala, end=" ")

#liður 4
tala = int(input("Sláðu inn tala sem er hærri en hunndrað"))
tala = tala + 1
while tala > 100:
    tala = tala - 1
    print(tala, end=" ")

#liður 5
tala = int(input("SLáðu inn tölu: "))
niu = tala // 9
fimm = tala // 5
print("9 gengur upp í töluna",niu,"sinnum")
print("5 gengur upp í töluna",fimm,"sinnum")
"""
#liður 6
print("sláðu inn tvær kommutölur")
tala = float(input("SLáðu inn fyrri töluna hér: "))
tala_1 = float(input("SLáðu inn seinni töluna hér: "))
svar = tala_1 / tala
(svar,round(5))
print(svar)


# Charlotte Capitanchik
# 25.04.18
# Based on fastCLIP tRNA strategy
# input is SAM of tRNA mapping
# and .fai of what is was mapped to

import pysam
from collections import defaultdict
import operator
from collections import Counter
import sys
import statistics
import random
import re

# input #
#trna_fai = sys.argv[1] # "/Volumes/lab-luscomben/working/charlotte/miCLIP_UN1710_CS2_delta2/summarising_trna_mapping_style/permanent_files/UCSC_hg38_tRNAs_mature_PTM-deduplicated.fa.fai" # #"/Users/capitac/PhD/trna_index/hg38-mature-tRNAs-deduplicated.fa.fai"
reads = sys.argv[1] #"../../results/maturetRNA_mapped/DDX3_WT_iCLIP_rep2.Aligned.out.sorted.deduplicated.bam"  # #/Volumes/lab-luscomben/working/charlotte/miCLIP_UN1710_CS2_delta2/summarising_trna_mapping_style/results/m1a_786o_noq_15min_20180413.Aligned.out.sorted.deduplicated.bam" # #"/Users/capitac/PhD/trna_index/m1a_HEK293_beta_1.Aligned.out.sorted.bam"
dir_and_samp = sys.argv[2] #"testing_DDX3_WT_iCLIP_rep2" #  #"/Volumes/lab-luscomben/working/charlotte/testing" #  #"/Users/capitac/PhD/trna_index/m1a_HEK293_beta_1"
samp = sys.argv[3] #"testing_DDX3_WT_iCLIP_rep2" #
seed = sys.argv[4] #1 #
fraction_merge = float(sys.argv[5]) #0.9 #

# output #
# "_trna_summary_stats.tsv"
# "_trna_mixedup_anticodons.tsv"
# "_ambig_ac_position_dist.tsv"

## Code to parse the .fai file to get the name of all the amino acids and counts .etc ##
## Not required for main output
# trna_index = open(trna_fai).readlines()
# aminoacid_count = defaultdict(int)
# anticodon_list=[]
# aminoacid_list=[]
# trna_index
# for i in range(0,len(trna_index)):
#     linelist = trna_index[i].split("\t")
#     namelist = trna_index[i].split("-")
#     anticodon = "-".join(namelist[1:3])
#     amino_acid = namelist[1]
#     anticodon_list.append(anticodon)
#     aminoacid_list.append(amino_acid)
#     aminoacid_count[amino_acid] += 1
# anticodon_list = set(anticodon_list)
# aminoacid_list = set(aminoacid_list)
# # Find out the amino acid with the highest number of genes
# max_key = max(aminoacid_count.iterkeys(), key=(lambda key: aminoacid_count[key]))
# aminoacid_count[max_key]
#print(aminoacid_count[max_key])

## Now move on to the reads ##

samfile = pysam.AlignmentFile(reads,"rb")

## Summarise at three levels
read_dictionary_amino_acid = defaultdict(set)
read_dictionary_anticodon = defaultdict(set)
read_dictionary_start_positions = defaultdict(list)
read_gene_dict = defaultdict(list)
readname_list = []

for read in samfile.fetch():
    if read.is_unmapped == False:
        aa_tmp = samfile.get_reference_name(read.reference_id)
        aa_tmp = re.sub("^(nm-|nmt-)","", aa_tmp).split("-")
        aa = aa_tmp[1]
        anticodon = "-".join(aa_tmp[1:3])
        read_dictionary_amino_acid[read.query_name].add(aa)
        read_dictionary_anticodon[read.query_name].add(anticodon)
        read_dictionary_start_positions[read.query_name] += [read.reference_start]
        read_gene_dict[read.query_name] += [samfile.get_reference_name(read.reference_id)]
        readname_list += [read.query_name]
    else:
        continue

## #### For a file where anticodons are not merged ######
FINAL_read_dictionary_notmerge = defaultdict(set)
for key, value in read_dictionary_anticodon.items():
    if len(value) == 1:
        FINAL_read_dictionary_notmerge[key] = value

final_cdna_counts_dict_notmerge = defaultdict(list)
random.seed(seed)
# assign positions, use mode if possible, but if no mode then select randomly
# from multimapped positions using random seed
for key,value in FINAL_read_dictionary_notmerge.items():
    try:
        x = statistics.mode(read_dictionary_start_positions[key])
        final_cdna_counts_dict_notmerge["".join(value)] += [x]
    except:
        final_cdna_counts_dict_notmerge["".join(value)] += [random.choice(read_dictionary_start_positions[key])]

for key,value in final_cdna_counts_dict_notmerge.items():
    final_cdna_counts_dict_notmerge[key] = Counter(value)

########### Unmerged ends here #################################
## Find out which anticodons are ambiguous ##
number_of_reads_multimapped_anticodon_dict = defaultdict(int)

for key, value in read_dictionary_anticodon.items():
    if len(value) > 1:
        number_of_reads_multimapped_anticodon_dict[";".join(sorted(value))] += 1

sorted_mixup_ranking = sorted(number_of_reads_multimapped_anticodon_dict.items(), key=operator.itemgetter(1), reverse=True)
total = sum(number_of_reads_multimapped_anticodon_dict.values())
fract_total = total * fraction_merge
print(total)
print(fract_total)

counting=0
thresholded_multimapped_anticodon_dict = defaultdict(int)

# Threshold mixed up anticodons based on % reads we want to recover
for i in sorted_mixup_ranking:
    counting += i[1]
    print(i)
    name=i[0]
    number=i[1]
    thresholded_multimapped_anticodon_dict[name] = number
    if counting >= fract_total:
        break

# This is the actual fraction of reads we are recovering
real_frac = float(sum(thresholded_multimapped_anticodon_dict.values())) / float(total)

# The anticodons we are going to merge
comparison_list = list(thresholded_multimapped_anticodon_dict.keys())

## Merge anticodons that are ambiguous################
FINAL_read_dictionary = defaultdict(set)

# If "readmap" has components in "checklist", then replace with the
# corresponsing entry in checklist ->
# If it doesn't - return readmap
# If it has conflicting entries - return readmap

def getIndex(readmap, checklist):
    """
    inputs:
        readmap - list of symbols
        checklist - list of strings(delineated by semicolons)
    """
    readmap = list(readmap) #type cast from set to list
    holder = []
    counter=0

    for i in readmap: # iterate over readmap
        switch=True
        for idx, j in enumerate(checklist):
            k = j.split(';')
            if i in k:
                holder.append(idx)
                switch=False
        if switch is True:
            holder.append(-1)

    if len(holder) < 1:
        holder = [-1]
    #evaluate holder values and return based on this
    holder = set(holder)
    holder = list(holder)

    if -1 in holder:
        return readmap
    elif len(holder) > 1:
        return readmap
    elif len(holder) == 1:
        return [checklist[holder[0]]] #checklist[list(holder)[0]]
    else:
        print('something went wrong')

tempdict = {}

for key, value in read_dictionary_anticodon.items():
    newval = getIndex(value, comparison_list)
    tempdict[key] = newval

# Filter for single mapping anticodon/merged anticodons
for key, value in tempdict.items():
    if len(value) == 1:
        FINAL_read_dictionary[key] = value

# assign possitions
final_cdna_counts_dict = defaultdict(list)
random.seed(seed)
for key,value in FINAL_read_dictionary.items():
    try:
        x = statistics.mode(read_dictionary_start_positions[key])
        final_cdna_counts_dict["".join(value)] += [x]
    except:
        final_cdna_counts_dict["".join(value)] += [random.choice(read_dictionary_start_positions[key])]

# count at positions
for key,value in final_cdna_counts_dict.items():
    final_cdna_counts_dict[key] = Counter(value)

# Write a file condensed to groups of anticodons
finalfile = open(dir_and_samp + "_tRNA.bed", "w")
for key,value in final_cdna_counts_dict.items():
    gene = key
    for k in value:
        finalfile.write(gene+"\t"+str(k)+"\t"+str(k+1)+"\t"+gene+"\t"+str(value[k])+"\t+\n")
finalfile.close()

# Write a file where only reads unambiguous for anticodon are used

finalfile = open(dir_and_samp + "_tRNA_unambig_AC.bed", "w")
for key,value in final_cdna_counts_dict_notmerge.items():
    gene = key
    for k in value:
        finalfile.write(gene+"\t"+str(k)+"\t"+str(k+1)+"\t"+gene+"\t"+str(value[k])+"\t+\n")
finalfile.close()

##### ESSENTIAL CODE STOPS HERE: JUST SUMMARY FILES FROM NOW ON  #####

# Some statistics for a summary file

total_mapped_reads = len(read_dictionary_amino_acid)
total_uniquely_mapped = 0
number_of_reads_multimapped_aa = 0 # number of reads that are ambigous for amino acid
number_of_reads_multimapped_anticodon = 0 # number of reads that are ambiguous for anticodon
number_of_reads_multimapped_anticodon_dict = defaultdict(int)

for key, value in read_dictionary_start_positions.items():
    if len(value) == 1:
        total_uniquely_mapped += 1

for key, value in read_dictionary_amino_acid.items():
    if len(value) > 1:
        number_of_reads_multimapped_aa += 1

reads_unambig_AC = defaultdict(str)
reads_unambig_pos = defaultdict(list)
for key, value in read_dictionary_anticodon.items():
    if len(value) > 1:
        number_of_reads_multimapped_anticodon += 1
        number_of_reads_multimapped_anticodon_dict[";".join(value)] += 1
    elif len(value) == 1:
        reads_unambig_AC[key] = "".join(value)
        reads_unambig_pos[key] = read_dictionary_start_positions[key]

# Sort the dictionary by highest to lowest
mixed_up_anticodons = open(dir_and_samp+"_tRNA_mixedup_anticodons.tsv", "w")
sorted_mixup_ranking = sorted(number_of_reads_multimapped_anticodon_dict.items(), key=operator.itemgetter(1), reverse=True)
for i in sorted_mixup_ranking:
    if i[1] > 200:
        mixed_up_anticodons.write(str(i[0])+"\t"+str(i[1])+"\t"+str(total_mapped_reads)+"\t"+samp+"\n")
mixed_up_anticodons.close()

summary_stats = open(dir_and_samp+"_tRNA_summary_stats.tsv", "w")
summary_stats.write("Total mapped reads to trna\t"+str(total_mapped_reads)+"\t"+samp+"\n")
summary_stats.write("Total single mapping reads to trna\t"+str(total_uniquely_mapped)+"\t"+samp+"\n")
summary_stats.write("Number of reads that are ambiguous at amino acid level\t"+str(number_of_reads_multimapped_aa)+"\t"+samp+"\n")
summary_stats.write("Percentage of reads used for amino acid summary\t"+str(round(float(total_mapped_reads-number_of_reads_multimapped_aa)/total_mapped_reads,3))+"\t"+samp+"\n")
summary_stats.write("Number of reads that are ambiguous at anticodon level\t"+str(number_of_reads_multimapped_anticodon)+"\t"+samp+"\n")
summary_stats.write("Percentage of reads used for anticodon summary\t"+str(round(float(total_mapped_reads-number_of_reads_multimapped_anticodon)/total_mapped_reads,3))+"\t"+samp+"\n")
summary_stats.write("The actual fraction of ambiguous reads that are recovered by merging anticodons\t"+str(real_frac)+"\n")
summary_stats.write("The anticodons that were merged are\t"+("|").join(comparison_list))

### AAlevel summary
aa_counts = defaultdict(int)
for key, value in read_dictionary_amino_acid.items():
    if len(value) == 1:
        aa_counts["".join(value)] += 1
### Write to file
aa_counts

### Anticodon level summary & also get names of reads that are unique to one anticodon
anticodon_counts = defaultdict(int)
unique_anticodon_read_names = set()
for key, value in read_dictionary_anticodon.items():
    if len(value) == 1:
        anticodon_counts["".join(value)] += 1
        unique_anticodon_read_names.add(key)
anticodon_counts

### How many reads that multimap within one anticodon have different starting positions? What is the maximum distance between them?

# list of reads that have ambiguous starting positions
# (here we assume that if there is a mode the starting position is "unambiguous")
# ie. if positions are 58 58 58 58 58 59 then we would take 58 as the "unambiguous" start

ambig_pos = 0
ambig_pos_reads = []
ambig_distances = []
for key,value in reads_unambig_pos.items():
    if len(set(value)) != 1:
        try:
            x = statistics.mode(value)
        except:
            ambig_pos += 1
            ambig_pos_reads.append(key)
            ambig_distances.append(max(value)-min(value))

anticodons_ambig_pos = dict((k, read_dictionary_anticodon[k]) for k in ambig_pos_reads if k in read_dictionary_anticodon)
ambig_pos_list = []
for key,value in anticodons_ambig_pos.items():
    ambig_pos_list.append("".join(value))

ac_affected = Counter(ambig_pos_list)

summary_stats.write("Number of reads that have an ambiguous position at anticodon level\t"+str(ambig_pos)+"\t"+samp+"\n")
if len(ambig_distances) != 0:
    summary_stats.write("Median distance between ambiguous positions at anticodon level\t"+str(statistics.median(ambig_distances))+"\t"+samp+"\n")
else:
    summary_stats.write("Median distance between ambiguous positions at anticodon level\t"+"NA"+"\t"+samp+"\n")
if len(ambig_distances) != 0:
    summary_stats.write("Biggest distance between ambiguous positions at anticodon level\t"+str(max(ambig_distances))+"\t"+samp+"\n")
else:
    summary_stats.write("Biggest distance between ambiguous positions at anticodon level\t"+"NA"+"\t"+samp+"\n")
summary_stats.close()

ambig_stats = open(dir_and_samp+"_tRNA_ambig_ac_position_dist.tsv", "w")
for key,value in ac_affected.items():
    ambig_stats.write(key+"\t"+str(value)+"\t"+str(ambig_pos)+"\n")
ambig_stats.close()


#1. Реализовать функцию, принимающую два числа (позиционные аргументы) и выполняющую их деление.
# Числа запрашивать у пользователя, предусмотреть обработку ситуации деления на ноль.

def div(*args):

    try:
        a1 = int(input("введите число которое нужно разделить "))
        a2 = int(input("введите число делитель "))
        res = a1 / a2
    except ValueError:
        return 'ошибка'
    except ZeroDivisionError:
        return "на ноль делить нельзя"

    return res

print(f'равно  {div()}')


from setuptools import setup, find_packages
from distutils.sysconfig import get_python_lib
from setuptools.command.install import install as _install
from setuptools.command.develop import develop as _develop
import os

def _post_install(libname, libpath):
    from JumpScale9 import j  # here its still the boostrap JumpScale9

    # remove leftovers
    for item in j.sal.fs.find("/usr/local/bin/", fileregex="js9*"):
        j.sal.fs.remove("/usr/local/bin/%s" % item)

    j.tools.executorLocal.initEnv()
    j.tools.jsloader.generate()

class install(_install):

    def run(self):
        _install.run(self)
        libname = self.config_vars['dist_name']
        libpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), libname)
        self.execute(_post_install, (libname, libpath), msg="Running post install task")

class develop(_develop):

    def run(self):
        _develop.run(self)
        libname = self.config_vars['dist_name']
        libpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), libname)
        self.execute(_post_install, (libname, libpath), msg="Running post install task")

long_description = ""
try:
    from pypandoc import convert
    long_description = convert('README.md', 'rst')
except ImportError:
    long_description = ""

setup(
    name='JumpScale9',
    version='9.4.0',
    description='Automation framework for cloud workloads',
    long_description=long_description,
    url='https://github.com/Jumpscale/core9',
    author='GreenItGlobe',
    author_email='info@gig.tech',
    license='Apache',
    packages=find_packages(),

    # IF YOU CHANGE ANYTHING HERE, LET DESPIEGK NOW (DO NOT INSTALL ANYTHING WHICH NEEDS TO COMPILE)
    install_requires=[
        'GitPython>=2.1.3',
        'click>=6.7',
        'colored_traceback',
        'colorlog>=2.10.0',
        'httplib2>=0.10.3',
        'ipython>=6.0.0',
        'libtmux>=0.7.1',
        'netaddr>=0.7.19',
        'path.py>=10.3.1',
        'pystache>=0.5.4',
        'python-dateutil>=2.6.0',
        'pytoml>=0.1.12',
        'toml',
        'redis>=2.10.5',
        'requests>=2.13.0',
        'future>=0.16.0',
        'watchdog',
        'msgpack-python',
        'npyscreen',
        'pyyaml',
        'pyserial>=3.0'
        'docker>=3',
        'fakeredis',
        'ssh2-python',
        'parallel_ssh>=1.4.0',
        'psutil>=5.4.3',
        'Unidecode>=1.0.22',

    ],
    cmdclass={
        'install': install,
        'develop': develop,
        'developement': develop
    },
    scripts=[
        'cmds/js9',
        'cmds/js9_code',
        'cmds/js9_docker',
        'cmds/js9_doc',
    ],
)


# Copyright (c) 2017 Computer Vision Center (CVC) at the Universitat Autonoma de
# Barcelona (UAB).
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.

import numpy as np
import math
import os

sldist = lambda c1, c2: math.sqrt((c2[0] - c1[0])**2 + (c2[1] - c1[1])**2)
flatten = lambda l: [item for sublist in l for item in sublist]

def get_colisions(selected_matrix, header):

    count_gen = 0
    count_ped = 0
    count_car = 0
    i = 1

    while i < selected_matrix.shape[0]:
        if (selected_matrix[i, header.index('collision_gen')]
                - selected_matrix[(i-10), header.index('collision_gen')]) > 40000:
            count_gen += 1
            i += 20
        i += 1

    i = 1
    while i < selected_matrix.shape[0]:
        if (selected_matrix[i, header.index('collision_car')]
                - selected_matrix[(i-10), header.index('collision_car')]) > 40000:
            count_car += 1
            i += 30
        i += 1

    i = 1
    while i < selected_matrix.shape[0]:
        if (selected_matrix[i, header.index('collision_ped')]
                - selected_matrix[i-5, header.index('collision_ped')]) > 30000:
            count_ped += 1
            i += 100
        i += 1

    return count_gen, count_car, count_ped

def get_distance_traveled(selected_matrix, header):

    prev_x = selected_matrix[0, header.index('pos_x')]
    prev_y = selected_matrix[0, header.index('pos_y')]

    i = 1
    acummulated_distance = 0
    while i < selected_matrix.shape[0]:

        x = selected_matrix[i, header.index('pos_x')]
        y = selected_matrix[i, header.index('pos_y')]
        # Here we defined a maximun distance in a tick, this case 8 meters or 288km/h
        if sldist((x, y), (prev_x, prev_y)) < 800:
            acummulated_distance += sldist((x, y), (prev_x, prev_y))

        prev_x = x
        prev_y = y

        i += 1

    return float(acummulated_distance)/float(100*1000)

def get_out_of_road_lane(selected_matrix, header):

    count_road = 0
    count_lane = 0

    i = 0

    while i < selected_matrix.shape[0]:
        # print selected_matrix[i,6]
        if (selected_matrix[i, header.index('sidewalk_intersect')]
                - selected_matrix[(i-10), header.index('sidewalk_intersect')]) > 0.3:
            count_road += 1
            i += 20
        if i >= selected_matrix.shape[0]:
            break

        if (selected_matrix[i, header.index('lane_intersect')]
                - selected_matrix[(i-10), header.index('lane_intersect')]) > 0.4:
            count_lane += 1
            i += 20

        i += 1

    return count_lane, count_road

def compute_summary(filename, dynamic_episodes):

    # Separate the PATH and the basename
    path = os.path.dirname(filename)
    base_name = os.path.basename(filename)

    f = open(filename, "rb")
    header = f.readline()
    header = header.split(',')
    header[-1] = header[-1][:-2]
    f.close()

    f = open(os.path.join(path, 'details_' + base_name), "rb")
    header_details = f.readline()
    header_details = header_details.split(',')
    header_details[-1] = header_details[-1][:-2]
    f.close()

    data_matrix = np.loadtxt(open(filename, "rb"), delimiter=",", skiprows=1)

    # Corner Case: The presented test just had one episode
    if data_matrix.ndim == 1:
        data_matrix = np.expand_dims(data_matrix, axis=0)

    tasks = np.unique(data_matrix[:, header.index('exp_id')])

    all_weathers = np.unique(data_matrix[:, header.index('weather')])

    reward_matrix = np.loadtxt(open(os.path.join(
        path, 'details_' + base_name), "rb"), delimiter=",", skiprows=1)

    metrics_dictionary = {'average_completion': {w: [0.0]*len(tasks) for w in all_weathers},
                          'intersection_offroad': {w: [0.0]*len(tasks) for w in all_weathers},
                          'intersection_otherlane': {w: [0.0]*len(tasks) for w in all_weathers},
                          'collision_pedestrians': {w: [0.0]*len(tasks) for w in all_weathers},
                          'collision_vehicles': {w: [0.0]*len(tasks) for w in all_weathers},
                          'collision_other': {w: [0.0]*len(tasks) for w in all_weathers},
                          'average_fully_completed': {w: [0.0]*len(tasks) for w in all_weathers},
                          'average_speed': {w: [0.0]*len(tasks) for w in all_weathers},
                          'driven_kilometers': {w: [0.0]*len(tasks) for w in all_weathers}
                          }

    for t in tasks:
        task_data_matrix = data_matrix[
            data_matrix[:, header.index('exp_id')] == t]
        weathers = np.unique(task_data_matrix[:, header.index('weather')])

        for w in weathers:
            t = int(t)

            task_data_matrix = data_matrix[np.logical_and(data_matrix[:, header.index(
                'exp_id')] == t, data_matrix[:, header.index('weather')] == w)]

            task_reward_matrix = reward_matrix[np.logical_and(reward_matrix[:, header_details.index(
                'exp_id')] == float(t), reward_matrix[:, header_details.index('weather')] == float(w))]

            km_run = get_distance_traveled(
                task_reward_matrix, header_details)

            metrics_dictionary['average_fully_completed'][w][t] = sum(
                task_data_matrix[:, header.index('result')])/task_data_matrix.shape[0]

            metrics_dictionary['average_completion'][w][t] = sum(
                (task_data_matrix[:, header.index('initial_distance')]
                 - task_data_matrix[:, header.index('final_distance')])
                / task_data_matrix[:, header.index('initial_distance')]) \
                / len(task_data_matrix[:, header.index('final_distance')])

            metrics_dictionary['driven_kilometers'][w][t]= km_run
            metrics_dictionary['average_speed'][w][t]= km_run/ \
                ((sum(task_data_matrix[:, header.index('final_time')]))/3600.0)

            if list(tasks).index(t) in set(dynamic_episodes):

                lane_road = get_out_of_road_lane(
                    task_reward_matrix, header_details)
                colisions = get_colisions(task_reward_matrix, header_details)

                metrics_dictionary['intersection_offroad'][
                    w][t] = lane_road[0]/km_run
                metrics_dictionary['intersection_otherlane'][
                    w][t] = lane_road[1]/km_run
                metrics_dictionary['collision_pedestrians'][
                    w][t] = colisions[2]/km_run
                metrics_dictionary['collision_vehicles'][
                    w][t] = colisions[1]/km_run
                metrics_dictionary['collision_other'][
                    w][t] = colisions[0]/km_run

    return metrics_dictionary


#!/usr/bin/env python
import sys
import matplotlib.pyplot as plt
import spada

def main():
    trackname = sys.argv[1]
    colors = 'rgbc'
    fig = plt.figure(figsize=(8,6))
    ax = fig.gca()
    for imass,mass in enumerate([1.0,2.0,3.0,4.0]):
        for ifeh,feh in enumerate([0.0,-0.2,-0.5,-0.8]):
            track = spada.SETrack(trackname,mass=mass,FeH=feh)
            if imass == 0:
                ax.plot(track.logTeff_lst,track.logL_lst,
                    color=colors[ifeh],ls='-',label='[Fe/H]=%+3.2f'%feh)
            else:
                ax.plot(track.logTeff_lst,track.logL_lst,
                    color=colors[ifeh],ls='-')
    t1,t2 = ax.get_xlim()
    ax.set_xlim(t2,t1)
    ax.set_xlabel('$\log{T_\mathrm{eff}}\,(\mathrm{K})$')
    ax.set_ylabel('$\log{\left(L/L_\odot\\right)}$')
    leg = ax.legend(loc='lower left')
    for t in leg.get_texts():
        t.set_fontsize(10)
    leg.get_frame().set_alpha(0.1)
    plt.show()

if __name__=='__main__':
    main()


# USAGE
# python train.py --dataset dataset

# import the necessary packages
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import VGG16
from tensorflow.keras.layers import AveragePooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix
from imutils import paths
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import cv2
import os
import code

def preprocess_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return cv2.resize(image, (224, 224))

def get_images():
    base_path = "multiclass_covid19_dataset/covid-chestxray-dataset/"
    metadata = pd.read_csv(base_path + "metadata.csv")
    label_names = [
        'COVID-19', 'ARDS', 'SARS',
        'Pneumocystis', 'Streptococcus',
        'No Finding', 'Pneumonia'
    ]
    label_encoding = list(range(len(label_names)))
    label_mapping = dict(zip(label_names, label_encoding))
    metadata["finding"] = metadata["finding"].map(label_mapping)
    data = []
    labels = []
    ignored = []
    for index, row in metadata.iterrows():
        try:
            image = preprocess_image(base_path + "images/" + row["filename"])
            data.append(image)
            labels.append(str(row["finding"]))
        except:
            continue
    return labels, data, label_mapping

def prepare_data(data, labels):
    """
    convert the data and labels to NumPy arrays while scaling the pixel
    intensities to the range [0, 255]
    """
    data = np.array(data) / 255.0
    labels = np.array(labels)
    labels = to_categorical(labels)
    return data, labels

def fit_model(trainX, trainY, testX, testY, INIT_LR = 1e-3, EPOCHS = 25):
    """
    load the VGG16 network, ensuring the head FC layer sets are left
    off

    place the head FC model on top of the base model (this will become
    the actual model we will train)
    """
    baseModel = VGG16(weights="imagenet", include_top=False,
                      input_tensor=Input(shape=(224, 224, 3)))

    # construct the head of the model that will be placed on top of the
    # the base model
    headModel = baseModel.output
    headModel = AveragePooling2D(pool_size=(4, 4))(headModel)
    headModel = Flatten(name="flatten")(headModel)
    headModel = Dense(64, activation="relu")(headModel)
    headModel = Dropout(0.5)(headModel)
    headModel = Dense(7, activation="softmax")(headModel)
    model = Model(inputs=baseModel.input, outputs=headModel)
    for layer in baseModel.layers:
        layer.trainable = False

    # compile our model
    print("[INFO] compiling model...")
    opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)
    model.compile(
        loss="categorical_crossentropy",
        optimizer=opt,
        metrics=["accuracy"]
    )
    print("[INFO] training head...")
    model_history = model.fit_generator(
        trainAug.flow(trainX, trainY, batch_size=batch_size),
        steps_per_epoch=len(trainX) // batch_size,
        validation_data=(testX, testY),
        validation_steps=len(testX) // batch_size,
        epochs=EPOCHS
    )

    return model, model_history

def get_classes(label_encoding):
    label_decoding = dict(
        zip(label_encoding.values(),
            label_encoding.keys()
        )
    )
    classes = []
    for label in testY.argmax(axis=1):
        classes.append(label_decoding[label])
    classes = set(classes)
    return classes

def get_sensitivity_specificity_accuracy(testY, pred_idxs):
    cm = confusion_matrix(testY.argmax(axis=1), pred_idxs)
    total = sum(sum(cm))
    acc = (cm[0, 0] + cm[1, 1]) / total
    sensitivity = cm[0, 0] / (cm[0, 0] + cm[0, 1])
    specificity = cm[1, 1] / (cm[1, 0] + cm[1, 1])
    print(cm)
    print("acc: {:.4f}".format(acc))
    print("sensitivity: {:.4f}".format(sensitivity))
    print("specificity: {:.4f}".format(specificity))

if __name__ == '__main__':
    plot = "plot.png"
    model_name = "covid19_2.model"
    batch_size = 8
    EPOCHS = 25
    labels, data, label_encoding = get_images()
    data, labels = prepare_data(data, labels)

    trainX, testX, trainY, testY = train_test_split(
        data, labels,
        test_size=0.20,
        stratify=labels,
        random_state=42
    )

    # initialize the training data augmentation object
    trainAug = ImageDataGenerator(rotation_range=15,
                                  fill_mode="nearest")

    model, model_history = fit_model(trainX, trainY, testX, testY, EPOCHS=EPOCHS)
    print("[INFO] evaluating network...")
    predIdxs = model.predict(testX, batch_size=batch_size)
    predIdxs = np.argmax(predIdxs, axis=1)

    classes = get_classes(label_encoding)
    print(classification_report(testY.argmax(axis=1), predIdxs,
                                target_names=classes))

    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, EPOCHS), model_history.history["loss"], label="train_loss")
    plt.plot(np.arange(0, EPOCHS), model_history.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, EPOCHS), model_history.history["accuracy"], label="train_acc")
    plt.plot(np.arange(0, EPOCHS), model_history.history["val_accuracy"], label="val_acc")
    plt.title("Training Loss and Accuracy on COVID-19 Dataset")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend(loc="lower left")
    plt.savefig(plot)
    model.save(model_name, save_format="h5")