#!/usr/bin/env python3

import argparse
import subprocess
import os
import sys
import pandas as pd
from yaml import dump

version = "1.3.4"


def create_config(configfile, args):
    # checks
    if not os.path.exists(args.sample_list):
        print("path to sample list " + args.sample_list + " does not exist")
        sys.exit(1)

    if not os.path.exists(args.kraken2db):
        print("path to kraken2 database " + args.kraken2db + " does not exist")
        sys.exit(1)

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

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

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

    print("You are using aquamis version" + version)

    config_dict = {
        'Created by AQUAMIS': version,
        'version'           : version,
        'workdir'           : args.working_directory,
        'samples'           : args.sample_list,
        'params'            : {
            'threads'    : args.threads_sample,
            'docker'     : args.docker,
            'run_name'   : args.run_name,
            'remove_temp': args.remove_temp,
            'fastp'      : {'length_required': args.min_trimmed_length},
            'confindr'   : {'database': args.confindr_database},
            'kraken2'    : {'db_kraken'         : args.kraken2db,
                            'read_length'       : args.read_length,
                            'taxonomic_qc_level': args.taxlevel_qc,
                            'taxonkit_db'       : args.taxonkit_db},
            'shovill'    : {'assembler'     : args.assembler,
                            'depth'         : args.shovill_depth,
                            'tmpdir'        : args.shovill_tmpdir,
                            'ram'           : args.shovill_ram,
                            'output_options': args.shovill_output_options,
                            'extraopts'     : args.shovill_extraopts,
                            'modules'       : args.shovill_modules},
            'mash'       : {'mash_refdb'     : args.mashdb,
                            'mash_kmersize'  : args.mash_kmersize,
                            'mash_sketchsize': args.mash_sketchsize},
            'mlst'       : {'scheme': args.mlst_scheme},
            'qc'         : {'thresholds': args.qc_thresholds},
            'json_schema': {'validation': args.json_schema,
                            'filter'    : args.json_filter},
        }
    }

    with open(configfile, 'w') as outfile:
        dump(config_dict, outfile, default_flow_style=False, sort_keys=False)


def run_snakemake(configfile, args, sample_list):
    if args.forceall:
        force = "--forceall"
    elif not args.force:
        force = ""
    else:
        force = "-f {args.force}"

    if args.no_assembly:
        force = "-f all_trimming_only"

    if args.unlock:
        unlock = "--unlock"
    else:
        unlock = ""

    if args.dryrun:
        dryrun = "-n"
    else:
        dryrun = ""

    if args.remove_temp:
        remove_temp = ""
    else:
        remove_temp = "--notemp"

    if args.conda_frontend:
        frontend = "conda"
    else:
        frontend = "mamba"

    if args.use_conda:
        call = "snakemake --conda-prefix {CondaPrefix} --keep-going --configfile {configfile} --config samples={sample_list} --snakefile {snakefile} {remove_temp} {force} {dryrun} {unlock} --cores {threads} --use-conda --conda-frontend {frontend}".format(
            CondaPrefix=args.condaprefix, snakefile=args.snakefile, configfile=configfile, sample_list=sample_list, remove_temp=remove_temp, force=force, dryrun=dryrun, unlock=unlock, threads=args.threads, frontend=frontend)
    else:
        call = "snakemake --keep-going --configfile {configfile} --config samples={sample_list} --snakefile {snakefile} {remove_temp} {force} {dryrun} {unlock} --cores {threads}".format(
            CondaPrefix=args.condaprefix, snakefile=args.snakefile, configfile=configfile, sample_list=sample_list, remove_temp=remove_temp, force=force, dryrun=dryrun, unlock=unlock, threads=args.threads)

    print(call)
    subprocess.call(call, shell=True)


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 + ".json")):
            present_set.add(s)

    print("Found the following samples: ", present_set)

    # subset pandas dataframe with present samples
    if len(present_set) == 1:
        samples_present = samples.loc[[present_set],]
    else:
        samples_present = samples.loc[present_set,]

    # write corrected sample sheet
    sample_list_corrected = os.path.join(args.working_directory, "samples_fixed.tsv")
    samples_present.to_csv(sample_list_corrected, sep="\t", index=True)

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

    failsfile.close()
    return sample_list_corrected


