content: #!python3
  """
  Predict protein pKa based on MCCE method.
  http://pka.engr.ccny.cuny.edu/
  
  Require MCCE 3.0 to work: https://anaconda.org/SalahSalah/mcce/files
  """
  import asyncio
  import glob
  import gzip
  import locale
  import logging
  import math
  import os
  import re
  import shutil
  import subprocess
  import sys
  import time
  from multiprocessing import Pool
  from urllib.request import urlopen
  
  import aioftp
  import pandas as pd
  import uvloop
  
  # Sapelo Locale is broken, quick fix
  locale.setlocale(locale.LC_ALL, "en_US.UTF-8")
  # Set working directory
  ROOTPATH = os.path.dirname(os.path.realpath(sys.argv[0]))
  os.chdir(ROOTPATH)
  # Log settings
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)
  handler = logging.FileHandler(f"./pKa_calculation_{__file__}.log")
  handler.setLevel(logging.INFO)
  formatter = logging.Formatter(
      "%(asctime)s\t%(levelname)s\t"
      "[%(filename)s:%(lineno)s -%(funcName)12s()]\t%(message)s"
  )
  handler.setFormatter(formatter)
  logger.addHandler(handler)
  
  
  class pdb:
      def __init__(self):
          self.all_ids = []
          self.download_ids = []  # Download -> Unzip -> Preprocess -> Calculate
          self.unzip_ids = []  # Unzip -> Preprocess -> Calculate
          self.preprocess_ids = []  # Preprocess -> Calculate
          self.ready_ids = []  # Calculate
          self.finished_ids = []  # Successfully calculated IDs
          self.error_ids = []  # Error in download, unzip, or calculation
          # IDs this script will work on (messy queue implementation)
          self.working_ids = []
  
      def load_id(self):
          """
          First try to get existing pKa values,
          then get the list of PDB files to download.
          """
          for folder in ["./pdb", "./annotation", "./results"]:
              try:
                  os.makedirs(folder)
              except OSError:
                  pass
          self.finished_ids = [id[-8:-4] for id in glob.glob("./results/*.pka")]
          logger.debug(f"{len(self.finished_ids)} finished files.")
          # Create file even at first run so that the results folder doesn't get deleted
          with open("./results/finished_ids.list", "a") as f:
              f.write("\n".join(self.finished_ids))
  
          self.ready_ids = list(set(
              [id[-12:-8].upper() for id in glob.glob("./pdb/*/*.pdb.bak")]) - set(self.finished_ids))
          logger.debug(f"{len(self.ready_ids)} files ready to be calculated.")
  
          self.preprocess_ids = list(set([id[-8:-4].upper() for id in glob.glob(
              "./pdb/*/*.pdb") if "out" not in id]) - set(self.finished_ids) - set(self.ready_ids))
          logger.debug(
              f"{len(self.preprocess_ids)} files ready to be preprocessed.")
  
          self.unzip_ids = [id[-11:-7].upper() for id in glob.glob("./*.ent.gz")]
          logger.debug(f"{len(self.unzip_ids)} files ready to be unzipped.")
  
          if not os.path.exists("./annotation/uniprot_id_mapping.dat"):
              with urlopen("ftp://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/idmapping/by_organism/HUMAN_9606_idmapping.dat.gz") as remotefile:
                  logger.debug(
                      "Saving UniProt ID mapping data since it doesn't exist...")
                  with open("./annotation/uniprot_id_mapping.dat.gz", "wb") as f:
                      f.write(remotefile.read())
              with gzip.open(
                  "./annotation/uniprot_id_mapping.dat.gz", "rb") as inFile, open(
                      "./annotation/uniprot_id_mapping.dat", "wb") as outFile:
                  shutil.copyfileobj(inFile, outFile)
              os.remove("./annotation/uniprot_id_mapping.dat.gz")
          else:
              logger.debug("UniProt ID mapping data exists.")
  
          logger.debug("Reading all possible PDB IDs...")
          annot = pd.read_csv("./annotation/uniprot_id_mapping.dat",
                              sep="\t", header=None,
                              names=["uniprot", "id", "value"])
          self.all_ids = annot.loc[annot.id == "PDB", "value"].tolist()
          self.download_ids = list(set(self.all_ids) - set(self.unzip_ids) - set(
              self.preprocess_ids) - set(self.ready_ids) - set(self.finished_ids))
          logger.info(
              f"{len(self.download_ids)} PDB files need to be downloaded.")
  
      def get_link(self, ids):
          """ Get PDB file links from:
              ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/pdb/ ,
              and create folders to store the files.
  
          Parameters
          ----------
              ids: list
                  The PDB IDs to download.
  
          Returns
          -------
              Links to download.
          """
          if isinstance(ids, list):
              ids = [id[:4].lower() for id in ids]  # pdb file IDs
              pdb_names = [f"{id}.ent.gz" for id in ids]  # pdb filenames
              # subdirectory of the pdb files
              pdbDirs = [id[1:3].lower() for id in ids]
              remoteaddr = [
                  f"ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/pdb/{pdbDir}/pdb{pdb_name}" for pdbDir, pdb_name in zip(pdbDirs, pdb_names)]
          else:
              raise TypeError(f"{id} is not a string or list.")
          return remoteaddr
  
      def make_dirs(self, ids):
          """Make sure the download directory exists."""
          for id in ids:
              try:
                  os.makedirs(os.path.join(ROOTPATH, "pdb", id.upper()))
              except OSError:
                  pass
  
      async def download_worker(self, session, url):
          """Download the given url to working directory."""
          url = url[len("ftp://ftp.wwpdb.org"):]
          logger.debug(f"Downloading {url}")
          try:
              await session.download(url)
              self.unzip_ids.append(url[-11:-7].upper())
          except Exception as e:
              self.error_ids.append(url[-11:-7].upper())
              logger.warning(f"Error when downloading {url}: {e}")
  
      async def download_session(self, sem, work_queue):
          """ Get urls from the queue and pass to worker.
  
          Parameters
          ----------
              sem: asyncio.Semaphore object
              work_queue: asyncio.Queue object
          """
          while not work_queue.empty():
              url = await work_queue.get()
              logger.debug(f"Got url from queue: {url}")
              async with sem:
                  async with aioftp.ClientSession("ftp.wwpdb.org") as session:
                      await self.download_worker(session, url)
  
      def download_queue(self, urls):
          """ Create a queue to download all the given urls.
  
          Parameters
          ----------
              urls: list
                  A list of urls to download.
  
          Returns
          -------
              Downloaded "*.ent.gz" files in working directory.
          """
          logger.debug(f"{len(urls)} urls to download.")
          loop = uvloop.new_event_loop()
          asyncio.set_event_loop(loop)
          q = asyncio.Queue()
          sem = asyncio.Semaphore(10)
          [q.put_nowait(url) for url in urls]
          tasks = [asyncio.ensure_future(self.download_session(sem, q))
                   for _ in range(len(urls))]
          loop.run_until_complete(asyncio.gather(*tasks))
          # Zero-sleep to allow underlying connections to close
          loop.run_until_complete(asyncio.sleep(0))
          loop.close()
  
      def check_mcce(self):
          """Check if MCCE 3.0 exists."""
          if not os.path.exists(os.path.join(ROOTPATH, "mcce3.0")):
              if not os.path.exists(os.path.join(ROOTPATH, "mcce3.0.tar.bz2")):
                  logger.debug("MCCE isn't downloaded yet. Retrieving...")
                  with urlopen("https://anaconda.org/SalahSalah/mcce/3.0/download/linux-64/mcce-3.0-0.tar.bz2") as remotefile:
                      with open("./mcce-3.0-0.tar.bz2", 'wb') as f:
                          f.write(remotefile.read())
              subprocess.run(["tar", "-xjf", "mcce-3.0-0.tar.bz2"])
              shutil.move("./info/recipe/mcce3.0", "./mcce3.0")
              shutil.rmtree(os.path.join(ROOTPATH, "info"), ignore_errors=True)
              shutil.rmtree(os.path.join(ROOTPATH, "bin"), ignore_errors=True)
          else:
              logger.info("MCCE 3.0 exists, proceeding to calculation...")
  
      def unzip(self, id):
          """Unzip downloaded *.ent.gz file."""
          try:
              saved_pdb = os.path.join(ROOTPATH, "pdb", id, f"{id}.pdb")
              with gzip.open(f"pdb{id.lower()}.ent.gz", "rb") as inFile, open(saved_pdb, "wb") as outFile:
                  shutil.copyfileobj(inFile, outFile)
              os.remove(f"pdb{id.lower()}.ent.gz")
              self.preprocess_ids.append(id)
          except Exception as e:
              self.error_ids.append(id)
              logger.warning(f"Unzip of {id} unsuccessful: {e}")
  
      def preprocess(self, id, backup=True):
          """
          This program will:
          1) strip lines other than ATOM and HETATM records
          2) keep the first model of an NMR structure
          3) delete H and D atoms
          4) MSE to MET residue
          5) keep only one atom alternate position
          6) keep defined chains, if chain ID(s) are given in command
          7) remove some cofactors and salt ions
  
          Parameters
          ----------
              id: str
                  The PDB ID to find the file.
              backup: bool, optional
                  Whether to backup the original file or not. Default is True,
                  and save to "original.bak".
  
          Returns
          -------
              Nothing, modify the file in place.
          """
          removable_res = [
              " ZN", "PCA", "XYP", " NA", " CL", " CA", " MG", " MN", "HOH"
          ]
          model_start = False
          newlines = []
          ID = id.upper()
          filepath = os.path.join(ROOTPATH, "pdb", ID, f"{ID}.pdb")
          if backup:
              shutil.copy2(filepath, f"{filepath}.bak")
          with open(filepath) as f:
              for line in f:
                  if line[:5] == "MODEL":
                      model_start = True
                  if model_start and line[:6] == "ENDMDL":
                      break
                  if line[:6] != "ATOM  " and line[:6] != "HETATM":
                      continue  # discard non ATOM records
                  if line[13] == "H" or line[12] == "H":
                      continue
                  if line[16] == "A":
                      line = f"{line[:16]} {line[17:]}"
                  elif line[16] != " ":
                      continue  # delete this line, alternative posion is not A or empty
                  if line[:6] == "HETATM" and line[17:20] == "MSE":
                      if line[12:15] == "SE ":
                          line = f"ATOM  {line[6:12]} SD{line[15:17]}MET{line[20:]}"
                      else:
                          line = f"ATOM  {line[6:17]}MET{line[20:]}"
                  res = line[17:20]
                  if res in removable_res:
                      continue
                  newlines.append(line.rstrip())
          with open(filepath, "w") as f:
              f.write("\n".join(newlines))
          logger.debug(f"{ID} preprocessing complete.")
  
      def set_params(self, id, quickrun=True):
          """
          Set the parameters for MCCE.
  
          Parameters
          ----------
              id: str
                  The PDB ID of the file.
              quickrun: bool, optional
                  Use "run.prm.quick" or "run.prm.default".
  
          Returns
          -------
              run.prm: a file describing the parameters that points to the PDB file.
          """
          pkgpath = os.path.join(ROOTPATH, "mcce3.0")
          ID = id.upper()
          filepath = os.path.join(ROOTPATH, "pdb", ID)
          newlines = []
          if quickrun:
              shutil.copy2(
                  os.path.join(pkgpath, "run.prm.quick"),
                  os.path.join(filepath, "run.prm")
              )
          else:
              shutil.copy2([
                  os.path.join(pkgpath, "run.prm.default"),
                  os.path.join(filepath, "run.prm")
              ])
          with open(os.path.join(filepath, "run.prm")) as f:
              for line in f:
                  line = line.rstrip()
                  if line.endswith("(INPDB)"):
                      line = re.sub(r"^[^\s]+", fr"{id}.pdb", line)
                  if line.endswith(("(DO_PREMCCE)", "(DO_ROTAMERS)",
                                    "(DO_ENERGY)", "(DO_MONTE)")):
                      line = re.sub(r"^f", r"t", line)
                  if line.endswith("(EPSILON_PROT)"):
                      line = re.sub(r"^[\d\.]+", r"8.0", line)
                  if line.startswith("/home/mcce/mcce3.0"):
                      line = re.sub(r"^/.*3\.0", pkgpath,
                                    line)
                  newlines.append(line)
          with open(os.path.join(filepath, "run.prm"), "w") as f:
              f.write("\n".join(newlines))
          self.ready_ids.append(ID)
          logger.debug(f"Parameters set for {ID}.")
  
      def split_ready_ids(self, num):
          """ A naive queue implementation for multiple scripts.
  
          Parameters
          ----------
              num: int
                  Which part of the IDs to work on.
  
          Returns
          -------
              A list of the actual IDs to work on, and save the lists of IDs for
              other scripts to work with if this is the first instance.
          """
          if os.path.isfile(os.path.join(ROOTPATH, "results", "working_ids.list")):
              with open(os.path.join(ROOTPATH, "results", f"working_ids.list{num}"), "r") as f:
                  self.working_ids = [line.strip() for line in f]
          else:
              n = math.ceil(len(self.ready_ids) / 10)
              self.working_ids = [self.ready_ids[i:i + n]
                                  for i in range(0, len(self.ready_ids), n)]
              metafile = []
              for i, ids in enumerate(self.working_ids):
                  metafile.append(os.path.join(
                      ROOTPATH, "results", f"working_ids.list{i}"))
                  with open(os.path.join(ROOTPATH, "results", f"working_ids.list{i}"), "w") as f:
                      f.write("\n".join(ids))
                  logger.debug(
                      f"Saved {len(ids)} IDs to file working_ids.list{i} .")
              with open(os.path.join(ROOTPATH, "results", "working_ids.list"), "w") as f:
                  f.write("\n".join(metafile))
              self.working_ids = self.working_ids[num]
  
      def calc_pka(self, id, clean=True):
          """ Calculate protein pKa values using MCCE.
              https://sites.google.com/site/mccewiki/home
  
          Parameters
          ----------
              id: str
                  The PDB ID of the protein calculated.
              clean: bool, optional
                  Only keep the PDB file, run log and pKa output.
  
          Returns
          -------
              A set of files in a subdirectory named after the ID.
              See user manual for detail.
          """
          id = id.upper()
          os.chdir(os.path.realpath(os.path.join(ROOTPATH, "pdb", id)))
          logger.info(f"{id} calculation started.")
          start = time.time()
          with open(f"{id}.run.log", "w") as f:
              subprocess.run(f"{ROOTPATH}/mcce3.0/mcce", stdout=f)
          with open(f"{id}.run.log", "rb") as f:
              last = f.readlines()[-1].decode().lstrip()
          if last.startswith(("Fatal", "FATAL", "WARNING", "STOP")):
              self.error_ids.append(id)
              logger.warning(
                  f"{id} calculation aborted after {time.time() - start}s, due to {last}")
          else:
              self.finished_ids.append(id)
              logger.info(
                  f"{id} calculation finished, used {time.time() - start}s.")
              shutil.move("pK.out", os.path.join(
                  ROOTPATH, "results", f"{id}.pka"))
          if clean:
              del_list = [i for i in os.listdir() if i not in (
                  "pK.out", f"{id}.run.log", f"{id}.pdb.bak")]
              [os.remove(item) for item in del_list]
  
  
  if __name__ == "__main__":
      x = pdb()
      x.load_id()
      urls = x.get_link(x.download_ids)
      x.make_dirs(x.all_ids)
      x.download_queue(urls)
  
      x.check_mcce()
      for id in x.unzip_ids:
          x.unzip(id)
      for id in x.preprocess_ids:
          try:
              x.preprocess(id)
              x.set_params(id)
          except Exception as e:
              x.error_ids.append(id)
              logger.warning(f"Preprocess of {id}: {e}")
      # subprocess.run(["find", ".", "-type", "d", "-empty", "-delete"])
  
      x.split_ready_ids(0)  # 0 - 9, run 0 first to generate other lists
  
      with Pool(os.cpu_count()) as p:
          p.map(x.calc_pka, x.working_ids)
  
      with open("./results/finished_ids.list", "a") as f:
          f.write("\n".join(x.working_ids))
  
      with open("./results/error_ids.list", "a") as f:
          f.write("\n".join(x.error_ids))
  