'''
This script obeys the file structure generated by TCGA for Level 3 Methylome data.
Currently, it is set up to read in values from platform: Illumina Human Methylation 27K.
This script must be reset for each cancer type.

Version: TCGA has two file formats that must be processed differently. User must identify version by inspection.
Version A has the following first line:
['barcode', 'probe name', 'beta value', 'gene symbol', 'chromosome', 'position']
Version B has the following first line:
['Hybridization REF', 'TCGA-BP-4335-01A-01D-1284-05', 'TCGA-BP-4335-01A-01D-1284-05', 'TCGA-BP-4335-01A-01D-1284-05', 'TCGA-BP-4335-01A-01D-1284-05']

Up_path: this is the file path in which manifest.txt is found

Out_path: this is the file path to which output files are sent

Use load_manifest() to generate a dictionary that maps each barcode to its associated file name.

Use build_probe_dict() to build the primary data object. It generates a dictionary with 
keys: barcode  
values: list of beta values in order given by probes_27.txt.

Use betas_to_list() to generate a single list of all raw beta values across all samples.

Use betas_to_list_by_island() to generate a 2 lists:
(1) of all raw beta values associated with CGIs across all samples 
(2) of all raw beta values not associated with CGIs across all samples

probes_27.txt
Object type: list of all probe names in sorted order. This is the primary indexing tool for accessing data in probe_dict objects

island_status.txt
Object type: list indexed by probes_27.txt that has value True if corresponding probe is found in a CGI and False otherwise

'''

import sys, os, time
import numpy as np

###################### MUST BE SET BEFORE USING OTHER METHODS ############################
version = "None"
up_path = "None"
out_path = "None"

# User must define file format version: "A" or "B"
def set_version(v):
	global version
	if v == "A" or v == "B":
		version = v
	else:
		print "Error: please enter valid file format version: A or B"
	return

# User must set file path in which manifest is found
def set_up_path(path):
	global up_path
	up_path = path
	return

# User must set file path in which output files should be placed
def set_out_path(path):
	global out_path
	out_path = path
	return

#################################### DIAGNOSTICS #########################################
def get_up_path():
	global up_path
	return up_path

def get_out_path():
	global out_path
	return out_path	

def get_barcodes(probe_dict,file_name):
	global out_path
	barcodes = []
	f_out = open(out_path+file_name,"w")
	f_out.write("\t".join(probe_dict.keys()))
	return probe_dict.keys()
		
################################# GLOBAL OBJECTS #########################################
probes_27 = open(str(os.getcwd()).replace("Scripts","/Annotation/probes_27.txt"),"r").readline().strip().split('\t')
island_status = open(str(os.getcwd()).replace("Scripts","/Annotation/island_status.txt"),"r").readline().strip().split('\t')

############################### DATA EXTRACTION METHODS ##################################
'''
Reads file manifest and returns a dictionary
Obeys file structure generated from TCGA data matrix
Keys: barcode
Values: file names
'''
def load_manifest():
	global up_path
	manifest = {}
	man_file = open(up_path+"file_manifest.txt","r").readlines()
	for line in man_file:
		spl = line.strip().split('\t')
		plat = spl[2]
		level = spl[3]
		if (plat == 'HumanMethylation27' and level == '3'):
			brc = spl[5]
			manifest[brc] = spl[6]
	return manifest

'''
Given a float raw beta value, returns a discretized float value.
'''
def discretize(beta):
	if (np.isnan(beta)):
		return np.nan
	elif (beta <= 0.2):
		return 0.0
	elif (beta > 0.6):
		return 2.0
	else:
		return 1.0

'''
For use with version A file type. 
Arguments:
fileName: file name without file path
discretize: True or False
Returns: dictionary of probe: float beta value entries
'''
def extract_betas_A(fileName,discrete):
	global up_path
	if discrete:
		file = open(up_path+"/DNA_Methylation/JHU_USC__HumanMethylation27/Level_3/"+fileName,"r").readlines()
		betas = {}
		for line in file:
			line = line.strip().split('\t')
			if line[2] == "beta value":
				continue
			else:
				if line[2] == "NA":
					betas[line[1]] = np.nan
				else:
					betas[line[1]] = discretize(float(line[2]))
		return betas
	elif not discrete:
		file = open(up_path+"/DNA_Methylation/JHU_USC__HumanMethylation27/Level_3/"+fileName,"r").readlines()
		betas = {}
		for line in file:
			line = line.strip().split('\t')
			if line[2] == "beta value":
				continue
			else:
				if line[2] == "NA":
					betas[line[1]] = np.nan
				else:
					betas[line[1]] = float(line[2])
		return betas
	else:
		sys.exit("Error: please enter valid value for discretize: True or False")

