<reponame>CedricMidoux/atlas
import logging

logger = logging.getLogger(__file__)

import multiprocessing
import os
import sys
import tempfile
from snakemake.io import load_configfile
from snakemake.utils import update_config as snakemake_update_config
from .default_values import *


def make_default_config():
    """generates a dict with all the default values, if they exist"""

    config = {}
    config["data_type"] = "metagenome"
    config["interleaved_fastqs"] = False
    config["tmpdir"] = tempfile.gettempdir()
    config["threads"] = 6

    config["simplejob_mem"] = 10
    config["simplejob_threads"] = 4
    config[
        "importqc_params"
    ] = "iupacToN=t touppercase=t qout=33 addslash=t trimreaddescription=t"

    config["mem"] = MEM
    config["large_mem"] = 250
    config["large_threads"] = 16
    config["preprocess_adapter_min_k"] = PREPROCESS_ADAPTER_MIN_K
    config["preprocess_minimum_base_quality"] = PREPROCESS_MINIMUM_BASE_QUALITY
    config[
        "preprocess_allowable_kmer_mismatches"
    ] = PREPROCESS_ALLOWABLE_KMER_MISMATCHES
    config[
        "preprocess_reference_kmer_match_length"
    ] = PREPROCESS_REFERENCE_KMER_MATCH_LENGTH
    config[
        "preprocess_minimum_passing_read_length"
    ] = PREPROCESS_MINIMUM_PASSING_READ_LENGTH
    config["preprocess_minimum_base_frequency"] = PREPROCESS_MINIMUM_BASE_FREQUENCY

    config["deduplicate"] = True
    config["error_correction_overlapping_pairs"] = True

    config["contaminant_max_indel"] = CONTAMINANT_MAX_INDEL
    config["contaminant_min_ratio"] = CONTAMINANT_MIN_RATIO
    config["contaminant_kmer_length"] = CONTAMINANT_KMER_LENGTH
    config["contaminant_minimum_hits"] = CONTAMINANT_MINIMUM_HITS
    config["contaminant_ambiguous"] = CONTAMINANT_AMBIGUOUS

    config["duplicates_only_optical"] = DUPLICATES_ONLY_OPTICAL
    config["duplicates_allow_substitutions"] = DUPLICATES_ALLOW_SUBSTITUTIONS

    config["normalize_reads_before_assembly"] = False
    config["normalization_kmer_length"] = NORMALIZATION_KMER_LENGTH
    config["normalization_target_depth"] = NORMALIZATION_TARGET_DEPTH
    config["normalization_minimum_kmer_depth"] = 5

    config["error_correction_before_assembly"] = True
    config["error_correction_remove_lowdepth"] = False
    config["error_correction_kmer"] = 31
    config["error_correction_lowdepth_fraction"] = 0.5
    config["error_correction_minimum_kmer_depth"] = 1
    config["error_correction_aggressive"] = False
    config["error_correction_minprob"] = 0.5  # for memory issues only , I think

    config["merge_pairs_before_assembly"] = True
    config["merging_k"] = MERGING_K
    config["merging_extend2"] = MERGING_EXTEND2
    config["merging_flags"] = MERGING_FLAGS

    config["assembler"] = "spades"
    config["assembly_memory"] = ASSEMBLY_MEMORY
    config["assembly_threads"] = ASSEMBLY_THREADS
    config["megahit_min_count"] = MEGAHIT_MIN_COUNT
    config["megahit_k_min"] = MEGAHIT_K_MIN
    config["megahit_k_max"] = MEGAHIT_K_MAX
    config["megahit_k_step"] = MEGAHIT_K_STEP
    config["megahit_merge_level"] = MEGAHIT_MERGE_LEVEL
    config["megahit_prune_level"] = MEGAHIT_PRUNE_LEVEL
    config["megahit_low_local_ratio"] = MEGAHIT_LOW_LOCAL_RATIO
    config["megahit_preset"] = "default"
    config["minimum_contig_length"] = MINIMUM_CONTIG_LENGTH
    config["spades_k"] = SPADES_K
    config["spades_use_scaffolds"] = True
    config["spades_preset"] = "meta"
    config["spades_extra"] = ""
    config["spades_skip_BayesHammer"] = False
    config["longread_type"] = None
    config["filter_contigs"] = True
    config["minimum_average_coverage"] = MINIMUM_AVERAGE_COVERAGE
    config["minimum_percent_covered_bases"] = MINIMUM_PERCENT_COVERED_BASES
    config["minimum_mapped_reads"] = MINIMUM_MAPPED_READS
    config["contig_trim_bp"] = CONTIG_TRIM_BP

    # config["translation_table"] = 11

    # map reads to contigs and to genes
    config["minimum_region_overlap"] = MINIMUM_REGION_OVERLAP
    config["feature_counts_allow_overlap"] = FEATURE_COUNTS_ALLOW_OVERLAP
    config["contig_count_multi_mapped_reads"] = CONTIG_COUNT_MULTI_MAPPED_READS
    config["contig_min_id"] = CONTIG_MIN_ID
    config["contig_map_paired_only"] = CONTIG_MAP_PAIRED_ONLY
    config["contig_max_distance_between_pairs"] = CONTIG_MAX_DISTANCE_BETWEEN_PAIRS

    config["maximum_counted_map_sites"] = MAXIMUM_COUNTED_MAP_SITES

    # gene cluster
    config["genecatalog"] = {
        "source": "genomes",
        "clustermethod": "linclust",
        "minlength": 100,
        "minid": 0.9,
        "coverage": 0.9,
        "extra": "",
        "SubsetSize": 500000,
    }

    config["eggNOG_use_virtual_disk"] = False
    config["virtual_disk"] = "/dev/shm"

    # binning
    config["perform_genome_binning"] = True

    config["final_binner"] = "DASTool"
    config["binner"] = ["metabat", "maxbin"]

    config["metabat"] = {"sensitivity": "sensitive", "min_contig_length": 1500}

    config["concoct"] = {
        "Nexpected_clusters": 200,  # important parameter
        "read_length": 100,  # change this parameter !
        "Niterations": 500,
        "min_contig_length": 1000,
    }

    config["maxbin"] = {
        "max_iteration": MAXBIN_MAX_ITERATION,
        "prob_threshold": MAXBIN_PROB_THRESHOLD,
        "min_contig_length": MAXBIN_MIN_CONTIG_LENGTH,
    }

    config["DASTool"] = {
        "search_engine": "diamond",
        "score_threshold": 0.5,
        "duplicate_penalty": 0.6,
        "megabin_penalty": 0.5,
    }

    config["cobining_min_contig_length"] = 2000
    config["cobining_min_bin_size"] = 200 * 1000
    config["semibin_options"] = " --recluster --max-node 1 --max-edges 200 "
    config["cobinning_separator"] = ":"

    config["annotations"] = ["gtdb_taxonomy", "checkm_taxonomy", "gtdb_tree"]
    config["rename_mags_contigs"] = True

    config["genome_dereplication"] = dict(
        filter=dict(noFilter=False, length=5000, completeness=50, contamination=10),
        score=dict(
            completeness=1,
            contamination=5,
            strain_heterogeneity=0,  # not in table
            N50=0.5,
            length=0,
        ),
        ANI=0.95,
        overlap=0.6,
        sketch_size=5000,
        opt_parameters="",
    )

    config["cat_range"] = 5
    config["cat_fraction"] = 0.3

    config["runtime"] = {"default": 5, "assembly": 24, "long": 12, "simplejob": 1}

    return config


