import sys
import os
import argparse

# This script will count all features for each sample. It can ignore stratification by species
# and also "un"-features if these options are set by the user.

def parse_arguments(args):
    """ 
    Parse the arguments from the user
    """
    parser = argparse.ArgumentParser(
        description= "Count features for each species\n",
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "-i", "--input",
        help="the feature table\n[REQUIRED]",
        metavar="<input.tsv>",
        required=True)
    parser.add_argument(
        "-o", "--output",
        help="file to write counts table\n[REQUIRED]",
        metavar="<output>",
        required=True)
    parser.add_argument(
        "--reduce-sample-name",
        help="remove the extra strings from the sample names",
        action='store_true')
    parser.add_argument(
        "--ignore-un-features",
        help="do not count UNMAPPED, UNGROUPED, or UNINTEGRATED features",
        action='store_true')
    parser.add_argument(
        "--ignore-stratification",
        help="only count the main feature not the stratified instances",
        action='store_true')
    parser.add_argument(
        "--include",
        help="only count features with this string included")
    parser.add_argument(
        "--filter",
        help="do not count features with this string included")

    return parser.parse_args()


def main():

    args=parse_arguments(sys)

    data=[]
    samples=[]
    with open(args.input) as file_handle:
        # remove RPKs from sample name if included
        if args.reduce_sample_name:
            samples=[sample.replace("_Abundance-RPKs","").replace("_genefamilies_Abundance","").replace("_Abundance","").replace("_taxonomic_profile","") for sample in file_handle.readline().rstrip().split("\t")[1:]]
        else:
            samples=file_handle.readline().rstrip().split("\t")[1:]

        for line in file_handle:
            if "|" in line and args.ignore_stratification:
                store=False
            else:
                store=True

            if "UNMAPPED" in line or "UNGROUPED" in line or "UNINTEGRATED" in line and args.ignore_un_features:
                store=False

            if args.include and not args.include in line:
                store=False

            if args.filter and args.filter in line:
                store=False

            if store:
                data.append(line.rstrip().split("\t")[1:])

    try:
        file_handle=open(args.output,"w")
    except EnvironmentError:
        sys.exit("Error: Unable to open output file: " + args.output)

    # write out the header
    file_handle.write("\t".join(["# samples","total features"])+"\n")

    # count the total features for each sample
    for i, sample in enumerate(samples):
        features=0
        for row in data:
            if float(row[i]) > 0:
                features+=1
        file_handle.write(sample+"\t"+str(features)+"\n")

    file_handle.close()

if __name__ == "__main__":
    main()

import sys
import os
import argparse

TOTAL_COUNT_TAG="reads; of these:"
NUCLEOTIDE_COUNT_TAG="Unaligned reads after nucleotide alignment:"
TRANSLATED_COUNT_TAG="Unaligned reads after translated alignment:"
SPECIES_COUNT_TAG="Total species selected from prescreen:"

def parse_arguments(args):
    """
    Parse the arguments from the user
    """
    parser = argparse.ArgumentParser(
        description= "Reads the HUMAnN2 logs and prints a table of read counts\n",
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "-i", "--input",
        help="the folder of log files\n[REQUIRED]",
        metavar="<input_folder>",
        required=True)
    parser.add_argument(
        "-o", "--output",
        help="file to write counts table\n[REQUIRED]",
        metavar="<output>",
        required=True)

    return parser.parse_args()

def main():
    # parse arguments from the user
    args = parse_arguments(sys)

    # search subfolders to find the log files
    log_files=[]
    for path, directories, files in os.walk(args.input):
        for file in filter(lambda x: x.endswith(".log"), files):
            log_files.append(os.path.join(path,file))

    try:
        file_handle=open(args.output,"w")
    except EnvironmentError:
        sys.exit("Error: Unable to open output file: " + args.output)

    # write out the header
    file_handle.write("\t".join(["# samples","total reads","total nucleotide aligned","total translated aligned","total species"])+"\n")

    for file in log_files:
        sample=os.path.basename(file).split(".log")[0]
        data=[sample,"NA","NA","NA","NA"]
        for line in open(file):
            if TOTAL_COUNT_TAG in line:
                print(line)
                print (line.split())
                data[1]=int(line.split()[7])
            elif NUCLEOTIDE_COUNT_TAG in line:
                try:
                    data[2]=int(data[1]*((100-float(line.split()[-2]))/100.0))
                except (ValueError, TypeError):
                    print("Warning: Unable to compute nucleotide reads aligned from log line: " + line)
            elif TRANSLATED_COUNT_TAG in line:
                try:
                    data[3]=int(data[1]*((100-float(line.split()[-2]))/100.0))
                except (ValueError, TypeError):
                    print("Warning: Unable to compute translated reads aligned from log line: " + line)
            elif SPECIES_COUNT_TAG in line:
                data[4]=line.split()[-1]

        file_handle.write("\t".join([str(i) for i in data])+"\n")

    file_handle.close()
    print("Read table written to file: " + args.output)

if __name__ == "__main__":
    main()

shell:
    """
    {SINGULARITY} kneaddata \
        --input {input.R1} \
        --input {input.R2} \
        --reference-db {BOWTIE_GRCH38_BUILD} \
        --bowtie2 {BOWTIE_DIR} \
        --trimmomatic {TRIMMOMATIC_DIR} \
        --output {params.directory}

    cat {params}/*_paired_*.fastq >> {output}

    """

shell:
    """
    {SINGULARITY} python {METAPHLAN_PY} {input} \
        --input_type fastq \
        --bowtie2_exe {BOWTIE2_EXEC} \
        --bowtie2_build {BOWTIE2_BUILD} \
        --samout output/metaphlan_analysis/$(basename {input}).sam \
        --output {output}
    """

shell:
    """
    {SINGULARITY} python {METAPHLAN_DIR}/utils/merge_metaphlan_tables.py {input} \
        > {output}
    """

shell:
    """
    {SINGULARITY} humann2 -i {input.fastq} \
        --bowtie2 {BOWTIE_DIR} \
        --metaphlan {METAPHLAN_DIR} \
        --taxonomic-profile {input.profile} \
        --nucleotide-database {NUCLEO_DB} \
        --protein-database {PROTEIN_DB} \
        --threads 4 \
        -o {params}
    """