import os
import sys
import numpy as np
import cPickle
import scipy
from scipy import spatial
from LoadTPLTGT import load_tpl as LoadTPL

InvalidDistance = -1
AdjacentCaCadist = 3.8

def Usage():
	print ' python TPL2PKL.py tplfile [ResDir]'
	print '	This script generates a .tpl.pkl file from an old .tpl file'
	print '	the .tpl file is generated by an old TPL_package '
	print '	the resultant file is named after proteinName.tpl.pkl and saved to ResDir if provided or current work directory'

if __name__ == "__main__":
	if len(sys.argv) < 2:
		Usage()
		exit(1)

	tplfile = sys.argv[1]

	ResDir=os.getcwd()
	if len(sys.argv)>=3:
		ResDir = sys.argv[2]
		if not os.path.isdir(ResDir):
			print 'ERROR: the folder for result saving does not exist: ', ResDir
			exit(1)

	tpl = LoadTPL(tplfile)

	### the old TPL file uses -999 to represent an invalid coordinate, which is not very good
	invalid = [ -999, -999, -999]

	## calculate CaCa distance
	CaCaDist = scipy.spatial.distance_matrix(tpl['Ca'], tpl['Ca']).astype(np.float16)
	CaValid = [ all( abs(ca-invalid)>0.1 ) for ca in tpl['Ca'] ]
	CaCaValid = np.outer(CaValid, CaValid)
	np.putmask(CaCaDist, CaCaValid==False, InvalidDistance)

	## set the Ca-Ca distance of the same residue to 0
	eps = 0.001
        np.fill_diagonal(CaCaDist, eps)

        ## set Ca-Ca distance of two sequentially adajcent residues to 3.8 if it is invalid
        seqLen = CaCaDist.shape[0]
        for i in np.arange(seqLen-1):
        	if CaCaDist[i, i+1] != InvalidDistance:
			continue
                CaCaDist[i, i+1] = AdjacentCaCadist
                CaCaDist[i+1, i] = AdjacentCaCadist


	## calculate CbCb distance
	CbCbDist = scipy.spatial.distance_matrix(tpl['Cb'], tpl['Cb']).astype(np.float16)
	CbValid = [ all( abs(cb-invalid)>0.1 ) for cb in tpl['Cb'] ]
	CbCbValid = np.outer(CbValid, CbValid)
	np.putmask(CbCbDist, CbCbValid==False, InvalidDistance)

	## set the Cb-Cb distance of the same residue to 0
	np.fill_diagonal(CbCbDist, eps)

	distMatrix = dict()
	distMatrix['CaCa'] = CaCaDist
	distMatrix['CbCb'] = CbCbDist

	tpl['atomDistMatrix'] = distMatrix

	savefile = os.path.join(ResDir, tpl['name'] + '.tpl.pkl')
	with open(savefile, 'w') as fh:
		cPickle.dump(tpl, fh, protocol=cPickle.HIGHEST_PROTOCOL)
