#!/usr/bin/env python

# python libraries
import os, sys, optparse

# dlcpar libraries
import dlcpar
from dlcpar import reconlib

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

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

def parse_args():
    parser = optparse.OptionParser()
    parser.add_option("-s", "--stree", dest="stree",
                      metavar="<species tree>",
                      help="species tree file in newick format")
    parser.add_option("-S", "--smap", dest="smap",
                      metavar="<species map>",
                      help="gene to species map")
    parser.add_option("--lmap", dest="lmap",
                      metavar="<locus map>",
                      help="gene to locus map (species-specific)")
    parser.add_option("-T", "--treeext", dest="treeext",
                      metavar="<tree file extension>",
                      default=".tree",
                      help="tree file extension (default: \".tree\")")
    parser.add_option("--by-fam", dest="by_fam", action="store_true")
    parser.add_option("--use-famid", dest="use_famid", action="store_true")
    parser.add_option("--explicit", dest="explicit",
                      action="store_true", default=False,
                      help="set to ignore extra lineages at implied speciation nodes")

    options, args = parser.parse_args()

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

    return options, args

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

def count_all_events(options, args):

    stree = treelib.read_tree(options.stree)
    gene2species = phylo.read_gene2species(options.smap)
    if options.lmap:
        gene2locus = phylo.read_gene2species(options.lmap)
    else:
        gene2locus = None

    treefiles = map(lambda line: line.rstrip(), util.read_strings(sys.stdin))
    gene_trees = []
    extras = []

    for treefile in treefiles:
        prefix = util.replace_ext(treefile, options.treeext, "")
        gene_tree, extra = reconlib.read_labeled_recon(prefix, stree)
        gene_trees.append(gene_tree)
        extras.append(extra)

    etree = reconlib.count_dup_loss_coal_trees(gene_trees, extras, stree, gene2species,
                                               gene2locus, implied=not options.explicit)

    # make table
    headers = ["genes", "dup", "loss", "coal", "appear"]
    ptable = treelib.tree2parent_table(etree, headers)

    # sort by post order
    lookup = util.list2lookup(x.name for x in stree.postorder())
    ptable.sort(key=lambda x: lookup[x[0]])

    ptable = [[str(row[0]), str(row[1]), float(row[2])] + row[3:]
              for row in ptable]

    tab = tablelib.Table(ptable,
                         headers=["nodeid", "parentid", "dist"] + headers)
    tab.write()

    return 0


def count_by_fam(options, args):

    stree = treelib.read_tree(options.stree)
    gene2species = phylo.read_gene2species(options.smap)
    if options.lmap:
        gene2locus = phylo.read_gene2species(options.lmap)
    else:
        gene2locus = None

    treefiles = map(lambda line: line.rstrip(), util.read_strings(sys.stdin))

    # write header
    lookup = util.list2lookup(x.name for x in stree.postorder())
    headers = ["genes", "dup", "loss", "coal", "appear"]
    print "\t".join(["famid", "nodeid", "parentid", "dist"] + headers)

    for treefile in treefiles:
        if options.use_famid:
            famid = os.path.basename(os.path.dirname(treefile))
        else:
            famid = treefile

        # read files and events
        prefix = util.replace_ext(treefile, options.treeext, "")
        gene_tree, extra = reconlib.read_labeled_recon(prefix, stree)

        etree = reconlib.count_dup_loss_coal_trees([gene_tree], [extra], stree, gene2species,
                                                   gene2locus, implied=not options.explicit)
        ptable = treelib.tree2parent_table(etree, headers)

        # sort by post order
        ptable.sort(key=lambda x: lookup[x[0]])

        # write table
        for row in ptable:
            print "\t".join(map(str, [famid] + row))

    return 0


def main():

    options, args = parse_args()

    if not options.by_fam:
        count_all_events(options, args)

    else:
        count_by_fam(options, args)


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