import os.path
import Bio.PDB.mmCIF.MMCIFlex
from UserDict import UserDict


class MMCIF2Dict(UserDict):
	# The token identifiers
	NAME=1
	LOOP=2
	DATA=3
	SEMICOLONS=4    
	DOUBLEQUOTED=5
	QUOTED=6
	SIMPLE=7

	def __init__(self, filename):
		# this dict will contain the name/data pairs 
		self.data={}
		# entry for garbage
		self.data[None]=[]
		if not os.path.isfile(filename):
			raise IOError, "File not found."
		Bio.PDB.mmCIF.MMCIFlex.open_file(filename)
		self._make_mmcif_dict()
		Bio.PDB.mmCIF.MMCIFlex.close_file()

	def _make_mmcif_dict(self):	
		# local copies
		NAME=self.NAME
		LOOP=self.LOOP
		DATA=self.DATA
		SEMICOLONS=self.SEMICOLONS
		DOUBLEQUOTED=self.DOUBLEQUOTED
		QUOTED=self.QUOTED
		SIMPLE=self.SIMPLE
		get_token=Bio.PDB.mmCIF.MMCIFlex.get_token
		# are we looping?
		loop_flag=0
		# list of names in loop
		temp_list=[]
		# last encountered name
		current_name=None
		# get first token/value pair
		token, value=get_token()
		# print token, value
		mmcif_dict=self.data
		# loop until EOF (token==0)
		while token:
			if token==NAME:
				if loop_flag:
					# Make lists for all the names in the loop
					while token==NAME:
						# create  a list for each name encountered in loop
						new_list=mmcif_dict[value]=[]
						temp_list.append(new_list)
						token, value=get_token()  
						# print token, value
					loop_flag=0			
					# nr of data items parsed
					data_counter=0
					# corresponding data name
					pos=0
					nr_fields=len(temp_list)
					# Now fill all lists with the data
					while token>3:
						pos=data_counter%nr_fields
						data_counter=data_counter+1
						temp_list[pos].append(value)
						token, value=get_token()  
						# print token, value
					if pos!=nr_fields-1:
						print "ERROR: broken name-data pair (data missing)!"
					# The last token was not used, so
					# don't set token to None! (this means the 
					# last parsed token goes through the loop again)
				else:	
					# simple name-data pair (no loop)
					# so next token should be the data
					next_token, data=get_token()  
					# print token, value
					mmcif_dict[value]=data
					if next_token<4:
						print "ERROR: broken name-data pair (name-non data pair)!"  
						# print token, value
					else:	
						# get next token
						token=None
			elif token==LOOP:
				loop_flag=1
				temp_list=[]
				# get next token
				token=None
			elif token==DATA:
				mmcif_dict[value[0:5]]=value[5:]
				token=None
			else:
				# we found some complete garbage
				print "ERROR: broken name-data pair (missing name)!"
				print token, value
				mmcif_dict[None].append(value)
				# get next token
				token=None
			if token==None:
				token, value=get_token()
				# print token, value

	def __getitem__(self, key):
		return self.data[key]


if __name__=="__main__":

	import sys

	if len(sys.argv)!=2:
		print "Usage: python MMCIF2Dict filename."

	filename=sys.argv[1]	

	mmcif_dict=MMCIF2Dict(filename)

	input=""
	print "Now type a key ('q' to end):"
	while(input!="q"):
		input=raw_input("MMCIF dictionary key ==> ") 	
		if input=="q":
			sys.exit()
		try:
			value=mmcif_dict[input]
			if type(value)==type([]):
				for item in value:
					print item
			else:
				print value
		except KeyError:
			print "No such key found."

