#!/usr/bin/env python

VERSION = "1.0.0"

# python libraries
import os, sys, optparse

# dlcpar libraries
import dlcpar
from dlcpar import common
from dlcpar import reconlib

# rasmus, compbio libraries
from rasmus import treelib, util
from compbio import phylo, phyloDLC

#==========================================================
# parser

def parse_args():
    """parse input arguments"""

    parser = optparse.OptionParser(
        usage = "usage: %prog [options] <gene tree> ...",

        version = "%prog " + VERSION,

        description =
        "%prog is a utility for converting dlcpar reconciliations " +
        "to dlcoal reconciliations.",

        epilog =
        "Written by Yi-Chieh Wu (yjw@cs.hmc.edu), Harvey Mudd College. " +
        "(c) 2012-2019. Released under the terms of the GNU General Public License.")

    grp_io = optparse.OptionGroup(parser, "Input/Output")
    grp_io.add_option("-s", "--stree", dest="stree",
                      metavar="<species tree>",
                      help="species tree file in newick format")
    grp_io.add_option("-S", "--smap", dest="smap",
                      metavar="<species map>",
                      help="gene to species map")
    parser.add_option_group(grp_io)

    grp_ext = optparse.OptionGroup(parser, "File Extensions")
    grp_ext.add_option("-I","--inputext", dest="inputext",
                       metavar="<input file extension>",
                       default=".tree",
                       help="input file extension (default: \".tree\")")
    grp_ext.add_option("-O", "--outputext", dest="outputext",
                       metavar="<output file extension>",
                       default="",
                       help="output file extension (default: \"\")")
    parser.add_option_group(grp_ext)

    grp_info = optparse.OptionGroup(parser, "Information")
    common.move_option(parser, "--version", grp_info)
    common.move_option(parser, "--help", grp_info)
    parser.add_option_group(grp_info)

    options, treefiles = parser.parse_args()

    #=============================
    # check arguments

    # input gene tree files
    if len(treefiles) == 0:
        parser.error("must specify input file(s)")

    # required options
    if not options.stree:
        parser.error("-s/--stree required")

    return options, treefiles

#==========================================================
# main

def main():
    """main"""

    # parse arguments
    options, treefiles = parse_args()

    # read species tree and species map
    stree = treelib.read_tree(options.stree)
    common.check_tree(stree, options.stree)
##    gene2species = phylo.read_gene2species(options.smap)

    # process genes trees
    for treefile in treefiles:

        # read dlcoal_recon files
        prefix = util.replace_ext(treefile, options.inputext, "")
        labeledrecon = reconlib.LabeledRecon()
        gene_tree, extra = labeledrecon.read(prefix, stree)

        # convert
        coal_tree, recon = reconlib.labeledrecon_to_recon(gene_tree, labeledrecon, stree)

        # output
        out = util.replace_ext(treefile, options.inputext, options.outputext)
        recon.write(out, coal_tree)


# main function
if __name__ == "__main__":
    sys.exit(main())