def make_config(
    database_dir,
    threads=8,
    assembler="spades",
    data_type="metagenome",
    interleaved_fastq=False,
    config="config.yaml",
):
    """
    Reads template config file with comments from ./template_config.yaml
    updates it by the parameters provided.

    Args:
        config (str): output file path for yaml
        database_dir (str): location of downloaded databases
        threads (int): number of threads per node to utilize
        assembler (str): either spades or megahit
        data_type (str): this is either metagenome or metatranscriptome
    """

    from ruamel.yaml import YAML  # used for yaml reading with comments

    yaml = YAML()
    # yaml.version = (1, 1)
    # yaml.default_flow_style = False

    template_conf_file = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "template_config.yaml"
    )

    with open(template_conf_file) as template_config:
        conf = yaml.load(template_config)

    conf["tmpdir"] = tempfile.gettempdir()
    conf["threads"] = multiprocessing.cpu_count() if not threads else threads
    conf["preprocess_adapters"] = os.path.join(database_dir, "adapters.fa")
    conf["contaminant_references"] = {
        "PhiX": os.path.join(database_dir, "phiX174_virus.fa")
    }

    if data_type == "metatranscriptome":
        conf["contaminant_references"]["rRNA"] = (
            os.path.join(database_dir, "silva_rfam_all_rRNAs.fa"),
        )

    conf["data_type"] = data_type
    conf["interleaved_fastqs"] = interleaved_fastq

    conf["assembler"] = assembler
    conf["database_dir"] = database_dir
    # conf["refseq_namemap"] = os.path.join(database_dir, "refseq.db")
    # conf["refseq_tree"] = os.path.join(database_dir, "refseq.tree")
    # conf["diamond_db"] = os.path.join(database_dir, "refseq.dmnd")

    if os.path.exists(config):
        logger.warning(
            f"Config file {config} already exists, I didn't dare to overwrite it. continue..."
        )
    else:

        with open(config, "w") as f:
            yaml.dump(conf, f)
        logger.info(
            "Configuration file written to %s\n"
            "        You may want to edit it using any text editor." % config
        )


