import os

shell.prefix("set +o pipefail; ")

singularity: "docker://zhuchcn/exce-rna-pipeline:latest"

configfile: os.path.join(workflow.basedir, "default_config.yml")
# configfile: "pipeline_config.yml"

if config["exogenous_mapping"]:
    from exceRNApipeline.includes.ensembl import  EnsemblFTP
    _BACTERIA_COLLECTION_INDICES = EnsemblFTP().get_bacteria_collection_numbers()

report: "reports/workflow.rst"

def all_input():
    """
    Returns the input files/diractories depends on whether the exogenous_mapping
    is set to tru.
    """
    inputs = [
        "output/results/endogenous/summary.png", 
        "output/results/qc/preprocess.tsv"
    ]
    inputs.extend(expand(
        "output/results/qc/{sample}_SE_fastqc/Images/per_base_quality.png",
        sample=config["samples"]
    ))
    if config["exogenous_mapping"]:
        inputs.extend(["output/results/silva/", "output/results/bacteria/"])
    return inputs

rule all:
    input: all_input()

# Step 1: preprocessing
# 
rule preprocess:
    input:
        sample=lambda wildcards: config['samples'][wildcards.sample]
    output: 
        fastq = temp("output/01-Preprocess/{sample}/{sample}_SE.fastq.gz"),
        log = "output/01-Preprocess/{sample}/{sample}.htsStats.log"
    params:
        adapter=config['adapter'],
        prefix="output/01-Preprocess/{sample}/{sample}",
        scratch = config.get('scratch')
    threads: 5
    shell: """
        task_preprocess.py \\
            --input-fq {input.sample} \\
            --output-fq {output.fastq} \\
            --sample-name {wildcards.sample} \\
            --prefix {params.prefix} \\
            --adapter {params.adapter} \\
            --log-file {output.log} \\
            --scratch-dir {params.scratch}
    """

rule preprocess_stats:
    input: 
        logs = expand("output/01-Preprocess/{sample}/{sample}.htsStats.log",
                      sample=config["samples"]),
    output: 
        tsv = "output/results/qc/preprocess.tsv",
        plot = report("output/results/qc/preprocess_barplot.png", 
                      caption="reports/preprocess_barplot.rst",
                      category="Preprocess")
    params:
        samples = list(config["samples"].keys())
    threads: 1
    shell: """
    task_preprocess_stats.py \\
        --log-file {input.logs} \\
        --sample-names {params.samples} \\
        --output-tsv {output.tsv} \\
        --output-plot {output.plot}
    """

rule fastqc:
    input:
        fastq = "output/01-Preprocess/{sample}/{sample}_SE.fastq.gz"
    output:
        report("output/results/qc/{sample}_SE_fastqc/Images/per_base_quality.png",
               caption="reports/qc_per_base_quality.rst",
               category="QC")
    params:
        outdir = "output/results/qc"
    threads: 1
    shell: """
    fastqc -o {params.outdir} {input.fastq}
    unzip {params.outdir}/{wildcards.sample}_SE_fastqc.zip \\
        -d {params.outdir}/
    rm -rf {params.outdir}/{wildcards.sample}_SE_fastqc.zip
    """

# Step 2: map to the UniVec library
rule univec_get_sequence:
    output: "genomes/UniVec_Core.fasta"
    params: url = config['urls']['univec']
    threads: 2
    shell: """
    wget {params.url} -O {output}
    """

