#import optparse
import os
import sys
import numpy
import MMTK

#from webnma.api.config import config
#from webnma.api.extract_modes import extract_modes


#class CorrelationMatrix():

 #   def __init__(self, mode_path):
 #       self.mode_filename = mode_path.split('/')[-1]
 #       self.protein_identifier = self.mode_filename.replace(".mode",         
def write_rawfiles(matrix,structname):
        filename = ('%s_correlation_matrix.dat' % 
                    (structname))
        vfile = open(filename, 'w')
        for i in xrange(1, len(matrix)+1):
                print >> vfile, str(i)+'',
        print >> vfile 
        for i2 in xrange(len(matrix)):
            print >> vfile, str(i2+1)+' ',
            matrix[i2].tofile(vfile, sep=" ")
            print >> vfile
        vfile.close()
        return filename

def calculate(structname):
        # Obtain the modes
	modes_original = MMTK.load(sys.argv[1])
        
        # Fetch the universe, protein and number of atoms
        natoms = len(modes_original[6])
        num_modes = len(modes_original) - 6
        
        # Use mode 6 - first non-trivial mode as a base
        #modes = extract_modes(modes_original)
        m = modes_original[6]

        # extracting frequencies
        freq = []
        for i in range(0, len(modes_original)):
            freq.append(modes_original.rawMode(i).frequency)

        # building initial matrix (first mode)
        correlation_matrix = (numpy.inner(m, m) *
                              1/(freq[6]**2))

        # Build Correlation matrix 
        for i in range(7, 6+num_modes):
            m = modes_original[i]
            correlation_matrix += (numpy.inner(m, m) * 
                                   1/(freq[i]**2))

        # Calculating basis for normalization
        a = [] # normalization factor
        for i in range(natoms):
            tmp = 0
            for j in range(6, 6+num_modes):
                m = modes_original[j]
                tmp += (numpy.inner(m[i], m[i]) *
                        1/(freq[j]**2))
            a.append(numpy.sqrt(tmp))

        # Normalizing correl.mat
        a = numpy.matrix(numpy.outer(a,a))
        correlation_matrix = numpy.matrix(correlation_matrix)/a

        # Write our cute correl mat to file
        matrixfilename = write_rawfiles(correlation_matrix,structname)
        return matrixfilename

calculate(sys.argv[2])