"""The snake make file that controls the process, imported by a python warper"""
import os
import glob
import pathlib
import pandas as pd
from join_asvbins.snake_functions import combine_mbstats_barrnap, \
    pullseqs_header_name_from_tab, filter_from_mbstats

search_tool = 'blast' if config['blast'] else 'mmseqs'
# TODO decide if this is neccicary
# output_name = config.get('output_name')
bins_path = config.get('bins')
asv_seqs_path = config.get('asv_seqs')
fasta_extention = config.get('fasta_extention')
generic_16s_path = config.get('generic_16s')
allow_empty= config.get('allow_empty')
s1_mmseqs_sensitivity = config.get('s1_mmseqs_sensitivity')
s2_mmseqs_sensitivity = config.get('s2_mmseqs_sensitivity')
min_len_with_overlap = config.get("min_len_with_overlap")
min_len_pct_no_overlap = config.get("min_len_pct_no_overlap")
s1_min_pct_id = config.get('s1_min_pct_id')
s2_min_pct_id = config.get('s2_min_pct_id')
s1_min_length = config.get('s1_min_length')
s2_min_length = config.get('s2_min_length')
s2_min_len_pct = config.get('min_len_pct')
s2_min_len_pct = config.get('min_len_pct')
s2_max_gaps = config.get('max_gaps')
s2_max_missmatch = config.get('max_missmatch')
verbosity= config.get('verbosity')
localy_combined_bins = "all_bins_combined"



# Set the appropriate output, This can be moved
if os.path.isdir(bins_path):
    bins_folder = bins_path
    path_to_combined_bins = localy_combined_bins
else:
    bins_folder = "None"
    path_to_combined_bins = bins_path


rule all:
    input:
        "candidate_statistics.tsv",
        "candidate_sequences.fna",
        "match_statistics.tsv",
        "match_sequences.fna"


rule search1_16s_bin_finds:
    input:
        "candidate_statistics.tsv",
        "candidate_sequences.fna"


rule search2_asv_bin_matches:
    input:
        "match_statistics.tsv",
        "match_sequences.fna"


# rule stage1_set_output_stats:
#     input:
#         mbstats_path = f"stage1_asvs_{search_tool}.tab",
#         barstats_raw_path = "barrnap_16S-gff.gff",
#         barstats_corrected_path = f"barrnap_corrected_stats_no_{search_tool}.tab"
#     output:
#         protected("16s_bin_finds_stats.tab")
#     run:
#         stage1_statistics(**input, output_path=output[0], search_tool=search_tool)


# rule stage2_set_output_fasta:
#     input:
#         f"barrnap_and_{search_tool}_asvs_matching_seqs.fna"
#     output:
#     shell:
#         """
#         cp {input} {output}
#         """


rule combine_barrnap_with_other: # where other is blast or mmseqs
    input:
        mbstats_fasta_path = f"stage1_asvs_{search_tool}_matches.fna",
        mbstats_stats_path = f"stage1_asvs_{search_tool}.tab",
        barrnap_fasta_path = "barrnap_fasta-16S.fna",
        barrnap_stats_path= "barrnap_16S-gff.gff"
    output:
        out_fasta_path = protected("candidate_sequences.fna"),
        out_stats_path = protected("candidate_statistics.tsv")
    run:
        combine_mbstats_barrnap(**input,
                                **output,
                                search_tool=search_tool,
                                allow_empty=allow_empty,
                                min_pct_id=s1_min_pct_id,
                                min_len_with_overlap=min_len_with_overlap,
                                min_len_pct_no_overlap=min_len_pct_no_overlap,
                                min_length=s1_min_length)


rule combine_input_fa:
    input:
        bins_folder
    output:
        temp(localy_combined_bins)
    run:
       input_list = glob.glob(os.path.join(input[0], f"*.{fasta_extention}"))
       if fasta_extention.endswith('gz'):
               shell(f"gzip -cd  {input_list} >> {{output}}")
       else:
               shell(f"cat {' '.join(input_list)} >> {{output}}")


rule mmseqs_stage1_search:
    input:
        path_to_combined_bins,
        generic_16s_path # Query
    output:
        temp(directory("mmseqs_stage1_db")),
        temp("stage1_asvs_mmseqs.tab")
    threads:
        workflow.cores
    params:
        sensitivity = s2_mmseqs_sensitivity,
        verbosity = verbosity if verbosity <= 3 else 3
    shell:
        """
        mkdir {output[0]}
        mmseqs createdb -v {params.verbosity} {input[0]} {output[0]}/target
        mmseqs createdb -v {params.verbosity} {input[1]} {output[0]}/query
        mmseqs search --search-type 3 \\
               -v {params.verbosity} \\
               -s {params.sensitivity} \\
               --threads {threads} \\
               {output[0]}/query \\
               {output[0]}/target \\
               {output[0]}/mmseqs_out \\
               temp
        mmseqs convertalis \\
               -v {params.verbosity} \\
               --format-output \'query,target,pident,alnlen,mismatch,gapopen,qstart,qend,tstart,tend,evalue,bits,qlen,tlen\' \\
               {output[0]}/query \\
               {output[0]}/target \\
               {output[0]}/mmseqs_out \\
               {output[1]}
        """


