import settings
import os, sys, string, re, time, math
import math
import util

from MMTK import *
from MMTK.PDB import PDBConfiguration
from MMTK.Proteins import Protein
from MMTK.NormalModes import NormalModes

from Numeric import sqrt, add, array, argmax, argsort, argmin
from Scientific import *
from Scientific.IO.ArrayIO import writeArray
from Scientific.Visualization import VMD


## display_helices.py - It really has alot of jobs to do. The name of the file is quite wrong... 
## Prime job is to identify the TM alpha helices. There is a large algorithm for doing this: FindBundle()
## We also want this file to generate a VMD file so the users can we their nice protein: DispalyHelices()
## There are some smaller functions as well. Basically, they are just helping the two functions described above.


## Function that starts it all, and returns the result of the algorithms
def HelixWrapper( items, helices, pdbfilename, hydro, residues, chains ):
 
    # Let us find the vectors and center of mass to represent the helices
    all_axes, all_cm =  FindHelixVectors( items, helices, pdbfilename, residues, chains )

    # Filter out all helices that are not in the TM bundle
    bundle, reason = FindBundle( items, helices, pdbfilename, all_axes, all_cm, hydro, residues )

    # Generate VMD and picture file of the protein with all the TM helices in it.
    result = DisplayHelices ( items, helices, pdbfilename, bundle, all_axes, all_cm, residues )

    # Let us just return the bundle (nevermind reason). The files generated here are stored on the local disk. Nevermind them.
    return bundle, reason


## Function for the filtering process
def FindBundle( items, helices, pdbfilename, all_axes, all_cm, hydro, residues ):

    length_threshold = 13
    scalar_threshold = 0.6
    distance_threshold = 5.5
    hydrophob_threshold = 0.0000
    
    # the array 'hydro' contains the hydropathy index to each helix found

    # reason[] will tell the user why the helix has been filtered out
    reason = ['']*len(helices)
        
    # bundle will contain the indices to the helices array which might be a part of the TM bundle
    bundle = []
    for i in range(0,len(helices)):
        bundle.append(i)

    popped = 0
    for i in range(0,len(bundle)):

        if len(helices[bundle[i-popped]]) < length_threshold:
            goaway = bundle.pop(i-popped)
            popped = popped +1
            reason[i-popped] = 'l'

    print 'bundle after helix length comparison: '
    print bundle
       

    # bundle_size[] is here an array that holds the number of
    # helices that are approximate parallel to helix number i
    # same index as all_axes[] is used
    bundle_size = len(helices)*[0.]

    for i in range( 0, len(bundle)):
        a = all_axes[bundle[i]]
        
        for j in range( 1, len(bundle)):
            b = all_axes[bundle[j]]

            scalar = a * b
                       
            if sqrt(scalar*scalar) > scalar_threshold:
                bundle_size[bundle[i]] = bundle_size[bundle[i]] + 1


    bundle_i = 0
    max_bundle = 0

    # finds which helix to use as "reference".
    # this is done to find the largest bundle in the protein
    for i in range(0, len(bundle)):
        if bundle_size[bundle[i]] > max_bundle:
            max_bundle = bundle_size[bundle[i]]
            bundle_i = bundle[i]
               
    popped = 0
    for i in range(0,len(bundle)):
        a = all_axes[bundle_i]
        b = all_axes[bundle[i-popped]]

        scalar = a * b

        if sqrt(scalar*scalar) < scalar_threshold:
            reason[bundle[i-popped]] = 's'
            bundle.pop(i-popped)
            popped = popped+1

    # this is the bundle after scalar comparison is done
    print 'bundle after scalar comparison: '
    print bundle

    avg_distance_array = len(helices)*[0.]

    # we filter out the helices that are "far" away from the other
    for i in range( 0, len(bundle)):
        a = all_cm[bundle[i]]
        ax = a.x()
        ay = a.y()
        az = a.z()

        distance = 0
        
        for j in range(1,len(bundle)):

            b = all_cm[bundle[j]]

            bx = b.x()
            by = b.y()
            bz = b.z()

            d = sqrt( (ax-bx)*(ax-bx) + (ay-by)*(ay-by) + (az-bz)*(az-bz) )
            distance = distance + d
            
        avg_distance_array[bundle[i]] = distance/(len(bundle)-1)

    # those helices that are more than 'distance_threshold' away from the other (in average)
    # is filtered out
    popped = 0
    for i in range(0,len(bundle)):
        d = avg_distance_array[bundle[i-popped]]

        if d > distance_threshold:
            reason[bundle[i-popped]] = 'd'
            goaway = bundle.pop(i-popped)
            popped = popped +1

    # this is the bundle after distance comparisonfor i in range(0, len(bundle)):

    print 'bundle after distance comparison: '
    print bundle

    # filtering out those helices that have negative hydropathy index
    popped = 0
    for i in range(0, len(bundle)):
        j = bundle[i-popped]

        if hydro[j] < hydrophob_threshold:
            reason[bundle[i-popped]] = 'h'
            go = bundle.pop(i-popped)
            popped = popped +1
            
    # the bundle after hydrophopbicity filtering
    print 'bundle after hydrophobicity filtering: '
    print bundle

    ## The bundle has hopefully been correctly identyfied, and can be returned.

    return bundle, reason


    
   
