'''Adapted from helix bundle rotation code written by Nathalie for TMM@ and some of Edvin's and Lars' old code'''
import MMTK
from MMTK import *
from MMTK.PDB import PDBConfiguration
from MMTK.Proteins import Protein
from MMTK.NormalModes import NormalModes
from numpy import sqrt, add, array, argmax, argsort, argmin,linalg, dot
from Scientific import *
from Scientific.IO.ArrayIO import writeArray
from Scientific.Visualization import VMD
import sys

def massWeightedNormalVector(v):
   return v/sqrt(v.massWeightedDotProduct(v))

def loadModes():
	filename=sys.argv[1].replace(".mode","")
	modes = MMTK.load(sys.argv[1]) #should have the input modes file for the TIM barrel structure in question
	universe = modes.universe
	protein = universe.protein
	protein.normalizeConfiguration('Ir')
	protein.writeToFile(sys.argv[2]+"_"+sys.argv[3]+"_from_modes.pdb")#sys.argv[2] should be filename, sys.argv[3] should be motion_type
	##sys.argv[4] should be mode type
	##the structure annotation part
	return modes,universe,protein,filename

def pickout(file1):
        name_list={}
        name=''
        first=0
        for line in file1:
                line=line.rstrip()
                row=line.split("\t")
                if len(row[0]) ==4:
                        name=row[0]
                        first = int(row[1])
                        name_list[(name,first)]=[]
			old_line=[]
                if len(row[0])<4 and "#" not in line:
                        start=int(row[0])-first
                        end=int(row[1])-first
                        name_list[(name,first)]=old_line
			old_line.append((start,end))

        return name_list

def make_collection(sse_annot,uni_prot):
	sse_dict = pickout(sse_annot)
	filename=uni_prot[3]
	protein=uni_prot[2]
	print filename
	for key in sse_dict.keys():
		if key[0] in filename:
			sse_list=[]
			chain1 = protein[0]
			for x in sse_dict[key]: 
			    sse_list.append(chain1[x[0]:x[1]+1])
			    print len(chain1[x[0]:x[1]+1])
			return Collection(sse_list)

def single_sse(atom_collection,universe):
	cm, mi = atom_collection.centerAndMomentOfInertia()
        ev, axes = mi.diagonalization()
        axis = axes[argmax(ev)].asVector().normal()      
        indiv_vector = ParticleVector(universe)
        return cm, axis, indiv_vector
        
def addToScene(graphics,scene,atom,indiv_vector,cm, sse_axis):
	scene.addObject(graphics.Arrow(atom.position(),atom.position()+1.*indiv_vector[atom],0.01))
        scene.addObject(graphics.Cylinder(cm-2.*Units.nm*sse_axis,
                                          cm+2.*Units.nm*sse_axis,
                                          0.05))
        return scene


