from __future__ import print_function
#for compatability between python 2.7 and python 3

#!/usr/bin/python
#nmwizConvertor.py by Nigel S. Michki
#usage: nmwizConverter filename1.pdb filename2.out proteinName
#returns: file in .nmd format which can be read by Normal Mode Wiz and used
#to generate a series of movies depicting motions for each normal mode in VMD.

#System imports
import sys

#other imports
import os
from tkinter import *
from tkinter.filedialog import *
from tkinter.simpledialog import askstring

#Local (nmwizConverter) imports
from mode import Mode
from atom import Atom

#isInt checks to see that a string translates to an integer.
#This function is needed to correctly parse to mode.out file.
def isInt(inString):
    try:
        int(inString)
        return True
    except ValueError:
        return False
        
#No longer need 4 arguments
##Check for correct number of command line arguments
#if len(sys.argv) != 4:
#    print("Error; wrong number of command line arguments.")
#    print("usage: nmwizConverter filename1.pdb filename2.out name_of_protein")

#Global variables 
#root things
if len(sys.argv) == 4:
    pdbFilename = sys.argv[1]
    modeFilename = sys.argv[2]
    proteinName = sys.argv[3]
    nmdFilename = proteinName + ".nmd"
else:
    root = Tk()
    root.withdraw()#delete the box or something

    #selecting files

    import ctypes  # An included library with Python install.
    ctypes.windll.user32.MessageBoxA(0, b"Please select .pdb file", b"PDB File" , 1)

    pdbFilename = askopenfilename(filetypes = (('pdb files','*.pdb'),('all files','*.*')))

    import ctypes  # An included library with Python install.
    ctypes.windll.user32.MessageBoxA(0, b"Please select .out file", b"OUT File" , 1)

    modeFilename = askopenfilename(filetypes = (('out files','*.out'),('all files','*.*')))
    proteinName = askstring("NMD Name","Enter desired name for NMD file")
    nmdFilename = proteinName + ".nmd"
	
import ctypes  # An included library with Python install.
ctypes.windll.user32.MessageBoxA(0, b"Please select output directory for .nmd file", b"NMD File", 1)


#ask for output directory
filedir = askdirectory()

atoms = []

###BEGIN HEAVY LIFTING###

#Read .pdb file and create Atom objects with corresponding data from file
with open(pdbFilename) as pdbFile:
    pdbLines = pdbFile.readlines()

for line in pdbLines:
   words = str.split(line)
   if len(words) < 1:
       pass
   elif words[0] == "ATOM":
       newAtom = Atom()
       newAtom.resID = words[1]
       newAtom.atomName = words[2]
       newAtom.resName = words[3]
       newAtom.coordinates = [words[5], words[6], words[7]]
       newAtom.bFactor = words[9]
       newAtom.tag = words[10]
       atoms.append(newAtom)


#Read .out file add corresponding mode data from file to existing Atom objects
with open(modeFilename) as modeFile:
    modeLines = modeFile.readlines()

currentModeNumber = 0
currentModeFrequency = 0.0
for line in modeLines:
    words = str.split(line)
    if len(words) < 1:
        pass
    elif words[0] == "VIBRATION":
        currentModeNumber = words[2]
        currentModeFrequency = words[4]
        
    elif words[0] == 'VIBRATION' and not isInt(words[2]):
        currentModeNumber = list(words[1])[4] + list(words[1])[5] +list(words[1])[6] + list(words[1])[7]
        currentModeFrequency = words[3]
        
    elif isInt(words[0]) and currentModeNumber != 0 and len(words) == 7: 
        newMode = Mode()
        newMode.number = currentModeNumber
        newMode.frequency = currentModeFrequency
        newMode.atomNumber = words[0]
        newMode.vector = [words[4], words[5], words[6]]
        atoms[int(newMode.atomNumber) - 1].modes.append(newMode)


#Assemble data into printable strings
printTheseLines = []

printTheseLines.append("nmwiz_load" + " " + proteinName + ".nmd")
printTheseLines.append("name" + " " + proteinName)

atomNames = ""
resNames = ""
resIDs = ""
chainIDs = ""
bFactors = ""
coordinates = ""
for atom in atoms:
    atomNames = atomNames + atom.atomName + " "
    resNames = resNames + atom.resName + " "
    resIDs = resIDs + atom.resID + " "
    chainIDs = chainIDs + "A" + " "
    bFactors = bFactors + atom.bFactor + " "
    coordinates = coordinates + atom.printCoordinates() + " "
printTheseLines.append("atomnames" + " " + atomNames)
printTheseLines.append("resnames" + " " + resNames)
printTheseLines.append("resids" + " " + resIDs)
printTheseLines.append("chainids" + " " + chainIDs)
printTheseLines.append("bfactors" + " " + bFactors)
printTheseLines.append("coordinates" + " " + coordinates)

numberOfLinesBeforeModes = len(printTheseLines)

numberOfModes = len(atoms[0].modes)
for i in range(0, numberOfModes):
    printTheseLines.append("mode" + " " + str(atoms[0].modes[i].number) + " "
                            + str(atoms[0].modes[i].frequency) + " ")
for atom in atoms:
    for mode in atom.modes:
        printTheseLines[numberOfLinesBeforeModes + int(mode.number) - 1] += ( 
            mode.printMode() + " ")

#import some more things for gui
#select output directory



#Write .nmd file with printable strings
with open(os.path.join(filedir,nmdFilename), "wt") as nmdFile:
    for line in printTheseLines:
        print(line, file=nmdFile)

import ctypes  # An included library with Python install.
ctypes.windll.user32.MessageBoxA(0, b"Success =]", b"NMD File", 1)


print("Output written to " + nmdFilename)

###END HEAVY LIFTING - phew###