# Copyright (C) 2012, 2013, 2016, 2019 by Brandon Invergo (b.invergo@gmail.com)
# This code is part of the Biopython distribution and governed by its
# license. Please see the LICENSE file that should have been included
# as part of this package.

"""Generate output for PAML tests."""

from __future__ import print_function

import os.path
import sys

from Bio._py3k import range


VERSIONS = ["4_1", "4_3", "4_4", "4_4c", "4_5", "4_6", "4_7", "4_8", "4_9a",
            "4_9i"]


def codeml(vers=None, verbose=False):
    from Bio.Phylo.PAML import codeml
    if vers is not None:
        versions = [vers]
    else:
        versions = VERSIONS
    tests = [("aa_model0", "aa_alignment.phylip", "species.tree"),
             ("aa_pairwise", "aa_alignment.phylip", "species.tree"),
             ("all_NSsites", "alignment.phylip", "species.tree"),
             ("branchsiteA", "alignment.phylip", "species.tree"),
             ("clademodelC", "alignment.phylip", "species.tree"),
             ("freeratio", "alignment.phylip", "species.tree"),
             ("ngene2_mgene02", "lysinYangSwanson2002.nuc", "lysin.trees"),
             ("ngene2_mgene34", "lysinYangSwanson2002.nuc", "lysin.trees"),
             ("pairwise", "alignment.phylip", "species.tree"),
             ("SE", "alignment.phylip", "species.tree"),
             ("m2a_rel", "alignment.phylip", "species.tree")]

    for test in tests:
        print(test[0])
        cml = codeml.Codeml()
        cml.working_dir = "temp"
        ctl_file = os.path.join("Control_files",
                                "codeml",
                                ".".join([test[0], "ctl"]))
        alignment = os.path.join("Alignments", test[1])
        tree = os.path.join("Trees", test[2])
        cml.read_ctl_file(ctl_file)
        cml.alignment = alignment
        cml.tree = tree
        for version in versions:
            # M2a_rel (NSsites 22) was introduced in PAML 4.6
            if test[0] == "m2a_rel" and int(version.split("_")[1][0]) < 6:
                continue
            print("\t{0}".format(version.replace("_", ".")))
            if test[0] in ["ngene2_mgene02", "ngene2_mgene34"] and \
               version == "4_6":
                cml.tree = ".".join([cml.tree, "4.6"])
            out_file = ".".join(["-".join([test[0], version]), "out"])
            cml.out_file = os.path.join("Results", "codeml", test[0], out_file)
            bin = "".join(["codeml", version])
            cml.run(command=bin, verbose=verbose, parse=False)


def baseml(vers=None, verbose=False):
    from Bio.Phylo.PAML import baseml
    if vers is not None:
        versions = [vers]
    else:
        versions = VERSIONS
    tests = [("model", list(range(0, 9))), ("nhomo", [1, 3, 4]),
             ("nparK", list(range(1, 5))), ("alpha1rho1", None), ("SE", None)]
    alignment = os.path.join("Alignments", "alignment.phylip")
    tree = os.path.join("Trees", "species.tree")
    for test in tests:
        print(test[0])
        bml = baseml.Baseml()
        for version in versions:
            print("\t{0}".format(version.replace("_", ".")))
            if test[1] is not None:
                for n in test[1]:
                    if (version in ["4_3", "4_4", "4_4c", "4_5"] and
                            test[0] == "nparK" and n in [3, 4]):
                        continue
                    print("\t\tn = {0}".format(n))
                    ctl_file = (os.path.join("Control_files", "baseml",
                                "{0}{1}.ctl".format(test[0], n)))
                    bml.read_ctl_file(ctl_file)
                    bml.alignment = alignment
                    bml.tree = tree
                    out_file = "{0}{1}-{2}.out".format(test[0], n, version)
                    bml.out_file = (os.path.join("Results", "baseml", test[0],
                                    out_file))
                    bin = "baseml{0}".format(version)
                    bml.run(command=bin, verbose=verbose, parse=False)
            else:
                if (version in ["4_3", "4_4", "4_4c", "4_5"] and
                        test[0] == "alpha1rho1"):
                    continue
                ctl_file = (os.path.join("Control_files", "baseml",
                            "{0}.ctl".format(test[0])))
                bml.read_ctl_file(ctl_file)
                bml.alignment = alignment
                bml.tree = tree
                out_file = "{0}-{1}.out".format(test[0], version)
                bml.out_file = (os.path.join("Results", "baseml", test[0],
                                out_file))
                bin = "baseml{0}".format(version)
                bml.run(command=bin, verbose=verbose, parse=False)


def yn00(vers=None, verbose=False):
    from Bio.Phylo.PAML import yn00
    if vers is not None:
        versions = [vers]
    else:
        versions = VERSIONS
    tests = ["yn00", "yn00_long", "yn00_dotted", "yn00_dottednum"]
    for test in tests:
        print(test)
        yn = yn00.Yn00()
        for version in versions:
            print("\t{0}".format(version.replace("_", ".")))
            ctl_file = (os.path.join("Control_files", "yn00",
                        "{0}.ctl".format(test)))
            yn.read_ctl_file(ctl_file)
            out_file = "{0}-{1}.out".format(test, version)
            yn.out_file = os.path.join("Results", "yn00", out_file)
            bin = "yn00{0}".format(version)
            yn.run(command=bin, verbose=verbose, parse=False)


def print_usage():
    versions = ", ".join(vers.replace("_", ".") for vers in VERSIONS)
    usage = """Usage: gen_results.py [-v] PROGRAM [VERSION]

Generate result files to be used in Bio.Phylo.PAML unit tests.

  -v         Use verbose output
  PROGRAM    codeml, baseml or yn00
  VERSION    %s

To use this, the PAML programs must be in your executable path and
they must be named programX_Y, where X and Y are the version numbers
(i.e. baseml4_5 or codeml4_4c). If VERSION is not specified, test
results will be generated for all versions listed above.
""" % (versions)
    sys.exit(usage)


if __name__ == "__main__":
    programs = ["codeml", "baseml", "yn00"]
    prog = None
    verbose = False
    vers = None
    if len(sys.argv) < 2:
        print_usage()
    for arg in sys.argv[1:]:
        if arg == "-v":
            verbose = True
        elif arg in programs:
            if prog is not None:
                print("Only one program at a time, please.")
                print_usage()
            prog = arg
        elif arg.replace(".", "_") in VERSIONS:
            if vers is not None:
                print("Only one version at a time, sorry.")
            vers = arg.replace(".", "_")
        else:
            print("Unrecognized argument")
            print_usage()
    if prog is None:
        print("No program specified")
        print_usage()
    if prog == "codeml":
        codeml(vers, verbose)
    elif prog == "baseml":
        baseml(vers, verbose)
    elif prog == "yn00":
        yn00(vers, verbose)
