#!/usr/bin/env python

import argparse, os
from phylogenetics.homologs import load_homologset
from phylogenetics.phyml import run_phyml
from phylogenetics.names import switch

def main():

    # ---------------------------------------------------
    # command-line arg parsing
    # ---------------------------------------------------

    parser = argparse.ArgumentParser(description="""
    Use PhyML to estimate the maximum likelihood tree for a set of aligned homologs.
    """)

    parser.add_argument("-i", "--input",
        help="Input pickle file containing homologset.",
        type=str)

    parser.add_argument("-o", "--output",
        help="Output pickle filename.",
        type=str)

    parser.add_argument("--dtype",
        help="Sequence datatype",
        type=str,
        default="aa")

    parser.add_argument("--labels",
        help="List of homolog attributes for labeling on tree.",
        nargs="+",
        type=str,
        default=["id"])

    parser.add_argument("--newick",
        help="Boolean to keep newick output. If False, removes the newick file containing the tree.",
        type=bool,
        default=True)

    # Parse the commandline arguments
    args = parser.parse_args()

    # ---------------------------------------------------
    # Run phyml
    # ---------------------------------------------------

    # Separate output filename and extension
    oname, oext = os.path.splitext(args.output)

    # load homolog set from pickle file
    hs = load_homologset(args.input)

    # Pass the homolog set to phyml
    hs2 = run_phyml(hs, oname, dtype=args.dtype, tree_file=args.newick, rm_tmp=True)

    # Change the names in the phyml output to specified labels.
    switch(oname+".nwk", hs2, "id",args.labels)

    # Write this out to a pickle file.
    hs2.write(args.output, format="pickle")

if __name__ == "__main__":
    main()
