#!/usr/bin/env python

import sys,os,go_lib, getopt

# ======================================================================

Usage="""

Usage:
	headgordon.py [-h] [--help] [-e eH] [-k key] [-n nrep] [--nrepeat=nrep] [--gmx] 
	[-D dslist] [--disulfides=dslist] [-p] [-H] file.pdb 

	in which:
		-e eH is the energy scale in kcal/mol
		-k key is a handy key for naming the output files (default = junk)
		-n nrep is the number of NC-linked repeats to make (default = 1)
		-d nongo_dist is the sigma to use for non-go contacts when gamma > 0.0
		--gmx will tell the script to build a force-field for gromacs
			in addition to that for CHARMM [still in testing!].
		-p : the file is in standard RCSB PDB format (default is 
		     to assume CHARMM pdb format

	Only the pdb file is a required argument. 

"""

# parse input
# ======================================================================

if (len(sys.argv)==1):
	print Usage
	sys.exit(0)

Tf = [ 350.0 ] 	# Kelvin
n_repeat = 1
key = "junk"
use_list = 0
use_siglist = 0
gamma = 0.0
average = 0
gmx = 0
symm = 0
water = 0
defbon = 0
scscale = 1.0
bbscale = 1.0
ngdist = 5.5	# default distance (in A) for non-native CA-CA contacts
geomh = 0
stata = 0
orland = 0
disulf_str = ""
sequence = ""
fixrep = 0
charmmpdb = 1
hessian = 0
threebodyexcl = 0
optlist, files = getopt.getopt(sys.argv[1:], 'hpt:k:n:g:d:l:r:w:sS:B:GAD:Os:H3:', \
		['help','tf=','key=','nrepeat=','gamma=','list=',\
		'symmetrize','nongodist=','gmx','water=','scscale=',\
		'scscale=','defbon','replist=','geomh','stat-angle','disulfides=',\
		'orland','seq=','fixrep','pdb'])
for opt, val in optlist:
	if opt in [ "-t", "--tf" ]:
		Tf = map(float, val.split(':'))
	elif opt in [ "-h", "--help" ]:
		print Usage
		sys.exit(0)
	elif opt in [ "-k", "--key" ]:
		key = val
	elif opt in [ "-p", "--pdb" ]:
		charmmpdb = 0
	elif opt in [ "-n", "--nrepeat" ]:
		n_repeat = int(val)
	elif opt in [ "-s", "--symmetrize" ]:
		symm = 1
	elif opt in [ "-g", "--gamma" ]:
		gamma = float(val)
	elif opt in [ "-d", "--nongodist" ]:
		ngdist = float(val)
	elif opt in [ "-l", "--list" ]:	
		go_list_file = val
		use_list = 1
	elif opt in [ "-r", "--replist" ]:	
		sig_list_file = val
		use_siglist = 1
	elif opt in [ "-w", "--water" ]:	
		water = int(val)
#	elif opt in [ "-a", "--average" ]:	
#		average = 1
	elif opt in [ "-S", "--scscale" ]:	
		scscale = float(val)
	elif opt in [ "-B", "--bbscale" ]:	
		bbscale = float(val)
	elif opt in [ "-G", "--geomh" ]:	
		geomh = 1
	elif opt in [ "-A", "--stat-angle" ]:	
		stata = 1
	elif opt in [ "-D", "--disulfides" ]:	
		disulf_str = val
	elif opt in ["--gmx" ]:	
		gmx = 1
	elif opt in ["--defbon" ]:	
		defbon = 1
	elif opt in ["-O", "--orland" ]:	
		orland = 1
	elif opt in ["--fixrep" ]:	
		fixrep = 1
	elif opt in ["-s", "--seq" ]:	
		sequence = val.upper()
	elif opt in ["-H" ]:	
		hessian = 1;
	elif opt in ["-3" ]:	
		threebodyexcl = int(val);
	else:
		print "\nUNKNOWN OPTION: %s\n\n" % opt
		print Usage
		sys.exit(1)
