#!/usr/bin/env python

VERSION = "1.0.0"

# python libraries
import os, sys, optparse
import StringIO

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

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

# yjw libraries
from yjw.bio import 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 dlcoal reconciliations " +
	     "to dlcpar reconciliations.",

        epilog =
             "Written by Yi-Chieh Wu (yjw@mit.edu), Massachusetts Institute of Technology. " +
             "(c) 2012. Released under the terms of the GNU General Public License.")

    grp_io = optparse.OptionGroup(parser, "Input/Output")
    common.add_common_options(grp_io, 
                              infiles=True,
			      stree=True, smap=True)
    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=".coal.tree",
	               help="input file extension (default: \".coal.tree\")")
    grp_ext.add_option("-O", "--outputext", dest="outputext",
                       metavar="<output file extension>",
                       default=".dlcpar",
                       help="output file extension (default: \".dlcpar\")")
    parser.add_option_group(grp_ext)

    grp_misc = optparse.OptionGroup(parser, "Miscellaneous")
    grp_misc.add_option("--use-locus-recon", dest="use_locus_recon",
                        default=False, action="store_true",
                        help="if set, use locus recon file rather than MPR")
    parser.add_option_group(grp_misc)

    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, args = parser.parse_args()
    if not options.input and len(args) == 0:
        parser.print_help()
	sys.exit(1)

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

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

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

    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, "")
	recon = phyloDLC.Recon()
	coal_tree, extra = recon.read(prefix, stree)
	
        # convert
	gene_tree, labeled_recon = reconlib.recon_to_labeledrecon(coal_tree, recon, stree, gene2species,
                                                                  locus_mpr=not options.use_locus_recon)
	
	# output
	out = util.replace_ext(treefile, options.inputext, options.outputext)
	labeled_recon.write(out, gene_tree)


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