#!/usr/bin/env python

# Copyright Amine Namouchi
# snpToolkit is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but without
# any warranty; without even the implied warranty of merchantability or fitness
# for a particular purpose. See the GNU General Public License for more details
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.


__licence__ = 'GPLv3'
__author__ = 'Amine Namouchi'
__author_email__ = 'amine.namouchi@gmail.com'
__version__ = '2.2.3'

if __name__ == "__main__":

    from argsLogger_snpToolkit import *
    import multiprocessing as mp
    import sys,os
    import subprocess
    from tqdm import tqdm
    import glob
    from pathlib import Path
    from datetime import datetime
    from subprocess import Popen, PIPE, STDOUT

    logger = setupLogger()

    import sys
    if len(sys.argv) == 1:
        logger.error('Run snpToolkit with -h option for instructions')
        sys.exit(0)

    args = get_options()
    cwd = os.path.dirname(__file__)


    if args.command == 'explore':
        cmd = ['python',cwd+'/explore_snpToolkit.py','-i', args.identifier]
        proc = subprocess.Popen(cmd, stderr=subprocess.STDOUT)
        proc.wait()

    elif args.command == 'annotate':
        from annotate_snpToolkit import parse_genbank_file
        from calls_snpToolkit import annotate
        try:
            from Bio import SeqIO
        except ImportError:
            logger.error('Please install Biopython before to continue')
            sys.exit(0)

        allowed_format = {".vcf", ".vcf.gz", ".vcf.zip"}
        FilesToProcess = [FILE for FILE in glob.glob('*' + args.identifier + '*') if FILE.endswith(tuple(allowed_format))]
       
        if len(FilesToProcess) == 0:
            logger.error(
                'No input file detected! Please check your(s) file(s) name(s)...')
            sys.exit(0)
        if args.genbank != None:
            if Path(args.genbank).exists():
                try:
                    annotationDB = parse_genbank_file(args.genbank)
                except BaseException as error:
                    logger.error(
                        'Something went wrong when trying to extract data from your genbank file: {}'.format(error))
                    sys.exit(0)
            else:
                logger.error('The file ' + args.genbank + ' does not exist')
                sys.exit(0)

        if args.excludeCloseSNPs == 0:
            logger.error('You must provide a value > 0 with the -f option')
            sys.exit(0)

        if args.exclude != None:   
            if Path(args.exclude).exists()==False:
                logger.error('The file ' + args.exclude + ' does not exist')
                sys.exit(0)
        snps_output_directory = 'snpToolkit_SNPs_output_' + datetime.now().strftime("%Y-%m-%d %H:%M:%S").replace(' ', '_at_').replace(':', '.')
        cmd = ['mkdir', snps_output_directory]
        proc = Popen(cmd, stdout=PIPE, stderr=STDOUT)
        proc.wait()
        indels_ouput_directory = 'snpToolkit_INDELS_output_' + datetime.now().strftime("%Y-%m-%d %H:%M:%S").replace(' ', '_at_').replace(':', '.')
        cmd = ['mkdir', indels_ouput_directory]
        proc = Popen(cmd, stdout=PIPE, stderr=STDOUT)
        proc.wait()
        options = args
        logger.info('{} CPUs requested out of {} detected on this machine'.format(options.processors,mp.cpu_count()))
        logger.info('snpToolkit is filtering and annotating your SNPs')
        pool = mp.Pool(processes=options.processors)
        pbar = tqdm(total=len(FilesToProcess))
  
        def update(*a):
            pbar.update()

        for inputFile in FilesToProcess:
            pool.apply_async(annotate,args=(options,inputFile,annotationDB,snps_output_directory,indels_ouput_directory),callback=update)
        
        pool.close()
        pool.join()
        logger.info('snpToolkit output files will be located in folders {} and {}'.format(snps_output_directory,indels_ouput_directory))


    elif args.command == 'combine':
        try:
            import pysam
        except ImportError:
            logger.error('Please install pysam before to continue')
            sys.exit()

        from combine_snpToolkit import *
        from calls_snpToolkit import combine
        
        combine(args)
    
    elif args.command == 'viz':
        cmd = ['python',cwd+'/plot_snpToolkit_output.py','--dir', args.directory]
        proc = subprocess.Popen(cmd, stderr=subprocess.STDOUT)
        proc.wait()

