#! /usr/bin/env python3

import sys,os,subprocess,webcolors,argparse,pathlib
from collections import OrderedDict

# Set up script arguments
parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    # Should probably expand this beyond just bw files in the future?
    description="Will take a set of bigWig files and turn them into a set of composite tracks,\
    one per directory. If no subdirs/dataset names provided, then will make a single composite.")
parser.add_argument('fileDir', type=str, help='name of directory where bw files are')
parser.add_argument("-d","--datasetList", type=str , help='dataset name or list of datasets. If list,\
    then dataset names should be enclosed in a set of quotes and space separated, e.g. "H3K27me3 \
    H3K36me3". These names should match the names of directories in fileDir. Will be used as composite track names.')
parser.add_argument("-c","--colors", type=str , help='file containing color associations. \
    Names in file should match those of the file names minus the file extension or match \
    those provided in -s/--shortLabels.')
parser.add_argument("-s","--shortLabels", type=str , help="if file names aren't what is desired \
    for track names, can provided a file containing file names in column 1 and desired track \
    names in column 2.")
parser.add_argument("-l","--longLabels", type=str , help="if file names aren't what is desired \
    for track names, can provided a file containing file names in column 1 and desired track \
    names in column 2.")
parser.add_argument("-f","--html", type=str , help="track description html file. Assumes same \
    one to be used by all.")
parser.add_argument("-b","--bbType", type=str , help="Indicates type (e.g. bigBed 12 +, bigBed 9, \
    etc) if hub contains bigBed tracks.")

args = parser.parse_args()

dsList = list()
# Handle case where people provide a list of track/dir names
if args.datasetList:
    dsList = args.datasetList.strip().split()
# Otherwise, just assume track names == dir names
else:
    for d in os.listdir(args.fileDir):
        if os.path.isdir(os.path.join(args.fileDir, d)):
            dsList.append(d)
# if dir has no subdirs, then just set dsList == input file dir
if len(dsList) == 0:
    dsList = [args.fileDir]

colors = dict()
if args.colors:
    # get suffix of colors file to know delimeter, i.e. tab-sep vs comma-sep
    # From https://stackoverflow.com/questions/541390/extracting-extension-from-filename-in-python
    cFtype = pathlib.Path(args.colors).suffix
    if cFtype == ".csv":
        csep = ","
    elif cFtype == ".tsv":
        csep = "\t"
    else:
        print("colors file must have tsv or csv file extension")
        exit(1)
    for line in open(args.colors, "r"):
        splitLine = line.strip().split(csep)
        colors[splitLine[0]] = splitLine[1]
shortLabels = dict()
if args.shortLabels:
    sFtype = pathlib.Path(args.shortLabels).suffix
    if sFtype == ".csv":
        ssep = ","
    elif sFtype == ".tsv":
        ssep = "\t"
    else:
        print("shortLabels file must have tsv or csv file extension")
        exit(1)
    for line in open(args.shortLabels, "r"):
        splitLine = line.strip().split(ssep)
        shortLabels[splitLine[0]] = splitLine[1]

longLabels = dict()
if args.longLabels:
    lFtype = pathlib.Path(args.longLabels).suffix
    if lFtype == ".csv":
        lsep = ","
    elif lFtype == ".tsv":
        lsep = "\t"
    else:
        print("longLabels file must have tsv or csv file extension")
        exit(1)
    #lfh = open(args.longLabels, "r")
    for line in open(args.longLabels, "r"):
        splitLine = line.strip().split(lsep)
        longLabels[splitLine[0]] = splitLine[1]

