#!/usr/bin/env python

import string, sys, os
import subprocess, shlex
import urllib


class getPDBFiles:

    def __init__(self):

        self.REFINE=False

        self.cif2mtzEXE=os.path.join(os.environ["CCP4"], "bin", "cif2mtz")
        self.uniqueifyEXE=os.path.join(os.environ["CCP4"], "bin", "uniqueify")
        self.refmacEXE=os.path.join(os.environ["CCP4"], "bin", "refmac5")

        #self.seqURL="http://www.rcsb.org/pdb/download/downloadFile.do?fileFormat=FASTA&compression=NO&structureId="
        self.seqURL="http://www.rcsb.org/pdb/download/downloadFastaFiles.do?structureIdList="
        self.pdbURL="http://www.rcsb.org/pdb/files/"
        #self.cifURL="http://www.rcsb.org/pdb/files/"
        self.cifURL="https://files.rcsb.org/download/"
       # https://files.rcsb.org/download/3C7T-sf.cif.gz
  
        self.scriptFILE=""
        self.pdbFILE=os.path.join(directory, code.upper(), code + ".pdb.gz")
        self.cifSFFILE=os.path.join(directory, code.upper(), code + ".cif.gz")
        self.fastaFILE=os.path.join(directory, code.upper(), code + ".fasta")
        
        self.code=""
        self.directory=""
        self.debug=False

    def get(self, code, directory="./"):
        """ Master function """

        self.code=code.lower()
        self.directory=directory

        if os.path.isdir(self.directory) == False:
            os.mkdir(self.directory)
        if os.path.isdir(os.path.join(self.directory, self.code.upper())) == False:
            os.mkdir(os.path.join(self.directory, self.code.upper()))
        
        urllib.urlretrieve(self.pdbURL + self.code.upper() + ".pdb.gz", self.pdbFILE)
        #urllib.urlretrieve(self.cifURL + "r" + self.code.upper() + "sf.ent.gz", self.cifSFFILE)
        urllib.urlretrieve(self.cifURL + self.code.upper() + "-sf.cif.gz", self.cifSFFILE)
        urllib.urlretrieve(self.seqURL + self.code.upper() + "&compressionType=uncompressed", self.fastaFILE)
        urllib.urlcleanup()
        
        currDIR=os.getcwd()
        os.chdir(os.path.join(self.directory, self.code.upper()))
        self.runJOB("gunzip -f " + self.code + ".pdb.gz")
        self.runJOB("gunzip -f " + self.code + ".cif.gz")
        self.runJOB(self.cif2mtzEXE + " hklin " + self.code + ".cif hklout " + self.code + ".mtz", key="end")
        if self.REFINE:
            self.runJOB(self.refmacEXE + " hklin " + self.code + ".mtz hklout " + self.code + \
                                         "_phased.mtz xyzin " + self.code + ".pdb xyzout " + self.code + "_phased.pdb", \
                                         key="ncyc 0\nend\n")
        if os.path.isfile(self.code + ".ent") and self.debug==False:
            os.remove(self.code + ".ent")
        if os.path.isfile(self.code + "_phased.pdb") and self.debug==False:
            os.remove(self.code + "_phased.pdb")
        os.chdir(currDIR)
        
    def runJOB(self, command_line, key="", logfile=""):
        """ Generic job runner """
    
        script = ""
        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()
    
        if logfile != "":
           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()
        else:
           while out:
              out=child_stdout_and_stderr.readline()
              sys.stdout.write(out)
           child_stdout_and_stderr.close()


if __name__ == "__main__":
   
    if len(sys.argv) != 3 and len(sys.argv) != 2:
        sys.stdout.write("Usage: getPDB <code> <directory>\n")
        sys.exit()
    if len(sys.argv) == 2: 
        directory="./"
    else:
        directory=sys.argv[2]
    code=sys.argv[1]
    
    os.environ["http_proxy"] = ""
    
    proxy=os.environ["http_proxy"]
    #if proxy == "":
    #    sys.stdout.write("Warning: Proxy not set!\n")

    G=getPDBFiles()
    G.get(code, directory=directory)