# 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."""


import os.path
import sys


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(f"\t{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(f"\t{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(f"\t\tn = {n}")
                    ctl_file = os.path.join(
                        "Control_files", "baseml", f"{test[0]}{n}.ctl"
                    )
                    bml.read_ctl_file(ctl_file)
                    bml.alignment = alignment
                    bml.tree = tree
                    out_file = f"{test[0]}{n}-{version}.out"
                    bml.out_file = os.path.join("Results", "baseml", test[0], out_file)
                    bin = f"baseml{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", f"{test[0]}.ctl")
                bml.read_ctl_file(ctl_file)
                bml.alignment = alignment
                bml.tree = tree
                out_file = f"{test[0]}-{version}.out"
                bml.out_file = os.path.join("Results", "baseml", test[0], out_file)
                bin = f"baseml{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(f"\t{version.replace('_', '.')}")
            ctl_file = os.path.join("Control_files", "yn00", f"{test}.ctl")
            yn.read_ctl_file(ctl_file)
            out_file = f"{test}-{version}.out"
            yn.out_file = os.path.join("Results", "yn00", out_file)
            bin = f"yn00{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)