#! /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 phmmer 
# Ronan Keegan - 10/9/14
#

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

class PHHit:

   def __init__(self):
      self.chainName=""
      self.pdbName=""
      self.chainID=""
      self.rank=""
      self.prob=0.0
      self.evalue=0.0
      self.pvalue=0.0
      self.score=0.0
      self.ndomains=0
      self.alignment=""
      self.alnRange=""
      self.cols=0
      self.alignments=dict([])

class Phmmer:

   def __init__(self):
      self.phmmerEXE=os.path.join(os.environ["CCP4"], "libexec", "phmmer")
      self.cpu=1
      self.seqin=""
      self.seqdb=""
      self.seqout=""
      self.logfile=""
      self.alnfile=""
      self.iterations=1
      self.termination=False
      self.workingDIR=""
      
      self.resultsList=[]
      self.resultsDict=dict([])
 
   def getAlignments(self):
       """ Get the alignments for each of the matches in the results list """
   
       if os.path.isfile(self.alnfile) == False:
           sys.stdout.write("Phmmer Error: can't find alignments file: \n  %s\n" % self.alnfile)
           sys.exit()
       
       f=open(self.alnfile, "r")
       lines=f.readlines()
       f.close()

       for line in lines:
           for i in self.resultsList:
               if i in line and "#=" not in line[0:2]:
                   self.resultsDict[i].alignment+=string.split(line)[-1]
                   self.resultsDict[i].alnRange=string.split(string.split(line, "/")[-1])[0]

   def runPhmmer(self, seqin, seqdb, cpu=1, iterations=1, logfile="", alnfile="", debug=False):
      """ Run HHblits """

      if logfile == "":
          self.logfile="phmmer.log"
      else:
          self.logfile=logfile
         
      if alnfile == "":
          self.alnfile="phmmerAlignment.log"
      else:
          self.alnfile=alnfile
         
      self.cpu=cpu
      self.iterations=iterations
      self.seqin=seqin
      self.seqdb=seqdb

      self.termination=False

      # Set the command line
      command_line=self.phmmerEXE + " -A %s %s %s" \
                    % (self.alnfile, self.seqin, self.seqdb)
      if debug == True:
         sys.stdout.write("Phmmer command line:\n  %s\n" % command_line)
         sys.stdout.write("\n")

      # Launch program
      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_stdout, child_stdin) = (p.stdout, p.stdin)

      # Write the input
      child_stdin.close()

      # Open the log file for writing
      log=open(self.logfile, "w")

      # Watch the output for successful termination
      out=child_stdout.readline()
      CAPTURE=False
      count=1

      while out:
         if debug == True:
            sys.stdout.write(out)

         if '[ok]' in out:
            self.termination=True

         m=dict([])
         if CAPTURE:
            hit=string.split(out)
            if len(hit) == 9:
               self.resultsList.append(hit[8])
               self.resultsDict[hit[8]]=PHHit()
               self.resultsDict[hit[8]].rank=count
               self.resultsDict[hit[8]].chainName=hit[8]
               self.resultsDict[hit[8]].pdbName=hit[8][0:4]
               if len(hit[8])>=6:
                  self.resultsDict[hit[8]].chainID=hit[8][5:]
               else:
                  self.resultsDict[hit[8]].chainID="A"
               self.resultsDict[hit[8]].score=float(hit[1])
               self.resultsDict[hit[8]].evalue=float(hit[0])
               self.resultsDict[hit[8]].ndomains=int(hit[7])
               #self.resultsDict[hit[1]].prob=float(out[35:40])
               #self.resultsDict[hit[1]].prob=float(hit[len(hit)-9])
               #self.resultsDict[hit[1]].score=float(hit[len(hit)-6])
               #self.resultsDict[hit[1]].evalue=float(hit[len(hit)-8])
               #self.resultsDict[hit[1]].pvalue=float(hit[len(hit)-7])
               #self.resultsDict[hit[1]].cols=float(hit[len(hit)-4])
               count=count+1
            else:
               CAPTURE=False
            
         if "E-value" in out and "score" in out and "bias" in out and "Sequence" in out:
            CAPTURE=True
            out=child_stdout.readline()

         log.write(out)
         log.flush()
         out=child_stdout.readline()

      child_stdout.close()

      # Get the alignemnts for each of the hits
      self.getAlignments()

      # Scan the output fasta file for the sequence alignments 

      #align=[]
      #seqDict=dict([])
      #seq=""
      #fastaFile=open(seqout, "r")
      #line=fastaFile.readline()
      #while line:
      #   if "# No" in line:
      #      if len(align) != 0:
      #         count=0
      #         for i in align:
      #            if ">" in i:
      #               if seq != "":
      #                  seqDict[title]=seq
      #               seq=""
      #               title=i
      #               if count==0:
      #                  title=title+"T"
      #               if count==1:
      #                  title=">Target_alignment"
      #               count=count+1
      #            else:
      #               seq=seq+i
      #         if seq != "":
      #            seqDict[title]=seq
      #         ModelName=string.strip(title).replace(">","")
      #         self.resultsDict[ModelName].alignments=seqDict
      #         seq=""
      #         count=count+1
      #      align=[]
      #      seqDict=dict([])
      #   else:
      #      align.append(string.strip(line))
      #   line=fastaFile.readline()
     # 
     # if len(align) != 0:
     ##    count=0
     #    for i in align:
     #       if ">" in i:
     #          if seq != "":
     #             seqDict[title]=seq
     #          seq=""
     #          title=i
     #          if count==0:
     #             title=title+"T"
     #          if count==1:
     #             title=">Target_alignment"
     #          count=count+1
     #       else:
     #          seq=seq+i
     #    if seq != "":
     #       seqDict[title]=seq
     #    ModelName=string.strip(title).replace(">","")
     #    self.resultsDict[ModelName].alignments=seqDict
     #    seq=""
     #    count=count+1


if __name__ == "__main__":

   seqin     = os.path.join(os.environ["CCP4"], "examples", "toxd", "toxd.seq")
   seqdb     = os.path.join(os.environ["CCP4"], "share", "mrbump", "data", "pdb_ATOMseqs.txt")

   ph=Phmmer()
   ph.runPhmmer(seqin=seqin, seqdb=seqdb)
  
   print ph.resultsList
   print ph.resultsDict["4bd9_B"].ndomains