#!/usr/bin/env ccp4-python
#
#     Copyright (C) 2012 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 in Ample for c-alpha tracing
#
# Ronan Keegan 04/05/2012

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

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 environment variables and required executables

if not "CCP4" in sorted(os.environ.keys()):
    raise RuntimeError('CCP4 not found')
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.phsoutFile=""

        self.phaserBestCC=0.0
        self.molrepBestCC=0.0
        self.phaserBestAvgChainLen=0.0
        self.molrepBestAvgChainLen=0.0
        self.phaserBestMaxChainLen=0
        self.molrepBestMaxChainLen=0
        self.phaserBestNoChains=0
        self.molrepBestNoChains=0
        self.runTime=0.0

        self.phaserShelxPDB="Not set"
        self.molrepShelxPDB="Not set"
        self.shelxScriptFile="shelx-script.sh"
        self.script=""

        self.labin=dict([])

        try:
            self.debug=eval(os.environ['AMPLE_DEBUG'])
        except:
            self.debug=False

        self.mtz2variousLogfile="mtz2various.log"
        self.shelxeLogfile="shelxe.log"
        self.workingDIR=""
        
        self._logParsed=False
        self._succeeded=False
        self.resultsDict = { 'SHELXE_CC' : None,
                            'SHELXE_ACL' : None,
                            'SHELXE_MCL' : None,
                            'SHELXE_NC' : None,
                            'SHELXE_time' : None,
                            'SHELXE_version' : None,
                            'SHELXE_wMPE' : None,
                            'SHELXE_os' : None, }

    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,ftype):
        """ 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 parseLog(self):
        if not os.path.isfile(self.shelxeLogfile):
            self._succeeded=False
            return False
        
        sp = parse_shelxe.ShelxeLogParser(self.shelxeLogfile)
        self.resultsDict["SHELXE_CC"] = sp.CC
        self.resultsDict["SHELXE_ACL"] = sp.avgChainLength
        self.resultsDict["SHELXE_MCL"] = sp.maxChainLength
        self.resultsDict["SHELXE_NC"] = sp.numChains
        self.resultsDict["SHELXE_time"] = sp.cputime
        self.resultsDict["SHELXE_version"] = sp.version
        self.resultsDict["SHELXE_wMPE"] = sp.wMPE
        self.resultsDict["SHELXE_os"] = sp.originShift
        
        if sp.CC and sp.CC >= 20 and sp.avgChainLength and sp.avgChainLength >= 8:
            self._succeeded=True
        
        self._logParsed=True
        return True

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

        self.hklinFile=os.path.splitext(self.mtzinFile)[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, resolution, mrProgram, pdbinFile, mtzinFile, pdboutFile="shelxe-output.pdb", phsoutFile="shelxe-output.phs", traceCycles=20, \
                  fp="FP", sigfp="SIGFP", free="FreeR_flag", shelxeEXE="shelxe", WEB_PATH_START=0, native=None, PHASES=False, invert_space_group=False):
        """ 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")

        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.phsoutFile=phsoutFile

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

        self.runMtz2Various(fp, sigfp, free)

        if resolution<=2.0:
            free_lunch=" -e1.0 -l5"
        else:
            free_lunch=""

        frac_solvent=solvent/100.0

        # 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

        # if we have phases we should use them in shelxe
        if PHASES:
            command_line=shelxeEXE + ' shelxe-input.pda shelxe-input_fa -a' + str(traceCycles) + ' -q -s' + str(frac_solvent) + ' -o -n -t4' + free_lunch
        else:
            command_line=shelxeEXE + ' shelxe-input.pda -a' + str(traceCycles) + ' -q -s' + str(frac_solvent) + ' -o -n -t4' + free_lunch

        if native != None:
            shutil.copyfile(native, "shelxe-input.ent")
            command_line=command_line+ ' -x'

        if invert_space_group:
            command_line=command_line+ ' -i'

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

        with open(self.shelxScriptFile, "w") as f:
            f.write(self.script)

        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()

        # Set the output file names
        if invert_space_group:
            outputPDBFile="shelxe-input_i.pdb"
            outputPHSFile="shelxe-input_i.phs"
            outputLSTFile="shelxe-input_i.lst"
            outputPHAFile="shelxe-input_i.pha"
            outputHATFile="shelxe-input_i.hat"
        else:
            outputPDBFile="shelxe-input.pdb"
            outputPHSFile="shelxe-input.phs"
            outputLSTFile="shelxe-input.lst"
            outputPHAFile="shelxe-input.pha"
            outputHATFile="shelxe-input.hat"

        pdbstatus=self.checkFileExist(outputPDBFile, "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(outputPDBFile, self.pdboutFile)
            script  = "mv %s %s\n" % (outputPDBFile, self.pdboutFile)

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

        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
            mcl=0
            noChains=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])
                noChains=int(string.split(line)[10])
            except:
                sys.stdout.write("Warning (shelxe_trace): Can't find Shelxe output score in output pdb\n")
            line=f.readline()
            currMin=1000000
            currMax=0
            maxLenFrag=0
            lenFrag=0
            first=True
            while line:
                if "ATOM" in line[0:4]:
                    resNo=int(string.split(line)[5])
                    if first==True:
                        currMin=resNo
                        first=False
                    if first == False:
                        if resNo < currMin:
                            currMin=resNo
                        if resNo > currMax:
                            currMax=resNo
                        if resNo < currMax:
                            lenFrag=currMax-currMin
                            if lenFrag > maxLenFrag:
                                maxLenFrag=lenFrag
                line=f.readline()
            f.close()
            if mrProgram.upper() == "PHASER":
                self.phaserBestCC=score
                self.phaserBestAvgChainLen=acl
                self.phaserBestMaxChainLen=maxLenFrag
                self.phaserBestNoChains=noChains
            if mrProgram.upper() == "MOLREP":
                self.molrepBestCC=score
                self.molrepBestAvgChainLen=acl
                self.molrepBestMaxChainLen=maxLenFrag
                self.molrepBestNoChains=noChains

        # Get the run time for SHELXE
        if os.path.isfile(outputLSTFile):
            lstfile=open(outputLSTFile, "r")
            line=lstfile.readline()
            while line:
                if "Total time: " in line and "finished at" in line:
                    try:
                        self.runTime=float(string.split(line)[-3])
                    except:
                        sys.stdout.write("SHELXE log file read error: can't read time taken\n")
                line=lstfile.readline()
            lstfile.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", outputPDBFile, \
                        outputPHAFile, outputHATFile, \
                        "shelxe-input.pda":
                if os.path.isfile(file):
                    os.remove(file)

        os.chdir(currDIR)

    def results(self, mrProgram):
        """  Compile a results log for this Shelx job """

        sys.stdout.write("SHELX>>> \n")
        sys.stdout.write("SHELX>>> ########################################################################\n")
        sys.stdout.write("SHELX>>> Shelx results for PDB: " + self.pdbinFile + "\n")
        sys.stdout.write("SHELX>>> \n")
        if mrProgram.upper() == "PHASER":
            sys.stdout.write("SHELX>>>    Best CC score for Phaser: " + str(self.phaserBestCC) + "\n")
            sys.stdout.write("SHELX>>>    Output PDB file: " + self.phaserShelxPDB + "\n")
        if mrProgram.upper() == "MOLREP":
            sys.stdout.write("SHELX>>>    Best CC score for Molrep: " + str(self.molrepBestCC) + "\n")
            sys.stdout.write("SHELX>>>    Output PDB file: " + self.molrepShelxPDB + "\n")
        sys.stdout.write("SHELX>>> \n")
        sys.stdout.write("SHELX>>> ########################################################################\n")
        sys.stdout.write("SHELX>>> \n")

    def succeeded(self):
        if not self._logParsed:
            self.parseLog()
        return self._succeeded


