from numpy import *
from pymol import cmd, stored
import pickle
import ReadAmberFiles as raf

#########################################################################
# GenNM - generate quasi-harmonic modes for visualization in pymol.     #
# Variables:                                                            #
#   obj: name of structure loaded into pymol. (usually the avg. coords) #
#   nmids: normal modes to visualize (ints separated by spaces)         #
#   mode: 1-start from mass weighted covariance matrix file             #
#         2-start from amber evecs file                                 #
#         3-start from pickle file of modes and frequencies generated by#
#                 previous run of GenNM                                 # 
#   inp: name of input file (either covar (mode1), evecs(mode2),        #
#             pickle(mode3)                                             #
#   amp: relative amplitude, default 1                                  #
#   nframes: # of frames to create for each mode visualization          #
#   dump_file: pickle modes and freq to this file                       #
#                                                                       #
# Output: creates a new pymol object for each mode specified in nmids   #
#                                                                       #
# Usage (commands in pymol terminal):                                   #
#    load na.pdb, NA                                                    #
#    run pymol_nmodes.py                                                #
#    gennm NA, 2, mode=1, inp=mwcvmat.dat                               #
#    gennm NA, 0 1 2, mode=2, inp=evecs.dat, dump_file=mypicklefile     #
#    gennm NA, 0 1 2, mode=3, inp=mypicklefile, amp=2, nframes=50       #
#                                                                       #
# This script has heavily relied on ideas from:                         #
#	Peter Mamonov (http://erg.biophys.msu.ru/wordpress/archives/197)    #
#=======================================================================#


#routine for reading evecs files
def ReadEvecs(inp):
	f=open(inp,'r')
	n=f.readline().split()
	n=int(n[n.index('nmodes')+1])
	natoms=int(f.readline().split()[0])
	w=zeros((n))
	v=zeros((n,natoms))
	n_vec=0
	n_partial=0
	width=len(f.readline().strip().split())
	while True:
		line=f.readline().strip().split()
		if "****" in line:
			w[n_vec]=f.readline().strip().split()[1]
			break
	while True:
		line=f.readline().strip().split()
		if not line: break
		if "****" in line:	
			n_vec+=1
			print n_vec
			n_partial=0
			w[n_vec]=f.readline().strip().split()[1]
		else:
			v[n_vec,n_partial:width+n_partial]=line
			n_partial+=width
	f.close()	
	return w,v.T

 
def GenNM(obj, nmids, mode, inp=None, amp=50., nframes=20, dump_file=False, \
          prmtop=None):
  nmids=[int(i) for i in nmids.split()]
  amp = float(amp)
  nframes=int(nframes)
  mode=int(mode)
  stored.NMA =  {}
  stored.NMA[obj] = {}
 
  if mode==1:
    if not inp:
      print "Please specify cpptraj covariance matrix file"
      return 0
    if (not prmtop):
      print "Please specify prmtop file for proper mass reweighing of vectors."
      return 0  
    print "Attempting to read mwcov matrix %s" %inp
    mwcvmat=genfromtxt(inp)
    print "Diagonalizing covariance matrix."
    w,v = linalg.eig(mwcvmat)
    #sort from largest eigenvvalue to smallest
    si=argsort(w)[::-1]
    w=w[si]
    v=v[:,si]
    #convert to cm-1
    for i in range(len(w)):
		if w[i]<0: w[i]= -108.587*sqrt(-0.6/w[i])
		elif w[i]>0: w[i]= 108.587*sqrt(0.6/w[i])
		else: w[i]=0
	#get masses and un-mass weigh the modes
    m=raf.prmtop(prmtop).Get_Masses()
    m=[[i]*3 for i in m]
    m=[item for sublist in m for item in sublist] 
    for vec in range(v.shape[1]):
      v[:,vec]=v[:,vec]/sqrt(m)  
    #store
    stored.NMA[obj]['w'] = w
    stored.NMA[obj]['v'] = v
  elif mode==2:
    if not inp:
      print "Please specify evecs frequency/eigenvector file"
      return 0 
    print "Attempting to read evecs file %s" %inp
    stored.NMA[obj]['w'], stored.NMA[obj]['v'] = ReadEvecs(inp)
  elif mode==3:
    if not inp:
      print "Please specify pickled normal mode object"
      return 0
    print "Attempting to read pickle file %s" %inp  
    stored.NMA[obj]=pickle.load( open(inp,"rb"))
    
  if dump_file:
    print "Dumping eigenvalues and modes to pickle file."	  
    pickle.dump(stored.NMA[obj], open(dump_file,"wb") )

 
  w = stored.NMA[obj]['w']
  v = stored.NMA[obj]['v']
  print "%7s%20s%20s" % ('NMID', 'f Hz', 'f cm^-1')
  for i in range(min(w.size,20)):
    print "%7d%20.5e\t%20.5f" % (i, w[i]*2.99792458e10, w[i]),
    if i in nmids:
      print "*"
    else:
      print ""
  for nmid in nmids:
    newobj = "%s-nm_%d" % (obj, nmid)
    cmd.delete(newobj)    
    nm = real(v[:,nmid])
    nm *= amp/sqrt((nm**2).sum())
    for ifr in range(nframes):
      cmd.create(newobj, obj, 1, ifr+1)
      stored.nm = nm*sin(2.*pi*ifr/nframes)
      cmd.alter_state(ifr+1, newobj, "x,y,z = x+stored.nm[(ID-1)*3], y+stored.nm[(ID-1)*3+1], z+stored.nm[(ID-1)*3+2]")
cmd.extend("gennm", GenNM)
