#!/opt/opt/bin/python
import sys
import numpy as np
from msmbuilder import Project, Conformation, Serializer
#import DihedralTools
import RefactoredDihedralTools

print("""USAGE:  GetProjectTorsions.py PDBFilename OutFilename


This will calculate all backbone torsion angles for an MSMBuilder project and save the result in OutFilename.
Notes:

1.  Assumes ProjectInfo.h5 is located in current directory, per MSMBuilder convention.
2.  Presently requires you to specify a 'reference' pdb filename.  KAB will fix this in a later version of MSMBuilder2.
3.  Obviously I need to add more error checking.

""")

ProjectFile="./ProjectInfo.h5"
P1=Project.Project.LoadFromHDF(ProjectFile)

PDBFilename=sys.argv[1]
C1=Conformation.Conformation.LoadFromPDB(PDBFilename)

OutFilename=sys.argv[2]

#Result=np.zeros((P1.GetNumTrajectories(),P1["TrajLengths"].max(),C1.GetNumberOfResidues()-1,2),dtype='float32')
for i in range(P1.GetNumTrajectories()):
	print(i)
	R1=P1.LoadTraj(i)
	ans1=RefactoredDihedralTools.GetTorsions(R1,C1,"Phi")
	ans2=RefactoredDihedralTools.GetTorsions(R1,C1,"Psi")
	#ans3=RefactoredDihedralTools.GetTorsions(R1,C1,"Chi")
	#ans=np.concatenate((ans1,ans2,ans3))[:,0]
	ans=np.concatenate((ans1,ans2))[:,0]
	print(ans.shape)
	if i==0:
		Result=np.zeros((P1.GetNumTrajectories(),P1["TrajLengths"].max(),ans.shape[0]),dtype='float32')
	print(Result.shape)
	Result[i,0:ans.shape[1]]=ans.transpose()

Serializer.SaveData(OutFilename,Result)