if len(files) < 1 and len(sequence) == 0:
	print Usage
	sys.exit(1)
elif len(files) > 1:
	average = 1

# create output directories and base file names
# ======================================================================

if gamma < 0.00001:
	if n_repeat == 1:
		#outp_base = "go_%s_tf%s" % (key,str(Tf))
		outp_base = "go_%s" % (key)
	else:
		outp_base = "go_%s_nr%i" % (key,n_repeat)
else:
	if n_repeat == 1:
		outp_base = "go_%s_ng%s" % (key,str(gamma))
	else:
		outp_base = "go_%s_ng%s_nr%i" % (key,str(gamma),n_repeat)

if symm:
	outp_base += "_sym"
if water>0:
	outp_base += "_w"

outp_dir = outp_base
if not os.path.exists(outp_dir):
	os.mkdir(outp_dir)

outp_base = outp_dir + "/" + outp_base

# print handy information to *_info.dat file and stdout
# ======================================================================

info = []
info.append("=========================================================\n")
command = reduce(lambda x, y: x+" "+y,sys.argv)
command = command[command.find("go_builder.py"):]
info.append("This go model generated with the command:\n\"%s\"\n" % (command))

if len(sequence) > 0:
	info.append("Building go model from sequence:\n")
	info.append("\t%s\n" % (sequence))
else:
	info.append("Building go model from PDB file %s\n" % (files[0]))
info.append("Target folding temperatures = ")
for t in Tf: 
	info.append("%8.3f " % (t))
info.append(" K\n")

info.append("Key for naming output files = %s\n" % (key))
info.append("Number of NC-linked repeats = %i\n" % (n_repeat))
info.append("Non-go weighting ('gamma') = %8.3f\n" % (gamma))
info.append("Default non-go contact distance = %8.3f\n" % (ngdist))
info.append("Scaling side-chain contacts by %8.3f\n" % (scscale))
info.append("Scaling backbone contacts by %8.3f\n" % (bbscale))
if (gmx):
	info.append("Writing GROMACS files\n")
if symm:
	info.append("Will attempt to symmetrize interactions in oligomer\n")
if average:
	info.append("Will average over multiple structures\n")
else:
	info.append("Will build model from single structure\n")
if water>0:
	info.append("Will place protein in water box of %ix%i cubes of\n"%(water,water))
	info.append("dimension %8.3f A containing %i pseudo waters\n"\
			%(go_lib.std_boxlen,go_lib.std_nwat))
	info.append("i.e. total box side = %8.3f A\n" % (float(water)*go_lib.std_boxlen))

if use_list:
	info.append("Taking go contacts from the following file = %s\n" \
			% (go_list_file))

info.append("Base for output file names = %s\n" % (outp_base))
info.append("=========================================================\n")

info_out = open(outp_base+"_info.dat","w")
for i in info:
	sys.stdout.write("%s" %(i))
	info_out.write("%s" %(i))
info_out.close()

# parse pdb and write out CA-only coords -- if there is a PDB file
# ======================================================================

if len(sequence)==0:	# we are building from pdb
	if average:
		meta_list = []
		xyz_list = []
		for file in files:
			meta_crd, xyz = go_lib.make_coords(open(file).readlines(),charmmpdb)
			meta_list.append(meta_crd)
			xyz_list.append(xyz)
	else:
		meta_crd, xyz = go_lib.make_coords(open(files[0]).readlines(),charmmpdb)
else:
	# if building from sequence ... make 'linear' coords
	# this allows us to continue as we normally would for PDB data
	meta_crd, xyz = go_lib.make_sequ_coords(sequence)

chains = meta_crd.keys()
chains.sort()
nchain = len(chains)

sys.stdout.write("Chains in system:\n")
for c in chains:
	sys.stdout.write("\t%s\n" % (c))

if len(chains) >1 and n_repeat>1:
	sys.stdout.write("\n******************************************************\n")
	sys.stdout.write("WARNING!\n")
	sys.stdout.write("There is more than one chain in your system and you are\n")
	sys.stdout.write("building an N-C linked repeat protein;\n")
	sys.stdout.write("only the first chain will be used!\n")
	sys.stdout.write("******************************************************\n")