def validate_config(config, workflow):
    conf = load_configfile(config)


#    validate_sample_defs(conf, workflow)
# could later add more validation steps


def update_config(config):
    """
    Populates config file with default config values.
    And made changes if necessary.

    """

    # in old version java_mem was used, new is mem
    if ("java_mem" in config) and (not ("mem" in config)):
        config["mem"] = config["java_mem"]

    # get default values and update them with values specified in config file
    default_config = make_default_config()
    snakemake_update_config(default_config, config)

    return default_config

import logging

logger = logging.getLogger(__file__)

import multiprocessing
import os
import sys
import tempfile
from snakemake.io import load_configfile
from snakemake.utils import update_config as snakemake_update_config
from .default_values import *


def make_default_config():
    """generates a dict with all the default values, if they exist"""

    config = {}
    config["data_type"] = "metagenome"
    config["interleaved_fastqs"] = False
    config["tmpdir"] = tempfile.gettempdir()
    config["threads"] = 6

    config["simplejob_mem"] = 10
    config["simplejob_threads"] = 4
    config[
        "importqc_params"
    ] = "iupacToN=t touppercase=t qout=33 addslash=t trimreaddescription=t"

    config["mem"] = MEM
    config["large_mem"] = 250
    config["large_threads"] = 16
    config["preprocess_adapter_min_k"] = PREPROCESS_ADAPTER_MIN_K
    config["preprocess_minimum_base_quality"] = PREPROCESS_MINIMUM_BASE_QUALITY
    config[
        "preprocess_allowable_kmer_mismatches"
    ] = PREPROCESS_ALLOWABLE_KMER_MISMATCHES
    config[
        "preprocess_reference_kmer_match_length"
    ] = PREPROCESS_REFERENCE_KMER_MATCH_LENGTH
    config[
        "preprocess_minimum_passing_read_length"
    ] = PREPROCESS_MINIMUM_PASSING_READ_LENGTH
    config["preprocess_minimum_base_frequency"] = PREPROCESS_MINIMUM_BASE_FREQUENCY

    config["deduplicate"] = True
    config["error_correction_overlapping_pairs"] = True

    config["contaminant_max_indel"] = CONTAMINANT_MAX_INDEL
    config["contaminant_min_ratio"] = CONTAMINANT_MIN_RATIO
    config["contaminant_kmer_length"] = CONTAMINANT_KMER_LENGTH
    config["contaminant_minimum_hits"] = CONTAMINANT_MINIMUM_HITS
    config["contaminant_ambiguous"] = CONTAMINANT_AMBIGUOUS

    config["duplicates_only_optical"] = DUPLICATES_ONLY_OPTICAL
    config["duplicates_allow_substitutions"] = DUPLICATES_ALLOW_SUBSTITUTIONS

    config["normalize_reads_before_assembly"] = False
    config["normalization_kmer_length"] = NORMALIZATION_KMER_LENGTH
    config["normalization_target_depth"] = NORMALIZATION_TARGET_DEPTH
    config["normalization_minimum_kmer_depth"] = 5

    config["error_correction_before_assembly"] = True
    config["error_correction_remove_lowdepth"] = False
    config["error_correction_kmer"] = 31
    config["error_correction_lowdepth_fraction"] = 0.5
    config["error_correction_minimum_kmer_depth"] = 1
    config["error_correction_aggressive"] = False
    config["error_correction_minprob"] = 0.5  # for memory issues only , I think

    config["merge_pairs_before_assembly"] = True
    config["merging_k"] = MERGING_K
    config["merging_extend2"] = MERGING_EXTEND2
    config["merging_flags"] = MERGING_FLAGS

    config["assembler"] = "spades"
    config["assembly_memory"] = ASSEMBLY_MEMORY
    config["assembly_threads"] = ASSEMBLY_THREADS
    config["megahit_min_count"] = MEGAHIT_MIN_COUNT
    config["megahit_k_min"] = MEGAHIT_K_MIN
    config["megahit_k_max"] = MEGAHIT_K_MAX
    config["megahit_k_step"] = MEGAHIT_K_STEP
    config["megahit_merge_level"] = MEGAHIT_MERGE_LEVEL
    config["megahit_prune_level"] = MEGAHIT_PRUNE_LEVEL
    config["megahit_low_local_ratio"] = MEGAHIT_LOW_LOCAL_RATIO
    config["megahit_preset"] = "default"
    config["minimum_contig_length"] = MINIMUM_CONTIG_LENGTH
    config["spades_k"] = SPADES_K
    config["spades_use_scaffolds"] = True
    config["spades_preset"] = "meta"
    config["spades_extra"] = ""
    config["spades_skip_BayesHammer"] = False
    config["longread_type"] = None
    config["filter_contigs"] = True
    config["minimum_average_coverage"] = MINIMUM_AVERAGE_COVERAGE
    config["minimum_percent_covered_bases"] = MINIMUM_PERCENT_COVERED_BASES
    config["minimum_mapped_reads"] = MINIMUM_MAPPED_READS
    config["contig_trim_bp"] = CONTIG_TRIM_BP

    # config["translation_table"] = 11

    # map reads to contigs and to genes
    config["minimum_region_overlap"] = MINIMUM_REGION_OVERLAP
    config["feature_counts_allow_overlap"] = FEATURE_COUNTS_ALLOW_OVERLAP
    config["contig_count_multi_mapped_reads"] = CONTIG_COUNT_MULTI_MAPPED_READS
    config["contig_min_id"] = CONTIG_MIN_ID
    config["contig_map_paired_only"] = CONTIG_MAP_PAIRED_ONLY
    config["contig_max_distance_between_pairs"] = CONTIG_MAX_DISTANCE_BETWEEN_PAIRS

    config["maximum_counted_map_sites"] = MAXIMUM_COUNTED_MAP_SITES

    # gene cluster
    config["genecatalog"] = {
        "source": "genomes",
        "clustermethod": "linclust",
        "minlength": 100,
        "minid": 0.9,
        "coverage": 0.9,
        "extra": "",
        "SubsetSize": 500000,
    }

    config["eggNOG_use_virtual_disk"] = False
    config["virtual_disk"] = "/dev/shm"

    # binning
    config["perform_genome_binning"] = True

    config["final_binner"] = "DASTool"
    config["binner"] = ["metabat", "maxbin"]

    config["metabat"] = {"sensitivity": "sensitive", "min_contig_length": 1500}

    config["concoct"] = {
        "Nexpected_clusters": 200,  # important parameter
        "read_length": 100,  # change this parameter !
        "Niterations": 500,
        "min_contig_length": 1000,
    }

    config["maxbin"] = {
        "max_iteration": MAXBIN_MAX_ITERATION,
        "prob_threshold": MAXBIN_PROB_THRESHOLD,
        "min_contig_length": MAXBIN_MIN_CONTIG_LENGTH,
    }

    config["DASTool"] = {
        "search_engine": "diamond",
        "score_threshold": 0.5,
        "duplicate_penalty": 0.6,
        "megabin_penalty": 0.5,
    }

    config["cobining_min_contig_length"] = 2000
    config["cobining_min_bin_size"] = 200 * 1000
    config["semibin_options"] = " --recluster --max-node 1 --max-edges 200 "
    config["cobinning_separator"] = ":"

    config["annotations"] = ["gtdb_taxonomy", "checkm_taxonomy", "gtdb_tree"]
    config["rename_mags_contigs"] = True

    config["genome_dereplication"] = dict(
        filter=dict(noFilter=False, length=5000, completeness=50, contamination=10),
        score=dict(
            completeness=1,
            contamination=5,
            strain_heterogeneity=0,  # not in table
            N50=0.5,
            length=0,
        ),
        ANI=0.95,
        overlap=0.6,
        sketch_size=5000,
        opt_parameters="",
    )

    config["cat_range"] = 5
    config["cat_fraction"] = 0.3

    config["runtime"] = {"default": 5, "assembly": 24, "long": 12, "simplejob": 1}

    return config


