lmp_processing

#!/usr/bin/env python

import os
import sys
import subprocess
import shutil
import stat

def is_exe(fpath):
  return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

def which(program):
  fpath, fname = os.path.split(program)
  if fpath:
    if is_exe(program):
      return program
  else:
    for path in os.environ["PATH"].split(os.pathsep):
      path = path.strip('"')
      exe_file = os.path.join(path, program)
      if is_exe(exe_file):
        return exe_file

  return None

def write_nc_chunk_script():
  """ Writes the nextclip chunk script to the nextclip dir """
  
  with open('nextclip/nextclip_chunk.sh', 'w') as f:
    f.write("#!/usr/bin/env bash\n\n")

    f.write("lines=`wc -l $1|awk '{print $1}'`\n")
    f.write("chunkseqs=$((($lines/4)/$4 + 1))\n\n")

    f.write("split -l $((chunkseqs * 4)) $1 $1.part_ &\n")
    f.write("split -l $((chunkseqs * 4)) $2 $2.part_ &\n")
    f.write("wait\n\n")

    f.write("for i in $1.part_*; do\n")
    f.write("  nextclip -i $i -j `echo $i|sed 's/'$1'/'$2'/'` -o $3.$i > $3.$i.report &\n")
    f.write("done\n")
    f.write("wait\n\n")

    f.write("cat $3.$1*A_R1.fastq > $3_A_R1.fastq &\n")
    f.write("cat $3.$1*A_R2.fastq > $3_A_R2.fastq &\n")
    f.write("cat $3.$1*B_R1.fastq > $3_B_R1.fastq &\n")
    f.write("cat $3.$1*B_R2.fastq > $3_B_R2.fastq &\n")
    f.write("cat $3.$1*C_R1.fastq > $3_C_R1.fastq &\n")
    f.write("cat $3.$1*C_R2.fastq > $3_C_R2.fastq &\n")
    f.write("cat $3.$1*D_R1.fastq > $3_D_R1.fastq &\n")
    f.write("cat $3.$1*D_R2.fastq > $3_D_R2.fastq &\n")
    f.write("wait\n\n")

    f.write("rm $1.part_* $2.part_* $3.$1.part_*.fastq\n")
  st=os.stat('nextclip/nextclip_chunk.sh')
  os.chmod('nextclip/nextclip_chunk.sh', st.st_mode | stat.S_IEXEC)
	
