#!/usr/bin/env python

# ------------------------------------------------------------------------------
#
# Work on a conformation ensemble in .pdb format and superimpose them all onto
# a user chosen .pdb file
#
# Thomas Coudrat, February 2014
# thomas.coudrat@gmail.com
#
# ------------------------------------------------------------------------------

import os
import sys
import glob
import shutil
import argparse
from os.path import basename
from subprocess import check_output, STDOUT, CalledProcessError
import socket
import json


def main():
    """
    Run the superimpose script
    """

    # Get arguments
    pdbDir, templatePath = parsing()

    # Get the paths corresponding to the platform where this is executed
    icm, script = getPaths()

    # Get paths of all .pdb files in this directory
    pdbPaths = sorted(glob.glob(pdbDir + "/*.pdb"))

    print
    print "Superimposing " + str(len(pdbPaths)) + " confs from dir: " + pdbDir
    print "Template used for superimposition: " + templatePath
    print

    superimpose(templatePath, pdbPaths, pdbDir, icm, script)


def parsing():
    """
    Parse arguments
    """
    descr = "Superimpose all .pdb conformations onto the provided template"
    descr_pdbDir = "Directory containing conformation ensemble"
    descr_templatePath = "Template onto which superimpose the conformation " \
        "ensemble"

    parser = argparse.ArgumentParser(description=descr)
    parser.add_argument("pdbDir", help=descr_pdbDir)
    parser.add_argument("templatePath", help=descr_templatePath)

    args = parser.parse_args()

    pdbDir = args.pdbDir
    templatePath = args.templatePath

    return pdbDir, templatePath


def getPaths():
    """
    Figure the paths to the ICM executable and the ICM superimpose script
    """

    # This Json file stores the ICM executable locations for each platform
    hostFiles_json = os.path.dirname(os.path.realpath(__file__)) + \
        "/hostFiles.json"

    # Read content of .json file
    with open(hostFiles_json, "r") as jsonFile:
        hostFiles_dict = json.load(jsonFile)

    # Get the hostname
    hostname = socket.gethostname()

    # Verify that the current hostname is present in the config file loaded from
    # json
    if hostname in hostFiles_dict.keys():
        icm = hostFiles_dict[hostname][0]
        script = hostFiles_dict[hostname][1]
    else:
        print "Error: hostname " + hostname + " filepaths not configured in " \
            " ./hostFiles.json"
        sys.exit()

    return icm, script


def superimpose(templatePath, pdbPaths, pdbDir, icm, script):
    """
    Get the template .pdb path as well as the paths to all the .pdb
    conformations
    Cycle through and for each of them run the ICM superimpose script.
    Make a temp copy of this script, and use sed to change for each execution
    """

    # Get the ICM name to be of the template
    templateName = basename(templatePath).replace(".pdb", "").replace("-", "_")
    # Temp directory name
    tempScript = "./" + os.path.dirname(pdbDir).lower() + ".icm"

    for pdbPath in pdbPaths:
        pdbName = basename(pdbPath).replace(".pdb", "")
        # This is to reflect the fact that ICM turns hyphens to underscores
        # when a file is loaded in
        pdbName = pdbName.replace("-", "_")

        # Copy the superimpose ICM script to a temp location
        shutil.copy(script, tempScript)

        # Make the sed changes
        os.system("sed -e 's|PDB_PATH_1|" + templatePath + "|g' " +
                  tempScript + " -i")
        os.system("sed -e 's|PDB_PATH_2|" + pdbPath + "|g' " +
                  tempScript + " -i")
        os.system("sed -e 's|PDB_NAME_1|" + templateName + "|g' " +
                  tempScript + " -i")
        os.system("sed -e 's|PDB_NAME_2|" + pdbName + "|g' " +
                  tempScript + " -i")

        print pdbName
        # Execute the temp script
        try:
            check_output(icm + " -s " + tempScript, stderr=STDOUT, shell=True)
        except CalledProcessError, e:
            print "\n Error during the ICM script execution"
            print e.output
            sys.exit()

        # Remove the temp script
        os.remove(tempScript)


if __name__ == "__main__":
    main()

@@ -868,28 +868,20 @@
    print
-%0A    print %22
+(%22%5Cn
 Superimp
@@ -910,16 +910,26 @@
 aths)) +
+%0A         
  %22 confs
@@ -945,24 +945,25 @@
 : %22 + pdbDir
+)
 %0A    print %22
@@ -960,17 +960,17 @@
    print
- 
+(
 %22Templat
@@ -1017,19 +1017,61 @@
 Path
+)%0A
 %0A    
-print%0A
+# Run the superimposition, using an ICM script
 %0A   
@@ -1129,16 +1129,299 @@
 cript)%0A%0A
+    # Finally, rename the directory to notify that the contained structures%0A    # were superimposed%0A    if pdbDir.endswith(%22_super/%22):%0A        print(%22%5CnDirectory not renamed: structures were already superimposed%22)%0A    else:%0A        os.rename(pdbDir, pdbDir.replace(%22/%22, %22_super/%22))%0A%0A
 %0Adef par
@@ -2193,30 +2193,30 @@
 orm%0A    
-hostFiles_json
+pathExecutable
  = os.pa
@@ -2257,20 +2257,46 @@
 __))
- + %5C
 %0A    
-   
+hostFiles_json = pathExecutable +
  %22/h
@@ -2742,17 +2742,17 @@
    print
- 
+(
 %22Error: 
@@ -2804,21 +2804,22 @@
 gured in
- 
 %22 
-%5C%0A
++%0A  
         
@@ -2841,16 +2841,17 @@
 es.json%22
+)
 %0A       
@@ -4221,24 +4221,25 @@
    print
- 
+(
 pdbName
+)
 %0A       
@@ -4390,17 +4390,19 @@
 essError
-,
+ as
  e:%0A    
@@ -4414,17 +4414,17 @@
    print
- 
+(
 %22%5Cn Erro
@@ -4457,16 +4457,17 @@
 ecution%22
+)
 %0A       
@@ -4480,17 +4480,17 @@
 rint
- 
+(
 e.output
 %0A   
@@ -4485,16 +4485,17 @@
 e.output
+)
 %0A       

Rename folder when superimposed, fixed for py3