def calculatingthings(sse_collection,suffix, motion_type, uni_prot,metric_file):
	modes = uni_prot[0]
	universe = uni_prot[1]
	protein = uni_prot[2]
	#print uni_prot
	
	initial = single_sse(sse_collection,universe)
	#print initial
	cm = initial[0]
	sse_axis = initial[1]
	vector = initial[2]
	p = []
	
	if motion_type == "indiv_translation":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			
			ele_initial=single_sse(element,universe)
			
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
			indiv_trans_vector = ele_initial[2]
			
			for atom in element.atomList():
				indiv_trans_vector[atom] = ele_axis
				scene=addToScene(graphics,scene,atom,indiv_trans_vector,cm,sse_axis)
			indiv_trans_vector=massWeightedNormalVector(indiv_trans_vector)
			for i in range(6, len(modes)): 
				mode = massWeightedNormalVector(modes[i])
				
	        		projection = mode.massWeightedDotProduct(indiv_trans_vector)
	        		#print projection 
				p.append((i+1,projection**2))
				
			p=array(p)
			writeArray(p, 'indiv_translation_%s_%s.pics.dat'% (suffix, counter+1))
			metric = 0
			
			for idx in range(len(p)):
				metric+=p[idx][1]*modes.omega_sq[idx]
			
			metric_file.write(str(metric)+"\n")
			pacc = add.accumulate(p)
			pacc[:,0] = p[:,0]
			writeArray(pacc, 'indiv_translation_%s_%s.esc.dat'%(suffix,counter+1))
			p = []	
			scene.writeToFile("indiv_translation_scene_%s.vmd"% suffix) 
			
			
	if motion_type=="indiv_slide":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			ele_initial=single_sse(element,universe)
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
                      	indiv_slide_vector = ele_initial[2]  #might be redundant
                      	
			indiv_rot_vector = sse_axis.cross(cm - element_cm)
			indiv_slide_v = sse_axis.cross(indiv_rot_vector)

			for atom in element.atomList():
				indiv_slide_vector[atom] = indiv_slide_v
				scene = addToScene(graphics,scene,atom,indiv_slide_vector,cm,sse_axis)
				
			indiv_slide_vector =massWeightedNormalVector(indiv_slide_vector)
			
			for i in range(6,len(modes)):		## changed from 0 to 6 to exclude first 
				mode = massWeightedNormalVector(modes[i])
				projection = mode.massWeightedDotProduct(indiv_slide_vector)
				p.append((i+1,projection**2))
				
			p=array(p)
			writeArray(p, 'indiv_slide_%s_%s.pics.dat'% (suffix, counter+1))
			metric = 0
			for idx in range(len(p)):
				metric+=p[idx][1]*modes.omega_sq[idx]
			metric_file.write(str(metric)+"\n")
			pacc = add.accumulate(p)
			pacc[:,0] = p[:,0]
			writeArray(pacc, 'indiv_slide_%s_%s.esc.dat'% (suffix, counter+1))
			p=[]
			scene.writeToFile("indiv_slide_scene_%s.vmd"% suffix)
			
	if motion_type=="indiv_tilt":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			ele_initial=single_sse(element,universe)
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
                      	indiv_tilt_vector = ele_initial[2]
                        
			indiv_rot_vector = sse_axis.cross(cm - element_cm)
			indiv_slide_v = sse_axis.cross(indiv_rot_vector)
			
			for atom in element.atomList():
				tilt_disp = atom.position()[1] - element_cm[1]
				indiv_tilt_v = indiv_slide_v * tilt_disp
				indiv_tilt_vector[atom] = indiv_tilt_v
				scene = addToScene(graphics, scene,atom,indiv_tilt_vector,cm,sse_axis)
				
			indiv_tilt_vector =massWeightedNormalVector(indiv_tilt_vector)
			
			for i in range(6,len(modes)):		## changed from 0 to 6 to exclude first 
				mode = massWeightedNormalVector(modes[i])
				projection = mode.massWeightedDotProduct(indiv_tilt_vector)
				p.append((i+1,projection**2))
			p=array(p)
			writeArray(p, 'indiv_tilt_%s_%s.pics.dat'% (suffix, counter+1))
			metric = 0
			for idx in range(len(p)):
				metric+=p[idx][1]*modes.omega_sq[idx]
			metric_file.write(str(metric)+"\n")
			pacc = add.accumulate(p)
			pacc[:,0] = p[:,0]
			writeArray(pacc, 'indiv_tilt_%s_%s.esc.dat'% (suffix, counter+1))
			p=[]	
			scene.writeToFile("indiv_tilt_scene_%s.vmd"% suffix)
			
	if motion_type=="indiv_bend_upper":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			ele_initial=single_sse(element,universe)
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
                      	indiv_bend_vector = ele_initial[2]
                      	
			indiv_rot_vector = sse_axis.cross(cm - element_cm)
			indiv_slide_vector = sse_axis.cross(indiv_rot_vector)
			
			for atom in element.atomList():
				bend_disp = (atom.position()[1] - element_cm[1])
				if bend_disp < 0:
					bend_disp = bend_disp * 0.000001
				indiv_bend_v = indiv_slide_vector * bend_disp
				
				indiv_bend_vector[atom] = indiv_bend_v
				scene = addToScene(graphics,scene,atom,indiv_bend_vector,cm,sse_axis)
				
			indiv_bend_vector =massWeightedNormalVector(indiv_bend_vector)
			
			for i in range(6,len(modes)):		## changed from 0 to 6 to exclude first 
				mode = massWeightedNormalVector(modes[i])
				projection = mode.massWeightedDotProduct(indiv_bend_vector)
				p.append((i+1,projection**2))
			p=array(p)
			writeArray(p, 'indiv_bend_upper_%s_%s.pics.dat'% (suffix, counter+1))
			metric = 0
			for idx in range(len(p)):
				metric+=p[idx][1]*modes.omega_sq[idx]
			metric_file.write(str(metric)+"\n")
			pacc = add.accumulate(p)
			pacc[:,0] = p[:,0]
			writeArray(pacc, 'indiv_bend_upper_%s_%s.esc.dat'% (suffix, counter+1))
			p=[]
			
			scene.writeToFile("indiv_bend_upper_scene_%s.vmd"% suffix)
	if motion_type=="indiv_bend_lower":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			ele_initial=single_sse(element,universe)
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
                      	indiv_bend_low_vector = ele_initial[2]
                      	
			indiv_rot_vector = sse_axis.cross(cm - element_cm)
			indiv_slide_vector = sse_axis.cross(indiv_rot_vector)
			#step = float(1)/len(element)
                        #weight=1
                        
			for atom in element.atomList():
				bend_low_disp = (atom.position()[1] - element_cm[1]) * -1
				if bend_low_disp < 0:
					bend_low_disp = bend_low_disp * 0.000001
				indiv_bend_low_v = indiv_slide_vector * bend_low_disp
				indiv_bend_low_vector[atom] = indiv_bend_low_v * bend_low_disp
                                #weight-=step
				scene = addToScene(graphics,scene,atom, indiv_bend_low_vector,cm,sse_axis)
				
			indiv_bend_alt_vector =massWeightedNormalVector(indiv_bend_low_vector)
			
			for i in range(6,len(modes)):		## changed from 0 to 6 to exclude first 
				mode = massWeightedNormalVector(modes[i])
				projection = mode.massWeightedDotProduct(indiv_bend_alt_vector)
				p.append((i+1,projection**2))
			p=array(p)
			writeArray(p, 'indiv_bend_low_%s_%s.pics.dat'% (suffix, counter+1))
			metric = 0
			for idx in range(len(p)):
				metric+=p[idx][1]*modes.omega_sq[idx]
			metric_file.write(str(metric)+"\n")
			
			pacc = add.accumulate(p)
			pacc[:,0] = p[:,0]
			writeArray(pacc, 'indiv_bend_low_%s_%s.esc.dat'% (suffix, counter+1))
			p=[]
			
			scene.writeToFile("indiv_bend_low_scene_%s.vmd"% suffix)
	if motion_type=="rotation":
		metric_file.write(suffix+"_"+motion_type+"\n")
		for atom in sse_collection.atomList():
			
                	vector[atom] = sse_axis.cross(atom.position()-cm)
    
		vector=massWeightedNormalVector(vector)

		for i in range(6, len(modes)): ##changed from 0 to 6 to exclude trivial modes
  			mode = massWeightedNormalVector(modes[i])
			projection = mode.massWeightedDotProduct(vector)
			
			p.append((i+1,projection**2))
		p = array(p)
	
		writeArray(p, 'rotation_%s.pics.dat' %suffix)
		metric = 0
		for idx in range(len(p)):
			metric+=p[idx][1]*modes.omega_sq[idx]
		metric_file.write(str(metric)+"\n")
		pacc = add.accumulate(p)
		pacc[:,0] = p[:,0]
		writeArray(pacc,'rotation_%s.esc.dat' %suffix)
	if motion_type == "translation":
		metric_file.write(suffix+"_"+motion_type+"\n")
		graphics = VMD
                scene = graphics.Scene(scale=1./Units.Ang)
                
		for atom in sse_collection.atomList():
                        vector[atom] = sse_axis
                        scene.addObject(graphics.Arrow(atom.position(),atom.position()+1.*vector[atom],0.01))
                        scene.addObject(graphics.Cylinder(cm-2.*Units.nm*sse_axis,
                                                  cm+2.*Units.nm*sse_axis,
                                                  0.05))
                vector=massWeightedNormalVector(vector)
                p = []
                pacc = 0
                for i in range(6,len(modes)): ##changed from 0
                        mode=massWeightedNormalVector(modes[i])
		        projection=mode.massWeightedDotProduct(vector)
			p.append((i+1, projection**2))
                p=array(p)
                writeArray(p, 'translation_%s.pics.dat' %suffix)
                metric=0
                for idx in range(len(p)):
			metric+=p[idx][1]*modes.omega_sq[idx]
		metric_file.write(str(metric)+"\n")
                pacc=add.accumulate(p)
                
                pacc[:,0] = p[:,0]
                writeArray(pacc, 'translation_%s.esc.dat' %suffix)
       
	if motion_type=="slide":
		graphics = VMD
		scene = graphics.Scene(scale=1./Units.Ang)
		metric_file.write(suffix+"_"+motion_type+"\n")
		for counter in xrange(len(sse_collection)):
			element=sse_collection[counter]
			ele_initial=single_sse(element,universe)
			element_cm = ele_initial[0]
			ele_axis = ele_initial[1]
                      	indiv_slide_vector = ele_initial[2]  #might be redundant
                      	
			indiv_rot_vector = sse_axis.cross(cm - element_cm)
			indiv_slide_v = sse_axis.cross(indiv_rot_vector)

			for atom in element.atomList():
				indiv_slide_vector[atom] = indiv_slide_v
				scene = addToScene(graphics,scene,atom,indiv_slide_vector,cm,sse_axis)
				
		slide_vector =massWeightedNormalVector(indiv_slide_vector)
			
		for i in range(6,len(modes)):		## changed from 0 to 6 to exclude first 
			mode = massWeightedNormalVector(modes[i])
			projection = mode.massWeightedDotProduct(slide_vector)
			p.append((i+1,projection**2))
				
		p=array(p)
		writeArray(p, 'slide_%s.pics.dat'% (suffix))
		metric = 0
		for idx in range(len(p)):
			metric+=p[idx][1]*modes.omega_sq[idx]
		metric_file.write(str(metric)+"\n")
		pacc = add.accumulate(p)
		pacc[:,0] = p[:,0]
		writeArray(pacc, 'slide_%s.esc.dat'% (suffix))
		p=[]
		scene.writeToFile("slide_scene_%s.vmd"% suffix)
       
