#!/usr/bin/env ccp4-python
#
#     Copyright (C) 2014 Ronan Keegan 
#
#     This code is distributed under the terms and conditions of the
#     CCP4 Program Suite Licence Agreement as a CCP4 Application.
#     A copy of the CCP4 licence can be obtained by writing to the
#     CCP4 Secretary, Daresbury Laboratory, Warrington WA4 4AD, UK.
#
# Class for running SHELXE for c-alpha tracing
#
# Ronan Keegan 21/11/2014

import sys, os, string
import shlex, subprocess
import shutil

# Update the python path with the mrbump folders
if not "CCP4" in sorted(os.environ.keys()):
    raise RuntimeError('CCP4 not found')
if os.path.isdir(os.path.join(os.environ["CCP4"], "share", "mrbump")):
    sys.path.append(os.path.join(os.environ["CCP4"], "share", "mrbump", "include", "file_info")) 
    sys.path.append(os.path.join(os.environ["CCP4"], "share", "mrbump", "include", "tools")) 
else:
    sys.stdout.write("Error: Can't find MrBUMP folders\n\n")
    sys.exit()
   
import MTZ_parse
import MRBUMP_phs2mtz


def which(program):
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

# Test for required executables
if os.name == "nt":
   if not which("mtz2various.exe"):
      raise RuntimeError('mtz2various.exe not found')
else:
   if not which("mtz2various"):
      raise RuntimeError('mtz2various not found')
#if not which("shelxe"):
#    raise RuntimeError('shelxe not found')

