# This is a python script  to extract N and NH 
# chemical shifts from BMRB files and 
# convert them to Sparky peak lists. Version 2 of this script 
# can also apply a systematic offset to 15N and 1H chemical shifts 
# Version # 3 (current) of this script can also apply a systematic 
# offset to residue numbers. 
# Version #4 has better exception handling of the output filenames

from string import split
from os import linesep
import os

print """\n README!!! README!!! README!!! README!!! README!!! README!!!


 1) When you enter filenames, please make sure that 
    filenames do not have spaces.
    
 2) This is a python script  to extract N and NH chemical shifts from BMRB files and
    convert them to Sparky peak lists.
 
    
     
 README END!!! README END!!! README END!!! README END !!!  """
 
from string import split
from os import linesep
import sys
import os

def openfile(phrase):
	done_input1_switch=0
	while not done_input1_switch:
		input_file0=raw_input(phrase)
		input_file0s=split(input_file0)
		input_file1=input_file0s[0]
		input_file1_exist=os.access(input_file1, os.F_OK)
		input_file1_read=os.access(input_file1, os.R_OK)
		if input_file1_exist==0:
			input_file1_exist_resolve=raw_input("""\n ERROR: File with this name does not exist.
 Please enter 1 to try again or 0 to exit ther program\n: """)
			if  input_file1_exist_resolve=='0':
				sys.exit(1)
			else:
				continue

		elif  input_file1_read==0:
			input_file1_read_resolve=raw_input("""\n ERROR: File with this name does not exist. 
 Please enter 1 to try again or 0 to exit ther program\n: """)
			if  input_file1_read_resolve=='0':
				sys.exit(1)
			else:
				continue

		elif (input_file1_exist==1) and (input_file1_read==1):
			done_input1_switch=1
	return input_file1
input_file1a=openfile("""\n Enter the name of BMRB file to be converted to Sparky peak list.\n : """)


done_output1_switch=0
while not done_output1_switch:
	output_file0=raw_input("""\n Enter the name of a Sparky peak file with N and NH chemical shifts 
 to be created from BMRB file.\n : """)
	output_file0s=split(output_file0)
	output_file1=output_file0s[0]
        output_dir1_exist=os.access(output_file1, os.F_OK)
        output_dir1_write=os.access(output_file1, os.W_OK)
        if output_dir1_exist!=0:
		output_dir1_exist_resolve_done=0
                while not output_dir1_exist_resolve_done:
			output_dir1_exist_resolve=raw_input("""\n ERROR: Directory or file with this name already exists.
 Please enter 1 to try again, 2 to overwrite it or 0 to exit the program\n: """)
        	        if  output_dir1_exist_resolve=='0':
				sys.exit(1)
        	        elif  (output_dir1_exist_resolve=='2') and (output_dir1_write==1):
				dir_out_test=os.path.isdir(output_file1)
				if dir_out_test==1:
						shutil.rmtree(output_file1)
				else:
					os.remove(output_file1)
				output_dir1_exist_resolve_done=1
				done_output1_switch=1
			elif (output_dir1_exist_resolve=='2') and (output_dir1_write==0):
				output_dir1_write_resolve_done=0
                		while not output_dir1_write_resolve_done:
                       			output_dir1_write_resolve=raw_input("""\n ERROR: You do not have permissions to overwrite it. Please enter 1 to try again, 
 2 if you have changed its permission and want to overwrite it or 0 to exit the program\n: """)
		                        if  output_dir1_write_resolve=='0':
		                                sys.exit(1)
		                        elif  output_dir1_write_resolve=='2':
		                                output_dir1_write=os.access(output_file1, os.W_OK)
		                                if output_dir1_write==0:
		                                        print "\n \n \n WARNING!!!: You still do not have persmission to overwrite the directory (or file)! Try again."
						else:
							dir_out_test=os.path.isdir(output_file1)
							if dir_out_test==1:
								shutil.rmtree(output_file1)
								output_dir1_exist_resolve_done=1
								done_output1_switch=1
								output_dir1_write_resolve_done=1
							else:
								os.remove(output_file1)
                        		        	        done_output1_switch=1
                                        			output_dir1_write_resolve_done=1
								output_dir1_exist_resolve_done=1
					elif output_dir1_write_resolve not in ["0","1","2"]:
						print ("""\n ERROR: Invalid input!!! Try again. \n: """)

	             			else:
 						output_dir1_write_resolve_done=1
						output_dir1_exist_resolve_done=1
			elif output_dir1_exist_resolve not in ["0","1","2"]:
				print ("""\n ERROR: Invalid input!!! Try again. \n: """)
			else:
				output_dir1_exist_resolve_done=1

	else:
		done_output1_switch=1

