import os
from pathlib import Path
from snakemake import rules
from snakemake.io import directory, protected, expand

configfile: "config/config.yaml"

container: "docker://condaforge/mambaforge:23.1.0-1"

rule create_database:
  input: config["file_names"]["bold_tsv"]
  output: 'results/temporary_database.db'
  params:
    marker=config["marker"],
    minlength=config["minlength"],
    log_level=config['log_level']
  conda: "envs/environment.yml"
  script:
      "scripts/bold_data_dump.py"

rule map_opentol:
  input: rules.create_database.output
  output: protected("results/databases/BOLD_{}_barcodes.db".format(config["marker"]))
  params:
    marker=config['marker'],
    log_level=config['log_level']
  conda: "envs/environment.yml"
  script:
      "scripts/map_opentol.py"

checkpoint family_fasta:
    input: rules.map_opentol.output
    output: directory(config["file_names"]["fasta_dir"])
    params:
      log_level=config['log_level'],
      fasta_dir=config["file_names"]["fasta_dir"],
      filter_level=config["fasta_filter"]["filter_level"],
      filter_name=config["fasta_filter"]["filter_name"],
      maxseq=config["maxseq"],
      minseq=config["minseq"]
    conda: "envs/environment.yml"
    script:
        "scripts/family_fasta.py"
def get_all_families(wildcards):
    checkpoint_output = checkpoints.family_fasta.get(**wildcards).output[0]
    families =  [
        f[:-6]  # remove the ".fasta"
        for f in os.listdir(checkpoint_output)
        if f [-6:] == ".fasta"
           and os.stat(os.path.join(checkpoint_output, f)).st_size != 0 # ignore empty files
           and "{" not in f ]  # ignore one file named "Erebidae-Plecoptera{Genus_moth}", it breaks the wildcard
    return expand(
        "results/raxml/{{family}}/{{family}}_{}.raxml.log".format(config['datatype']),
        family=families[:200]  # TODO remove the limit
    )

rule all:
    input:
        # "results/raxml/Cercopithecidae/Cercopithecidae_{}.raxml.log".format(config['datatype'])
        get_all_families

rule hmm_align:
    input:  "results/fasta/family/{family}.fasta"
    output: "results/fasta/alignment/{{family}}_{}.fasta".format(config['datatype'])
    conda: "envs/environment.yml"
    script:
        "scripts/msa_hmm.py"


rule get_opentree_zip:
    # FIXME the version "opentree13.4" is static here but the download URL is user-defined
    output: "results/opentree/opentree13.4_tree/labelled_supertree/labelled_supertree.tre"
    params:
        directory=lambda wildcards, output: Path(output[0]).parent.parent.parent
    conda: "envs/environment.yml"
    shell:
        """
        wget -P {params.directory} {config[opentol_url]} && 
        tar -xvzf {params.directory}/opentree13.4tree.tgz --directory {params.directory}
        """


rule megatree_loader:
    input: rules.get_opentree_zip.output
    output: protected("results/databases/opentree_nodes.db")
    conda: "envs/environment.yml"
    shell: "megatree-loader -i {input} -d {output}"


# rule replace_ids:
#     input:
#         rules.map_opentol.output,
#         rules.hmm_align.output
#     output: "results/fasta/alignment/{{family}}_{}_corrected.fasta".format(config['datatype'])
#     conda: "envs/environment.yml"
#     script: "scripts/replace_alignment_ids.py"
#
#
# rule unique_otts:
#     input: rules.replace_ids.output
#     output: "results/ott_lists/{family}_otts.txt"
#     conda: "envs/environment.yml"
#     shell: 'cat {input} | egrep ">" | cut -f1 -d "_" | tr -d ">" | sort -u > {output}'
#
#
# rule megatree_pruner:
#     input:
#         rules.megatree_loader.output,
#         rules.unique_otts.output
#     output: "results/newick/{family}_constraint_temp.nwk"
#     conda: "envs/environment.yml"
#     shell: "megatree-pruner -d {input[0]} -i {input[1]} > {output}"
#
#
# rule edit_newick:
#     input:
#         rules.replace_ids.output,
#         rules.megatree_pruner.output
#     output:  "results/newick/{family}_constraint.nwk"
#     conda: "envs/environment.yml"
#     script: "scripts/edit_constraint.py"
#
#
# rule raxml:
#     input:
#         rules.replace_ids.output,
#         rules.edit_newick.output
#     output: "results/raxml/{{family}}/{{family}}_{}.raxml.log".format(config['datatype'])
#     threads: 4
#     params:
#         directory=lambda wildcards, output: Path(output[0]).parent.parent / wildcards.family
#     conda: "envs/environment.yml"
#     shell:
#         """
#         mkdir -p {params.directory} |
#         raxml-ng --msa {input[0]} --model GTR+G --tree-constraint {input[1]} --threads {threads}
#         --prefix {params.directory}/{wildcards.family}_{config[datatype]}
#         """