def make_config(
    database_dir,
    threads=8,
    assembler="spades",
    data_type="metagenome",
    interleaved_fastq=False,
    config="config.yaml",
):
    """
    Reads template config file with comments from ./template_config.yaml
    updates it by the parameters provided.

    Args:
        config (str): output file path for yaml
        database_dir (str): location of downloaded databases
        threads (int): number of threads per node to utilize
        assembler (str): either spades or megahit
        data_type (str): this is either metagenome or metatranscriptome
    """

    from ruamel.yaml import YAML  # used for yaml reading with comments

    yaml = YAML()
    # yaml.version = (1, 1)
    # yaml.default_flow_style = False

    template_conf_file = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "template_config.yaml"
    )

    with open(template_conf_file) as template_config:
        conf = yaml.load(template_config)

    conf["tmpdir"] = tempfile.gettempdir()
    conf["threads"] = multiprocessing.cpu_count() if not threads else threads
    conf["preprocess_adapters"] = os.path.join(database_dir, "adapters.fa")
    conf["contaminant_references"] = {
        "PhiX": os.path.join(database_dir, "phiX174_virus.fa")
    }

    if data_type == "metatranscriptome":
        conf["contaminant_references"]["rRNA"] = (
            os.path.join(database_dir, "silva_rfam_all_rRNAs.fa"),
        )

    conf["data_type"] = data_type
    conf["interleaved_fastqs"] = interleaved_fastq

    conf["assembler"] = assembler
    conf["database_dir"] = database_dir
    # conf["refseq_namemap"] = os.path.join(database_dir, "refseq.db")
    # conf["refseq_tree"] = os.path.join(database_dir, "refseq.tree")
    # conf["diamond_db"] = os.path.join(database_dir, "refseq.dmnd")

    if os.path.exists(config):
        logger.warning(
            f"Config file {config} already exists, I didn't dare to overwrite it. continue..."
        )
    else:

        with open(config, "w") as f:
            yaml.dump(conf, f)
        logger.info(
            "Configuration file written to %s\n"
            "        You may want to edit it using any text editor." % config
        )


