import torch


def load_corrupt(BASE_DIR, date_to_load):
	"""
	Loads corrupted genome vectors from file
	
	Arguments:
		BASE_DIR (str) -- path to working directory
		date_to_load (str) -- name of file to load
		
	Returns:
		corrupted_train (tensor) -- corrupted training data. Rows are genomes, columns are genes. 1 denotes gene encoded by genome, 0 denotes absence.
		c_train_genomes (list) -- tnum corresponding to each row (genome) of corrupted_train
		train_input_mods (list of lists) -- lists of the mods that were retained during the corruption process (in same order as genome rows / c_train_genomes)
		corrupted_test (tensor) -- corrupted test data. Rows are genomes, columns are genes. 1 denotes gene encoded by genome, 0 denotes absence.
		c_test_genomes (list) -- tnum corresponding to each row (genome) of corrupted_test
		test_input_mods (list) -- lists of the mods that were retained during the corruption process (in same order as genome rows / c_test_genomes)
	"""
	print("Loading corrupted genome vectors from "+date_to_load)
	corrupted_train = torch.load(BASE_DIR+"corrupted_train_"+date_to_load+".pt")
	c_train_genomes = torch.load(BASE_DIR+"c_train_genomes_"+date_to_load+".pt")
	corrupted_test = torch.load(BASE_DIR+"corrupted_test_"+date_to_load+".pt")
	c_test_genomes = torch.load(BASE_DIR+"c_test_genomes_"+date_to_load+".pt")
	train_input_mods = torch.load(BASE_DIR+"train_input_mods_"+date_to_load+".pt")
	test_input_mods = torch.load(BASE_DIR+"test_input_mods_"+date_to_load+".pt")
	
	return corrupted_train, c_train_genomes, train_input_mods, corrupted_test, c_test_genomes, test_input_mods