#!/usr/bin/env python3

from argparse import ArgumentParser
from pathlib import Path
import sys
import platform
import os
import subprocess

## This is a launch script for FastQC.  It simply parses the command line options
## given to the program and creates a java launch command which does the actual 
## work.

# We need to add resolve() in case this script is being accessed via a symlink
base_dir = Path(__file__).resolve().parent
options = None

def main():
    ## Often people try to download and run the source distribution of fastqc.  We'll
    ## test for this specifically.

    if (base_dir / "uk/ac/babraham/FastQC/FastQCApplication.java").exists():
        exit_with_error("You have downloaded the sourcecode of FastQC, rather than the compiled version. You need to download the compiled version to be able to run the program.")

    global options
    options = get_options()

    modify_classpath()

    java_bin = find_java_bin()

    java_args = get_java_args()

    returncode = run_fastqc(java_bin,java_args)

    sys.exit(returncode)

def run_fastqc(java_bin,java_args):

    command = [java_bin]
    command.extend(java_args)

    command.append("uk.ac.babraham.FastQC.FastQCApplication")
    command.extend(options.files)


    write_debug("Launch command sections are "+str(command))


    completed = subprocess.run(command)

    return completed.returncode

def get_java_args():
    java_args = []

    # Make the program do something sensible if it runs out of memory
    java_args.append('-XX:+ExitOnOutOfMemoryError')

    # Add OSX specific decorations if we're on a mac
    if platform.system() == "Darwin":
        java_args.append('-Xdock:name=FastQC')
        java_args.append("-Xdock:icon=$RealBin/../Resources/seqmonk.icns")
        java_args.append('-Dapple.laf.useScreenMenuBar=true')

    if options.version:
        java_args.append("-Dfastqc.show_version=true")


    if options.outdir:
        java_args.append(f"-Dfastqc.output_dir={options.outdir}")
        
    if options.contaminants:
        java_args.append(f"-Dfastqc.contaminant_file={options.contaminants}")

    if options.adapters:
        java_args.append(f"-Dfastqc.adapter_file={options.adapters}")

    if options.limits:
        java_args.append(f"-Dfastqc.limits_file={options.limits}")

    if options.dir:
        java_args.append(f"-Djava.io.tmpdir={options.dir}")

    if options.min_length:
        java_args.append(f"-Dfastqc.min_length={options.min_length}")

    if options.max_length:
        java_args.append(f"-Dfastqc.max_length={options.max_length}")

    if options.dup_length:
        java_args.append(f"-Dfastqc.dup_length={options.dup_length}")


    # For memory we multiply their memory amount by the number of files
    # which will be processed in parallel, this will be based on the number
    # of files supplied and the number of threads selected.
    used_memory = options.memory

    if len(options.files) > 1:
        file_multipler = len(options.files)

        thread_multiplier = max(1,options.threads // 4)

        used_memory *= min(file_multipler, thread_multiplier)


    java_args.insert(0,f"-Xmx{used_memory}m")

    if options.quiet:
        java_args.append("-Dfastqc.quiet=true")

    if options.fast5:
        java_args.append("-Dfastqc.nano=true")

    if options.phred64:
        java_args.append("-Dfastqc.phred64=true")

    if options.nogroup:
        java_args.append("-Dfastqc.nogroup=true")

    if options.extract:
        java_args.append("-Dfastqc.unzip=true")

        if options.delete:
            java_args.append("-Dfastqc.delete=true")

    if options.threads:
        java_args.append(f"-Dfastqc.threads={options.threads}");

    if options.format:
        java_args.append(f"-Dfastqc.sequence_format={options.format}")

    if options.png:
        java_args.append("-Dfastqc.svg=false")

    # We've found that on systems with large numbers of CPUs, and especially on systems
    # where IO contention can occur, that the program can consume large amounts of CPU
    # during parallel garbage collection.  Since interactive performance isn't really
    # important when running non-interactively we'll limit the GC to using only a single
    # thread if running with a list of filenames

    if options.files:
        java_args.append("-XX:ParallelGCThreads=1")

    if options.files or options.version:
        java_args.append("-Djava.awt.headless=true")
    
    else:
        # This is an attempt to work around a known issue with XQuartz on
        # OSX displaying programs with a black background.  We'll only set
        # this if we appear to be running interactively.  We don't do much 
        # interaction so losing xrender is hopefully not too onerous
        java_args.append("-Dsun.java2d.xrender=false")


    write_debug("Java args are "+str(java_args))

    return java_args
    

def find_java_bin():

    # We need to find the interpreter

    # If they've specifically specified a java path then we use that
    if options.java:
        return options.java


    # Let's start by assuming that we're just using something in the path
    java_bin = "java"

    # They may have a jre bundled with the distribution. We'll use that if it's there

    # Windows first
    if (base_dir/"jre/bin/java.exe").exists():
        java_bin = (base_dir/"jre/bin/java.exe")

    # Linux
    elif(base_dir/"jre/bin/java").exists():
        java_bin = (base_dir/"jre/bin/java")

    # OSX
    elif(base_dir/"/jre/Contents/Home/bin/java").exists():
        java_bin = (base_dir/"/jre/Contents/Home/bin/java")


    write_debug("Using java binary at "+str(java_bin))

    return java_bin



def modify_classpath():
    # We need to add some of the jars for fastqc to the classpath.  We may have an 
    # existing classpath set which we need to extend, or we might need to start from
    # scratch.
    # 
    # Annoyingly the classpath delimiter is different on windows to unix OSs so we 
    # need to figure that out.

    classpath_delimiter = ":"

    if platform.system() == "Windows":
        classpath_delimiter = ";"

    write_debug("Setting CLASSPATH delimiter to "+classpath_delimiter)

    new_elements = [base_dir,base_dir/"htsjdk.jar",base_dir/"commons-io-2.22.0.jar",base_dir/"commons-compress-1.28.0.jar",base_dir/"cisd-jhdf5.jar"]

    # Convert to non-relative full path strings
    new_elements = [str(x.resolve()) for x in new_elements]

    if "CLASSPATH" in os.environ:
        write_debug("Found existing CLASSPATH containing "+os.environ["CLASSPATH"])
        new_elements.insert(0,os.environ["CLASSPATH"])

    classpath = classpath_delimiter.join(new_elements)

    write_debug("Setting CLASSPATH to "+classpath)
    
    os.environ["CLASSPATH"] = classpath


def write_warning(message):
    print("[WARNING] "+message)


def write_debug(message):
    if options.debug:
        print("[DEBUG] "+message)

def exit_with_error(message):
    print("[ERROR] "+message, file=sys.stderr)
    sys.exit(1)

def get_options():
    parser = ArgumentParser(description="A high throughput sequence QC analysis tool")

    parser.add_argument("-v","--version", action="store_true", help="Print FastQC version")
    parser.add_argument("--outdir", type=str, help="Create all output files in the specified output directory. Please note that this directory must exist as the program will not create it.  If this option is not set then the output file for each sequence file is created in the same directory as the sequence file which was processed.")
    parser.add_argument("-c","--contaminants", type=str, help="Specifies a non-default file which contains the list of contaminants to screen overrepresented sequences against. The file must contain sets of named contaminants in the form name[tab]sequence.  Lines prefixed with a hash will be ignored.")
    parser.add_argument("-a","--adapters", type=str, help="Specifies a non-default file which contains the list of adapter sequences which will be explicity searched against the library. The file must contain sets of named adapters in the form name[tab]sequence.  Lines prefixed with a hash will be ignored.")
    parser.add_argument("-l","--limits", type=str, help="Specifies a non-default file which contains a set of criteria which will be used to determine the warn/error limits for the various modules.  This file can also be used to selectively remove some modules from the output all together.  The format needs to mirror the default limits.txt file found in the Configuration folder.")
    parser.add_argument("-d","--dir", type=str, help="Selects a directory to be used for temporary files written when generating report images. Defaults to system temp directory if not specified.")
    parser.add_argument("--min_length", type=int, help="Ignore reads shorter than this length (bp)")
    parser.add_argument("--max_length", type=int, help="Ignore reads longer than this length (bp)")
    parser.add_argument("--dup_length", type=int, help="Sets a length to which the sequences will be truncated when defining them to be duplicates, affecting the duplication and overrepresented sequences plot.  This can be useful if you have long reads with higher levels of miscalls, or contamination with adapter dimers containing UMI sequences.")
    parser.add_argument("--memory", type=int, help="Sets the base amount of memory, in Megabytes, used to process each file.  Defaults to 512MB. You may need to increase this if you have a file with very long sequences in it.", default=512)
    parser.add_argument("-t","--threads", type=int, help="Total number of CPU threads to use. Processing a single file can use up to 4 threads. Defaults to 4, capped at the number of available CPUs.", default=4)
    parser.add_argument("-q","--quiet", action="store_true", help="Supress all progress messages")
    parser.add_argument("--fast5", help="Files come in fast5 format. In this mode you can pass in directories to process and the program will take in all fast5 files within those directories and produce a single output file from the sequences found in all files.")
    parser.add_argument("--phred64", action="store_true", help="Quality scores are encoded in the obsolete Phred64 format.")
    parser.add_argument("-j","--java",type=str,help="Provides the full path to the java binary you want to use to launch FastQC. If not supplied then the program will look for an embedded JRE in a folder called jre in the root of the FastQC installation, or java will be assumed to be in your path.")
    parser.add_argument("--nogroup", action="store_true", help="Disable grouping of bases for reads >50bp. All reports will show data for every base in the read.  WARNING: Using this option will cause fastqc to crash and burn if you use it on really long reads, and your plots may end up a ridiculous size. You have been warned!")
    parser.add_argument("--png", action="store_true", help="Save the graphs in the report in PNG format. Without this option graphs will be SVG format by default")
    parser.add_argument("--extract", action="store_true", help="If set then the zipped output file will be uncompressed in the same directory after it has been created. If --delete is also specified then the zip file will be removed after the contents are unzipped.")
    parser.add_argument("--delete", action="store_true", help="If used with --extract then the zip file of metrics will be deleted after being extracted")
    parser.add_argument("-f","--format", type=str, help="Bypasses the normal sequence file format detection and forces the program to use the specified format.  Valid formats are bam, sam,bam_mapped, sam_mapped and fastq.")
    parser.add_argument("--debug",action="store_true", help="Write debug information from launch script")
    parser.add_argument("files", type=str, nargs="*", help="Files to process")

    options = parser.parse_args()

    # Now we do some sanity checks

    # Does outdir exist
    if options.outdir:
        if not Path(options.outdir).exists():
            exit_with_error(f"Output dir {options.outdir} does not exist")

        if not Path(options.outdir).is_dir():
            exit_with_error(f"Output dir {options.outdir} isn't a directory")


    # Does temp dir exist
    if options.dir:
        if not Path(options.dir).exists():
            exit_with_error(f"Temp dir {options.dir} does not exist")

        if not Path(options.dir).is_dir():
            exit_with_error(f"Temp dir {options.dir} isn't a directory")

    # Does the comtaminant file exist
    if options.contaminants:
        if not Path(options.contaminants).exists():
            exit_with_error(f"Contaminants file {options.contaminants} does not exist")

        if not Path(options.contaminants).is_file():
            exit_with_error(f"Contaminants file {options.contaminants} isn't a file")

    # Does the adapters file exist
    if options.adapters:
        if not Path(options.adapters).exists():
            exit_with_error(f"Adapters file {options.adapters} does not exist")

        if not Path(options.adapters).is_file():
            exit_with_error(f"Adapters file {options.adapters} isn't a file")

    # Does the limits file exist
    if options.limits:
        if not Path(options.limits).exists():
            exit_with_error(f"Limits file {options.limits} does not exist")

        if not Path(options.limits).is_file():
            exit_with_error(f"Limits file {options.limits} isn't a file")

    # Check min length
    if options.min_length:
        if options.min_length < 1:
            exit_with_error("Min length value must be positive")

    # Check min length
    if options.max_length:
        if options.max_length < 1:
            exit_with_error("Max length value must be positive")

    if options.min_length and options.max_length:
        if options.max_length < options.min_length:
            exit_with_error("Max length value was less than min length")

    # Check dup length
    if options.dup_length:
        if options.dup_length < 1:
            exit_with_error("Dup length value must be positive")

    # Check memory
    if options.memory:
        if options.memory < 100 or options.memory > 25600:
            exit_with_error("Memory must be in the range 100 - 25600MB")

    if options.threads:
        if options.threads < 1:
            exit_with_error("Number of threads must be positive")

        if options.threads > len(options.files)*4:
            # We don't emit a warning if they're running interactively 
            # as we don't know what they're going to load.
            if len(options.files) > 0:
                write_warning("You can't use more than 4 threads per file - setting this many is a waste")

        if options.threads > 4 and options.threads % 4 != 0:
            write_warning("Above 4 threads you should use a multiple of 4 (so 4, 8, 12 etc) for maximum benefit")

    # Check format
    if options.format:
        valid_options = ["bam", "sam","bam_mapped", "sam_mapped", "fastq"]
        if not options.format in valid_options:
            exit_with_error("Invalid format specified, allowed values are "+",".join(valid_options))


    # Check java bin
    if options.java:
        if not Path(options.java).exists():
            exit_with_error("Specified java path doesn't exist")

    # Check the input files exist
    found_stdin = False
    for file in options.files:

        if file.startswith("stdin:"):

            if len(options.files) > 1:
                exit_with_error("You can't specify multiple files if reading from stdin")

            break

        if not Path(file).exists():
            exit_with_error(f"File {file} doesn't exist")

    return options


main()