if (n_repeat==1):
	oxyz = xyz
	ochains = chains
else:
	oxyz = { chains[0]: go_lib.make_CA_nclink(xyz,n_repeat) }
	ochains = [ chains[0] ]

go_lib.write_CA_pdb(outp_base + ".pdb", oxyz, ochains)

if (gmx):
	go_lib.write_CA_gro(outp_base + ".gro", oxyz, ochains)

if symm:
	# all chains must AT LEAST be the same length!
	l0 = len(xyz[chains[0]])
	for chaini in chains:
		li = len(xyz[chaini])
		for chainj in chains:
			lj = len(xyz[chainj])
			if li != lj:
				sys.stderr.write("  !!! ERROR !!!  \n")
				sys.stderr.write("Cannot symmetrize contacts!\n")
				sys.stderr.write("Length of chain %s = %i\n"\
						% ( chaini, li ))
				sys.stderr.write("Length of chain %s = %i\n"\
						% ( chainj, lj ))
				sys.exit(1)


# solvate protein and write coords if necc.
# ======================================================================

if water > 0:
	# (i) centre protein coordinates
	xyz = go_lib.center_coords(xyz)
	#go_lib.write_CA_pdb(outp_base + "_center.pdb", xyz, chains)
	# (ii) create BIG water box!
	wbox = go_lib.make_wbox(water)
	# (iii) attempt to solvate by deleting waters
	go_lib.solvate(xyz,wbox)
	go_lib.write_CA_pdb(outp_base + "_solv.pdb", xyz, xyz.keys())


# write out sequences
# ======================================================================
seqs = {}
if symm:
	for c in chains:
		seqs[c] = map(lambda x: "G%s"%(x), range(1,len(xyz[c])+1))
	nres_tot = len(xyz[chains[0]])
else:
	off = 0
	nres_tot = 0
	for c in chains:
		seqs[c] = map(lambda x: "G%s"%(x), range(1+off,len(xyz[c])+1+off))
		off += len(xyz[c])
		nres_tot += len(xyz[c])

if water == 0:
	nsolv = 0
else:
	nsolv = len(xyz["SOLV"])

go_lib.write_sequences(outp_base+"_seq.inp", seqs, disulf_str, key, nsolv)

# write out topology
# ======================================================================
if symm:
	chain = chains[0]
	nres = len(xyz[chain])
	residues = []
	for i in range(1,nres+1):
		name = meta_crd[chain][i]["name"]
		residues.append(("G%i"%(i),name))
else:
	residues = []
	idx = 0
	for c in chains:
		nres = len(xyz[c])
		for i in range(1,nres+1):
			idx += 1
			name = meta_crd[c][i]["name"]
			residues.append(("G%i"%(idx),name))

go_lib.write_topology(outp_base+"_top.inp",residues,key)

# calculate parameters
# ======================================================================

if n_repeat > 1:
	link = 1
else:
	link = 0

ntf = 1+nchain*(nchain+1)/2
tfmap = {}
tfbb = 350.0

print nchain
print Tf
print len(Tf),ntf
if len(Tf) == ntf:
	idx = 0
	tfbb = Tf[idx]
	idx += 1 
	for i in range(nchain):
		chaini = chains[i]
		tfmap[chaini] = {}
		tfmap[chaini][chaini] = Tf[idx]
		idx+=1
		for j in range(i+1,nchain):
			chainj = chains[j]
			tfmap[chaini][chainj] = Tf[idx]
			idx+=1
else:
	if len(Tf) == 1:
		tfbb = Tf[0]
	else:
		print "wrong number of Tf's specified!"
		print len(Tf),ntf
		sys.exit(1)
	for i in range(nchain):
		chaini = chains[i]
		tfmap[chaini] = {}
		tfmap[chaini][chaini] = tfbb
		for j in range(i+1,nchain):
			chainj = chains[j]
			tfmap[chaini][chainj] = tfbb