def option(question1):
	option_switch_done=0
	while not option_switch_done:
		option_switch=raw_input(question1)
		if option_switch not in ["1","0"]:
			print "\n \n Wrong choice! Try again. \n \n"
			continue
		else: 
			option_switch_done=1
	return option_switch

corr_switch=option("""\n Do you want to apply systematic correction to chemical shift values?.\n 
 If yes, enter 1.\n If not, enter 0 \n : """)

def enter_number(phrase2):
	enter_done=0
	while not enter_done:
		number=raw_input(phrase2)
		try:
			float(number)
		except ValueError:
			print """ \n ERROR: You did not enter a number!!! You entered '%s'. Try again. \n """ % (number)
			continue	
		else:
		        number_int=eval(number)
			enter_done=1
			print number_int

	return number_int

if corr_switch == "1":
	corr_n_int=enter_number("""\n Enter a value of systematic correction of 15N chemical shift. \n : """)
	corr_h_int=enter_number("""\n Enter a value of systematic correction of 1H chemical shift. \n : """)


res_corr_switch=option("""\n Do you want to apply systematic correction to residue numbers?.\n 
 If yes, enter 1.\n If not, enter 0 \n : """)
if res_corr_switch == "1":
	corr_res_int=enter_number("""\n Enter a value of systematic correction of residue number. \n : """)


input1=open(input_file1a, 'r')
output1 = open(output_file1, 'w+' )
unassigned_out2="%s_%s" %(output_file1,"unassigned")
output2 = open(unassigned_out2, 'w+' )

aLine_1 = input1.readline()
cells_1=split(aLine_1)
residue_num=cells_1[1]
input1.close()
N_shift='999.000'
HN_shift='999.000'

input1=open(input_file1a, 'r')
done_1=0
while not done_1:
	aLine_1 = input1.readline()
	if aLine_1 !="":
	  	cells_1=split(aLine_1)
		if cells_1[1]==residue_num:
			if cells_1[3]=="N":
				N_shift=cells_1[5]
				residue_name=cells_1[2]
				residue_num=cells_1[1]
			elif cells_1[3]=="H":
				HN_shift=cells_1[5]				
				residue_name=cells_1[2]
				residue_num=cells_1[1]
		if cells_1[1]!=residue_num:
			if (N_shift!='999.000') and (HN_shift!='999.000'):
				res_name_out="%s%sN-H" %(residue_name,residue_num)
				if corr_switch == "1":
					N_shift_int=eval(N_shift)
					N_shift_corr=N_shift_int+corr_n_int
					N_shift="%.3f" %(N_shift_corr)
					HN_shift_int=eval(HN_shift)
					HN_shift_corr=HN_shift_int+corr_h_int
					HN_shift="%.3f" %(HN_shift_corr)
				cells_1_out=res_name_out+' '+N_shift+' '+HN_shift+linesep
				cells_1_write=[cells_1_out]
				output1.writelines(cells_1_write)
				N_shift='999.000'
				HN_shift='999.000'
				
			else:
				res_name_out="B%sN-H" %(residue_num)
				cells_2_out=res_name_out+' '+N_shift+' '+HN_shift+linesep
				cells_2_write=[cells_2_out]
				output2.writelines(cells_2_write)
			residue_num=cells_1[1]
			if cells_1[3]=="N":
				N_shift=cells_1[5]
				residue_name=cells_1[2]
				residue_num=cells_1[1]
			elif cells_1[3]=="H":
				HN_shift=cells_1[5]				
				residue_name=cells_1[2]
				residue_num=cells_1[1]
	else:
		done_1=1

if (N_shift!='999.000') and (HN_shift!='999.000'):
	res_name_out="%s%sN-H" %(residue_name,residue_num)
	if corr_switch == "1":
		N_shift_int=eval(N_shift)
		N_shift_corr=N_shift_int+corr_n_int
		N_shift="%.3f" %(N_shift_corr)
		HN_shift_int=eval(HN_shift)
		HN_shift_corr=HN_shift_int+corr_h_int
		HN_shift="%.3f" %(HN_shift_corr)
	cells_1_out=res_name_out+' '+N_shift+' '+HN_shift+linesep
	cells_1_write=[cells_1_out]
	output1.writelines(cells_1_write)
else:
	res_name_out="B%sN-H" %(residue_num)
	cells_2_out=res_name_out+' '+N_shift+' '+HN_shift+linesep
	cells_2_write=[cells_2_out]
	output2.writelines(cells_2_write)

print 'Done!'