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

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)
	ans=DihedralTools.GetRama(R1,C1)
	a2=np.transpose(ans,(1,0,2))
	Result[i,0:len(a2)]=a2

Serializer.SaveData(OutFilename,Result)