# new ----------------------------------------------
bonds = {}
angles = {}
dihedrals = {}
nonbonded = {}
threebody ={}

NSFAC = 1.0
disulfides = disulf_str.split(":")
for i in range(nchain):
	chaini = chains[i]
	bb_bonds = go_lib.calc_bonds(meta_crd[chaini],link)
	if (defbon==1 or len(sequence)>0):
		bb_bonds = map(lambda x: (x[0],x[1],go_lib.std_bond), bb_bonds)

	bb_angles = go_lib.calc_angles(meta_crd[chaini],link,stata)
	bb_dihedrals = go_lib.setup_dihe(meta_crd[chaini],link)
	nonbonded[chaini] = {}
	#nonbonded[chaini][chaini],tmp_nsfac = go_lib.calc_ncon_intra(meta_crd[chaini],tfmap[chaini][chaini],ngdist,gamma,geomh,fixrep)
	nonbonded[chaini][chaini] = go_lib.calc_ncon_intra(meta_crd[chaini],ngdist,gamma,geomh,fixrep)
	#
	# include disulfides if necc...
	clen = len(chaini)
	my_disulfides = map(lambda x: x[clen:], filter(lambda x: x.find(chaini)==0,disulfides))
	#
	if len(my_disulfides)>0:
		ds_bonds, ds_angles, ds_dihedrals = go_lib.calc_disulf_bonded(meta_crd[chaini],tfbb,my_disulfides)
		bonds[chaini] = bb_bonds+ds_bonds
		angles[chaini] = bb_angles+ds_angles
		dihedrals[chaini] = bb_dihedrals+ds_dihedrals
	else:
		bonds[chaini] = bb_bonds
		angles[chaini] = bb_angles
		dihedrals[chaini] = bb_dihedrals
	#
	# do 3 body if asked
	if threebodyexcl > 0:
		nresi = len(meta_crd[chaini])
		#nnatpair = len(nonbonded[chaini][chaini]['nc'].keys())+\
		#		len(nonbonded[chaini][chaini]['hb'].keys())
		threebody[chaini] = {}
		inte3 = tfmap[chaini][chaini]*go_lib.big_fat_fudge_factor*float(nresi)
		threebody[chaini][chaini] = go_lib.calc_threebody_intra(meta_crd[chaini],\
				inte3,threebodyexcl)
		print len(threebody[chaini][chaini].keys())
	#
	# inter-subunit contacts
	for j in range(i+1,nchain):
		chainj = chains[j]
		#nonbonded[chaini][chainj] = go_lib.calc_ncon_inter(meta_crd[chaini],meta_crd[chainj],\
		#		tfmap[chaini][chainj],ngdist,gamma,geomh,fixrep)
		nonbonded[chaini][chainj] = go_lib.calc_ncon_inter(meta_crd[chaini],meta_crd[chainj],\
				ngdist,gamma,geomh,fixrep)

if symm:
	bonds = go_lib.average_bonded_list(bonds)
	angles = go_lib.average_bonded_list(angles)
	nonbonded =go_lib.average_nonbonded_matrix(nonbonded)
	dihedrals = dihedrals[chains[0]]
else:
	bonds = go_lib.merge_bonded(bonds)
	angles = go_lib.merge_bonded(angles)
	dihedrals = go_lib.merge_dihe(dihedrals)
	if len(Tf) == 1:
		nonbonded,NSFAC =go_lib.merge_nonbonded_matrix(nonbonded,xyz,Tf)
	else:
		nonbonded,NSFAC =go_lib.merge_nonbonded_matrix(nonbonded,xyz,Tf[1:])

if threebodyexcl > 0:
	threebody =go_lib.merge_threebody_matrix(threebody,xyz)

if len(sequence) > 0:
	nonbonded = go_lib.calc_ncon_intra_helical(meta_crd[chains[0]],gamma)

if use_list:
	nonbonded['nc'] = go_lib.make_nc_from_list(meta_crd[chains[0]],go_list_file)
	nonbonded['hb'] = {}
