#!/usr/bin/env python3

# Python Standard Packages
import argparse
import subprocess
import os
import sys

# Other Packages
import pandas as pd
from yaml import dump as yaml_dump


# %% Functions

def create_config(configfile, args):
    #region # checks
    if not os.path.exists(args.kraken2db):
        print(f"ERROR: path to kraken2 database {args.kraken2db} does not exist")
        sys.exit(1)

    if not os.path.exists(args.taxonkit_db):
        print(f"ERROR: path to taxonkit_db database {args.taxonkit_db} does not exist")
        sys.exit(1)

    if not os.path.exists(args.mashdb):
        print(f"ERROR: path to mash database {args.mashdb} does not exist")
        sys.exit(1)

    if args.taxlevel_qc != "S" and args.taxlevel_qc != "G":
        print(f"ERROR: Please specify flag --taxlevel_qc as either S or G. Your choice: {args.taxlevel_qc} is invalid")
        sys.exit(1)
    else:
        print(f"Your taxlevel_qc choice: {args.taxlevel_qc} is valid")

    #endregion # checks
    #region # settings
    # adapt execution arguments to single-end input
    if args.single_end or args.iontorrent:
        aquamis_mode = "single-end"
        if args.snakefile == os.path.join(os.path.dirname(os.path.realpath(__file__)), "Snakefile"):
            args.snakefile = os.path.join(os.path.dirname(os.path.realpath(__file__)), "Snakefile_SingleEnd.smk")
        args.shovill_modules = args.shovill_modules.replace("--noreadcorr", "")  # remove incompatible option
    elif args.shovill_modules == "":
        aquamis_mode = "paired-end"
        args.shovill_modules = "--noreadcorr"  # add as default for paired-end

    if args.iontorrent:
        args.shovill_extraopts = (args.shovill_extraopts + " --iontorrent").strip() if not 'iontorrent' in args.shovill_extraopts else args.shovill_extraopts

    #endregion # settings
    #region # create config dictionary
    config_dict = {
        'pipeline'   : pipeline,
        'version'    : version_fixed,
        'git_version': version_git,
        'workdir'    : args.working_directory,
        'samples'    : args.sample_list,
        'smk_params' : {
            'log_directory'  : args.logdir,
            'mode'           : aquamis_mode,
            'profile'        : args.profile,
            'rule_directives': args.rule_directives,
            'threads'        : args.threads_sample,
            'thread_factor'  : args.thread_factor,
            'remove_temp'    : args.remove_temp,
            'ephemeral'      : args.json if args.json else args.ephemeral,
            'batch'          : args.batch,
            'docker'         : args.docker
        },
        'params'     : {
            'run_name'   : args.run_name,
            'fastp'      : {'length_required': args.min_trimmed_length},
            'confindr'   : {'database' : args.confindr_database,
                            'use_rmlst': args.confindr_use_rmlst
                            },
            'kraken2'    : {'db_kraken'         : args.kraken2db,
                            'read_length'       : args.read_length,
                            'taxonomic_qc_level': args.taxlevel_qc,
                            'taxonkit_db'       : args.taxonkit_db
                            },
            'shovill'    : {'depth'         : args.shovill_depth,
                            'output_options': args.shovill_output_options,
                            'tmpdir'        : args.shovill_tmpdir,
                            'ram'           : args.shovill_ram,
                            'assembler'     : args.assembler,
                            'kmers'         : args.shovill_kmers,
                            'extraopts'     : args.shovill_extraopts,
                            'modules'       : args.shovill_modules
                            },
            'mash'       : {'mash_refdb'     : args.mashdb,
                            'mash_kmersize'  : args.mash_kmersize,
                            'mash_sketchsize': args.mash_sketchsize,
                            'mash_protocol'  : args.mash_protocol
                            },
            'mlst'       : {'scheme': args.mlst_scheme},
            'quast'      : {'min_contig'  : args.quast_min_contig,
                            'min_identity': args.quast_min_identity
                            },
            'qc'         : {'thresholds': args.qc_thresholds},
            'json_schema': {'validation': args.json_schema,
                            'filter'    : args.json_filter
                            }
        }
    }

    # endregion # create config dictionary

    # write dictionary to yaml
    with open(configfile, 'w') as outfile:
        yaml_dump(config_dict, outfile, default_flow_style = False, sort_keys = False)
    print(f"Config written to {configfile}")