class Shelxe:

   def __init__(self):
      self.pdbinFile=""
      self.mtzinFile=""
      self.hklinFile=""
      self.pdboutFile=""
      self.mtzoutFile=""
      self.phsoutFile=""

      self.shelxScriptFile="shelx-script.sh"
      self.script=""

      self.labin=dict([])
      self.debug=True

      self.mtz2variousLogfile="mtz2various.log"
      self.shelxeLogfile="shelxe.log"
      self.workingDIR=""
     
   def setDebug(self, debugValue):
      self.debug=debugValue

   def setShelxScriptFile(self, filename):
      self.shelxScriptFile=filename

   def setShelxeLogFile(self, filename):
      self.shelxeLogfile=filename

   def setShelxWorkingDIR(self, dirname):
      self.workingDIR=dirname

   def checkFileExist(self, filename, program, type):
      """ Check to see if the PDB file has been outout from Refmac in MrBUMP """

      if os.path.isfile(filename):
         status = 0
      else:
         status = 1
         sys.stdout.write("Warning: no file output from "+ program +"  for this job\n")
      
      return status

   def runMtz2Various(self, fp="FP", sigfp="SIGFP", free="FreeR_flag"):
      """ Run mtz2various to convert the MTZ file for Shelxe """

      self.hklinFile=os.path.splitext(os.path.split(self.mtzinFile)[-1])[0] + ".hkl"

      command_line='mtz2various HKLIN %s HKLOUT %s' % (self.mtzinFile, self.hklinFile)

      self.script += command_line + " << eof\n"

      key  = "LABIN FP=%s SIGFP=%s FREE=%s\n" % (fp, sigfp, free)
      key += "OUTPUT SHELX\n"
      key += "FSQUARED\n"
      key += "END\n"
   
      self.script += key + "eof\n\n"

      if os.name == "nt":
         process_args = shlex.split(command_line, posix=False)
         p = subprocess.Popen(process_args, bufsize=0, shell="False", stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
      else:
         process_args = shlex.split(command_line)
         p = subprocess.Popen(process_args, stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)


      #process_args = shlex.split(command_line)
      #p = subprocess.Popen(process_args, stdin = subprocess.PIPE,
      #                            stdout = subprocess.PIPE, stderr=subprocess.STDOUT)
  
      (child_stdin, child_stdout_and_stderr) = (p.stdin, p.stdout)
  
      # Write the keyword input
      child_stdin.write(key)
      child_stdin.close()
  
      # Watch the output for successful termination
      out=child_stdout_and_stderr.readline()
  
      log=open(self.mtz2variousLogfile, "w")

      while out:
         out=child_stdout_and_stderr.readline()
         log.write(out)
         if self.debug:
            sys.stdout.write(out)
  
      child_stdout_and_stderr.close()
      log.close()

      status=self.checkFileExist(self.hklinFile, "mtz2various", "hkl")
      if status==1:
         sys.stdout.write("Error: No output file from mtz2various\n")

   def runMtz2hkl(self, rfree_flag=""):
      """ Run mtz2hkl to convert the MTZ file for Shelx """

      if rfree_flag != "":
         command_line='mtz2hkl -2 orig -r %s' % rfree_flag
      else:
         command_line='mtz2hkl -2 orig'

      self.script += command_line + "\n"

      if os.name == "nt":
         process_args = shlex.split(command_line, posix=False)
         p = subprocess.Popen(process_args, bufsize=0, shell="False", stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
      else:
         process_args = shlex.split(command_line)
         p = subprocess.Popen(process_args, stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)

#      process_args = shlex.split(command_line)
#      p = subprocess.Popen(process_args, stdin = subprocess.PIPE,
#                                  stdout = subprocess.PIPE, stderr=subprocess.STDOUT)
  
      (child_stdin, child_stdout_and_stderr) = (p.stdin, p.stdout)
  
      # Write the keyword input
      child_stdin.close()
  
      # Watch the output for successful termination
      out=child_stdout_and_stderr.readline()
  
      while out:
         sys.stdout.write(out)
         out=child_stdout_and_stderr.readline()
  
      child_stdout_and_stderr.close()

      status=self.checkFileExist("orig.hkl", "mtz2hkl", "hkl")
      if status==1:
         sys.stdout.write("Error: No output file from mtz2hkl\n")

   def runShelxe(self, solvent, pdbinFile, mtzinFile, workingDIR, pdboutFile="shelxe-output.pdb", \
                 mtzoutFile="shelxe-output.mtz", phsoutFile="shelxe-output.phs", traceCycles=15, \
                 dmCycles=20, shelxeEXE="shelxe", WEB_PATH_START=0):
      """ Run shelxe """

      # Give a header output
      sys.stdout.write("############################################################\n")
      sys.stdout.write("Running Shelxe to do phase improvement and c-alpha trace...\n")
      sys.stdout.write("############################################################\n")
      sys.stdout.write("\n")

      self.workingDIR=workingDIR
      if os.path.isdir(self.workingDIR) == False:
          try:
              os.mkdir(self.workingDIR)
          except:
              sys.stderr.write("Error: Can't create working directory:\n  %s\n" % self.workingDIR)  
              sys.exit()
 
      sys.stdout.write("Running directory is:\n  %s \n" % self.workingDIR)
      sys.stdout.write("\n")

      self.mtzinFile=mtzinFile
      self.pdbinFile=pdbinFile
      self.pdboutFile=pdboutFile
      self.mtzoutFile=mtzoutFile
      self.phsoutFile=phsoutFile

      currDIR=os.getcwd()
      os.chdir(self.workingDIR)

      mtz = MTZ_parse.MTZ_parse()
      mtz.run_mtzdmp(self.mtzinFile)

      self.runMtz2Various(mtz.F, mtz.SIGF, mtz.FreeR_flag)
   
      if mtz.resolution<=2.0:
          free_lunch=" -e1.0 -l5"
      else:
          free_lunch=""

      if solvent > 1.0:
          frac_solvent=solvent/100.0
      else:
          frac_solvent=solvent

      command_line_args=' -a' + str(traceCycles) + ' -m' + str(dmCycles) + \
                        ' -q -s' + str(frac_solvent) + ' -o -n -t4' + free_lunch

      # Set up input files for shelxe
      if os.path.isfile("shelxe-input.hkl"):
         os.remove("shelxe-input.hkl")
      shutil.copyfile(self.hklinFile, "shelxe-input.hkl")
      self.script += "cp %s shelxe-input.hkl\n" % self.hklinFile
      if os.path.isfile("shelxe-input.pda"):
         os.remove("shelxe-input.pda")
      shutil.copyfile(self.pdbinFile, "shelxe-input.pda")
      self.script += "cp %s shelxe-input.pda\n\n" % self.pdbinFile

      command_line=shelxeEXE + ' shelxe-input.pda' + command_line_args
      self.script += command_line + "\n\n"

      file=open(self.shelxScriptFile, "w")
      file.write(self.script) 
      file.close()

      if os.name == "nt":
         process_args = shlex.split(command_line, posix=False)
         p = subprocess.Popen(process_args, bufsize=0, shell="False", stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
      else:
         process_args = shlex.split(command_line)
         p = subprocess.Popen(process_args, stdin = subprocess.PIPE,
                                stdout = subprocess.PIPE, stderr = subprocess.STDOUT)

      (child_stdin, child_stdout_and_stderr) = (p.stdin, p.stdout)
  
      # Write the keyword input
      child_stdin.close()
  
      # Watch the output for successful termination
      out=child_stdout_and_stderr.readline()
  
      cycle=1
      CAPTURE=False
      #capture_count=0
      #frag_count=0
      #GET_FRAGS=True
      #acl=0.0
      log=open(self.shelxeLogfile, "w")
      while out:
         log.write(out)
         if "residues left after pruning" in out and "divided into chains" in out:
            CAPTURE=True
            sys.stdout.write("---------------------------------\n")
            sys.stdout.write("SHELXE tracing cycle: %d\n" % cycle)
            sys.stdout.write("---------------------------------\n")
            sys.stdout.write("\n")
            cycle=cycle+1
         if self.debug == False and CAPTURE:
            sys.stdout.write("  " + string.strip(out) + "\n")
         #if CAPTURE and capture_count >= 1 and string.strip(out) != "" and GET_FRAGS:
         #   for i in string.split(out):
         #      if ":" not in i:
         #         frag_count=frag_count+1
         #         acl=(acl+float(i))/float(frag_count)
         #if CAPTURE and acl > 0.0 and string.strip(out) == "":
         #   GET_FRAGS=False
         if "CC for partial structure against native data" in out:
            CAPTURE=False
         #   frag_count=0
         #   acl=0
         #   capture_count=0
            sys.stdout.write("\n")
         if self.debug==False and "Best trace (cycle " in out:
            sys.stdout.write("  " + string.split(string.strip(out),"was")[0] + "\n")
            sys.stdout.write("\n")
         if self.debug:
            sys.stdout.write(out)
         #if CAPTURE:
         #   capture_count=capture_count+1
         out=child_stdout_and_stderr.readline()
 
      child_stdout_and_stderr.close()
      log.close()

      pdbstatus=self.checkFileExist("shelxe-input.pdb", "shelxe", "pdb")
      if pdbstatus==1:
         sys.stdout.write("Error: No output pdb file from shelxe\n")
         script = "# No output from SHELXE\n"
      else:
         shutil.move("shelxe-input.pdb", self.pdboutFile)
         script  = "mv shelxe-input.pdb %s\n" % self.pdboutFile

      phsstatus=self.checkFileExist("shelxe-input.phs", "shelxe", "phs")
      if phsstatus==1:
         sys.stdout.write("Error: No output phs file from shelxe\n")
      else:
         shutil.move("shelxe-input.phs", self.phsoutFile)
         script += "mv shelxe-input.phs %s\n\n" % self.phsoutFile

         c=MRBUMP_phs2mtz.PHS2MTZ()
         if mtz.FreeR_flag != "":
             c.phs2mtz(self.phsoutFile, self.pdboutFile, self.mtzoutFile, os.getcwd(), hklref=self.mtzinFile, freeLabel=mtz.FreeR_flag, resolution=mtz.resolution, debug=self.debug)
         else:
             c.phs2mtz(self.phsoutFile, self.pdboutFile, self.mtzoutFile, os.getcwd(), resolution=mtz.resolution, debug=self.debug)
       

      scriptFile=open(self.shelxScriptFile, "a")
      scriptFile.write(script)
      scriptFile.close()
 
      # Set the best scores
      if os.path.isfile(self.pdboutFile):
         score=0.0
         acl=0.0
         f=open(self.pdboutFile, "r")
         line=f.readline()
         try:
            score=float(string.split(line)[6].replace("%",""))
            acl=float(string.split(line)[7])/float(string.split(line)[10])
         except: 
            sys.stdout.write("Warning (shelxe_trace): Can't find Shelxe output score in output pdb\n")
         f.close()

      sys.stdout.write("Log file written to:\n  %s \n" % self.shelxeLogfile[WEB_PATH_START:])
      sys.stdout.write("\n")

      if pdbstatus != 1:
         sys.stdout.write("Output PDB file written to:\n  %s \n" % self.pdboutFile[WEB_PATH_START:])
         sys.stdout.write("\n")

      if self.debug == False:
        for file in self.hklinFile, "shelxe-input.hkl", "shelxe-input.pdb", \
                    "shelxe-input.pha", "shelxe-input.hat", \
                    "shelxe-input.pda": 
           if os.path.isfile(file):
              os.remove(file) 

      os.chdir(currDIR)

if __name__ == "__main__":

   if  len(sys.argv) == 7:
      pdbinFile=sys.argv[1]
      mtzinFile=sys.argv[2]
      solvent=float(sys.argv[3])
      workingDIR=sys.argv[4]
      pdboutFile=sys.argv[5]
      mtzoutFile=sys.argv[6]
   elif len(sys.argv) == 9:
      pdbinFile=sys.argv[1]
      mtzinFile=sys.argv[2]
      solvent=float(sys.argv[3])
      workingDIR=sys.argv[4]
      pdboutFile=sys.argv[5]
      mtzoutFile=sys.argv[6]
      traceCycles=sys.argv[7]
      dmCycles=sys.argv[8]
   else:
      sys.stdout.write("Usage: python ccp4iSHELXE.py <pdbin> <mtzin> <solvent> <working directory> <pdbout> <mtzout> [optional: <no. of auto tracing cycles> <dm cycles>]\n")
      sys.stdout.write("\n")
      sys.stdout.write("       e.g. python mrSHELXE.py foo.pdb foo.mtz 45.33 /tmp/test bar.pdb bar.mtz 15 20\n")
      sys.stdout.write("\n")
      sys.exit()
     
   SHELX=Shelxe()

   status=SHELX.checkFileExist(pdbinFile, "mrbump", "pdb")
   if status==1:
      sys.exit(1)

   status=SHELX.checkFileExist(mtzinFile, "mtzfile", "mtz")
   if status==1:
      sys.exit(1)

   SHELX.runShelxe(solvent, pdbinFile, mtzinFile, workingDIR, pdboutFile=pdboutFile, mtzoutFile=mtzoutFile, traceCycles=traceCycles, dmCycles=dmCycles)