#!/opt/conda/conda-bld/vitap_1776377915507/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placeho/bin/python

import argparse
from VITAP_assignment import assignment
from VITAP_upd import upd


BANNER = (
    "*******************************************************************\n"
    "The VITAP (Viral Taxonomic Assignment Pipeline) v.1.12\n"
    "*******************************************************************"
)

CITATION = (
    "Please cite: Zheng, K. et al. VITAP: a high precision tool for DNA and RNA "
    "viral classification based on meta-omic data. Nat Commun 16, 2226 (2025)."
)


def main():
    parser = argparse.ArgumentParser(
        prog="VITAP",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=BANNER,
        epilog=(
            "Usage: VITAP <command> [options]\n"
            "Commands:\n"
            "  assignment    VITAP taxonomic assignment pipeline.\n"
            "  upd           Update program of VITAP.\n"
            "                The VITAP will update database based on ICTV VMR "
            "(or other formated source) and UniRef90 database.\n\n"
            f"{CITATION}\n"
            "For command-specific help, run: VITAP <command> -h"
        ),
    )

    subparsers = parser.add_subparsers(
        dest="command",
        required=True
    )

    # ================= assignment =================
    parser_assign = subparsers.add_parser(
        "assignment",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=BANNER,
        help="VITAP taxonomic assignment pipeline.",
        epilog="Example: VITAP assignment -i target_genome.fasta -d ./DB/ -o VITAP_result"
    )
    parser_assign.add_argument("-i", "--fasta", required=True, help="Input genome FASTA")
    parser_assign.add_argument("-d", "--db", required=True, help="Database directory")
    parser_assign.add_argument("-p", "--cpu", default=1, help="Threads for DIAMOND (ICTV blast)")
    parser_assign.add_argument("-o", "--out", required=True, help="Output directory")
    parser_assign.add_argument(
        "--low_conf",
        action="store_true",
        help="Export low-confidence assignments (lineage_score < 0.1). Default: discard."
    )
    parser_assign.set_defaults(func=assignment)

    # ================= upd =================
    parser_upd = subparsers.add_parser(
        "upd",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=BANNER,
        help="Update program of VITAP.",
        epilog=(
            "Example: VITAP upd --vmr raw_VMR.csv -o VMR_reformat.csv -d VMR-MSL\n\n"
            "The VITAP will update database based on ICTV VMR "
            "(or other formated source) and UniRef90 database."
        )
    )
    parser_upd.add_argument("--vmr", required=True, help="Input raw ICTV VMR source file in csv format.")
    parser_upd.add_argument("-o", "--out", required=True, help="Reformat ICTV VMR source file for VITAP utilization")
    parser_upd.add_argument("-d", "--db", required=True, help="You need to assign a name for a new DB")
    parser_upd.set_defaults(func=upd)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
