preprocessor.py

import pandas as pd
import os
from collections import defaultdict
import logging
import argparse
import shutil
import multiprocessing
import psutil
import yaml
import sys
import json
import re
from pathlib import Path


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
LOG_FILE = os.path.join(Path.cwd(), "app.log")
file_sink = logging.FileHandler(LOG_FILE)
file_sink.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
file_sink.setFormatter(formatter)
logging.getLogger().addHandler(file_sink)
logging.info(f"logging with MIN Level >>> INFO at {LOG_FILE}")

PATTERNS = {
    "Novagen": re.compile(r"""
            (?P<sample_name>.+)-   
            (?P<sample_id>.+)       
            (?P<read_num>[12])     
            \.(?P<ext>fastq\.gz|fastq|fq\.gz|fq)$
            """, re.VERBOSE),
        "illumina": re.compile(r"""
            (?P<sample_name>.+)     
            _(?P<sample_id>S\d+)    
            _(?P<lane>L00\d)        
            _(?P<read_num>R[12]|read[12]) 
            _(?P<tail>00\d)         
            \.(?P<ext>fastq\.gz|fastq|fq\.gz|fq)$
            """, re.VERBOSE),

        "SRR": re.compile(r"""
            (?P<sample_name>SRR\d+) 
            (?P<sep>[_.])           
            (?P<read_num>[12]|R[12]) 
            \.(?P<ext>fastq\.gz|fastq|fq\.gz|fq)$
            """, re.VERBOSE),

        "general": re.compile(r"""
            (?P<sample_name>.+)    
            (?P<sep>[_.])        
            (?P<read_num>[12]|R[12]) 
            \.(?P<ext>fastq\.gz|fastq|fq\.gz|fq)$
            """, re.VERBOSE)
}


# colors
RED = "\033[;31;1m"
RED_ = "\033[;31;4m"
GRE = "\033[;32;1m"
YEL = "\033[;33;1m"
BLU = "\033[;34;1m"
PRP = "\033[;35;1m"
CYN = "\033[;36;1m"
BLD = "\033[;37;1m"
GRY = "\033[;30;1m"
NC = "\033[;39;0m"
NC_ = "\033[;39;4m"


def check_extension(df): # takes pandas df and returns string
    # checks all files have same extension from pandas df, to use in generete sample table function
    uniques = df['ext'].unique()
    if len(uniques) > 1:
        logging.fatal(f"Your input directory has multiple fastq file extensions, please check directory.")

    else:
        return uniques[0]


def check_PE(df): # takes pandas df and returns string
    """checks all files either single or paried ended from pandas df, to use in generete sample table function"""
    uniques = df['PE'].unique()
    if len(uniques) > 1:
        logging.fatal(f"Your input directory has both Paired and single files, please check directory.")
    else:
        return uniques[0]


def check_R(df): # takes pandas df and returns string
    """checks all files have same naming patterns from pandas df, to use in generete sample table function"""
    logging.info("Checking sample file patterns is consistent and unique")
    uniques = df['read_num'].unique()
    if len(uniques) > 1:
        logging.error(f"Your input directory has multiple fastq file naming patterns, please check directory.")
    else:
        return uniques[0].replace("1","")

def check_R_pattern(df): # takes pandas df and returns string
    """checks all files have same naming patterns from pandas df, to use in generete sample table function"""
    logging.info("Checking sample file patterns is consistent and unique")
    uniques = df['R_pattern'].unique()
    if len(uniques) > 1:
        logging.fatal(f"Your input directory has multiple fastq file naming patterns, please check directory.")
    else:
        return uniques[0].replace("1","")

def check_pattern(df): # takes pandas df and returns a string 
    """checks all files have same naming patterns from pandas df, to use in generete sample table function"""
    logging.info("Checking sample file patterns is consistent and unique")
    uniques = df['matched_pattern'].unique()
    if len(uniques) > 1:
        logging.fatal(f"Your input directory has multiple fastq file naming patterns, please check directory.")
    else:
        return uniques[0]
def recognize_pattern(filename):
    logging.info(f"Extracting samples from raw_data >>> {filename}")
    logging.info("Recognizing file naming patterns")
    matched_data = None
    matched_pattern = None
    for pattern_name, regex in PATTERNS.items():
        match = regex.match(filename)
        if match:
            matched_pattern = pattern_name
            matched_data = match.groupdict()
            break
    
    if not matched_data:
        logging.error(f"Unrecognized file naming pattern detected. Only illumina, Novagen, SRR and General Sample file naming are supported  >>> {filename}")
        return None
    
    sample_info = {
        "file_name": filename,
        "matched_pattern": matched_pattern,
        "sample_name": None,
        "sample_id": None,
        "read_num": None,
        "lane": None,
        "tail": None,
        "ext": None,
        "unit": None,
    }
    sample_info.update(matched_data)
    if matched_pattern == "SRR" and not sample_info["sample_id"]:
        sample_info["sample_id"] = sample_info["sample_name"]
    if not sample_info["sample_id"]:
        sample_info["sample_id"] = sample_info["sample_name"]
        
    return sample_info
            
