#!/usr/bin/env python

"""
Script for taking IGV screenshots. Run igv_plotter -h for more info.
"""

import configargparse
import os
import re

import igv_api

parser = configargparse.ArgParser(description="This script takes IGV snapshots.",
                                  default_config_files=["~/.igv_plotter", "igv_plotter.yaml"],
                                  formatter_class=configargparse.ArgumentDefaultsRawHelpFormatter)

main_args = parser.add_argument_group("core arguments")
main_args.add("-o", "--output", help="Output image file prefix", default="igv_")
main_args.add("-L", "--locus-file", dest="locus_file", metavar="LOCUS_FILE", help="Text file with one locus per line "
    "formatted like X:12345 or chr1:12345-54321. A snapshot will be taken at each locus.")
main_args.add("files_or_loci", metavar="files or loci", nargs="+", help="Files to load into IGV and/or loci at which to "
    "take snapshots (as an alternative to -L). Locus format can be X:12345 or chr1:12345-54321. Files and loci "
    "will be processed in the order they are specified, so if the command line order is 'file1.bed 1:12345 "
    "file2.bed', file1 will be loaded, then a snapshot will be taken, then file2 will be loaded. "
    "For a list of accepted IGV file formats, see https://www.broadinstitute.org/igv/RecommendedFileFormats")
main_args.add("-v", "--verbose", action="store_true", help="Show log messages from igv_plotter and IGV.")

extra_args = parser.add_argument_group("additional options (YMMV - sometimes IGV ignores these)")
extra_args.add("--width", type=int, help="IGV window width.", default=800)
extra_args.add("--height", type=int, help="IGV window height.", default=600)
extra_args.add("--max-panel-height", type=int, help="IGV max panel height.", default=600)

extra_args.add("-c", "--config-file", help="Config file path", is_config_file=True)
extra_args.add("-m", "--max-memory", help="IGV max memory limit (eg. 2G)", default="1024M")
extra_args.add("-g", "--genome", help="Either the local path of a .genome file, or a genome id (eg. hg19) chosen from "
               "the right-most column of http://igv.broadinstitute.org/genomes/genomes.txt", default="hg19")
extra_args.add("-p", "--preference", action="append", metavar="KEY=VALUE", help="Changes one of the IGV settings "
    "(eg. -p DEFAULT_FONT_SIZE=10). For a full list of keys & example values that can be set, see: "
    "https://github.com/broadinstitute/IGV/blob/master/src/org/broad/igv/PreferenceManager.java#L928")
extra_args.add("--view-as-pairs", action="store_true", help="In .bam tracks, view reads as pairs")
main_args.add("--igv-jar-path", help="Path to igv.jar", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "igv.jar"))
g1 = extra_args.add_mutually_exclusive_group()
g1.add("--expand", dest="compactness_command", action="store_const", const="expand", help="Expand all tracks")
g1.add("--squish", dest="compactness_command", action="store_const", const="squish", help="Squish all tracks")
g1.add("--collapse", dest="compactness_command", action="store_const", const="collapse", help="Collapse all tracks")

#g2 = extra_args.add_mutually_exclusive_group()
#g2.add("--hide-igv-window", dest="hide_igv_window", action="store_true", help="To run in the foreground and show a "
#   "window, IGV needs to know whether the computer it's running on has a display monitor attached. igv_plotter "
#	"attempts to detect this automatically by (checking whether the DISPLAY env. var is set). If this detection "
#	"fails, it may cause IGV to crash when there's no monitor, or to hide the window unnecessarily when there is a "
#	"monitor. This option can be used to override the automatic display detection mechanism.")
#g2.add("--show-igv-window", dest="hide_igv_window", action="store_false", help="See help for --hide-igv-window")

args = parser.parse()

# process command line args
#if args.hide_igv_window is None:
#	# attempt to automatically detect whether a display monitor is connected to this computer
#	has_display_monitor = "DISPLAY" in os.environ
#	hide_igv_window = not has_display_monitor
#else:
#	hide_igv_window = args.hide_igv_window

hide_igv_window = True  # for now always hide the window since it opens and closes too fast to be useful anyway

# helper funcitons
def is_valid_locus_or_region(locus_or_region):
    return igv_robot._match_locus_string(locus_or_region) or igv_robot._match_region_string(locus_or_region)

snapshot_counter = 0
def take_snapshot(locus_or_region):
    global snapshot_counter
    snapshot_counter += 1
    locus_or_region_tag = re.sub("[:-]", "_", locus_or_region)  # if os.name != "nt" 
    snapshot_filename = "%ss%d__%s.png" % (args.output, snapshot_counter, locus_or_region_tag)
    igv_robot.goto(locus_or_region)
    igv_robot.snapshot(snapshot_filename)


# create igv_robot
args.igv_jar_path = os.path.abspath(os.path.expanduser(args.igv_jar_path))
if not os.path.isfile(args.igv_jar_path) or not args.igv_jar_path.endswith(".jar"):
    parser.error("Invalid IGV jar path: " + args.igv_jar_path)

igv_robot = igv_api.IGVCommandLineRobot(args.igv_jar_path,
    igv_window_width=args.width,
    igv_window_height=args.height,
    max_memory=args.max_memory,
    hide_igv_window=hide_igv_window,
    verbose=args.verbose)


# global settings
if args.genome:
    igv_robot.genome(args.genome)
if args.view_as_pairs:
    igv_robot.view_all_tracks_as_pairs()
if args.preference:
    for p in args.preference:
        key_value = p.split("=")
        if len(key_value) != 2:
            parser.error("Invalid preference arg: %(p)s. It should be of the form -p key=value" % locals())
        igv_robot.preference(key_value[0], key_value[1])

# files, loci, regions
for f in args.files_or_loci:
    if os.path.isfile(f):
        if f.endswith(".bam") and not (os.path.isfile(f[:-1]+"i") or os.path.isfile(f+".bai")):
            parser.error(".bai index not found for bam file: " + f)
        igv_robot.load(f)
        if args.compactness_command is not None:
            igv_robot.command(args.compactness_command)
        if args.max_panel_height is not None:
            igv_robot.max_panel_height(args.max_panel_height)
    elif is_valid_locus_or_region(f):
        take_snapshot(locus_or_region=f)
    else:
        parser.error("File doesn't exist: " + f)

if args.locus_file:
    if not os.path.isfile(args.locus_file):
        parser.error("File doesn't exist: " + args.locus_file)

    with open(args.locus_file) as f:
        # parse file that contains loci or regions
        for i, line in enumerate(f):
            line = line.strip()
            if not is_valid_locus_or_region(line):
                parser.error("%(file_path)s line %(i)s: invalid region or locus: %(line)s" % locals())
            take_snapshot(locus_or_region=line)

igv_robot.exit_igv()

# execute queued commands
if snapshot_counter == 0:
    parser.exit("Finished - 0 loci given, so no snapshots will be taken")

try:
    igv_robot.execute()
except igv_api.IGVException as e:
    parser.exit("IGV crashed: %s" % e)