## Function for identifying all alpha helical axes
def FindHelixVectors( items, helices, pdbfilename, residues, chains ):
    ## Let us store all axes in a list. Well, let us do the same with the center of masses as well.

    all_axes = []
    all_cm = []

    ## We open the mode file, and load the universe. Then have our protein... 
    modes = load(items['modefile_name'])
    universe = modes.universe
    protein = universe[0]
    
    num_helices = len(helices)

    ## For each helix we have found, we have have to find the axis of inertia
    for i in range (0, num_helices):

        chain_nr = util.get_chain_nr(helices[i][0], chains)
        chain = protein[chain_nr]
        
        start_res = helices[i][1]
        start_res = util.findIndex(start_res, residues, chain_nr )
        h_len = len(helices[i])-1
        end_res = helices[i][h_len]
        end_res = util.findIndex(end_res, residues, chain_nr)

        helix = chain[start_res:end_res]
        #print str(start_res) + ':' + str(end_res) + '-'+  str(chain_nr)
        cm, mi = helix.centerAndMomentOfInertia()

        ev, axes = mi.diagonalization()
        helix_axis = axes[argmin(ev)].asVector().normal()

        all_axes.append( helix_axis )
        all_cm.append( cm )

    ## we can just return the result again.
    return all_axes, all_cm
  

## Function that is used for building the VMD file and generating a image used to display the protein.
## Nothing exciting to see here..
def DisplayHelices ( items, helices, pdbfilename, bundle, all_axes, all_cm, residues ):

    graphics = VMD
    scene = graphics.Scene(scale=1./Units.Ang)

    pngfiles = [ items['display_0'], items['display_1'] ]
    
 
        
    print pdbfilename
    scene.addObject(graphics.Molecules(pdbfilename))
    vmdfilename = 'scratch/' + items['jobid'] + '.vmd'
    scene.writeToFile(vmdfilename)

    vmdfile1 = open(vmdfilename, 'r').readlines()

    for i in range(0,1):
        
        array = []

        for line in vmdfile1:
            array.append(line)

        for j in range( 0, len(helices)):
        
            h_len = len(helices[j])
            helix_axis = all_axes[j]
            cm = all_cm[j]

            pos1 = cm-(h_len*0.075)*Units.nm*helix_axis
            pos2 = cm+(h_len*0.075)*Units.nm*helix_axis

            color = 0
            for k in bundle:
                if k == j:
                    color = 1
                    
            array.append('graphics 0 color ')
            array.append(str(color) +'\n')
            
            array.append('graphics top cylinder {')
            array.append(str(pos1[0]*10) + ' ')
            array.append(str(pos1[1]*10) + ' ')
            array.append(str(pos1[2]*10) + '} {')

            array.append(str(pos2[0]*10) + ' ')
            array.append(str(pos2[1]*10) + ' ')
            array.append(str(pos2[2]*10) + '} radius 0.89999999999999991 filled yes\n')

            array.append('graphics 0 color ')
            array.append('8 \n')

            array.append('graphics top text {')
            array.append(str(pos1[0]*10) + ' ')
            array.append(str(pos1[1]*10) + ' ')
            array.append(str(pos1[2]*10))
            array.append('} "' + str(j) + '" size 1.5\n')
            
        array.append('mol delrep 0 top\n')
        array.append('mol representation Trace 0.300000 6.000000\n')
        array.append('mol color Occupancy\n')
        array.append('mol selection {all}\n')
        array.append('mol material Opaque\n')
        array.append('mol addrep top\n')
        array.append('scale by 1.65\n')
        array.append('axes location off\n')

        if i == 0:
            xrotation = 0
            yrotation = 0
            zrotation = 0
        elif i == 1:
            xrotation = 90
            yrotation = 90
            zrotation = 0
        else:
            xrotation = 280
            yrotation = 0
            zrotation = 0
            
        array.append('display projection orthographic\n')
        array.append('rotate x by '+str(xrotation)+'\n')
	array.append('rotate z by '+str(yrotation)+'\n')
	array.append('rotate y by '+str(zrotation)+'\n')

        pngfile = pngfiles[i]
        povrayfile = pngfile.replace('png', 'povray')
        
        #array.append('render POV3 ' + povrayfile + ' ' + settings.POVRAYPATH + ' +H300 +W250 -I'+povrayfile+' -O'+pngfile+' +D +X +A +FN\n')
        #array.append('quit\n')
    
        vmdfile = open(vmdfilename,'w')

        c=0
        while c < len(array):
            vmdfile.write(array[c])
            c=c+1
        vmdfile.close()
        
        #os.system(settings.VMDPATH + ' -e ' + vmdfilename + ' -dispdev none')
                
        print pngfile
    
    return 1

def DisplayHelices2 ( items, helices, pdbfilename, vmdfilename, residues ):

    starttime = time.clock()

    num_helices = len(helices)

    f = os.popen('pymol -c -p', 'w')
    f.write('load ' + pdbfilename + '\n' )
    f.write('hide all\n')
    f.write('show cartoon\n')

    for i in range (0, 1):
        
        start_res = helices[i][0]
        h_len = len(helices[i])-1
        end_res = helices[i][h_len]
        
        print h_len

        f.write('color red, resi ' + str(start_res) + '-' + str(end_res) + '\n')
        f.write('ray 400,400\n')
        f.write('png ' + 'scratch/' + items['jobid'] + '.' + str(i) + '.pymol.png'  + '\n') 
        f.write('color green, all' + '\n')
        
    f.write('quit\n')

    stoptime = time.clock()

    return stoptime - starttime