from asyncio import run
from argparse import ArgumentParser, Namespace
from time import time

from locimend.model import infer_sequence, train_model


def parse_arguments() -> Namespace:
    parser = ArgumentParser()
    subparsers = parser.add_subparsers(dest="task")
    parser_train = subparsers.add_parser("train")
    parser_infer = subparsers.add_parser("infer")
    parser_train.add_argument(
        "data_file", help="FASTQ file containing the sequences with errors"
    )
    parser_train.add_argument(
        "label_file", help="FASTQ file containing the sequences without errors"
    )
    parser_infer.add_argument("sequence", help="DNA sequence with errors")
    return parser.parse_args()


async def execute_task(args):
    if args.task == "train":
        start_time = time()
        train_model(data_file=args.data_file, label_file=args.label_file)
        end_time = time()
        print(f"Training time: {end_time - start_time}")
    else:
        prediction = await infer_sequence(sequence=args.sequence)
        print(f"Error-corrected sequence: {prediction}")


def main() -> None:
    args = parse_arguments()
    run(execute_task(args))


if __name__ == "__main__":
    main()