def run_snakemake(configfile, args, sample_list):
    # force mode
    force_mode = "--forceall" if args.forceall else ""
    force_mode = "--force" if args.force else force_mode

    # user-specified force target; force_targets are mutually exclusive
    force_target = args.forceall if args.forceall else ""
    force_target = args.force if args.force else force_target

    # all-rule force target, add new all rules here
    force_target = "all_trimming_only" if args.no_assembly else force_target
    force_target = "all_ephemeral" if args.ephemeral else force_target
    force_target = "all_json" if args.json else force_target
    force = f"{force_mode} {force_target}"

    # other workflow options
    remove_temp = "" if args.remove_temp else "--notemp"
    batch_rule = "delete_ephemeral_dirs" if args.ephemeral else "write_report"
    batch = f"--batch {batch_rule}={args.batch}" if args.batch else ""
    threads = "--cores" if args.profile == "none" else ""  # use all available cores, if no profile is specified
    threads = f"--cores {args.threads}" if args.threads > 0 else threads  # override cores by commandline arg if provided
    dryrun = "--dryrun --quiet" if args.dryrun else ""
    unlock = "--unlock" if args.unlock else ""

    call = f"snakemake --profile {smk_profile} --configfile {configfile} --config samples={sample_list} " \
           f"--snakefile {args.snakefile} {remove_temp} {batch} {force} {threads} {dryrun} {unlock}"

    if args.use_conda:
        frontend = "conda" if args.conda_frontend else "mamba"
        call = call + f" --use-conda --conda-prefix {args.condaprefix} --conda-frontend {frontend}"

    print(call)
    subprocess.call(call, shell = True, cwd = args.logdir)


def find_missing(args):
    samples = pd.read_csv(args.sample_list, index_col = "sample", sep = "\t")
    samples.index = samples.index.astype('str', copy = False)  # in case samples are integers, need to convert them to str

    # find present samples
    present_set = set()
    for s in samples.index:
        if os.path.isfile(os.path.join(args.working_directory, "json", "post_assembly", s + ".aquamis.json")):
            present_set.add(s)

    # write fails to file
    fail_counter = 0
    with open(os.path.join(args.working_directory, "Assembly", "fails.txt"), "w") as failsfile:
        for sample in samples.index:
            if sample not in present_set:
                fail_counter += 1
                failsfile.write(sample + "\n")
    print(f"Number of FAILED samples: {fail_counter}")

    # exit if none completed
    if len(present_set) == 0:
        print("ERROR: None of the samples completed the pipeline")
        sys.exit(1)
    else:
        print(f"Set of COMPLETE samples: {present_set}")

    # subset pandas dataframe with present samples and write corrected sample sheet
    sample_list_corrected = os.path.join(args.working_directory, "samples_fixed.tsv")
    samples_present = samples.loc[present_set,]
    samples_present.to_csv(sample_list_corrected, sep = "\t", index = True)

    return sample_list_corrected


# %% Main