rule blast_stage1_search:
    input:
        path_to_combined_bins,
        generic_16s_path # Query
    output:
        temp(directory("blast_stage1_db")),
        temp("stage1_asvs_blast.tab")
    run:
       shell("mkdir {output[0]}")
       shell("makeblastdb -dbtype nucl -in {input[0]} -out {output[0]}/blast_db")
       shell("blastn -db {output[0]}/blast_db -out {output[1]} -query {input[1]} -outfmt \"6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore qlen slen\"")


rule pullseq_header_name:
    input:
        path_to_combined_bins,
        "{level}_asvs_{tool}.tab"
    output:
        temp("{level}_asvs_{tool}_matches.fna")
    run:
       pullseqs_header_name_from_tab(in_fasta_path=input[0],
                                     out_fasta_path=output[0],
                                     tab_file_path=input[1],
                                     header_column='sseqid')


rule mmseqs_stage2_search:
    input:
       "candidate_sequences.fna", # Target
       asv_seqs_path # Query
    output:
       temp(directory("mmseqs_stage2_db")),
       temp("stage2_asvs_mmseqs.tab")
    threads:
        workflow.cores
    params:
        sensitivity = s1_mmseqs_sensitivity,
        verbosity = verbosity if verbosity <= 3 else 3
    shell:
        # TODO Split out the db creation, into other rules if it makes sense seeing as you may need to limit cores
        """
        mkdir {output[0]}
        mmseqs createdb -v {params.verbosity} \\
                        {input[0]} {output[0]}/target
        mmseqs createdb -v {params.verbosity} \\
                        {input[1]} {output[0]}/query
        mmseqs search --search-type 3 \\
               {output[0]}/query \\
               {output[0]}/target \\
               {output[0]}/mmseqs_out \\
               temp \\
               -s {params.sensitivity} \\
               -v {params.verbosity}
        mmseqs convertalis  \\
               -v {params.verbosity} \\
               --format-output \'query,target,pident,alnlen,mismatch,gapopen,qstart,qend,tstart,tend,evalue,bits,qlen,tlen\' \\
               {output[0]}/query \\
               {output[0]}/target \\
               {output[0]}/mmseqs_out \\
               {output[1]}
        """


rule stage2_filtering:
    input:
       stats_file_in = f"stage2_asvs_{search_tool}.tab",
       fasta_file_in = "candidate_sequences.fna"
    output:
        fasta_file_out = protected("match_sequences.fna"),
        stats_file_out = protected("match_statistics.tsv")
    run:
       filter_from_mbstats(
                       **input,
                       **output,
                       min_pct_id=s2_min_pct_id,
                       min_length=s2_min_length,
                       min_len_pct=s2_min_len_pct,
                       max_gaps=s2_max_gaps,
                       max_missmatch=s2_max_missmatch,
                       search_tool=search_tool
                       )


rule blast_stage2_search:
    input:
       "barrnap_and_blast_asvs.fna", # Target
        asv_seqs_path, # Query
    output:
        temp(directory("blast_stage2_db")),
        temp("stage2_asvs_blast.tab")
    shell:
        """
        mkdir {output[0]}
        makeblastdb -dbtype nucl -in {input[0]} -out {output[0]}/blast_db
        blastn -db {output[0]}/blast_db -out {output[1]} -query {input[1]} \
        -outfmt \"6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore qlen slen\"
        """


rule run_barrnap_barrnap:
    input:
        path_to_combined_bins
    output:
        temp("barrnap_rrna.gff")
    threads:
        workflow.cores
    params:
        verbosity = "--quiet" if verbosity < 3 else ""
    shell:
        "barrnap --threads {threads} {params.verbosity} {input} > {output[0]}"


rule run_barrnap_16s_gtff:
    input:
        "barrnap_rrna.gff"
    output:
        temp("barrnap_16S-gff.gff")
    shell:
        "grep \"16S\" {input} > {output}"


rule run_barrnap_fasta_filter:
    input:
        path_to_combined_bins,
        "barrnap_16S-gff.gff"
    output:
        temp("barrnap_fasta_raw.fna"),
        temp(f"{path_to_combined_bins}.fai")
    shell:
        "bedtools getfasta -fi {input[0]} -bed {input[1]} -fo {output[0]}"


rule run_barrnap_headers:
    input:
        "barrnap_fasta_raw.fna"
    output:
        temp("barrnap_16S-id.txt")
    shell:
        "grep \">\" {input} | sed 's/>//g' > {output}"


rule run_barrnap_fasta_trim:
    input:
        "barrnap_fasta_raw.fna",
        "barrnap_16S-id.txt"
    output:
        temp("barrnap_fasta-16S.fna"),
        temp("barrnap_fasta_raw.fna.fai")
    shell:
        "xargs samtools faidx {input[0]} < {input[1]} > {output[0]}"