def validate_config(config, workflow):
    conf = load_configfile(config)


#    validate_sample_defs(conf, workflow)
# could later add more validation steps


def update_config(config):
    """
    Populates config file with default config values.
    And made changes if necessary.

    """

    # in old version java_mem was used, new is mem
    if ("java_mem" in config) and (not ("mem" in config)):
        config["mem"] = config["java_mem"]

    # get default values and update them with values specified in config file
    default_config = make_default_config()
    snakemake_update_config(default_config, config)

    return default_config
generates a dict with all the default values, if they exist # for memory issues only , I think # config["translation_table"] = 11 # map reads to contigs and to genes # gene cluster # binning # important parameter # change this parameter ! # not in table Reads template config file with comments from ./template_config.yaml     updates it by the parameters provided.      Args:         config (str): output file path for yaml         database_dir (str): location of downloaded databases         threads (int): number of threads per node to utilize         assembler (str): either spades or megahit         data_type (str): this is either metagenome or metatranscriptome # used for yaml reading with comments # yaml.version = (1, 1) # yaml.default_flow_style = False # conf["refseq_namemap"] = os.path.join(database_dir, "refseq.db") # conf["refseq_tree"] = os.path.join(database_dir, "refseq.tree") # conf["diamond_db"] = os.path.join(database_dir, "refseq.dmnd") #    validate_sample_defs(conf, workflow) # could later add more validation steps Populates config file with default config values.     And made changes if necessary. # in old version java_mem was used, new is mem # get default values and update them with values specified in config file