#!python

"""
 Copyright 2020, Conner Kojima <cykojima@usc.edu>

 This file is the main driver of rrap, a bionformatics pipeline that
 generates bam.stats files and rpkm visualization given reference 
 genome and metagenome directories.

"""

import os
import argparse
import subprocess
from pathlib import Path
from rrap import concatenator
from rrap import indexer
from rrap import read_recruiter
from rrap import visualizer


def main():
    c = Controller()
    c.run()


class Controller:
    def __init__(self):
        self.p = argparse.ArgumentParser(prog="RRAP",
                                         description="Run read recruitment on a set of cleaned fna files")

        # argument groups
        self.inputs = None
        self.outputs = None
        self.optional = None
        self.subcommands = None
        self.arg_groups = []

        # args
        self.args = None

        # pipes
        self.concatenator = None
        self.indexer = None
        self.read_recruiter = None
        self.visualizer = None

        # intermediate products
        self.cat_file_path = None
        self.index_dir_path = None

        self.rpkm_heater_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "rpkm_heater.py")
        self.stats_dir_path = None

    def run(self):
        self.add_arguments()

        print("---------making output dir if needed-------------")
        self.set_output_dir()

        if self.args.crg:
            self.cat_file_path = self.args.crg
        else:
            print("---------concatenating reference genomes-------------")
            self.concatenator = concatenator.Concatenator(self.args)
            self.cat_file_path = self.concatenator.concatenate()

        if not self.args.index_pass:
            print("---------indexing reference genomes-------------")
            self.indexer = indexer.Indexer(self.args, self.index_dir_path, self.cat_file_path)
            self.indexer.index()

        if not self.args.rr_pass:
            print("---------read recruitment and data transform-------------")
            self.read_recruiter = read_recruiter.ReadRecruiter(self.args, self.index_dir_path,
                                                            self.cat_file_path, self.stats_dir_path)
            self.read_recruiter.read_recruit()

        if not self.args.vis_pass:
            print("---------rpkm_heater-------------")
            self.visualizer = visualizer.Visualizer(self.args, self.rpkm_heater_path, self.stats_dir_path)
            self.visualizer.visualize()

        # TODO V2
        # self.fastq_dump()
        # self.clean()

    def add_arguments(self):
        # TODO specify argument groups
        self.inputs = self.p.add_argument_group("## input arguments")
        self.outputs = self.p.add_argument_group("## output arguments")
        self.optional = self.p.add_argument_group("## options")

        self.arg_groups.extend([self.inputs, self.outputs, self.optional])

        # Add the arguments
        self.inputs.add_argument('-i', help='text file of all dir paths that contain cleaned metaG fna files', 
                                 required=True, metavar='path of input txt file')
        self.inputs.add_argument('-crg', help=' path for concatenated reference genome fa file', 
                                 required=False, metavar='path of concatenated genome fa file')
        self.inputs.add_argument('-rg', help='input directory for reference genomes', 
                                 required=True, metavar='input reference genome dir path')
        self.outputs.add_argument('-o', help='output directory path', 
                                  required=True, metavar='output dir path')
        self.inputs.add_argument('-n', help='name of the project', required=True, metavar='project name')
        self.inputs.add_argument('-sort_gen', help='txt file of sorted genomes', required=False, 
                                 metavar='genome sorted txt file')
        self.inputs.add_argument('-sort_samples', help='txt file of sorted samples', required=False,
                                 metavar='sample sorted txt file')
        self.inputs.add_argument("--threads", help='number of available threads', required=False,
                                 metavar='thread number')

        # specify optional args
        self.options.add_argument("--merge-contigs", default=False, dest='contig_merge',
                                  action='store_true', help="Concatenate contigs under individual organisms")
        self.options.add_argument("--skip-indexing", default=False, dest='index_pass',
                                  action='store_true', 
                                  help='Specify if the indexing step has already been completed and can be skipped. \
                                        If this flag is used, please check that bowtie2 index files exist e.g. \
                                        <output_dir_path>/index_dir/<project_name>.x.bt2')
        self.options.add_argument("--skip-rr", default=False, dest='rr_pass',
                                  action='store_true',
                                  help='Specify if the read recruitment step has already been completed and can be' \
                                       'skipped. If this flag is used, please check that read recruitment files exist' \
                                       'e.g. <output_dir_path>/')
        self.options.add_argument("--skip-vis", default=False, dest='vis_pass',
                                  action='store_true',
                                  help='Specify if the visualization step can be skipped')

        self.args = self.p.parse_args()
        

    def set_output_dir(self):
        if self.args.o:
            # create output dir and create inner stats and index dir
            self.stats_dir_path = os.path.join(self.args.o, "stats_dir_" + self.args.n)
            self.index_dir_path = os.path.join(self.args.o, "index_dir", self.args.n)
            if not os.path.isdir(self.args.o):
                subprocess.run("mkdir " + self.args.o, shell=True)
            if not os.path.isdir(self.stats_dir_path):
                subprocess.run("mkdir " + self.stats_dir_path, shell=True)
            if not os.path.isdir(self.index_dir_path):
                subprocess.run("mkdir " + os.path.join(self.args.o, "index_dir"), shell=True)


if __name__ == "__main__":
    main()
