quackers_stages.py

import os
import re
import sys
import time
from datetime import datetime as dt
import quackers_commands as q_com
import MetaPro_utilities_v2 as mp_util


class q_stage:
    def __init__(self, out_path, path_obj, dir_obj, args_pack):
        self.path_obj = path_obj
        self.dir_obj = dir_obj
        self.bin_tools = ["cct", "mbat2", "mbin2"]
        self.start_s_path   = args_pack["s_path"]
        self.start_f_path  = args_pack["p1_path"]
        self.start_r_path  = args_pack["p2_path"]
        self.job_control = mp_util.mp_util(out_path, self.path_obj.bypass_log_name)
        self.op_mode = args_pack["op_mode"]
        self.quality_encoding = "64"
        if(self.op_mode == "single"):
            self.quality_encoding = self.job_control.determine_encoding(self.start_s_path)
        else:
            self.quality_encoding = self.job_control.determine_encoding(self.start_f_path)

        
        self.command_obj = q_com.command_obj(path_obj, dir_obj, self.quality_encoding)
        self.mspades_contig_fail = False

        if(os.path.exists(self.dir_obj.assembly_alt_contigs)):
            print(dt.today(), "metaspades contigs overrided with megahit-backup")
            self.dir_obj.assembly_contigs = self.dir_obj.assembly_alt_contigs
        
        
        

        self.hosts_bypassed = False

        

    def check_host_bypass(self):
        #checks if there are hosts. if not, trigger the bypass
        list_of_hosts = sorted(self.path_obj.config["hosts"].keys())

        if(len(list_of_hosts) == 0):
            #skip host-cleaning. move data
            self.hosts_bypassed = True 
            print(dt.today(), "no hosts used: bypassing host filter")
            if(self.op_mode == "single"):
                self.dir_obj.host_final_s = self.dir_obj.clean_dir_final_s
                #self.dir_obj.host_final_s = self.dir_obj.start_s
                
            elif(self.op_mode == "paired"):
                self.dir_obj.host_final_f = self.dir_obj.clean_dir_final_f
                self.dir_obj.host_final_r = self.dir_obj.clean_dir_final_r
                self.dir_obj.host_final_s = self.dir_obj.clean_dir_final_s
                #self.dir_obj.host_final_f = self.dir_obj.start_f
                #self.dir_obj.host_final_r = self.dir_obj.start_r
    
    def low_quality_filter(self):
        if (os.path.exists(self.dir_obj.clean_dir_mkr)):
            print(dt.today(), "skipping cleaning step")
        else:
             
            command = self.command_obj.adapterremoval_command(self.quality_encoding, self.dir_obj.clean_dir_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.clean_dir_job, command)
            

        self.job_control.wait_for_mp_store()
        
        if(self.dir_obj.check_mkr_clean):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.clean_dir)
        else:
            message_line = str(dt.today()) + " broken at: "+ self.path_obj.clean_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")

    def host_filter(self):
        #launch for each new ref path
        #walk through each host and launch a bwa job
        list_of_hosts = sorted(self.path_obj.config["hosts"].keys())
        list_of_mkrs = list()
        if(len(list_of_hosts) == 0):
            #skip host-cleaning. move data
            self.hosts_bypassed = True 
            print(dt.today(), "no hosts used: bypassing host filter")
            if(self.op_mode == "single"):
                self.dir_obj.host_final_s = self.dir_obj.clean_dir_final_s
                #self.dir_obj.host_final_s = self.dir_obj.start_s
                
            elif(self.op_mode == "paired"):
                self.dir_obj.host_final_f = self.dir_obj.clean_dir_final_f
                self.dir_obj.host_final_r = self.dir_obj.clean_dir_final_r
                self.dir_obj.host_final_s = self.dir_obj.clean_dir_final_s
                #self.dir_obj.host_final_f = self.dir_obj.start_f
                #self.dir_obj.host_final_r = self.dir_obj.start_r

        else:

            
            for host_key in list_of_hosts:
                
                
                host_ref_path = self.path_obj.hosts_path_dict[host_key]
                ref_basename = os.path.basename(host_ref_path)
                ref_basename = ref_basename.split(".")[0]
                host_bt2_marker_path_p = os.path.join(self.dir_obj.host_dir_top, ref_basename + "_p_host_bt2_mkr")
                host_bt2_marker_path_s = os.path.join(self.dir_obj.host_dir_top, ref_basename + "_s_host_bt2_mkr")
                list_of_mkrs.append(host_bt2_marker_path_p)
                list_of_mkrs.append(host_bt2_marker_path_s)
                if(os.path.exists(host_bt2_marker_path_p)):
                    print("skipping Host filter P:", ref_basename)
                else:
                
                    #command = self.command_obj.clean_reads_bwa_simple_p(host_ref_path, ref_basename, self.dir_obj.clean_dir_final_f, self.dir_obj.clean_dir_final_r, host_bwa_marker_path)
                    command = self.command_obj.host_rem_bt2_p(host_ref_path, ref_basename, self.dir_obj.clean_dir_final_f, self.dir_obj.clean_dir_final_r, host_bt2_marker_path_p)
                    #command = self.command_obj.clean_reads_bwa_simple_p(host_ref_path, ref_basename, self.dir_obj.start_f, self.dir_obj.start_r, host_bwa_marker_path)
                    #self.job_control.launch_and_create_v2_with_mp_store(script_path, command)
                    script_path = os.path.join(self.dir_obj.host_dir_top, "host_filter_p_" + ref_basename + ".sh")
                    self.job_control.launch_and_create_v2_with_mp_store(script_path, command)
                
                if(os.path.exists(host_bt2_marker_path_s)):
                    print(dt.today(), "skipping host filter S:", ref_basename)
                else:
                    #command = self.command_obj.clean_reads_bwa_simple_s(host_ref_path, ref_basename, self.dir_obj.clean_dir_final_s, host_bwa_marker_path)
                    command = self.command_obj.host_rem_bt2_s(host_ref_path, ref_basename, self.dir_obj.clean_dir_final_s, host_bt2_marker_path_s)
                    #command = self.command_obj.clean_reads_bwa_simple_s(host_ref_path, ref_basename, self.dir_obj.start_s, host_bwa_marker_path)
                    #self.job_control.launch_and_create_v2_with_mp_store(script_path, command)
                    script_path = os.path.join(self.dir_obj.host_dir_top, "host_filter_s_" + ref_basename + ".sh")
                    self.job_control.launch_and_create_v2_with_mp_store(script_path, command)
            
            self.job_control.wait_for_mp_store()


            
            for host_key in list_of_hosts:
                host_ref_path = self.path_obj.hosts_path_dict[host_key]
                ref_basename = os.path.basename(host_ref_path)
                ref_basename = ref_basename.split(".")[0]
                sam_sift_marker_path_p = os.path.join(self.dir_obj.host_dir_top, ref_basename + "_sift_mkr_p")
                sam_sift_marker_path_s = os.path.join(self.dir_obj.host_dir_top, ref_basename + "_sift_mkr_s")
                list_of_mkrs.append(sam_sift_marker_path_p)
                list_of_mkrs.append(sam_sift_marker_path_s)
                sam_path = ""
                score_out_path = ""
                sam_name = ref_basename
                sam_path_p = os.path.join(self.dir_obj.host_dir_sam, sam_name + "_p.sam")
                sam_path_s = os.path.join(self.dir_obj.host_dir_sam, sam_name + "_s.sam")
                
                score_out_path_s = os.path.join(self.dir_obj.host_dir_sam, sam_name + "_s_score_bt2.out")
                score_out_path_p = os.path.join(self.dir_obj.host_dir_sam, sam_name + "_p_score_bt2.out")

                
                if(os.path.exists(sam_sift_marker_path_p)):
                    print(dt.today(), "skipping sam sift on P:", ref_basename)
                else:
                    #if(os.path.exists(score_out_path_p)):
                    command = self.command_obj.sift_bt2_sam_command(sam_path_p, score_out_path_p, sam_sift_marker_path_p)
                    script_path = os.path.join(self.dir_obj.host_dir_top, "sam_sift_p_" + ref_basename + ".sh")
                    self.job_control.launch_and_create_v2_with_mp_store(script_path, command)
                if(os.path.exists(sam_sift_marker_path_s)):
                    print(dt.today(), "skipping SAM SIFT on s")
                else:
                    command = self.command_obj.sift_bt2_sam_command(sam_path_s, score_out_path_s, sam_sift_marker_path_s)
                    script_path = os.path.join(self.dir_obj.host_dir_top, "sam_sift_s_" + ref_basename + ".sh")
                    self.job_control.launch_and_create_v2_with_mp_store(script_path, command)

            self.job_control.wait_for_mp_store()

            if(os.path.exists(self.dir_obj.host_recon_mkr_p)):
                print(dt.today(), "Skipping host read reconcile P")
            else:    
                list_of_mkrs.append(self.dir_obj.host_recon_mkr_p)
                command = self.command_obj.clean_reads_reconcile(self.dir_obj.host_dir_sam, self.dir_obj.host_dir_end, "None", self.dir_obj.clean_dir_final_f, self.dir_obj.clean_dir_final_r, self.dir_obj.host_recon_mkr_p)
                self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.host_recon_job_p, command)

            if(os.path.exists(self.dir_obj.host_recon_mkr_s)):
                print(dt.today(), "skipping host read reconcile S")
            else:
                if(os.path.getsize(self.dir_obj.clean_dir_final_s) > 1000):
                    list_of_mkrs.append(self.dir_obj.host_recon_mkr_s)
                    print("FILE:", self.dir_obj.clean_dir_final_s, " is not empty:", os.path.getsize(self.dir_obj.clean_dir_final_s))
                    print(dt.today(), "working on clean reads reconcile: S")
                    command = self.command_obj.clean_reads_reconcile(self.dir_obj.host_dir_sam, self.dir_obj.host_dir_end, self.dir_obj.clean_dir_final_s, "None", "None", self.dir_obj.host_recon_mkr_s)
                    self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.host_recon_job_s, command)
                else:
                    print("size of clean dir final s:", os.path.getsize(self.dir_obj.clean_dir_final_s))
                    print("file:", self.dir_obj.clean_dir_final_s)
                    time.sleep(10)
                    print(dt.today(), "No singletons for host filtering.  Skip")
            #command = self.command_obj.clean_reads_reconcile(self.dir_obj.host_dir_data, self.dir_obj.host_dir_end, self.dir_obj.start_s, self.dir_obj.start_f, self.dir_obj.start_r, self.dir_obj.host_recon_mkr)
            
            print(dt.today(), "working on clean reads reconcile")
            self.job_control.wait_for_mp_store()

            if(self.dir_obj.check_mkr_host(list_of_mkrs)):
                self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.host_dir)
            else:
                message_line = str(dt.today()) + " broken at: " + self.path_obj.host_dir
                self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
                sys.exit("bad_stage")
        
    def assembly(self):
        
        if(not os.path.exists(self.dir_obj.assembly_mkr)):
            command = ""
            if(("megahit" in self.path_obj.contig_tool) or ("MEGAHIT" in self.path_obj.contig_tool)):
                print(dt.today(), "choosing MEGAHIT for contig generation")
                time.sleep(3)
                if(self.op_mode == "single"):
                    command = self.command_obj.megahit_command_s(self.dir_obj.host_final_s, self.dir_obj.assembly_alt_dir_data, self.dir_obj.assembly_mkr)
                    self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_mhit_s_job, command)
                    print("command:", command)
                else:
                    print("working in contigs paired")
                    time.sleep(3)
                    if(os.path.getsize(self.dir_obj.clean_dir_final_s) > 1000):
                        command = self.command_obj.megahit_command_p(self.dir_obj.host_final_f, self.dir_obj.host_final_r, self.dir_obj.host_final_s, self.dir_obj.assembly_alt_dir_data, self.dir_obj.assembly_mkr)
                        print("command:", command)
                    else:
                        command = self.command_obj.megahit_command_p(self.dir_obj.host_final_f, self.dir_obj.host_final_r, "None", self.dir_obj.assembly_alt_dir_data, self.dir_obj.assembly_mkr)
                        
                    self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_mhit_p_job, command)
                self.job_control.wait_for_mp_store()
                #set new contigs
                self.dir_obj.assembly_contigs = self.dir_obj.assembly_alt_contigs
            else:
                print(dt.today(), "MetaSPADES discontinued.  Use MegaHit")
                sys.exit("Stop")
            
        else:
            print(dt.today(), "skipping: contig assembly")


        if(not os.path.exists(self.dir_obj.assembly_bt2_idx_mkr)):
            print(dt.today(), "indexing contigs for BT2")
            print("using:", self.dir_obj.assembly_bt2_idx)
            command = self.command_obj.bowtie2_index_ref_command(self.dir_obj.assembly_contigs, self.dir_obj.assembly_bt2_idx, self.dir_obj.assembly_bt2_idx_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_bwa_idx_job, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping BT2 contig indexing")
        
        if(not os.path.exists(self.dir_obj.assembly_pp_mkr)):
            command = ""
            if(self.op_mode == "single"):
                command = self.command_obj.clean_reads_bowtie2_command_s(self.dir_obj.assembly_bt2_idx, self.dir_obj.host_final_s, self.dir_obj.assembly_pp_mkr)
            else:
                command = self.command_obj.clean_reads_bowtie2_command_p(self.dir_obj.assembly_bt2_idx, self.dir_obj.host_final_f, self.dir_obj.host_final_r, self.dir_obj.assembly_pp_mkr) 
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_pp_job, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping BT2 align contigs")

        if(not os.path.exists(self.dir_obj.assembly_scan_sam_mkr)):
            command = self.command_obj.bt2_scan_sam(self.dir_obj.assembly_scan_sam_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_scan_sam_job, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping BT2 sam/bam business")
                        



        if(not os.path.exists(self.dir_obj.assembly_reconcile_mkr_p)):
            command = self.command_obj.contig_reconcile(self.dir_obj.assembly_score_out, "None", self.dir_obj.host_final_f, self.dir_obj.host_final_r, "None", self.dir_obj.assembly_final_f, self.dir_obj.assembly_final_r, self.dir_obj.assembly_reconcile_mkr_p)
            #command = self.command_obj.contig_reconcile(self.dir_obj.assembly_score_out, self.dir_obj.assembly_dir_data, self.dir_obj.host_final_s, self.dir_obj.host_final_f, self.dir_obj.host_final_r, self.dir_obj.assembly_reconcile_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_recon_job_p, command)
        else:
            print(dt.today(), "skipping contig-read reconciliation P")
        if(not os.path.exists(self.dir_obj.assembly_reconcile_mkr_s)):
            #command = self.command_obj.contig_reconcile(self.dir_obj.assembly_score_out, self.dir_obj.assembly_dir_data, "None", self.dir_obj.host_final_f, self.dir_obj.host_final_r, self.dir_obj.assembly_reconcile_mkr)
            command = self.command_obj.contig_reconcile(self.dir_obj.assembly_score_out, self.dir_obj.host_final_s, "None", "None", self.dir_obj.assembly_final_s, "None", "None", self.dir_obj.assembly_reconcile_mkr_s)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.assembly_recon_job_s, command)
            
        else:
            print(dt.today(), "skipping contig-read reconciliation S")
        self.job_control.wait_for_mp_store()    
    
        if(self.dir_obj.check_mkr_assembly()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.assembly_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.assembly_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
        #sys.exit("bad_stage")


    
    def concoct_binning(self):
        #print(dt.today(), "temp holder")
        #requires adjusting headers and junk before sending off to concoct.
        #code just removes the dangling portion of the ID in each contig
        if(not os.path.exists(self.dir_obj.cct_prep_mkr)):
            command = self.command_obj.concoct_prep_command(self.dir_obj.cct_prep_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.cct_prep_job_path, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping concoct prep")

        if(not os.path.exists(self.dir_obj.cct_mkr)):
            command = self.command_obj.concoct_command(self.dir_obj.cct_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.cct_job_path, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping concoct binning")

        if(not os.path.exists(self.dir_obj.cct_checkm_mkr)):
            print(dt.today(), "running checkm")
            command = self.command_obj.checkm_command(self.dir_obj.cct_checkm_mkr)
            self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.cct_checkm_job_path, command)
            self.job_control.wait_for_mp_store()
        else:
            print(dt.today(), "skipping checkm")
            print(self.dir_obj.cct_checkm_mkr)

        if(self.dir_obj.check_mkr_bin_cct()):    
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.cct_bin_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.cct_bin_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")
        

    def metabat2_binning(self):
        print(dt.today(), "running metawrap-binning: metabat2")
        command = self.command_obj.metabat2_bin_command(self.op_mode, self.hosts_bypassed, self.dir_obj.mbat2_mkr)
        self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.mbat2_job, command)
        self.job_control.wait_for_mp_store()
    
        if(self.dir_obj.check_mkr_bin_mbat2()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.mbat2_bin_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.mbat2_bin_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")

    def maxbin2_binning(self):
        print(dt.today(), "running metawrap-binning: maxbin2")
        command = self.command_obj.maxbin2_bin_command(self.op_mode, self.hosts_bypassed, self.dir_obj.mbin2_mkr)
        self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.mbin2_job, command)
        self.job_control.wait_for_mp_store()
        if(self.dir_obj.check_mkr_bin_mbin2()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.mbin2_bin_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.mbin2_bin_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")

    def metawrap_bin_refine(self):
        print("running metawrap bin_refinement")
        command = self.command_obj.metawrap_bin_refinement_command(self.dir_obj.mwrap_bin_r_mkr)
        self.job_control.launch_and_create_v2_with_mp_store(self.dir_obj.mwrap_bin_r_job, command)
        self.job_control.wait_for_mp_store()
        if(self.dir_obj.check_mkr_mwrap_bin_r()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.mwrap_bin_r_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.mwrap_bin_r_dir
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")


    def gtdbtk_classify(self):
        print("running GTDB-tk classify")
        
        marker_path = self.dir_obj.gtdbtk_mkr
        job_path = self.dir_obj.gtdbtk_job
        bin_choice = "refined"
            
        command = self.command_obj.gtdbtk_command(bin_choice, marker_path)
        self.job_control.launch_and_create_v2_with_mp_store(job_path, command)
        
        self.job_control.wait_for_mp_store()
        if(self.dir_obj.check_mkr_gtdbtk()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.gtdbtk_class_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.gtdbtk_class_dir 
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")

    def metawrap_quant(self):
        print(dt.today(), "running metawrap quant bin")
        marker_path = self.dir_obj.mwrap_quant_mkr
        job_path = self.dir_obj.mwrap_quant_job
        bin_choice = "refined"
        command = self.command_obj.metawrap_quantify_command(bin_choice, self.dir_obj.host_final_f, self.dir_obj.host_final_r, self.dir_obj.host_final_s, marker_path)
        self.job_control.launch_and_create_v2_with_mp_store(job_path, command)
        
        
        self.job_control.wait_for_mp_store()
        if(self.dir_obj.check_mkr_mwrap_quant()):
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, self.path_obj.mwrap_quant_dir)
        else:
            message_line = str(dt.today()) + " broken at: " + self.path_obj.mwrap_quant_dir 
            self.job_control.write_to_bypass_log(self.path_obj.bypass_log, message_line)
            sys.exit("bad_stage")