rule univec_index:
    input:
        genome = "genomes/UniVec_Core.fasta"
    output: 
        temp(directory("genomes/star_index_univec"))
    params:
        extra = "--genomeSAindexNbases 8 --genomeChrBinNbits 16",
        scratch = config.get("scratch"),
        mem = config.get("univec_index_mem_gb")
    threads: config["univec_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """

rule univec_mapping:
    input:
        genome="genomes/star_index_univec",
        sample="output/01-Preprocess/{sample}/{sample}_SE.fastq.gz"
    output:
        bam=temp("output/02-UniVec/{sample}/Aligned.out.bam"),
        unmapped=temp("output/02-UniVec/{sample}/Unmapped.fastq.gz")
    params:
        prefix="output/02-UniVec/{sample}/",
        scratch = config.get("scratch")
    threads: config["univec_align_cpu"]
    shell: """
    task_star_align.py \\
        -i {input.sample} \\
        -g {input.genome} \\
        -n {wildcards.sample} \\
        -b {output.bam} \\
        -u {output.unmapped} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    """

# Step 3: map to the rRNA sequences.
rule rrna_make_sequence:
    input: config['rRNA']
    output: "genomes/RiboRNA_homo_sapiens.fasta"
    params: 
        rna5s1=config['urls']['rna5s1']
    threads: 2
    shell: """
    curl \
        -o genomes/rna5s1.fasta \
        -H 'Content-type:text/x-fasta' \
        {params.rna5s1}
    cat genomes/rna5s1.fasta {input} \
        > {output}
    """

rule rrna_index:
    input:
        genome = "genomes/RiboRNA_homo_sapiens.fasta"
    output: 
        temp(directory("genomes/star_index_rrna_human"))
    params:
        extra = "--genomeSAindexNbases 8 --genomeChrBinNbits 16",
        scratch = config.get('scratch'),
        mem = config.get("rrna_index_mem_gb")
    threads: config["rrna_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """ 

rule rrna_mapping:
    input:
        genome = "genomes/star_index_rrna_human",
        sample = "output/02-UniVec/{sample}/Unmapped.fastq.gz"
    output:
        bam=temp("output/03-rRNA/{sample}/Aligned.out.bam"),
        unmapped=temp("output/03-rRNA/{sample}/Unmapped.fastq.gz")
    params:
        prefix="output/03-rRNA/{sample}/",
        scratch = config.get('scratch')
    threads: config["rrna_align_cpu"]
    shell: """
    task_star_align.py \\
        -i {input.sample} \\
        -g {input.genome} \\
        -n {wildcards.sample} \\
        -b {output.bam} \\
        -u {output.unmapped} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    """

# Step 4: map to the human genome. 
rule genome_download:
    output: "genomes/human_genome_primary_assembly.fa"
    params: url = config['urls']['human']
    threads: 2
    shell: """
    wget {params.url} -O {output}.gz
    gunzip -c {output}.gz > {output}
    """

rule genome_index:
    input:
        genome = "genomes/human_genome_primary_assembly.fa"
    output: 
        directory("genomes/star_index_human_genome")
    params:
        extra = "--genomeSAindexNbases 14 --genomeChrBinNbits 18",
        scratch = config.get('scratch'),
        mem = config["endogenous_index_ram_gb"]
    threads: config["endogenous_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """ 

rule genome_mapping:
    input:
        genome = "genomes/star_index_human_genome",
        sample = "output/03-rRNA/{sample}/Unmapped.fastq.gz"
    output:
        bam="output/04-Genome/{sample}/Aligned.out.bam",
        unmapped=temp("output/04-Genome/{sample}/Unmapped.fastq.gz")
    params:
        prefix="output/04-Genome/{sample}/",
        scratch = config.get('scratch')
    threads: config["endogenous_align_cpu"]
    shell: """
    task_star_align.py \\
        -i {input.sample} \\
        -g {input.genome} \\
        -n {wildcards.sample} \\
        -b {output.bam} \\
        -u {output.unmapped} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    """

rule make_genome_annotation:
    output: 
        gencode = 'genomes/gencode_human_annotation.gtf',
        tRNA = 'genomes/tRNA.gtf',
        piRNA = 'genomes/piRNA.gtf'
    params: 
        output_dir = 'genomes',
        tRNA = config['urls']['tRNA'],
        piRNA = config['urls']['piRNA'],
        gencode = config['urls']['gencode']
    threads: 2
    script: """
    task_make_human_genome_gtf.py \\
        -c {params.genecode} \\
        -r {params.tRNA} \\
        -i {params.piRNA} \\
        -g {output.gencode} \\
        -t {output.tRNA} \\
        -p {output.piRNA} \\
        -o {params.output_dir}
    """

# count endogenous genes
rule endogenous_count:
    input:
        gencode = 'genomes/gencode_human_annotation.gtf',
        tRNA = 'genomes/tRNA.gtf',
        piRNA = 'genomes/piRNA.gtf',
        bam="output/04-Genome/{sample}/Aligned.out.bam"
    output: 
        gencode = "output/04-Genome/{sample}/ReadsPerGene_gencode.txt",
        tRNA = "output/04-Genome/{sample}/ReadsPerGene_tRNA.txt",
        piRNA = "output/04-Genome/{sample}/ReadsPerGene_piRNA.txt"
    threads: config["endogenous_count_cpu"]
    shell: """
    htseq-count -f bam -s no -i gene_id \\
        --additional-attr gene_name --additional-attr gene_type \\
        {input.bam} \\
        {input.gencode} \\
        > {output.gencode}
    htseq-count -f bam -s no -i gene_id \\
        --additional-attr gene_name --additional-attr gene_type \\
        {input.bam} \\
        {input.tRNA} \\
        > {output.tRNA}
    htseq-count -f bam -s no -i gene_id \\
        --additional-attr gene_name --additional-attr gene_type \\
        {input.bam} \\
        {input.piRNA} \\
        > {output.piRNA}
    """

rule summarize_counts:
    input:
        gencode = expand("output/04-Genome/{sample}/ReadsPerGene_gencode.txt", sample=config["samples"]),
        tRNA = expand("output/04-Genome/{sample}/ReadsPerGene_tRNA.txt", sample=config["samples"]),
        piRNA = expand("output/04-Genome/{sample}/ReadsPerGene_piRNA.txt", sample=config["samples"])
    output:
        gencode = "output/04-Genome/ReadsPerGene_gencode.txt",
        tRNA = "output/04-Genome/ReadsPerGene_tRNA.txt",
        piRNA = "output/04-Genome/ReadsPerGene_piRNA.txt",
        summary = "output/04-Genome/ReadsPerGene_summary.txt"
    params:
        samples = list(config['samples'].keys())
    threads: 4
    shell: '''
    task_summarize_counts.py \\
        --input-gencode {input.gencode} \\
        --input-tRNA {input.tRNA} \\
        --input-piRNA {input.piRNA} \\
        --output-gencode {output.gencode} \\
        --output-tRNA {output.tRNA} \\
        --output-piRNA {output.piRNA} \\
        --output-summary {output.summary} \\
        --sample-names {params.samples}
    '''

rule organize_endogenous_counts:
    input:
        gencode = "output/04-Genome/ReadsPerGene_gencode.txt",
        tRNA = "output/04-Genome/ReadsPerGene_tRNA.txt",
        piRNA = "output/04-Genome/ReadsPerGene_piRNA.txt",
        summary = "output/04-Genome/ReadsPerGene_summary.txt"
    output:
        gencode = "output/results/endogenous/ReadsPerGene_gencode.txt",
        tRNA = "output/results/endogenous/ReadsPerGene_tRNA.txt",
        piRNA = "output/results/endogenous/ReadsPerGene_piRNA.txt",
        summary = "output/results/endogenous/ReadsPerGene_summary.txt"
    threads: 1
    shell: """
    cp {input.gencode} {output.gencode}
    cp {input.tRNA} {output.tRNA}
    cp {input.piRNA} {output.piRNA}
    cp {input.summary} {output.summary}
    """

rule endogenous_count_barplot:
    input: "output/results/endogenous/ReadsPerGene_summary.txt"
    output: 
        report("output/results/endogenous/summary.png",
               caption="reports/endogenous_summary.rst",
               category="Endogenous Genome")
    threads: 1
    shell: '''
    task_endogenous_count_barplot.py \\
        -i {input} \\
        -o {output}
    '''

# Step 5: map to repetitive elements
rule download_repeat_masker:
    output: temp("genomes/repeat_masker.fa.out")
    params: url = config['urls']['repeat_masker']
    threads: 1
    shell: """
    wget {params.url} -O {output}.gz
    zcat {output}.gz > {output}
    """

rule make_re_fasta:
    input: 
        annotation = "genomes/repeat_masker.fa.out",
        genome = "genomes/GRCh38.primary_assembly.genome.fa"
    output: "genomes/RepeatMasker.fa"
    threads: 24
    shell: """
    anno_to_fasta.py \\
        -a {input.annotation} -g {input.genome} -o {output} \\
        -c 5 -s 6 -e 7 -n 10 11 5 6 7 -k 3
    """

rule re_index:
    input: genome = "genomes/RepeatMasker.fa"
    output: 
        temp(directory("genomes/star_index_re"))
    params:
        extra = "--genomeSAindexNbases 14 --genomeChrBinNbits 8",
        scratch = config.get('scratch'),
        mem = config["re_index_ram_gb"]
    threads: config["re_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """ 

rule re_mapping:
    input:
        genome="genomes/star_index_re",
        sample="output/04-Genome/{sample}/Unmapped.fastq.gz"
    output:
        bam=temp("output/05-RE/{sample}/Aligned.out.bam"),
        unmapped=temp("output/05-RE/{sample}/Unmapped.fastq.gz")
    params:
        prefix="output/05-RE/{sample}/",
        scratch = config.get('scratch')
    threads: config["re_align_cpu"]
    shell: """
    task_star_align.py \\
        -i {input.sample} \\
        -g {input.genome} \\
        -n {wildcards.sample} \\
        -b {output.bam} \\
        -u {output.unmapped} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    """

# Step 6: map to SILVA
rule download_silva:
    output:
        silva_lsu = temp("genomes/silva_lsu.fa.gz"),
        silva_ssu = temp("genomes/silva_ssu.fa.gz"),
        taxmap_lsu = temp("genomes/silva_lsu_taxmap.txt"),
        taxmap_ssu = temp("genomes/silva_ssu_taxmap.txt"),
    params:
        lsu = config['urls']['silva_lsu'],
        ssu = config['urls']['silva_ssu'],
        taxmap_lsu = config['urls']['silva_taxmap_lsu'],
        taxmap_ssu = config['urls']['silva_taxmap_ssu']
    shell: """
    wget {params.lsu} -O genomes/silva_lsu.fa.gz
    wget {params.ssu} -O genomes/silva_ssu.fa.gz
    wget {params.taxmap_lsu} -O genomes/silva_taxmap_lsu.txt.gz
    zcat genomes/silva_taxmap_lsu.txt.gz > {output.taxmap_lsu}
    wget {params.taxmap_ssu} -O genomes/silva_taxmap_ssu.txt.gz
    zcat genomes/silva_taxmap_ssu.txt.gz > {output.taxmap_ssu}
    """

rule filter_silva:
    input:
        silva_lsu = "genomes/silva_lsu.fa.gz",
        silva_ssu = "genomes/silva_ssu.fa.gz"
    output:
        silva_lsu = temp("genomes/silva_lsu_filter.fa.gz"),
        silva_ssu = temp("genomes/silva_ssu_filter.fa.gz"),
        silva_lsu_namelist = temp("genomes/silva_lsu_namelist.txt"),
        silva_ssu_namelist = temp("genomes/silva_ssu_namelist.txt")
    threads: 4
    shell: """
    zcat {input.silva_lsu} | grep '^>' |\\
        grep -e ' Bacteria;' -e ' Archaea;' \\
             -e ' Eukaryota;Opisthokonta;Nucletmycea;Fungi;' |\\
        sed 's/^>//g' > {output.silva_lsu_namelist}
    zcat {input.silva_ssu} | grep '^>' |\\
        grep -e ' Bacteria;' -e ' Archaea;' \\
             -e ' Eukaryota;Opisthokonta;Nucletmycea;Fungi;' |\\
        sed 's/^>//g' > {output.silva_ssu_namelist}
    task_extract_fastx.py \\
        --input-file {input.silva_lsu} \\
        --output-file {output.silva_lsu} \\
        --namelist-file {output.silva_lsu_namelist} \\
        --seq-format fasta \\
        --verbose
    task_extract_fastx.py \\
        --input-file {input.silva_ssu} \\
        --output-file {output.silva_ssu} \\
        --namelist-file {output.silva_ssu_namelist} \\
        --seq-format fasta \\
        --verbose
    """

rule solve_silva_taxa:
    input: 
        fasta = "genomes/silva_{su}_filter.fa.gz",
        taxmap = "genomes/silva_{su}_taxmap.txt"
    output: temp("genomes/silva_{su}_filter_ncbiTaxa.fa.gz")
    threads: 4
    shell: """
    task_solve_silva_taxa.py \\
        -f {input.fasta} \\
        -t {input.taxmap} \\
        -s {wildcards.su} \\
        -o {output}
    """

rule make_silva:
    input: expand("genomes/silva_{su}_filter_ncbiTaxa.fa.gz", su=["ssu", "lsu"])
    output:
        combine = temp("genomes/silva_filter_ncbiTaxa.fasta"),
        silva = temp(expand(
            "genomes/silva_tax/silva_{silva_ind}.fasta",
            silva_ind=range(1, config["silva_align_parallel"] + 1)
        ))
    params:
        parallel = config["silva_align_parallel"],
        output_prefix = "genomes/silva_tax/silva_"
    threads: 4
    shell: """
    zcat {input} | sed -E 's/ .+$/:SSU/' > {output.combine}
    task_split_fasta.py \\
        --input-file  {output.combine} \\
        --output-prefix {params.output_prefix} \\
        --n-batch {params.parallel}
    """


rule silva_index:
    input:
        genome = "genomes/silva_tax/silva_{silva_ind}.fasta"
    output:
        temp(directory("genomes/silva_tax/star_index_silva_{silva_ind}"))
    wildcard_constraints:
        silva_ind="[0-9]+"
    params:
        extra = "--genomeSAindexNbases 10 --genomeChrBinNbits 7",
        scratch = config.get('scratch'),
        mem = config["silva_index_ram_gb"]
    threads: config["silva_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """ 

rule silva_mapping:
    input:
        sample = "output/05-RE/{sample}/Unmapped.fastq.gz",
        genome = "genomes/silva_tax/star_index_silva_{silva_ind}"
    output: temp("output/06-SILVA/{sample}/Aligned_{silva_ind}.txt.gz")
    wildcard_constraints:
        silva_ind="[0-9]+"
    params:
        prefix="output/06-SILVA/{sample}/Aligned_{silva_ind}/",
        scratch = config.get('scratch')
    threads: config["silva_align_cpu"]
    shell: '''
    task_star_align_silva.py \\
        -i {input.sample} \\
        -n {wildcards.sample} \\
        -c {wildcards.silva_ind} \\
        -g {input.genome} \\
        -o {output} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    '''

rule silva_combine:
    input: 
        expand(
            "output/06-SILVA/{{sample}}/Aligned_{silva_ind}.txt.gz",
            silva_ind=range(1, config["silva_align_parallel"] + 1)
        )
    output: "output/06-SILVA/{sample}/Aligned_combine.txt.gz"
    threads: 3,
    shell: """
    zcat {input} | 
    gzip -c > {output}
    """

rule download_ncbi_taxdump:
    output: directory("genomes/taxdump")
    params:
        url = config["urls"]["ncbi_taxdump"]
    threads: 1
    shell: """
    wget {params.url} -O genomes/taxdump.tar.gz
    tar -zxf genomes/taxdump.tar.gz -c {output} 
    """

rule silva_count_taxa:
    input: 
        sample = "output/06-SILVA/{sample}/Aligned_combine.txt.gz",
        taxdump = "genomes/taxdump"
    output: 
        expand(
            "output/06-SILVA/{{sample}}/exogenousAligned_SILVA_{tax}.txt",
            tax=config["tax_levels"]
        )
    params: 
        prefix = "output/06-SILVA/{sample}/exogenousAligned_SILVA_"
    threads: 4
    shell: """
    task_count_exogenous_taxa.py \\
        --input-file {input.sample} \\
        --output-prefix {params.prefix} \\
        --nodes-dmp  {input.taxdump}/nodes.dmp \\
        --names-dmp {input.taxdump}/names.dmp
    """

rule silva_count_combine:
    input: 
        expand(
            "output/06-SILVA/{sample}/exogenousAligned_SILVA_{{tax}}.txt",
            sample=config["samples"]
        )
    output: "output/06-SILVA/SILVA_count_{tax}.txt"
    threads: 4
    run: 
        from pandas import read_csv

        counts = None
        for sample in config["samples"].keys():
            path = 'output/06-SILVA/' + sample + '/exogenousAligned_SILVA_' + \
                   wildcards.tax + '.txt'
            new_counts = read_csv(
                path, sep='\t', index_col=0, names=[sample]
            )
            if counts is None:
                counts = new_counts
            else:
                counts = counts.join(new_counts, how='outer')
        counts.to_csv(output[0], sep='\t')
        
rule silva_extract_unmapped:
    input:
        aligned = "output/06-SILVA/{sample}/Aligned_combine.txt.gz",
        fastq = "output/05-RE/{sample}/Unmapped.fastq.gz"
    output: 
        unmapped = "output/06-SILVA/{sample}/Unmapped.fastq.gz",
    params:
        namelist_path = "output/06-SILVA/{sample}/aligned_reads.txt",
        scratch = config.get("scratch")
    threads: 3
    shell: '''
    task_silva_extract_unmapped.py \\
        -a {input.aligned} \\
        -u {input.fastq} \\
        -o {output.unmapped} \\
        -l {params.namelist_path} \\
        -s {params.scratch}
    '''


# Step 7: map to Bacteria

rule bacteria_download_genomes:
    output: 
        # FIXME: The bacteria genomes are not marked as temp for now, but
        # should be removed in the end of the day.
        directory(expand(
            "genomes/bacteria/fasta/bacteria_{bacteria_ind}_collection",
            bacteria_ind=_BACTERIA_COLLECTION_INDICES
        ))
    params:
        version = config.get("version") or "current",
        output_dir = "genomes/bacteria/fasta",
        scratch = config.get('scratch')
    threads: config["bacteria_download_cpu"]
    shell: '''
    task_bacteria_download_genome.py \\
        -o {params.output_dir} \\
        -v {params.version} \\
        -s {params.scratch} \\
        -t {threads}
    '''

def bacteria_index_prev(wildcards):
    ind = int(wildcards.bacteria_ind)
    jobs = config["bacteria_align_parallel"]
    indices = _BACTERIA_COLLECTION_INDICES
    inds_each_job = [indices[start::jobs] for start in range(jobs)]
    cuts = [len(x) for x in inds_each_job]
    cuts = [sum(cuts[:i]) for i in range(len(cuts))]
    if ind in cuts:
        return f"genomes/bacteria/fasta/bacteria_{ind}_collection"
    return expand(
        f"output/07-Bacteria/{{sample}}/Aligned_{ind-1}.txt.gz",
        sample=config["samples"]
    )

rule bacteria_combine_genome:
    input:
        genome_collection = "genomes/bacteria/fasta/bacteria_{bacteria_ind}_collection",
        prev = bacteria_index_prev
    output:
        temp("genomes/bacteria/fasta/bacteria_{bacteria_ind}_collection.fasta")
    threads: 2
    shell: """
    for fa in $(find "{input.genome_collection}" -name "*.dna.toplevel.fa.gz")
    do
        taxa=$(basename -- $fa |cut -d '.' -f1)
        zcat $fa | sed "s/^>/>$taxa:/g" >> {output}
    done
    """

rule bacteria_index:
    input:
        genome = "genomes/bacteria/fasta/bacteria_{bacteria_ind}_collection.fasta" 
    output: 
        temp(directory("genomes/bacteria/star_index/bacteria_{bacteria_ind}_collection"))
    params:
        extra = "--genomeSAindexNbases 13 --genomeChrBinNbits 16",
        scratch = config.get('scratch'),
        mem = config["bacteria_index_ram_gb"]
    threads: config["bacteria_index_cpu"]
    shell: """
    task_star_index.py \\
        -i {input.genome} \\
        -o {output} \\
        -t {threads} \\
        -m {params.mem} \\
        -a "{params.extra}" \\
        -s {params.scratch}
    """ 

rule bacteria_mapping:
    input:
        genome = "genomes/bacteria/star_index/bacteria_{bacteria_ind}_collection",
        sample = "output/06-SILVA/{sample}/Unmapped.fastq.gz"
    output: 
        temp("output/07-Bacteria/{sample}/Aligned_{bacteria_ind}.txt.gz")
    params:
        prefix = lambda wildcards: f"output/06-SILVA/{wildcards.sample}/Aligned_{wildcards.bacteria_ind}/",
        scratch = config.get('scratch')
    threads: config["bacteria_align_cpu"]
    shell: '''
    task_star_align_bacteria.py \\
        -i {input.sample} \\
        -n {wildcards.sample} \\
        -c {wildcards.bacteria_ind} \\
        -g {input.genome} \\
        -o {output} \\
        -p {params.prefix} \\
        -t {threads} \\
        -s {params.scratch}
    '''

rule bacteria_combine_result:
    input: 
        expand(
            "output/07-Bacteria/{{sample}}/Aligned_{bacteria_ind}.txt.gz",
            bacteria_ind = _BACTERIA_COLLECTION_INDICES
        )
    output: "output/07-Bacteria/{sample}/Aligned.txt.gz"
    threads: 2
    shell: """
    zcat {input} | \\
    awk -F '\t' '{{split($2, taxa, ":"); print($1"\t"taxa[1])}}' | \\
    uniq | gzip \\
    > {output}
    """

rule bacteria_count_taxa:
    input:
        sample = "output/07-Bacteria/{sample}/Aligned.txt.gz",
        taxdump = "genomes/taxdump"
    output:
        expand(
            "output/07-Bacteria/{{sample}}/exogenousAligned_bacteria_{tax}.txt",
            tax=config["tax_levels"]
        )
    params:
        prefix = "output/07-Bacteria/{sample}/exogenousAligned_bacteria_"
    threads: 4
    shell: """
    task_count_exogenous_taxa.py \\
        --input-file {input.sample} \\
        --output-prefix {params.prefix} \\
        --nodes-dmp  {input.taxdump}/nodes.dmp \\
        --names-dmp {input.taxdump}/names.dmp
    """

rule bacteria_count_combine:
    input: 
        expand(
            "output/07-Bacteria/{sample}/exogenousAligned_bacteria_{{tax}}.txt",
            sample=config["samples"]
        )
    output: "output/07-Bacteria/bacteria_count_{tax}.txt"
    threads: 4
    run: 
        from pandas import read_csv

        counts = None
        for sample in config["samples"].keys():
            path = 'output/07-Bacteria/' + sample + '/exogenousAligned_bacteria_' + \
                   wildcards.tax + '.txt'
            new_counts = read_csv(
                path, sep='\t', index_col=0, names=[sample]
            )
            if counts is None:
                counts = new_counts
            else:
                counts = counts.join(new_counts, how='outer')
        counts.to_csv(output[0], sep='\t')


rule organize_exogenous:
    input:
        silva = expand("output/06-SILVA/SILVA_count_{tax}.txt",tax=config["tax_levels"]),
        bacteria = expand("output/07-Bacteria/bacteria_count_{tax}.txt",tax=config["tax_levels"])
    output:
        silva = directory("output/results/silva/"),
        bacteria = directory("output/results/bacteria/")
    threads: 1
    shell: """
    cp -r {input.silva} {output.silva}
    cp -r {input.bacteria} {output.bacteria} 
    """
    