#!/export/ginkgo/bin/python

# Building up the system and calculating the normal modes
# input: PDB 'string', forcefield, nb of modes to be calculated
# (default is number_of_atoms/5)
#

import os,sys,string,re,time
from MMTK import *
from DomainFinderCalc import DomainAnalysisData
#from WriteHTMLTable import WriteHTMLTable
from MMTK.PDB import PDBConfiguration

from Scientific.Visualization.Color import ColorByName, ColorScale

'''CalcModes uses MMTK to calculate the normal modes for structures.  '''

def CalcModes(pdbfilename, modesfilename, nmodes):
    
    import Numeric 
    N = Numeric
    starttime = time.clock()
    
    # Construct the system
    data = DomainAnalysisData()
    data.sequence = PDBConfiguration(pdbfilename).peptide_chains
    natoms = data.universe.numberOfCartesianCoordinates()
    #print 'atoms: ' + str(natoms)
    #nmodes = max(natoms/2, min(200, 3*natoms))
    #nmodes = 3*natoms/2

    nmodes_calc = int(float(nmodes)*3*natoms)
    #nmodes_calc = 3*natoms
    #nmodes_calc = 200
    print 'number of modes to calculate: ' + str(nmodes_calc)
    
    data.nmodes_calc = nmodes_calc
    data.nmodes_keep = nmodes_calc
    
    #data.universe.writeToFile(modesfilename.replace('modes', 'pdb'), None, 'pdb')
    
    # making a modes file and saving it
    modes=data.modes
    print modes
    save(modes,modesfilename)
        
    # calculating deformation energies for all modes from 7
    mode_deformations = data.mode_deformations
    average_energies = 6*[0.]

    # an array to store the mode deformations. to be used in further calculations
    #mode_def_array = []
    
    #print mode_deformations[6].array #prints mode def energies for mode 7

    for i in range(6, len(mode_deformations)):
        #mode_def_array.append(mode_deformations[i].array)
        d_mean = N.add.reduce(mode_deformations[i].array)/natoms
        average_energies.append(d_mean)

    # we're finished!
    stoptime = time.clock()
    elapsed = stoptime - starttime

    # write table to an html file. All of this should be moved up to the Zope layer. 
    #web_table=WriteHTMLTable(average_energies,6,20,'Normal mode index','Deformation Energy')
    #firststring = '<br>'+'<h3>'+'Deformation energies'+'</h3>'
    #fourthstring = '<br>Elapsed time (normal modes calculations) ' + str(elapsed) + ' seconds'
    #htmlstring = web_table + fourthstring
    
    return 1





if __name__ == '__main__':
    CalcModes(sys.argv[1], sys.argv[2], sys.argv[3])