if use_siglist:
	nonbonded['sig'] = go_lib.make_sig_from_list(meta_crd[chains[0]],sig_list_file)

go_map,nongo_map, qlist = go_lib.make_go_nongo_map(nonbonded, \
		nres_tot,scscale,bbscale)
go_and_nongo = go_lib.mergemap(go_map,nongo_map)
GOMODEL_nbfix = go_lib.gotozero(go_map,nongo_map,nres)

if water > 0:
	winte = go_lib.calc_nbond_w(residues)
else:
	winte = {}

# write linear pdb for each chain.
# these will probably overlap, so will need to be translated/rotated
# before energy/force evaluation
# ======================================================================

if len(sequence) == 0:
	for chain in chains:
		go_angles = go_lib.calc_angles(meta_crd[chain],link,0,0)
		bb_bonds = go_lib.calc_bonds(meta_crd[chain],link)
		tmpxyz = { chain: go_lib.make_CA_linear(bb_bonds,go_angles) }
		go_lib.write_CA_pdb(outp_base+"_linear_%s.pdb"%(chain),tmpxyz, [chain])

# write out parameters
# ======================================================================
go_lib.write_parms(outp_base+"_parm.inp",key,bonds,angles,dihedrals,
		nonbonded['sig'],go_and_nongo,NSFAC,tfbb,winte)
go_lib.write_parms(outp_base+"_gomodel_parm.inp",key,bonds,angles,dihedrals,
		nonbonded['sig'],GOMODEL_nbfix,NSFAC,tfbb,winte)
go_lib.write_golist(outp_base+"_gomodel_golist.dat", go_map,
		nres, n_repeat)

if threebodyexcl > 0:
	go_lib.write_threebody("%s_3body_excl%i.dat"%(outp_base,threebodyexcl),threebody)

# calculate and save Hessian if required
# ======================================================================

if hessian:
	sys.stdout.write("Dumping dVdr matrix\n")
	go_lib.dump_dVdr(outp_base+"_dVdr.dat",bonds,go_and_nongo,tfbb)



# WRITE Q-LISTS
# ======================================================================
# global qlist
go_lib.write_qlist(outp_base+"_qlist.dat",qlist,nres,n_repeat)
sys.exit(0)
# for AFM concatemers
if n_repeat>1:
	# write inter-module qlists
	go_lib.write_qlist(outp_base+"_qlist.dat",qlist,nres,n_repeat)
	go_lib.write_qlist_intra(outp_base+"_intra_qlist.dat",qlist,nres,n_repeat)
	go_lib.write_qlist_inter(outp_base+"_inter_qlist.dat",qlist,nres,n_repeat)
	for r in range(n_repeat):
		go_lib.write_module_qlist("%s_r%i_qlist.dat"%(outp_base,r),qlist,nres,r)

# for multiple chains
if nchain>1:
	for i in range(nchain):
		chaini = chains[i]
		for j in range(i,nchain):
			chainj = chains[j]
			go_lib.dump_qlist("%s_%s-%s_qlist.dat"%(outp_base,chaini,chainj), \
				qlists[chaini][chainj])

# write GROMACS topology
# ======================================================================
if (gmx):
	go_lib.write_gro_top(outp_base+"_gromacs.top",key,residues,bonds,
			angles,dihedrals,nonbonded['sig'],
			go_and_nongo,NSFAC,tfbb,seqs,symm,winte,"12_6")
	go_lib.write_gro_top(outp_base+"_gromacs.top",key,residues,bonds,
			angles,dihedrals,nonbonded['sig'],
			go_and_nongo,NSFAC,tfbb,seqs,symm,winte,"12_10")
	go_lib.write_gro_top(outp_base+"_gromacs.top",key,residues,bonds,
			angles,dihedrals,nonbonded['sig'],
			go_and_nongo,NSFAC,tfbb,seqs,symm,winte,"12_10_6")

# write 'qdetails' file
# ======================================================================
go_lib.write_qdetails(outp_base+"_qdetails.dat",meta_crd[chains[0]],nonbonded['details'],
		go_map)