def main():
    parser = argparse.ArgumentParser()

    # path arguments
    parser.add_argument('-l', '--sample_list', help='List of samples to assemble, format as defined by ...',
                        required=True, type=os.path.abspath)
    parser.add_argument('-d', '--working_directory', help='Working directory',
                        type=os.path.abspath, required=True)
    parser.add_argument('-s', '--snakefile', help='Path to Snakefile of bakcharak pipeline; default: path to Snakefile in same directory',
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "Snakefile"), type=os.path.abspath, required=False),
    parser.add_argument('-r', '--run_name', 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)  # TODO: working dir
    # 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='Definition of thresholds in JSON file; default: ' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "thresholds", "AQUAMIS_thresholds.json"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "resources", "AQUAMIS_thresholds.json"), type=os.path.abspath, required=False)
    # json schema arguments
    parser.add_argument('--json_schema', help='JSON schema used for validation; default: ' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "resources", "AQUAMIS_schema_v20210226.json"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "resources", "AQUAMIS_schema_v20210226.json"), type=os.path.abspath, required=False)
    parser.add_argument('--json_filter', help='Definition of thresholds in JSON file; default: ' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "thresholds", "AQUAMIS_schema_filter_v20210226.json"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "resources", "AQUAMIS_schema_filter_v20210226.json"), type=os.path.abspath, required=False)
    # mash arguments
    parser.add_argument('-m', '--mashdb', help='Path to reference mash database; default:' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "mash", "mashDB.msh"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "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)
    # kraken arguments
    parser.add_argument('--kraken2db', help='Path to kraken2 database; default: ' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "kraken2"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "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, required=False, type=int)
    parser.add_argument('--taxonkit_db', help='Path to taxonkit_db; default:' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "taxonkit"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "taxonkit"), type=os.path.abspath, required=False)
    # confindr database
    parser.add_argument('--confindr_database', help='Path to confindr databases; default:' + os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "confindr"),
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "reference_db", "confindr"), type=os.path.abspath, required=False)
    # assembly options
    parser.add_argument('--min_trimmed_length', help='Minimum length of a read to keep; default: 15',
                        default=15, 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_output_options', help='Extra output options for shovill; default: ""',
                        default="", 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('--shovill_ram', help='Limit amount of RAM provided to shovill; default: 16',
                        default=16, type=int, required=False)
    parser.add_argument('--shovill_tmpdir', help='Fast temporary directory; default: "/tmp/shovill"',
                        default="/tmp/shovill", type=os.path.abspath, required=False)
    parser.add_argument('--shovill_extraopts', help='Extra options for shovill; default: ""',
                        default="", required=False)
    parser.add_argument('--shovill_modules', help='Module options for shovill; choose from --noreadcorr --trim --nostitch --nocorr --noreadcorr; default: "--noreadcorr"',
                        default="--noreadcorr", required=False)
    # mlst options
    parser.add_argument('--mlst_scheme', help='Extra options for MLST; default: ""',
                        default="", required=False)
    # cpu arguments
    parser.add_argument('-t', '--threads', help='Number of Threads to use. Ideally multiple of 10; default: 10',
                        default=10, type=int, required=False)
    parser.add_argument('--threads_sample', help='Number of Threads to use per sample; default: 1',
                        default=1, type=int, required=False)
    parser.add_argument('-c', '--condaprefix', help='Path of default conda environment, enables recycling built environments; default: "<AQUAMIS>/conda_env"',
                        default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "conda_env"), required=False)
    parser.add_argument('-n', '--dryrun', help='Snakemake dryrun. Only calculate graph without executing anything',
                        default=False, action='store_true', required=False)
    parser.add_argument('--forceall', help='Snakemake force. Force recalculation of all steps',
                        default=False, action='store_true', required=False)
    parser.add_argument('-f', '--force', help='Snakemake force. Force recalculation of output (rule or file) speciefied here',
                        default=False, type=str, 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('--unlock', help='Unlock a snakemake execution folder if it had been interrupted',
                        default=False, action='store_true', required=False)
    parser.add_argument('--no_assembly', help='Only trimming and kraken analysis',
                        default=False, action='store_true', required=False)
    # parser.add_argument('--no_conda', help='Do not utilize "--useconda" option, i.e. all software dependencies are available in a single env',
    #                     default=False, action='store_true', required=False)
    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('--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('--logdir', help='Path to directory where .snakemake output is to be saved',
                        default='NA', required=False)

    args = parser.parse_args()

    if not os.path.exists(args.working_directory):
        os.makedirs(args.working_directory)

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

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

    if args.fix_fails and not args.dryrun:
        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)


if __name__ == '__main__':
    main()