def get_flash_stats(lib_prefix):
  """ Gets the useful stats from the FLASH and dedup log files """

  flash_log = os.path.join("flash", "{0}_flash.log".format(lib_prefix))
  dedup_log = os.path.join("flash", "{0}_dedup.log".format(lib_prefix))

  p = subprocess.Popen(['tail', flash_log], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  flash_text, err = p.communicate()
  if p.returncode != 0:
    raise IOError(err)
 
  f_fields = flash_text.replace("[FLASH]", "").split("\n")

  with open(dedup_log, 'r') as f:
    lines = f.readlines()

  lines = [x.strip() for x in lines]

  return "{0}:\n{1}\n{2}\n{3}\n{4}\n\nAfter de-duplication: {5}\n".format(lib_prefix, f_fields[1], f_fields[2], f_fields[3], f_fields[4], lines[5])

def get_nextclip_stats(lib_prefix):
  """ Gets stats from the .nc_counts files """

  nc_log = os.path.join("nextclip", "{0}.nc_counts".format(lib_prefix))

  with open(nc_log, 'r') as f:
    lines = f.readlines()

  lines = [x.strip() for x in lines]
  dict = {}
  for x in lines:
    (cat, count) = x.split()
    dict[cat] = count

  total = int(dict["A"]) + int(dict["B"]) + int(dict["C"])
  str = "{0}:\n     catA: {1}\n     catB: {2}\n     catC: {3}\n     catD: {4}\nTotal reads after processing: {5}\n".format(lib_prefix, dict["A"], dict["B"], dict["C"], dict["D"], total)

  return str

if __name__ == '__main__':

  if len(sys.argv) != 3 or sys.argv[1] == "--help":
    print "\n#### w2rap LMP processing ####\n"
    print "Usage: {0} libs_list ncpus\n".format(sys.argv[0])
    print "libs_list is a text file containing a list of your FASTQ LMP read files, eg.\n" \
    "/path/to/LIB1_R1.fastq\n/path/to/LIB1_R2.fastq\n/path/to/LIB2_R1.fastq\n/path/to/LIB2_R2.fastq\n\n" \
    "FASTQ read files must be uncompressed and end in _R1.fastq or _R2.fastq\n"
    sys.exit()

  libs_list = sys.argv[1]
  ncpus = sys.argv[2]

  print "\n#### w2rap LMP processing ####\n"

  # can we find the input file and does it look ok
  if not os.path.exists(libs_list):
    print "Cannot find libraries file {0}.\n".format(libs_list)
    sys.exit()

  # get the directory where this script is running
  #abspath = os.path.abspath(__file__)	
  #dname = os.path.dirname(abspath)	
  #bin_dir = os.path.join(dname, "bin")

  # if the bin directory exists, add it to the path
  #if os.path.exists(bin_dir):
  #  os.environ["PATH"] += os.pathsep + bin_dir
  
  # read env var and add to path
  bin_dir = os.environ['W2RAP_PATH']
  if bin_dir == None:
    print "No path to binaries, please set W2RAP_PATH environment variable"
    sys.exit()
  else:
    os.environ["PATH"] += os.pathsep + bin_dir

  # check the required executables exist in the path
  FLASH_PATH = which("flash")
  if FLASH_PATH == None:
    print "ERROR: Cannot find FLASH in PATH.\n"
    sys.exit()

  DEDUP_PATH = which("dedup_fastq")
  if DEDUP_PATH == None:
    print "ERROR: Cannot find dedup_fastq in PATH.\n"
    sys.exit()

  NXCLIP_PATH = which("nextclip")
  if NXCLIP_PATH == None:
    print "ERROR: Cannot find Nextclip in PATH.\n"
    sys.exit()

  print "FLASH found: {0}".format(FLASH_PATH)
  print "dedup_fastq found: {0}".format(DEDUP_PATH)
  print "Nextclip found: {0}".format(NXCLIP_PATH)

  cwd = os.getcwd()

  libraries = []
  with open(libs_list, 'r') as f:
    while True:
      r1_in = f.readline().strip()
      r2_in = f.readline().strip()
      if r2_in:
        r1 = os.path.abspath(r1_in)
        r2 = os.path.abspath(r2_in)
        if r1 == r2:
          print "Read 1 and read 2 files are the same - {0}.".format(r1)
          sys.exit()

        if not os.path.exists(r1):
          print "Read file {0} does not exist.".format(r1)
          sys.exit()

        if not os.path.exists(r2):
          print "Read file {0} does not exist.".format(r2)
          sys.exit()       

        # get the read length from r1 (required for flash)
        fp = open(r1)
        for i, line in enumerate(fp):
          if i == 1:
            break
        fp.close()
        read_length = len(line.strip("\n"))

        # get the library prefix
        base=os.path.basename(r1)
	prefix = os.path.splitext(base)[0].replace("_R1", "")

        libraries.append({"r1": r1, "r2": r2, "read_length": read_length, "prefix": prefix})
      else: break  # EOF

  print "Number of libraries to process: {0}".format(len(libraries))
  for lib in libraries:
    print lib["r1"], lib["r2"]

  print "\nRunning FLASH and de-duplicating combined reads..."
  if not os.path.exists("flash"):
    os.makedirs("flash")

  # write the script for flash stage
  with open('flash/run.sh', 'w') as f:
    f.write("#!/usr/bin/env bash\n")

    for lib in libraries:
      f.write("flash -t {0} -M {1} -o {2} {3} {4} > {2}_flash.log &\n".format(ncpus, lib["read_length"], lib["prefix"], lib["r1"], lib["r2"]))
    f.write("wait\n\n")

    for lib in libraries:
      f.write("dedup_fastq -i {0}.extendedFrags.fastq -o {0}.extendedFrags.dedup > {0}_dedup.log &\n".format(lib["prefix"]))
    f.write("wait\n\n")

    for lib in libraries:
      f.write("cat {0}.notCombined_1.fastq {0}.extendedFrags.dedup.fastq > {0}_extended_R1.fastq\n".format(lib["prefix"]))
      f.write("cat {0}.notCombined_2.fastq {0}.extendedFrags.dedup_rc.fastq > {0}_extended_R2.fastq\n".format(lib["prefix"]))
    f.write("wait\n\n")

    for lib in libraries:
      f.write("rm {0}.*\n".format(lib["prefix"]))
    f.write("wait\n")

  # run the script in the flash directory
  os.chdir("flash")
  st=os.stat('./run.sh')
  os.chmod('./run.sh', st.st_mode | stat.S_IEXEC)
  subprocess.call(['./run.sh'])
  os.chdir(cwd)

  # for each library, get the stats from flash and dedup
  for lib in libraries:
    print get_flash_stats(lib["prefix"])    

  print "Running Nextclip..."
  if not os.path.exists("nextclip"):
    os.makedirs("nextclip")

  # symlink the flashed files here
  for file in os.listdir("flash"):
    if file.endswith(".fastq"):
      os.symlink(os.path.join(cwd, "flash", file), os.path.join(cwd, "nextclip", file))

  # write the nextclip script
  write_nc_chunk_script()

  # write the script for nextclip stage
  with open('nextclip/run.sh', 'w') as f:
    f.write("#!/usr/bin/env bash\n")

    for lib in libraries:
      f.write("./nextclip_chunk.sh {0}_extended_R1.fastq {0}_extended_R2.fastq {0}_nc {1}\n".format(lib["prefix"], ncpus))
      f.write("rm *.report\n")

    for lib in libraries:
      f.write("wcA=`wc -l {0}_nc_A_R1.fastq | cut -d' ' -f 1`\n".format(lib["prefix"]))
      f.write("wcB=`wc -l {0}_nc_B_R1.fastq | cut -d' ' -f 1`\n".format(lib["prefix"]))
      f.write("wcC=`wc -l {0}_nc_C_R1.fastq | cut -d' ' -f 1`\n".format(lib["prefix"]))
      f.write("wcD=`wc -l {0}_nc_D_R1.fastq | cut -d' ' -f 1`\n".format(lib["prefix"]))
      f.write("let wcA=wcA/4\n")
      f.write("let wcB=wcB/4\n")
      f.write("let wcC=wcC/4\n")
      f.write("let wcD=wcD/4\n")
      f.write("printf \"A $wcA\\nB $wcB\\nC $wcC\\nD $wcD\\n\" > {0}.nc_counts\n".format(lib["prefix"]))

    for lib in libraries:
      f.write("cat {0}_nc_[ABC]_R1.fastq > {0}_nc_ABC_R1.fastq &\n".format(lib["prefix"]))
      f.write("cat {0}_nc_[ABC]_R2.fastq > {0}_nc_ABC_R2.fastq &\n".format(lib["prefix"]))
    f.write("wait\n\n")
    f.write("rm *_nc_?_R?.fastq *extended*\n")

  # run the script in the nextclip directory
  os.chdir("nextclip")
  st=os.stat('./run.sh')
  os.chmod('./run.sh', st.st_mode | stat.S_IEXEC)
  subprocess.call(['./run.sh'])
  os.chdir(cwd)

  for lib in libraries:
    print get_nextclip_stats(lib["prefix"]) 

  # get rid of the temporary flash dir and extra bits in the nextclip dir
  shutil.rmtree(os.path.join(cwd, "flash"))
  os.remove(os.path.join(cwd, "nextclip", "run.sh"))
  os.remove(os.path.join(cwd, "nextclip", "nextclip_chunk.sh"))
  for lib in libraries:
    os.remove(os.path.join(cwd, "nextclip", "{0}.nc_counts".format(lib["prefix"])))

print "DONE."