if __name__ == "__main__":

    if  len(sys.argv) == 7:
        pdbinFile=sys.argv[1]
        mtzinFile=sys.argv[2]
        nmol=int(sys.argv[3])
        solvent=float(sys.argv[4])
        resolution=float(sys.argv[5])
        mrProgram=sys.argv[6]
        FP="FP"
        SIGFP="SIGFP"
        FREE="FREE"
    elif  len(sys.argv) == 10:
        pdbinFile=sys.argv[1]
        mtzinFile=sys.argv[2]
        nmol=int(sys.argv[3])
        solvent=float(sys.argv[4])
        resolution=float(sys.argv[5])
        mrProgram=sys.argv[6]
        FP=sys.argv[7]
        SIGFP=sys.argv[8]
        FREE=sys.argv[9]
    else:
        sys.stdout.write("Usage: python MRBUMP_Shelxe.py <pdb> <mtz> <nmol> <solvent> <mtz resolution> <MR program> [<FP> <SIGFP> <FREE>]\n")
        sys.stdout.write("\n")
        sys.stdout.write("       e.g. python MRBUMP_Shelxe.py foo.pdb foo.mtz 2 45.33 1.543 PHASER FP SIGFP FREE\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.runMtz2Various(fp=FP, sigfp=SIGFP, free=FREE)
    SHELX.runShelxe(solvent, resolution, mrProgram, pdbinFile, traceCycles=1)
    SHELX.results(mrProgram)