#!/usr/bin/env python

import sys
import argparse

parser = argparse.ArgumentParser(description='This script is for creating a standard iToL "label" and/or "branch" color file when given the IDs of the genomes you want to color.')

required = parser.add_argument_group('required arguments')

required.add_argument("-g", "--target_genomes", help='Single-column file with the genomes to color (need to match the IDs in the tree file, with no "">")', action="store", dest="target_genomes", required=True)
parser.add_argument("-w", "--what_to_color", help='What to color, must be: "branches", "labels", or "both" (default: "both")', action="store", dest="to_color", default="both")
parser.add_argument("-c", "--color", help='Color to use of either: "blue", "green", or "red" (default: "blue", of course, \'cause it\'s the best)', action="store", dest="color", default="blue")
parser.add_argument("-o", "--output_file", help='Output file for iToL (default: "iToL-colors.txt")', action="store", dest="output_file", default="iToL-colors.txt")

if len(sys.argv)==1:
    parser.print_help(sys.stderr)
    sys.exit(0)

args = parser.parse_args()

if args.color == "blue":
    col = "#0000ff"
elif args.color == "green":
    col = "#00a33f"
elif args.color == "red":
    col = "#a30000"
else:
    print("\n\tSorry, we're not prepared to handle \"" + str(args.color) + "\" as the color... :(\n")
    parser.print_help(sys.stderr)
    sys.exit(1)

if args.to_color not in ["both", "branches", "labels"]:
    print("\n\tSorry, we're not prepared to handle \"" + str(args.to_color) + "\" as the argument for what to color... :(\n")
    parser.print_help(sys.stderr)
    sys.exit(1)

target_list = []

with open(args.target_genomes, "r") as target_genomes:
    for genome in target_genomes:
        target_list.append(genome.strip())

out_file = open(args.output_file, "w")

out_file.write("TREE_COLORS\nSEPARATOR TAB\nDATA\n\n")

# writing lines for coloring labels if needed
if args.to_color in ["both", "labels"]:

    for target in target_list:
        out_file.write(str(target) + "\tlabel\t" + str(col) + "\tbold\n")

# writing lines for coloring branches if needed
if args.to_color in ["both", "branches"]:

    for target in target_list:
        out_file.write(str(target) + "\tbranch\t" + str(col) + "\tnormal\t1.5\n")

out_file.close()
