#!/usr/bin/env python

import argparse
from phylogenetics.homologs import load_homologset
from phylogenetics.msaprobs import run_msaprobs

def main():

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

    parser = argparse.ArgumentParser(description="""
    Construct a multiple sequence alignment using MSAProbs
    """)

    # command line arguments
    parser.add_argument("-i", "--input",
            help="Input pickle file containing homologset.",
            type=str,
            required=True)

    parser.add_argument("-o", "--output",
            help="Filename to save pickled homologset.",
            type=str,
            required=True)

    parser.add_argument("--fasta",
            help="Boolean; if set to False, do not output an alignment fasta",
            type=bool,
            default=True)

    # Parse the arguments
    args = parser.parse_args()

    # ---------------------------------------------------
    # Main code for MSA
    # ---------------------------------------------------

    # Open homologset and run msaprobs to align sequences.
    hs = load_homologset(args.input)

    # Run MSAProbs
    hs2 = run_msaprobs(hs)

    # Write to homologset pickle
    hs2.write(args.output, format="pickle")

    # Split output name from extension
    oname, oext = os.path.splitext(args.output)

    # Print alignment to fasta file.
    if args.fasta:
        hs2.write(oname + '.fasta', format="fasta", aligned=True)

if __name__ == "__main__":
    main()