# Store everything in an ordered dict
output=OrderedDict()
for ds in dsList:
    dirName = ds.lower()

    # Store everything in a OrderedDict, because output order matters
    # I think 'track' always needs to be first
    output[dirName] = OrderedDict()
    output[dirName]["track"] = dirName
    output[dirName]["compositeTrack"] = "on"
    output[dirName]["shortLabel"] = ds
    output[dirName]["longLabel"] = ds
    output[dirName]["visibility"] = "dense"
    output[dirName]["autoScale"] = "group"
    # type is empty for now, will be filled in later by the step that checks types of child tracks
    output[dirName]["type"] = ""
    # Use track html if provided
    if args.html:
        output[dirName]["html"] = args.html

    if dsList[0] == args.fileDir:
        # Case where the current dir is the only dataset/composite
        fdir = args.fileDir
    else:
        # Otherwise, we make file paths out of the fileDir and dirName
        fdir = os.path.join(args.fileDir,dirName)

    #Store stanzas for child tracks in an ordered dict
    children = OrderedDict()
    for f in os.listdir(fdir):
        # Basically take everything before final suffix to use as track name since we need something unique
        cluster = f.rsplit(".",1)[0]
        ftype = pathlib.Path(f).suffix
        tname = dirName + "_" + cluster
        children[tname] = OrderedDict()
        children[tname]["track"] = tname
        children[tname]["parent"] = dirName + " on"

        # If short labels are specified, use those
        if len(shortLabels) != 0:
            # need try/except here in cases where not every track has a specified short label
            # otherwise you get a KeyError and the script just crashes
            try:
                shortLabel = shortLabels[f]
            except KeyError:
                shortLabel = cluster
        # If no shortLabel provided just use the file name
        else:
            shortLabel = cluster

        children[tname]["shortLabel"] = shortLabel
        # Again, if someone specifies long labels, use those
        if len(longLabels) != 0:
            # need try/except here in cases where not every short label has a corresponding long label
            # otherwise you get a KeyError and the script just crashes
            try:
                longLabel = shortLabel + " - " + longLabels[shortLabel]
            except KeyError:
                longLabel = shortLabel
        # If no long labels specified, just use the short label
        else:
            longLabel = shortLabel
        if dsList[0] == args.fileDir:
            children[tname]["longLabel"] = longLabel
        else:
            children[tname]["longLabel"] = ds + " - " + longLabel
        fpath = os.path.join(fdir,f)
        # Do a bunch of things specific to bigWig files
        if ftype in {".bw",".bigWig",".bigwig"}:
            output[dirName]["type"] = "bigWig"
            # Calculate min/max via bigWigInfo since that's required for the type line
            cmd = ["bigWigInfo", fpath]
            p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            cmdout, cmderr = p.communicate()
            results = cmdout.decode("ASCII").split("\n")
            resultsDict = dict()
            for item in results:
                splitItem = item.replace(" ","").split(":")
                if len(splitItem) > 1:
                    resultsDict[splitItem[0]] = splitItem[1]
            children[tname]["type"] = "bigWig" + " " + resultsDict["min"] + " " + resultsDict["max"]
            # autoScale group: http://genome.ucsc.edu/goldenPath/help/trackDb/trackDbHub.html#autoScale
            children[tname]["autoScale"] = "group"
        # Set type to args.bbType, otherwise error
        elif ftype in {".bb", ".bigBed", ".bigbed"}:
            output[dirName]["type"] = "bigBed"
            if args.bbType:
                children[tname]["type"] = args.bbType
            else:
                print("Your hub contains bigBed tracks, please specify -b/-bbType\n")
                #parser.print_help(sys.stderr)
                exit(1)
        else:
            print("makeCbHub supports bigWig, bigBed and other big* tracks at the moment.\nEnsure that your files end with one of the following extentsions:\n\nbw, bigWig, bigwig, bb, bigBed, bigbed\n")
            exit(1)
        # If someone specifies colors, use those
        if len(colors) != 0:
            if shortLabel in colors.keys():
                # Assumes hexcolors that start with '#'
                if colors[shortLabel].startswith("#"):
                    rgbTuple = webcolors.hex_to_rgb(colors[shortLabel])
                    rgbStr = str(rgbTuple.red) + "," + str(rgbTuple.green) + "," + str(rgbTuple.blue)
                    children[tname]["color"] = rgbStr
                # otherwise, assume they're rgb
                else:
                    children[tname]["color"] = colors[shortLabel]
        children[tname]["bigDataUrl"] = fpath
        children[tname]["visibility"] = "dense"
    # After processing all child tracks, add it to original dict for that dataset
    output[dirName]["children"] = children

for parent in output:
    for setting in output[parent]:
        if setting != "children":
            print(setting, output[parent][setting])
    print("")
    # Need to have space between parent stanza and child stanzas
    # So process parent settings first, then move on to printing child tracks
    for child in output[parent]["children"]:
        for child_setting in output[parent]["children"][child]:
            print("    ", child_setting, output[parent]["children"][child][child_setting])
        print("")