'''	
For use with version B file type.	
Arguments:
fileName: file name without file path
discretize: True or False
Returns: dictionary of probe: beta value entries
'''
def extract_betas_B(fileName,discrete):
	global platform, up_path
	if discrete:
		file = open(up_path+"/DNA_Methylation/JHU_USC__HumanMethylation27/Level_3/"+fileName,"r").readlines()
		betas = {}
		for line in file:
			line = line.strip().split('\t')
			if line[0] == "Hybridization REF" or line[0] == "Composite Element REF":
				continue
			else:
				if line[1] == "NA":
					betas[line[0]] = np.nan
				else:
					betas[line[0]] = discretize(float(line[1]))
		return betas
	elif not discrete:
		file = open(up_path+"/DNA_Methylation/JHU_USC__HumanMethylation27/Level_3/"+fileName,"r").readlines()
		betas = {}
		for line in file:
			line = line.strip().split('\t')
			if line[0] == "Hybridization REF" or line[0] == "Composite Element REF":
				continue
			else:
				if line[1] == "NA":
					betas[line[0]] = np.nan
				else:
					betas[line[0]] = float(line[1])
		return betas
	else:
		sys.exit("Error: please enter valid value for discretize: True or False")

'''
Arguments:
manifest: dictionary
file_name: file name without file path
discrete: True or False; discretize or not
Returns: A dictionary with keys: barcode, values: list of float beta values indexed by probes_27.txt

'''
def build_probe_dict(manifest, file_name, discrete):
	
	global out_path, probes_27, version
	samples = sorted(manifest.keys())
	probe_dict = {}
	
	for s in samples:
		if version == "A":
			betas = extract_betas_A(manifest[s],discrete)
			sorted_betas = []
			for p in probes_27:
				if p in betas:
					sorted_betas.append(betas[p])
				else:
					sorted_betas.append(np.nan)
			probe_dict[s] = np.array(sorted_betas)
		else:
			betas = extract_betas_B(manifest[s],discrete)
			sorted_betas = []
			for p in probes_27:
				if p in betas:
					sorted_betas.append(betas[p])
				else:
					sorted_betas.append(np.nan)
			probe_dict[s] = np.array(sorted_betas)
	if file_name == "NA":
		return probe_dict
	else:
		f_out = open(out_path+file_name,"w")
		f_line = ["Samples"]+probes_27
		f_out.write('\t'.join(f_line)+'\n')
		barcodes = sorted(probe_dict.keys())
		for b in barcodes:
			f_out.write(b+'\t')
			f_out.write('\t'.join(str(a) for a in probe_dict[b]))
			f_out.write('\n')
		f_out.close()
		return probe_dict
	
	

def remove_SNPs(probe_dict,file_name):
	global out_path, probes_27
	
	# Convert dictionary to pandas dataframe
	pd_1 = pd.DataFrame.from_dict(probe_dict,orient="index")
	pd_1.columns = probes_27
	
	# Remove any columns where all values are NA
	pd_1 = pd_1.dropna(axis=1,how='all')
	
	if file_name == "NA":
		return pd_1
	else:
		# Write pandas object to txt 
		f_out = open(out_path+file_name,"w")
		f_line = ["Samples"] + pd_1.columns.tolist()
		f_out.write('\t'.join(f_line)+'\n')
		barcodes = sorted(pd_1.index.tolist())
		for sample in barcodes:
			f_out.write(sample+'\t')
			f_out.write('\t'.join(str(a) for a in pd_1.loc[sample]))
			f_out.write('\n')
		f_out.close()
		return pd_1