def parse_samples(inpath): # takes path return contains fastq files, returns df contains sample information
    ## takes input path
    ## gets the file names containg fastq and fq
    ## performs the recogize_pattern function to 
    ## capture sample information and stores it in 
    ## pandas df
    # input path to absolute path
    logging.info(f"Parsing samples out of raw data path '{inpath}'")
    path = os.path.abspath(inpath)
    # list all files
    all_files = os.listdir(path)
    samples = defaultdict(dict)
    # takes fastq files only
    for file_name in all_files:
        if os.path.isfile(path + "/" + file_name) and ("fastq" in file_name or "fq" in file_name):
            # Captures the file path and name
            filename, file_extension = os.path.splitext(file_name)
            if "fastq" in filename or "fq" in filename:
                filename, new_ext = os.path.splitext(filename)
                file_extension = new_ext + file_extension
            # recogize_pattern function returns a dictitionary with sample names, id, and read information
            sample_info = recognize_pattern(file_name)
            logging.info(f"Sample pattern info recognized from file: '{file_name}' >>> {sample_info}")

            # get only forward reads and replace the read number to get R2
            # appends sample information to a dict of dicts
            
            if "1" in sample_info["read_num"]:
                read_2 = sample_info["read_num"].replace("1","2")
                if sample_info["matched_pattern"] == "illumina":
                    read_1 = f"{sample_info['read_num']}_{sample_info['tail']}.f"
                    read_2 = f"{read_2}_{sample_info['tail']}.f"
                else:
                    read_1 = f"{sample_info['read_num']}.f"
                    read_2 = f"{read_2}.f"

                f2 = file_name.replace(read_1, read_2)
                if f2 in all_files:
                    sample_info["file2"] = f2
                    sample_info["PE"] = True
                    samples[sample_info["sample_name"]] = sample_info

                else:
                    sample_info["file2"] = ""
                    sample_info["PE"] = False

    # converts the dict to pandas df and returns the df
    m_samples = samples
    print(f"Samples parsed out to be converted to data frame {m_samples}")
    logging.info(f"Samples parsed out to be converted to data frame {m_samples}")
    samples= pd.DataFrame(samples).T
    samples = samples.sort_values(by=['sample_id'])

    return samples



if __name__=="__main__":
    logging.info("Hello.................")
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", required=True, type=os.path.abspath, help="Path to the input or raw data directory")
    parser.add_argument("-o", "--output", required=True, type=os.path.abspath, help="Path to the output directory")
    parser.add_argument("--overwrite", help="Should wexygen clear out everything in the output directory before processing?", action="store_true")
    args = parser.parse_args()
    
    logging.info("Welcome to the preprocessor.py...Let's create necessary files and configurations for this workflow to run successfully!")
    
    outpath = os.path.abspath(args.output)
    path = os.path.abspath(args.input)
    
    
    all_threads = multiprocessing.cpu_count()
    all_mem = int( (psutil.virtual_memory().total ) / 1000000000 )

    # create output path if doesn't exsit 
    if not os.path.exists(outpath):
        logging.info(f"Creating output directory >>> '{outpath}'")
        os.mkdir(outpath)
    else:        
        if args.overwrite:
            logging.info(f"Purging output directory  before processing >>> '{outpath}'")
            shutil.rmtree(outpath)
            os.mkdir(outpath)             
    logging.info(f"Output directory ready for processing >>> '{outpath}'")    
    all_args = vars(args) #there will only be input and output here, I'll had the rest collected from cpp later...
    samples = parse_samples(args.input)
    ext = str(check_extension(samples))
    PE = bool(check_PE(samples))
    R = str(check_R(samples))
    # R_pattern = str(check_R_pattern(samples))
    R_pattern = ""
    compressed = False
    EXT = ext
    pattern = str(check_pattern(samples))
    tail = "001" if pattern == "illumina" else ""

    # to perform gunzipping 
    if ".gz" in ext:
        compressed = True
        EXT = ext.replace(".gz","")
    
    # check if analysis run before and created sample table 
    sample_file = os.path.join(outpath, "samples.tsv")
    logging.info(f"Exporting sample data frame into '{sample_file}'")
    if os.path.exists(sample_file):
        logging.warning(f"Found an exsiting sample.tsv file in output directory, WILL OVERWRITE!!")
        samples.to_csv(sample_file,sep='\t',index=False)
    else:
        samples.to_csv(sample_file,sep='\t',index=False)  
        
    
      
    CWD = Path.cwd()
    extra_info = {
    "path": str(path),
    "working_dir": str(outpath),
    "ext": ext,
    "tail": tail,
    "R": R,
    "naming_pattern": pattern,
    "R_pattern": R_pattern,
    "compressed" : compressed,
    "total_mem": all_mem,
    "wexygen_DIR": str(CWD),
     #they should be able to specify these
    "common_rules": f"{CWD}/workflows/common/rules/",
    "snakemake_rules_dir": f"{CWD}/workflows/snakemake/",
    "nextflow_dir": f"{CWD}/workflows/nextflow/"
    
}

    if "decompress" not in all_args:
        all_args.update({"decompress":False})
    all_args.update(extra_info)
    args_json_path = os.path.join(outpath, "all_args.json")
    logging.info(f"Writing config to json in output path >>> '{args_json_path}'")
    with open(args_json_path, "w", encoding="utf8") as f:
        json.dump(all_args, f, ensure_ascii=False, indent=4)
    # create config file 
    config_yaml_path = os.path.join(outpath, "config.yaml")
    with open(config_yaml_path, "w") as yaml_file:
        yaml.safe_dump(all_args, yaml_file, default_flow_style=False, sort_keys=False)

    # if len(sys.argv) > 1:
    #     command = " ".join(sys.argv[1:])
        
    #     # Write the command to a file to be deserialized
    #     with open(f"{outpath}/command.txt", "w") as f:
    #         f.write(command)