vstrains

#!/usr/bin/env python3

import argparse
import sys
import os
import platform
import numpy
import logging
import time
from datetime import date

from utils import VStrains_SPAdes

__author__ = "Runpeng Luo"
__copyright__ = "Copyright 2022-2025, VStrains Project"
__credits__ = ["Runpeng Luo", "Yu Lin"]
__license__ = "MIT"
__version__ = "1.1.0"
__maintainer__ = "Runpeng Luo"
__email__ = "John.Luo@anu.edu.au"
__status__ = "Production"


def run(args, logger):
    numpy.seterr(all="raise")
    RUNNER = {
        "spades": VStrains_SPAdes.run,
    }
    RUNNER[args.assembler](args, logger)


def main():
    parser = argparse.ArgumentParser(
        prog="VStrains",
        description="""Construct full-length viral strains under de novo approach 
        from contigs and assembly graph, currently supports SPAdes""",
    )

    parser.add_argument(
        "-a",
        "--assembler",
        dest="assembler",
        type=str,
        required=True,
        choices=["spades"],
        help="name of the assembler used. [spades]",
    )

    parser.add_argument(
        "-g",
        "--graph",
        dest="gfa_file",
        type=str,
        required=True,
        help="path to the assembly graph, (.gfa format)",
    )

    parser.add_argument(
        "-p",
        "--path",
        dest="path_file",
        type=str,
        required=False,
        help="contig file from SPAdes (.paths format), only required for SPAdes. e.g., contigs.paths",
    )

    parser.add_argument(
        "-mc",
        "--minimum_coverage",
        dest="min_cov",
        default=None,
        type=int,
        help=argparse.SUPPRESS,
        # (
        #     "minimum node coverage cutoff [default: auto]"
        # ),
    )

    parser.add_argument(
        "-ml",
        "--minimum_contig_length",
        dest="min_len",
        default=None,
        type=int,
        help=argparse.SUPPRESS,
        # ("minimum initial contig length [default: 250]"),
    )

    parser.add_argument(
        "-r",
        "--reference_fa",
        dest="ref_file",
        default=None,
        type=str,
        help=argparse.SUPPRESS,
    )

    parser.add_argument(
        "-o",
        "--output_dir",
        dest="output_dir",
        default="acc/",
        type=str,
        help="path to the output directory [default: acc/]",
    )

    parser.add_argument(
        "-d",
        "--dev_mode",
        dest="dev",
        action="store_true",
        default=False,
        help=argparse.SUPPRESS,
    )

    parser.add_argument(
        "-fwd",
        "--fwd_file",
        dest="fwd",
        required=True,
        default=None,
        type=str,
        help="paired-end sequencing reads, forward strand (.fastq format)",
    )

    parser.add_argument(
        "-rve",
        "--rve_file",
        dest="rve",
        required=True,
        default=None,
        type=str,
        help="paired-end sequencing reads, reverse strand (.fastq format)",
    )

    args = parser.parse_args()

    # parsing arguments, sanity check
    if (not args.gfa_file) or (not os.path.exists(args.gfa_file)):
        print("\nPath to the assembly graph is required, (.gfa format)")
        print("Please ensure the path is correct")
        print("\nExiting...\n")
        sys.exit(1)

    args.assembler = args.assembler.lower()

    if args.assembler.lower() == "spades":
        if (not args.path_file) or (not os.path.exists(args.path_file)):
            print(
                "\nPath to Contig file from SPAdes (.paths format) is required for SPAdes assmbler option. e.g., contigs.paths"
            )
            print("\nExiting...\n")
            sys.exit(1)
    else:
        print("\nPlease make sure to provide the correct assembler type (SPAdes).")
        print("\nExiting...\n")
        sys.exit(1)

    if args.min_len != None:
        if args.min_len < 0:
            print(
                "\nPlease make sure to provide the correct option (invalid value for min_len or min_cov)."
            )
            print("\nExiting...\n")
            sys.exit(1)
    else:
        args.min_len = 250

    if args.min_cov != None:
        if args.min_cov < 0:
            print(
                "\nPlease make sure to provide the correct option (invalid value for min_len or min_cov)."
            )
            print("\nExiting...\n")
            sys.exit(1)

    if args.output_dir[-1] == "/":
        args.output_dir = args.output_dir[:-1]

    # initialize output directory
    os.makedirs(args.output_dir, exist_ok=True)
    try:
        os.makedirs(args.output_dir + "/gfa/")
        os.makedirs(args.output_dir + "/tmp/")
        os.makedirs(args.output_dir + "/paf/")
        os.makedirs(args.output_dir + "/aln/")
    except OSError as _:
        print("\nCurrent output directory is not empty")
        print("Please empty/re-create the output directory: " + str(args.output_dir))
        print("\nExiting...\n")
        sys.exit(1)

    if os.path.exists(args.output_dir + "/vstrains.log"):
        os.remove(args.output + "/vstrains.log")

    # Setup logger
    # -----------------------
    logger = logging.getLogger("VStrains %s" % __version__)
    logger.setLevel(logging.DEBUG if args.dev else logging.INFO)

    consoleHeader = logging.StreamHandler()
    consoleHeader.setLevel(logging.INFO)
    consoleHeader.setFormatter(logging.Formatter("%(message)s"))
    logger.addHandler(consoleHeader)

    fileHandler = logging.FileHandler(args.output_dir + "/vstrains.log")
    fileHandler.setLevel(logging.DEBUG if args.dev else logging.INFO)
    fileHandler.setFormatter(logging.Formatter("%(message)s"))
    logger.addHandler(fileHandler)

    logger.info("Welcome to VStrains!")
    logger.info(
        "VStrains is a strain-aware assembly tools, which constructs full-length "
    )
    logger.info("virus strain with aid from de Bruijn assembly graph and contigs.")
    logger.info("")
    logger.info("System information:")
    try:
        logger.info("  VStrains version: " + str(__version__).strip())
        logger.info("  Python version: " + ".".join(map(str, sys.version_info[0:3])))
        logger.info("  OS: " + platform.platform())
    except Exception:
        logger.info("  Problem occurred when getting system information")

    logger.info("")
    start_time = time.time()

    logger.info("Input arguments:")
    logger.info("Assembly type: " + args.assembler)
    logger.info("Assembly graph file: " + args.gfa_file)
    logger.info("Forward read file: " + args.fwd)
    logger.info("Reverse read file: " + args.rve)
    if args.assembler == "spades":
        logger.info("Contig paths file: " + args.path_file)
    logger.info("Output directory: " + os.path.abspath(args.output_dir))
    if args.dev:
        logger.info("*DEBUG MODE is turned ON")
    logger.info("\n\n")
    logger.info(
        "======= VStrains pipeline started. Log can be found here: "
        + os.path.abspath(args.output_dir)
        + "/vstrains.log\n"
    )

    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    consoleHeader.setFormatter(formatter)
    fileHandler.setFormatter(formatter)

    # all good
    run(args, logger)

    elapsed = time.time() - start_time

    consoleHeader.setFormatter(logging.Formatter("%(message)s"))
    fileHandler.setFormatter(logging.Formatter("%(message)s"))

    logger.info("")
    logger.info("Thanks for using VStrains")
    logger.info(
        "Result is stored in {0}/strain.fasta".format(os.path.abspath(args.output_dir))
    )
    logger.info(
        "You can visualise the path stored in {0}/strain.paths via {0}/gfa/graph_L0.gfa".format(
            os.path.abspath(args.output_dir)
        )
    )
    logger.info("Finished: {0}".format(date.today().strftime("%B %d, %Y")))
    logger.info("Elapsed time: {0}".format(elapsed))
    logger.info("Exiting...")
    logger.removeHandler(fileHandler)
    logger.removeHandler(consoleHeader)

    return 0


if __name__ == "__main__":
    main()