'''
Arguments:
manifest: dictionary
file_out: file name without file path
Returns: A list of raw float beta values associated with all probes across all samples
Pickles output
'''
def betas_to_list(manifest, file_out):
	global out_path, probes_27, version
	samples = sorted(manifest.keys())
	all_betas = [np.nan]*(len(samples)*len(probes_27))
	index = 0
	for s in samples:
		if version == "A":
			betas = extract_betas_A(manifest[s],False)
			betas = betas.values()
			for i in range(0,len(betas)):
				if betas[i] == "NA":
					all_betas[index+i] = np.nan
				else:
					all_betas[index+i] = float(betas[i])
			index += len(probes_27)
		else:
			betas = extract_betas_B(manifest[s],False)
			betas = betas.values()
			for i in range(0,len(betas)):
				if betas[i] == "NA":
					all_betas[index+i] = np.nan
				else:
					all_betas[index+i] = float(betas[i])
			index += len(probes_27)
	if file_out == "NA":
		return all_betas
	else:
		#pickle.dump(all_betas,open(out_path+file_out,"wb"))
		return all_betas

'''
Arguments:
manifest: dictionary
file_out1: file name without file path for betas associated with CGIs
file_out2: file name without file path for betas not associated with CGIs
Returns: Nothing
Pickles both lists
'''
def betas_to_list_by_island(manifest, file_out1, file_out2):
	global out_path, probes_27, version, island_status
	num_island = sum(np.array(island_status) == "True")
	num_nonisland = sum(np.array(island_status) == "False")
	samples = sorted(manifest.keys())
	all_betas_island = [np.nan]*(len(samples)*num_island)
	all_betas_nonisland = [np.nan]*(len(samples)*num_nonisland)
	
	ind_island = 0
	ind_nonisland = 0
	for s in samples:
		if version == "A":
			betas = extract_betas_A(manifest[s],False)
			betas = betas.values()
			for i in range(0,len(betas)):
				if island_status[i] == "True":
					if betas[i] == "NA":
						all_betas_island[ind_island] = np.nan
					else:
						all_betas_island[ind_island] = float(betas[i])
					ind_island += 1
				else:
					if betas[i] == "NA":
						all_betas_nonisland[ind_nonisland] = np.nan
					else:
						all_betas_nonisland[ind_nonisland] = float(betas[i])
					ind_nonisland += 1
		else:
			betas = extract_betas_B(manifest[s],False)
			betas = betas.values()
			for i in range(0,len(betas)):
				if island_status[i] == "True":
					if betas[i] == "NA":
						all_betas_island[ind_island] = np.nan
					else:
						all_betas_island[ind_island] = float(betas[i])
					ind_island += 1
				else:
					if betas[i] == "NA":
						all_betas_nonisland[ind_nonisland] = np.nan
					else:
						all_betas_nonisland[ind_nonisland] = float(betas[i])
					ind_nonisland += 1
	
	# Write lists to txt files
	f_out = open(file_out1,"w")
	f_out.write('\t'.join(str(i) for i in all_betas_island))
	f_out.close()
	f_out = open(file_out2,"w")
	f_out.write('\t'.join(str(i) for i in all_betas_nonisland))
	f_out.close()
	return 

#################################### MAIN METHOD #########################################

main_up_path = os.getcwd().replace("Scripts","Input/")
main_out_path = os.getcwd().replace("Scripts","Output/")

cancers = ["brca","coad","gbm","kirc","luad","lusc","ov","ucec"]
cancers_A = ["brca","coad","gbm","lusc","ov"]
cancers_B = ["kirc","luad","ucec"]

for cancer in cancers_A:
	print "Started: "+str(cancer)
	
	set_version("A")
	set_up_path(main_up_path+str(cancer)+"/")
	print get_up_path()
	set_out_path(main_out_path+str(cancer)+"/")
	print get_out_path()
	
	manifest = load_manifest()
	probe_dict = build_probe_dict(manifest,"all_raw_betas.txt",False)
	probe_dict_r = remove_SNPs(probe_dict,"probe_dict_dis.txt")
	betas_to_list_by_island(manifest,"betas_raw_island.txt","betas_raw_nonisland.txt")
	
	print "Finished: "+str(cancer)

for cancer in cancers_B:
	print "Started: "+str(cancer)
	
	set_version("B")
	set_up_path(main_up_path+str(cancer)+"/")
	set_out_path(main_out_path+str(cancer)+"/")
	
	manifest = load_manifest()
	probe_dict = build_probe_dict(manifest,"all_raw_betas.txt",True)
	probe_dict_r = remove_SNPs(probe_dict,"probe_dict_dis.txt")
	betas_to_list_by_island(manifest,"betas_raw_island.txt","betas_raw_nonisland.txt")

	print "Finished: "+str(cancer)