def main():
    # set global vars
    global pipeline
    global repo_path
    pipeline = "AQUAMIS"

    # set paths
    repo_path = os.path.dirname(os.path.realpath(__file__))

    #region # parse arguments
    parser = argparse.ArgumentParser()

    # path arguments
    parser.add_argument('--sample_list', '-l', help = '[REQUIRED] List of samples to analyze, as a three column tsv file with columns sample and fastq paths. Can be generated with provided script create_sampleSheet.sh',
                        default = "/", type = os.path.abspath, required = False)
    parser.add_argument('--working_directory', '-d', help = '[REQUIRED] Working directory where results are saved',
                        default = "/", type = os.path.abspath, required = False)
    parser.add_argument('--snakefile', '-s', help = f'Path to Snakefile of {pipeline}; default: {os.path.join(repo_path, "Snakefile")}',
                        default = os.path.join(repo_path, "Snakefile"), type = os.path.abspath, required = False),
    # docker arguments
    parser.add_argument('--docker', help = 'Mapped volume path of the host system',
                        default = "", required = False)
    # qc tresholds arguments
    parser.add_argument('--qc_thresholds', help = f'Definition of thresholds in JSON file; default: {os.path.join(repo_path, "resources", "AQUAMIS_thresholds.json")}',
                        default = os.path.join(repo_path, "resources", "AQUAMIS_thresholds.json"), type = os.path.abspath, required = False)
    # json schema arguments
    parser.add_argument('--json_schema', help = f'JSON schema used for validation; default: {os.path.join(repo_path)}/resources/AQUAMIS_schema_v20230906.json',
                        default = os.path.join(repo_path, "resources", "AQUAMIS_schema_v20230906.json"), type = os.path.abspath, required = False)
    parser.add_argument('--json_filter', help = f'Definition of thresholds in JSON file; default: {os.path.join(repo_path, "thresholds", "AQUAMIS_schema_filter_v20230906.json")}',
                        default = os.path.join(repo_path, "resources", "AQUAMIS_schema_filter_v20230906.json"), type = os.path.abspath, required = False)
    # fastp arguments
    parser.add_argument('--min_trimmed_length', help = 'Minimum length of a read to keep; default: 15',
                        default = 15, type = int, required = False)
    # mash arguments
    parser.add_argument('--mashdb', help = f'Path to reference mash database; default: {os.path.join(repo_path, "reference_db", "mash", "mashDB.msh")}',
                        default = os.path.join(repo_path, "reference_db", "mash", "mashDB.msh"), type = os.path.abspath, required = False)
    parser.add_argument('--mash_kmersize', help = 'kmer size for mash, must match size of database; default: 21',
                        default = 21, type = int, required = False)
    parser.add_argument('--mash_sketchsize', help = 'sketch size for mash, must match size of database; default: 1000',
                        default = 1000, type = int, required = False)
    parser.add_argument('--mash_protocol', help = 'Transfer protocol for reference retrieval, choose between https or ftp; default: "https"',
                        default = "https", required = False)
    # kraken arguments
    parser.add_argument('--kraken2db', help = f'Path to kraken2 database; default: {os.path.join(repo_path, "reference_db", "kraken2")}',
                        default = os.path.join(repo_path, "reference_db", "kraken"), type = os.path.abspath, required = False)  # Note: minikraken2.tar.gz extracts to "../kraken/" not kraken2
    parser.add_argument('--taxlevel_qc', help = 'Taxonomic level for kraken2 classification quality control. Choose S for species or G for genus; default: "G"',
                        default = "G", required = False)
    parser.add_argument('--read_length', help = 'Read length to be used in bracken abundance estimation; default: 150',
                        default = 150, type = int, required = False)
    parser.add_argument('--taxonkit_db', help = f'Path to taxonkit_db; default: {os.path.join(repo_path, "reference_db", "taxonkit")}',
                        default = os.path.join(repo_path, "reference_db", "taxonkit"), type = os.path.abspath, required = False)
    # confindr database
    parser.add_argument('--confindr_database', help = f'Path to ConFindR databases; default: {os.path.join(repo_path, "reference_db", "confindr")}',
                        default = os.path.join(repo_path, "reference_db", "confindr"), type = os.path.abspath, required = False)
    parser.add_argument('--confindr_use_rmlst', help = f'Restrict ConFindR to use only rMLST-derived databases, even when cgMLST-derived exist; default: False',
                        default = False, type = os.path.abspath, required = False)
    # assembly options
    parser.add_argument('--shovill_tmpdir', help = f'Fast temporary directory; default: /tmp/shovill_$USER',
                        default = os.path.join(os.sep, "tmp", "shovill_" + str(os.getenv("USER"))), type = os.path.abspath, required = False)
    parser.add_argument('--shovill_ram', help = 'Limit amount of RAM (in GB, integer) provided to shovill; default: 16',
                        default = 16, type = int, required = False)
    parser.add_argument('--shovill_depth', help = 'Sub-sample --R1/--R2 to this depth. Disable with --depth 0; default: 100',
                        default = 100, type = int, required = False)
    parser.add_argument('--assembler', help = 'Assembler to use in shovill, choose from megahit velvet skesa spades; default: spades',
                        default = "spades", required = False)
    parser.add_argument('--shovill_kmers', help = 'K-mers to use <blank=AUTO>; default: ""',
                        default = "", required = False)
    parser.add_argument('--shovill_extraopts', help = 'Extra assembler options in quotes and equal sign notation e.g. spades: --shovill_extraopts="--iontorrent"; default: ""',
                        default = "", required = False)
    parser.add_argument('--shovill_modules', help = 'Module options for shovill; choose from --noreadcorr --trim --nostitch --nocorr; default: --noreadcorr; Note: add choices as string in quotation marks',
                        default = "", required = False)
    parser.add_argument('--shovill_output_options', help = 'Extra options for shovill; default: ""',
                        default = "", required = False)
    # mlst options
    parser.add_argument('--mlst_scheme', help = 'Extra option for MLST; default: ""',
                        default = "", required = False)
    # quast options
    parser.add_argument('--quast_min_contig', help = 'Extra option for QUAST; default: 500',
                        default = 500, type = int, required = False)
    parser.add_argument('--quast_min_identity', help = 'Extra option for QUAST; default: 80',
                        default = 80, type = float, required = False)
    # aquamis arguments
    parser.add_argument('--run_name', '-r', help = 'Name of the sequencing run for all samples in the sample list, e.g. "210401_M02387_0709_000000000-HBXX6", or a self-chosen analysis title',
                        default = "", required = False)
    parser.add_argument('--no_assembly', help = 'Perform only trimming and contamination assessment on reads',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--single_end', help = 'Use only single-end reads as input',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--iontorrent', help = 'Use single-end Ion Torrent reads as input',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--ephemeral', help = 'Remove most intermediate data. Analysis results are reduced to JSONs and reports',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--json', help = 'Only keep assemblies and JSONs (implies --ephemeral). High-Throughput mode',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--fix_fails', help = 'Re-run Snakemake after failure removing failed samples',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--rule_directives', help = f'Process DAG with rule grouping and/or rule prioritization via Snakemake rule directives YAML; default: {os.path.join(repo_path, "profiles", "smk_directives.yaml")} equals no directives',
                        default = os.path.join(repo_path, "profiles", "smk_directives.yaml"), type = os.path.abspath, required = False)
    # cpu arguments
    parser.add_argument('--use_conda', help = 'Utilize the Snakemake "--useconda" option, i.e. Smk rules require execution with a specific conda env',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--conda_frontend', help = 'Do not use mamba but conda as frontend to create individual conda environments',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--condaprefix', '-c', help = f'Path of default conda environment, enables recycling of built environments; default: {os.path.join(repo_path, "conda_env")}',
                        default = os.path.join(repo_path, "conda_env"), required = False)
    parser.add_argument('--remove_temp', help = 'Remove large temporary files. May lead to slower re-runs but saves disk space',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--threads', '-t', help = f'Number of Threads/Cores to use. This overrides the "<{pipeline}>/profiles" settings',
                        default = 0, type = int, required = False)
    parser.add_argument('--threads_sample', '-p', help = 'Number of Threads to use per sample in multi-threaded rules; default: 1',
                        default = 1, type = int, required = False)
    parser.add_argument('--thread_factor', help = 'Factor to increase threads_sample for shovill and quast; default: 3',
                        default = 3, type = int, required = False)
    parser.add_argument('--batch', '-b', help = 'Compute sample list in batches. Please state a fraction string, e.g. 1/3',
                        default = "", type = str, required = False)
    parser.add_argument('--force', '-f', help = 'Snakemake force. Force recalculation of output (rule or file) specified here',
                        default = False, type = str, required = False)
    parser.add_argument('--forceall', help = 'Snakemake force. Force recalculation of all steps',
                        default = False, type = str, required = False)
    parser.add_argument('--profile', help = f'(A) Full path or (B) directory name under "<{pipeline}>/profiles" of a Snakemake config.yaml with Snakemake parameters, e.g. available CPUs and RAM. Default: "bfr.hpc"',
                        default = "none", type = str, required = False)
    parser.add_argument('--dryrun', '-n', help = 'Snakemake dryrun. Only calculate graph without executing anything',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--unlock', help = 'Unlock a Snakemake execution folder if it had been interrupted',
                        default = False, action = 'store_true', required = False)
    parser.add_argument('--logdir', help = 'Path to directory where .snakemake output is to be saved',
                        default = "/", type = os.path.abspath, required = False)
    parser.add_argument('--version', '-V', help = 'Print program version',
                        default = False, action = 'store_true', required = False)

    args = parser.parse_args()

    #endregion # parse arguments
    #region # eval arguments

    # version
    global version_fixed
    global version_git

    version_fixed = open(os.path.join(repo_path, 'VERSION'), 'r').readline().strip()
    try:
        version_git = subprocess.check_output(["git", "describe", "--always"], stderr = subprocess.DEVNULL, cwd = repo_path).decode('ascii').strip()
    except:
        version_git = version_fixed

    if args.version:
        print(f"{pipeline} version\t{version_fixed}")
        if not version_fixed == version_git:
            print(f"commit  version\t{version_git.lstrip('v')}")
        sys.exit(0)
    else:
        print(f"You are using {pipeline} version {version_fixed} ({version_git})")

    # checks
    if args.sample_list == "/":
        print("ERROR: the following argument is required: -l/--sample_list")
        sys.exit(1)
    elif not os.path.exists(args.sample_list):
        print(f"ERROR: path to sample list {args.sample_list} does not exist")
        sys.exit(1)

    if args.working_directory == "/":
        print("ERROR: the following argument is required: -d/--working_directory")
        sys.exit(1)
    elif not os.path.exists(args.working_directory):
        os.makedirs(args.working_directory)

    args.logdir = args.working_directory if args.logdir == "/" else args.logdir
    print(f"Snakemake logs are saved in {args.logdir}/.snakemake/log")

    # snakemake profile
    global smk_profile
    if os.path.exists(os.path.join(repo_path, "profiles", args.profile)):
        smk_profile = os.path.join(repo_path, "profiles", args.profile)
    elif os.path.exists(os.path.realpath(args.profile)):
        smk_profile = os.path.realpath(args.profile)
    elif os.path.exists(os.path.join(repo_path, os.environ.get('SNAKEMAKE_PROFILE', '$'))):
        smk_profile = os.path.join(repo_path, os.environ.get('SNAKEMAKE_PROFILE'))
    elif os.path.exists(os.path.realpath(os.environ.get('SNAKEMAKE_PROFILE', '$'))):
        smk_profile = os.path.realpath(os.environ.get('SNAKEMAKE_PROFILE'))
    elif os.path.exists(os.path.join(repo_path, "profiles", "bfr.hpc")):
        smk_profile = os.path.join(repo_path, "profiles", "bfr.hpc")
    else:
        print(f"ERROR: path to Snakemake config.yaml {args.profile} does not exist")
        sys.exit(1)

    #endregion # eval arguments

    # write config
    configfile = os.path.join(args.working_directory, "config_aquamis.yaml")
    create_config(configfile, args)

    run_snakemake(configfile, args, sample_list = args.sample_list)

    # business logic: rerun with filtered sample set if snakemake failed
    if args.fix_fails and not args.dryrun and not args.unlock:
        if not os.path.isfile(os.path.join(args.working_directory, "reports", "assembly_report.html")):
            print("Re-evaluating samples")
            sample_list_corrected = find_missing(args)
            run_snakemake(configfile, args, sample_list = sample_list_corrected)


# %% Main Call

if __name__ == '__main__':
    main()