#		
	return


#if "indiv" in sys.argv[3]:
	#uni_prot = loadModes()
	#ALPHA_annot = open("ALPHA_from_pdb.txt","r")
	#BETA_annot = open("BETA_from_pdb.txt","r")
	#metric_file = open("METRIC_for_indiv_overlaps_BETAonly.txt","a")
	#alpha_collection = make_collection(ALPHA_annot,uni_prot)
	#calculatingthings(alpha_collection, sys.argv[2]+"_ALPHA_"+sys.argv[4], sys.argv[3],uni_prot,metric_file)
	#beta_collection = make_collection(BETA_annot, uni_prot)
	#calculatingthings(beta_collection, sys.argv[2]+"_BETA_"+sys.argv[4], sys.argv[3],uni_prot,metric_file)
	#metric_file.close()
	
metric_file=open("metric_for_bundle_movements.txt","a")

ALPHA_annot = open("ALPHA_from_pdb.txt","r")
BETA_annot = open("BETA_from_pdb.txt","r")
uni_prot = loadModes()
print "ALPHA"
alpha_collection=make_collection(ALPHA_annot,uni_prot)
calculatingthings(alpha_collection, sys.argv[2]+"_ALPHA_"+sys.argv[4], sys.argv[3],uni_prot,metric_file)
#make_collection(ALPHA_annot,uni_prot)
print "BETA"
beta_collection=make_collection(BETA_annot,uni_prot)
calculatingthings(beta_collection, sys.argv[2]+"_BETA_"+sys.argv[4], sys.argv[3],uni_prot,metric_file)
#make_collection(BETA_annot,uni_prot)
metric_file.close()