#!/usr/bin/python

import sys

def getvars(arguments):
	switches={}

	version='0.1'
	
	try:
		if "-i" in arguments:
			switches['in_one']=arguments[arguments.index('-i')+1]
			print 'Input: %s '% (switches['in_one'])
		else:
			arguments="--help";
	except:
		arguments="--help";
		
	try:
		if "-o" in arguments:
			switches['o']=arguments[arguments.index('-o')+1].lower()
			print 'Output: %s'% switches['o']
		else:
			arguments="--help";
	except:
		arguments="--help";

	try:
		if "-w" in arguments:
			switches['w']=float(arguments[arguments.index('-w')+1])
			print 'Weighting: %i'% switches['w']
		else:
			print 'Assuming no weighting'
			switches['w']=1.0;
	except:
		switches['w']=1.0;

	doexit=0
	try:
		if ("-h" in arguments) or ("--help" in arguments):
			print '\t\t nebinterpolate version %s' % version
			print 'Creates interpolated structures'
			print 'from an multixyz file'
			print '--input \t-i\t multi-xyz file to morph'
			print '--output\t-o\t output file'
			print '--weight\t-w\t weight first structure vs second one (1=average; 0=start; 2=end)'
			print 'Exiting'
			doexit=1
	except:
		a=0 # do nothing
	if doexit==1:
		sys.exit(0)

	return switches

def getcmpds(switches):
	
	cmpds={}
	
	g=open(switches['in_one'],'r')	
	n=0
	xyz=[]
	atoms=[]
	structure_id=1

	for line in g:

		try:
			if len(xyz)==cmpds['atoms_'+str(structure_id)]:
				cmpds['coords_'+str(structure_id)]=xyz
				cmpds['elements_'+str(structure_id)]=atoms			
				structure_id+=2
				n=0
				atoms=[]
				xyz=[]
		except:
			pass

		n+=1
		line=line.rstrip('\n')
#		print structure_id, line
			
		if n==1:
			cmpds['atoms_'+str(structure_id)]=int(line)
		elif n==2:
			cmpds['title_'+str(structure_id)]=line
		else:
			line=line.split(' ')
			line=filter(None,line)
			xyz+=[[float(line[1]),float(line[2]),float(line[3])]]
			atoms+=[line[0].capitalize()]
	g.close

	cmpds['coords_'+str(structure_id)]=xyz
	cmpds['elements_'+str(structure_id)]=atoms			
	cmpds['w']=switches['w']
	cmpds['structures']=(structure_id)
	
	return cmpds

def morph(cmpds,structure_id):
	coords_one=cmpds['coords_'+str(structure_id)]
	coords_two=cmpds['coords_'+str(structure_id+2)]
	
	coords_morph=[]
	coords_diff=[]
	for n in range(0,cmpds['atoms_'+str(structure_id)]):
		print coords_one[n][0], coords_two[n][0],(coords_two[n][0]-coords_one[n][0])/2.0
		morph_x=coords_one[n][0]+cmpds['w']*(coords_two[n][0]-coords_one[n][0])/2.0
		morph_y=coords_one[n][1]+cmpds['w']*(coords_two[n][1]-coords_one[n][1])/2.0
		morph_z=coords_one[n][2]+cmpds['w']*(coords_two[n][2]-coords_one[n][2])/2.0
		diff_x=coords_two[n][0]-coords_one[n][0]
		diff_y=coords_two[n][1]-coords_one[n][1]
		diff_z=coords_two[n][2]-coords_one[n][2]
		coords_morph+=[[morph_x,morph_y,morph_z]]
		coords_diff+=[[diff_x,diff_y,diff_z]]

	cmpds['atoms_'+str(structure_id+1)]=cmpds['atoms_'+str(structure_id)]
	cmpds['elements_'+str(structure_id+1)]=cmpds['elements_'+str(structure_id)]
	cmpds['title_'+str(structure_id+1)]='structure '+str(structure_id+1)
	cmpds['coords_'+str(structure_id+1)]=coords_morph

	return cmpds

def genxyzstring(coords,element):
	x_str='%10.5f'% coords[0]
	y_str='%10.5f'% coords[1]
	z_str='%10.5f'% coords[2]
	
	xyz_string=element+(3-len(element))*' '+10*' '+\
	(8-len(x_str))*' '+x_str+10*' '+(8-len(y_str))*' '+y_str+10*' '+(8-len(z_str))*' '+z_str+'\n'
	
	return xyz_string

def writemorph(cmpds,outfile):
	g=open(outfile,'w')	

	for m in range(1,cmpds['structures']+1):
		g.write(str(cmpds['atoms_'+str(m)])+'\n')
		g.write(str(cmpds['title_'+str(m)])+'\n')
		for n in range(0,cmpds['atoms_'+str(m)]):
			coords=cmpds['coords_'+str(m)][n]
			g.write(genxyzstring(coords, cmpds['elements_'+str(m)][n]))
	g.close

	return 0

if __name__=="__main__":
	arguments=sys.argv[1:len(sys.argv)]
	switches=getvars(arguments)
	cmpds=getcmpds(switches)

#check that the structures are compatible
	for n in range(1,cmpds['structures'],2):

		if cmpds['atoms_'+str(n)]!=cmpds['atoms_'+str(n+2)]:
			print 'The number of atoms differ. Exiting'
			sys.exit(1)
		elif cmpds['elements_'+str(n)]!=cmpds['elements_'+str(n+2)]:
			print 'The types of atoms differ. Exiting'
			sys.exit(1)
		cmpds=morph(cmpds,n)
	success=writemorph(cmpds,switches['o'])
	if success==0:
		print 'Conversion seems successful'