#! /usr/bin/env python3

import argparse
import sys
from pathlib import Path
from ribomala import index, qc, analysis


def get_arguments():
    parser = argparse.ArgumentParser(
        prog="ribomala",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""

        --◉--◎--◉--◎--◉--◎--◉--

Ribomala is a tool to analyze ribosome profiling data.

It contains three modes:
- Index: Indexes the transcriptome fasta
- QC: Checks for ribosome reading frame and periodicity
- Analysis: Calculates E-, P- and A-site ribosome occupancy

""",
        epilog="""

If you find ribomala useful for your research please cite as:
https://github.com/ps-puneetsharma/ribomala

        --◉--◎--◉--◎--◉--◎--◉--

""",
        add_help=True,
    )

    # Create subparsers for each mode
    subparsers = parser.add_subparsers(
        dest="mode", required=True, help="Mode of operation"
    )

    # Index mode parser
    index_parser = subparsers.add_parser("index", help="Index the transcriptome fasta")

    index_parser.add_argument(
        "-f",
        "--fasta",
        type=Path,
        required=True,
        help="Path to transcriptome FASTA file",
    )

    # QC mode parser
    qc_parser = subparsers.add_parser(
        "qc", help="Check for ribosome reading frame and periodicity"
    )

    qc_parser.add_argument(
        "-s",
        "--samples",
        required=True,
        help="A text file (.txt) containing samples to be processed",
    )

    qc_parser.add_argument(
        "-n",
        "--min",
        type=int,
        metavar="int",
        required=False,
        default=28,
        help="Minimum ribosome protected fragment size. DEFAULT: 28",
    )

    qc_parser.add_argument(
        "-x",
        "--max",
        type=int,
        metavar="int",
        required=False,
        default=33,
        help="Maximum ribosome protected fragment size. DEFAULT: 33",
    )

    qc_parser.add_argument(
        "-i",
        "--input",
        type=Path,
        required=True,
        help="Input directory containing BAM files",
    )

    qc_parser.add_argument(
        "-o", "--output", type=Path, required=True, help="Output directory for QC files"
    )

    # Analysis mode parser
    analysis_parser = subparsers.add_parser(
        "analysis", help="Calculates E-, P- and A-site ribosome occupancy"
    )

    analysis_parser.add_argument(
        "-c",
        "--txcsv",
        type=Path,
        required=True,
        help="Path to transcriptome CSV file generated by QC module",
    )

    analysis_parser.add_argument(
        "-s",
        "--samples",
        required=True,
        help="A (tab-separated) text file (.tsv or .txt) containing samples to be processed",
    )

    analysis_parser.add_argument(
        "-a",
        "--exclstart",
        type=int,
        metavar="int",
        required=False,
        default=60,
        help="Number of nucleotides to exclude from CDS start (must be multiple of 3). DEFAULT: 60",
    )

    analysis_parser.add_argument(
        "-z",
        "--exclend",
        type=int,
        metavar="int",
        required=False,
        default=60,
        help="Number of nucleotides to exclude from CDS end (must be multiple of 3). DEFAULT: 60",
    )

    analysis_parser.add_argument(
        "-u",
        "--upstream",
        type=int,
        metavar="int",
        required=False,
        default=30,
        help="The number of upstream codons to check from the provided offset position. DEFAULT: 30",
    )

    analysis_parser.add_argument(
        "-d",
        "--downstream",
        type=int,
        metavar="int",
        required=False,
        default=30,
        help="The number of downstream codons to check from the provided offset position. DEFAULT: 30",
    )

    analysis_parser.add_argument(
        "-m",
        "--codon",
        type=str,
        metavar="str",
        required=False,
        default="AAA",
        help="Codon(s) for which mean relative enrichment at provided offsets needs to be checked. Provide multiple codons separated with comma",
    )

    analysis_parser.add_argument(
        "-i",
        "--input",
        type=Path,
        required=True,
        help="Input directory containing BAM files",
    )

    analysis_parser.add_argument(
        "-o",
        "--output",
        type=Path,
        required=True,
        help="Output directory for analysis files",
    )

    return parser.parse_args()


def main():
    args = get_arguments()

    if args.mode == "index":
        index.run(args)
    elif args.mode == "qc":
        qc.run(args)
    elif args.mode == "analysis":
        analysis.run(args)
    else:
        args.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
