'''
A "pli" file is an arbitary XML format invented by Barbara Engelhardt
to store GO annotations for proteins.

This parser is here for backwards compatibility. The more common format
is "GAF".
'''
import xml.etree.ElementTree as xml
import json

def annotated_protein(protein_id, evidence_set):
	return {'protein_id':json.dumps(protein_id),
			'evidence_set':json.dumps(evidence_set)}
	
def parser(evidence_file, evidence_constraints):
	'''
	Returns dictionary "evidence_set"
	keys are protein identifiers
	values are dictionaries as created by "annotated_protein(...)" above.
	
	For example, pprint(evidence_set) looks like:
	{'ADA_BOVIN': {'evidence_set': '[["0004000", "TAS"]]',
               'protein_id': '{"uniprot_id": "ADA_BOVIN", "uniprot_acc": "P56658"}'},
	 'ADA_HUMAN': {'evidence_set': '[["0004000", "TAS"]]',
	               'protein_id': '{"uniprot_id": "ADA_HUMAN", "uniprot_acc": "P00813"}'},
	 'ADA_MOUSE': {'evidence_set': '[["0004000", "IMP"]]',
	               'protein_id': '{"uniprot_id": "ADA_MOUSE", "uniprot_acc": "P03958"}'},
	 'ADA_RAT': {'evidence_set': '[["0004000", "IDA"]]',
	             'protein_id': '{"uniprot_id": "ADA_RAT", "uniprot_acc": "Q920P6"}'},
	 'ADA_YEAST': {'evidence_set': '[["0000034", "IMP"]]',
	               'protein_id': '{"uniprot_id": "ADA_YEAST", "uniprot_acc": "P53909"}'},
	 'ADD_ECO57': {'evidence_set': '[["0004000", "TAS"]]',
	               'protein_id': '{"uniprot_id": "ADD_ECO57", "uniprot_acc": "Q8X661"}'},
	 'ADD_ECOL6': {'evidence_set': '[["0004000", "TAS"]]',
	               'protein_id': '{"uniprot_id": "ADD_ECOL6", "uniprot_acc": "Q8FH99"}'},
	 'ADD_ECOLI': {'evidence_set': '[["0004000", "TAS"], ["0005515", "IPI"], ["0005515", "IPI"],
	 								["0005515", "IPI"], ["0005515", "IPI"], ["0005515", "IPI"],
	 								["0005515", "IPI"], ["0005515", "IPI"]]',
	               'protein_id': '{"uniprot_id": "ADD_ECOLI", "uniprot_acc": "P22333"}'},
	 'ADD_STRVG': {'evidence_set': '[["0004000", "TAS"]]'}
	''' 
	evidence_set = {}
	
	tree = xml.parse(evidence_file)
	
	# Get family name
	#family = tree.getroot()
	#FamilyID = tree.getroot().find("FamilyID");
	#FamilyID = FamilyID.text
	
	# Iterate over all proteins
	proteins = tree.getroot().findall("Protein")
	for pro in proteins:
		"""
		# This gets the raw XML of the protein
		old_stdout = sys.stdout
		sys.stdout = mystdout = StringIO()
		xml.dump(pro)
		sys.stdout = old_stdout
		xml_str = mystdout.getvalue()
		"""
		
		# Get identifier(s)
		uniprot_acc_xml = pro.find("ProteinNumber")
		if (uniprot_acc_xml == None):
			continue
		uniprot_id_xml = pro.find("ProteinName")
		if (uniprot_id_xml == None):
			continue 
		
		# Get evidence GO codes if they exist
		go_xml = pro.find("GONumber")
		moc_xml = pro.find("MOC")
		
		p_ev_set = []
		if go_xml is not None \
		and moc_xml is not None:
			go_codes = go_xml.text[1:-1].replace(' ', '').split(',')
			moc_codes = moc_xml.text[1:-1].replace(' ', '').split(',')
			
			p_ev_set = [('GO:'+go_codes[i], moc_codes[i]) \
						for i in range(len(go_codes)) \
							if evidence_constraints.has_key(moc_codes[i]) \
							and evidence_constraints[moc_codes[i]] > 0]
		
		# Only add protein to evidence set if it has annotation after filtering. 
		if len(p_ev_set) > 0:
			evidence_set[uniprot_id_xml.text] = \
							annotated_protein(protein_id={'uniprot_id':uniprot_id_xml.text,
												  		  'uniprot_acc':uniprot_acc_xml.text},
											  evidence_set=p_ev_set)
	
	return evidence_set
