#! /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.
#
# A MrBUMP wrapper for phs2mtz - convert shelxe phs file to mtz format 
# Ronan Keegan - 10/11/14
#

import os, string, sys, shlex, subprocess

if not "CCP4" in sorted(os.environ.keys()):
    raise RuntimeError('CCP4 not found')


class PHS2MTZ:

   def __init__(self):

      if os.name == "nt":
          self.sftoolsEXE=os.path.join(os.environ["CCP4"], "bin", "sftools.exe")
          self.cadEXE=os.path.join(os.environ["CCP4"], "bin", "cad.exe")
          self.f2mtzEXE=os.path.join(os.environ["CCP4"], "bin", "f2mtz.exe")
      else:
          self.sftoolsEXE=os.path.join(os.environ["CCP4"], "bin", "sftools")
          self.cadEXE=os.path.join(os.environ["CCP4"], "bin", "cad")
          self.f2mtzEXE=os.path.join(os.environ["CCP4"], "bin", "f2mtz")

      self.pdbin=""
      self.cell=""
      self.symm=""
      self.resolution=0.0

      self.F2hkin=""
      self.F2hklout=""
      self.F2logfile=""

      self.CADhklin1=""
      self.CADhklin2=""
      self.CADhklout=""
      self.CADlogfile=""

      self.SFTOOLSlogfile=""

      self.script="#!/bin/bash\n\n"
      self.scriptFILE=""

      self.workingDIR=""
      self.debug=False

   def phs2mtz(self, phsin, pdbin, hklout, workingDIR, hklref="", f="", sigf="", freeLabel="", resolution=0.0,debug=False):
      """ run phs2mtz """
 
      try:
          self.debug=eval(os.environ["MRBUMP_DEBUG"])
      except:
          self.debug=debug

      self.workingDIR=workingDIR

      self.F2hklout   = os.path.join(self.workingDIR, "f2mtz_phs.mtz")
      self.F2logfile  = os.path.join(self.workingDIR, "f2mtz.log")
      self.SFTOOLSlogfile = os.path.join(self.workingDIR, "sftools.log")
      self.scriptFILE = os.path.join(self.workingDIR, "phs2mtz.sh")

      self.getPDBinfo(pdbin)
      self.F2MTZ(phsin, hklout=self.F2hklout)
      self.CAD(self.F2hklout, os.path.join(self.workingDIR, "cad_out.mtz"),\
               os.path.join(self.workingDIR, "cad.log"), hklin2=hklref, \
               f=f, sigf=sigf, freeLabel=freeLabel, resolution=resolution)
      self.SFTOOLS(self.CADhklout, hklout, self.SFTOOLSlogfile)

      if self.debug:
          self.cleanUP()
          f=open(self.scriptFILE, "w")
          f.write(self.script)
          f.close()

   def getPDBinfo(self, pdbin):
       """ Extract PDB stuff """

       self.pdbin=pdbin
   
       file=open(self.pdbin, "r")
       line=file.readline()
       line=file.readline()
       file.close()
       
       self.cell="  ".join(string.split(line)[1:7])
       self.symm="".join(string.split(line)[7:])

   def F2MTZ(self, hklin, hklout="", cell="", symm=""):
      """ Run f2mtz to do stuff """

      if cell != "":
         self.cell=cell

      if symm != "":
         self.symm=symm

      self.F2hklin=hklin
      if hklout != "":
         self.CADhklout=hklout

      command_line="%s hklin %s hklout %s" % (self.f2mtzEXE, hklin, hklout)

      # sftools seems to have a 66-character limit on file names so need to use relative path
      hklin=os.path.relpath(hklin)
      hklout=os.path.relpath(hklout)

      key  = "TITLE   f2mtz\n"
      key += "cell    %s\n" % self.cell
      key += "symm    %s\n" % self.symm
      key += "labout  H K L F FOM PHI SIGF\n"
      key += "CTYPOUT H H H F W P Q\n"
      key += "pname   f2mtz\n"
      key += "dname   f2mtz\n"
      key += "END\n"

      self.runJOB(command_line, key, self.F2logfile)

   def CAD(self, hklin1, hklout, logfile, hklin2="", f="FP", sigf="SIGFP", freeLabel="", resolution=0.0):
      """ Run cad to do stuff """
 
      if resolution != 0.0:
         self.resolution=resolution

      # file setup
      self.CADhklin1=hklin1
      if hklin2 != "":
         self.CADhklin2=hklin2
      self.CADhklout=hklout
      self.CADlogfile=logfile
 
      if hklin2 != "":
         command_line='%s hklin1 %s hklin2 %s hklout %s' % (self.cadEXE, self.CADhklin1, self.CADhklin2, self.CADhklout)
      else:
         command_line='%s hklin1 %s hklout %s' % (self.cadEXE, self.CADhklin1, self.CADhklout)

      if hklin2 != "":
         if resolution != 0.0:
            key  = "LABIN FILE 1 E1=PHI E2=FOM\n"
            key += "LABIN FILE 2 E1=%s E2=%s E3=%s\n" % (f, sigf, freeLabel)
            key += "LABOUT FILE 1 E1=PHI_SHELXE E2=FOM_SHELXE\n"
            key += "LABOUT FILE 2 E1=F E2=SIGF E3=%s\n" % freeLabel
            key += "RESOLUTION OVERALL 200.0 %.2lf\n" % resolution
            key += "END\n"
         else:
            key  = "LABIN FILE 1 E1=PHI E2=FOM\n"
            key += "LABIN FILE 2 E1=%s E2=%s E3=%s\n" % (f, sigf, freeLabel)
            key += "LABOUT FILE 1 E1=PHI_SHELXE E2=FOM_SHELXE\n"
            key += "LABOUT FILE 2 E1=F E2=SIGF E3=%s\n" % freeLabel
            key += "END\n"
      
      else:
         if resolution != 0.0:
            key  = "LABIN FILE 1 E1=F E2=SIGF E3=PHI E4=FOM\n"
            key += "LABOUT FILE 1 E1=F E2=SIGF E3=PHI_SHELXE E4=FOM_SHELXE\n"
            key += "RESOLUTION OVERALL 200.0 %.2lf\n" % resolution
            key += "END\n"
         else:
            key  = "LABIN FILE 1 E1=F E2=SIGF E3=PHI E4=FOM\n"
            key += "LABOUT FILE 1 E1=F E2=SIGF E3=PHI_SHELXE E4=FOM_SHELXE\n"
            key += "END\n"
   
      self.runJOB(command_line, key, self.CADlogfile)

   def SFTOOLS(self, hklin, hklout, logfile, FWT="FWT", PHWT="PHWT"):
      """ Run sftools to generate map coefficients """

      command_line=self.sftoolsEXE

      # sftools seems to have a 66-character limit on file names so need to use relative path
      hklin=os.path.relpath(hklin)
      hklout=os.path.relpath(hklout)

      key  = "mode batch\n"
      key += 'read "%s" mtz\n' % hklin
      key += "CALC F COL FWT = COL F COL FOM_SHELXE *\n"
      key += "CALC P COL PHWT = COL PHI_SHELXE 0 +\n"
      key += 'write "%s" mtz\n' % hklout
      key += "EXIT\n"
      key += "YES\n"

      self.runJOB(command_line, key, self.F2logfile)

   def runJOB(self, command_line, key, logfile):
      """ Generic job runner """
 
      self.script += command_line + " << eof\n" + 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)


      (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(logfile, "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()

   def cleanUP(self):
       """ Clean up after the job """
 
       if os.path.isfile(self.F2logfile):
           os.remove(self.F2logfile)
       if os.path.isfile(self.CADlogfile):
           os.remove(self.CADlogfile)
       if os.path.isfile(self.F2hklout):
           os.remove(self.F2hklout)
        
if __name__ == "__main__":
 
    if len(sys.argv) == 5:
       phsin=sys.argv[1]
       pdbin=sys.argv[2]
       hklout=sys.argv[3]
       workingDIR=sys.argv[4]
       c=PHS2MTZ()
       c.phs2mtz(phsin,pdbin,hklout,workingDIR,hklref="1ttz.mtz",freeLabel="FREE",debug=True)
    elif len(sys.argv) == 6:
       phsin=sys.argv[1]
       pdbin=sys.argv[2]
       hklin=sys.argv[3]
       workingDIR=sys.argv[4]
       resolution=float(sys.argv[5])
       c=PHS2MTZ()
       c.phs2mtz(phsin,pdbin,hklout,workingDIR,resolution)
    else:
       sys.stdout.write("Usage: python MRBUMP_phs2mtz.py <shelxe_phs> <shelxe_pdb> <mtz out file> <working directory> [optional <resolution>]\n")
       sys.exit()