#!/usr/bin/env python

'''
Author: Shohei Kojima @ RIKEN
Copyright (c) 2020 RIKEN
All Rights Reserved
See file LICENSE for details.
'''


import os,sys,datetime,argparse,glob,shutil,logging


# version
version='v1.0.2 2022/01/09'


# args
parser=argparse.ArgumentParser(description='')
parser.add_argument('-merge_mei', help='Specify a file containing paths to output directories when you merge MEIs.', action='store_true')
parser.add_argument('-merge_absent_me', help='Specify a file containing paths to output directories when you merge absent MEs.', action='store_true')
parser.add_argument('-f', metavar='str', type=str, help='Required. Specify paths to the vcf files to be merged. One line should contain one path to a vcf file.')
parser.add_argument('-fa', metavar='str', type=str, help='Required. Specify reference genome which are used when input reads were mapped. Example: GRCh38DH.fa')
parser.add_argument('-rep', metavar='str', type=str, help='Required. Specify RepBase file used for repeatmasking. Example: humrep.ref')
parser.add_argument('-repremove', metavar='str', type=str, help='Optional. Specify full path to a file containing the names of non-ME repeat class. Default: /path/to/MEGAnE/docs/human_non_ME_rep_headers.txt')
parser.add_argument('-chr', metavar='str', type=str, help='Optional. Specify name(s) of the chromosome(s) to be analyzed with comma as a delimiter (e.g. chr1,chr2). By default (i.e. without this flag), MEGAnE analyzes all chromosomes.')
parser.add_argument('-male_sex_chr', metavar='str', type=str, help='Optional. Specify name(s) of the male-specific chromosome(s). Default: chrY,Y', default='chrY,Y')
parser.add_argument('-female_sex_chr', metavar='str', type=str, help='Optional. Specify name(s) of the chromosome(s) that is diploid in female. Default: chrX,X', default='chrX,X')
parser.add_argument('-no_sex_chr', help='Optional. Specify if the input species does not have sex chromosome (e.g. environmental sex determination).', action='store_true')
parser.add_argument('-outdir', metavar='str', type=str, help='Optional. Specify output directory. Default: ./jointcall_out', default='./jointcall_out')
parser.add_argument('-cohort_name', metavar='str', type=str, help='Optional. Specify a cohort name. This will be used for the variant names as well the output file name. Default: YYYY-MM-DD-HHMMSS')
parser.add_argument('-make_scaffold', help='Optional. Specify if you only generate a scaffold VCF. This option is for joint calling of massive samples (e.g. >10,000).', action='store_true')
parser.add_argument('-input_scaffold', metavar='str', type=str, help='Optional. Specify a scaffold VCF generated by the "-make_scaffold" option. This option is for joint calling of massive samples (e.g. >10,000).')
parser.add_argument('-chunk_f', metavar='str', type=str, help='Optional. Specify a file with chunked samples. Please use with the "-input_scaffold" option. This option is for joint calling of massive samples (e.g. >10,000).')
parser.add_argument('-chunk_vcf_list', metavar='str', type=str, help='Optional. Specify file a file if you merge multiple VCF files generated by the "-input_scaffold" and "-chunk_f" flags. This option is for joint calling of massive samples (e.g. >10,000).')
parser.add_argument('-pybedtools_tmp', metavar='str', type=str, help='Optional. Specify directory for temporary bedtools files, e.g. /dev/shm')
parser.add_argument('-do_not_overwrite', help='Optional. Specify if you do NOT overwrite previous results.', action='store_true')
parser.add_argument('-p', metavar='int', type=int, help='Optional. Number of threads. 3 or more is recommended. Default: 2', default=2)
parser.add_argument('-v', '--version', action='version', version='MEGAnE %s %s' % (os.path.basename(__file__), version))
parser.add_argument('-mmap', action='store_true', help=argparse.SUPPRESS)
args=parser.parse_args()
args.version=version


# start
import init
init.init_jointcall(args, version)


# logging
import log
if args.merge_mei is True:
    args.logfilename='for_debug_jointcall_ins.log'
else:
    args.logfilename='for_debug_jointcall_abs.log'
if os.path.exists(os.path.join(args.outdir, args.logfilename)) is True:
    os.remove(os.path.join(args.outdir, args.logfilename))
log.start_log(args)
log.logger.debug('Logging started.')


# initial check
import initial_check
log.logger.debug('This is %s version %s' % (__file__, version))
print()
log.logger.info('Initial check started.')
initial_check.check_merge_vcf(args, sys.argv)


# set up
import setup
setup.setup_merge_vcf(args, init.base)
params=setup.params
args.fai=setup.fai_path
args.rep_headers_to_be_removed=setup.rep_headers_to_be_removed


# output file names
import utils
filenames=utils.empclass()

filenames.repdb           =os.path.join(args.outdir, 'repdb')
filenames.rep_unknown_fa  =os.path.join(args.outdir, 'rep_unknown.fa')
filenames.blast_tmp_res   =os.path.join(args.outdir, 'blastn_tmp.txt')
filenames.reshaped_rep    =os.path.join(args.outdir, 'reshaped_repbase.fa')

filenames.merged_vcf_ins  =os.path.join(args.outdir, '%s_MEI_scaffold.vcf.gz' % args.cohort_name)
filenames.filled_vcf_ins  =os.path.join(args.outdir, '%s_MEI_jointcall.vcf.gz' % args.cohort_name)
filenames.merged_vcf_abs  =os.path.join(args.outdir, '%s_MEA_scaffold.vcf.gz' % args.cohort_name)
filenames.filled_vcf_abs  =os.path.join(args.outdir, '%s_MEA_jointcall.vcf.gz' % args.cohort_name)
filenames.scaffold_info_i =os.path.join(args.outdir, '%s_scaffold_info_MEI.txt' % args.cohort_name)
filenames.scaffold_info_a =os.path.join(args.outdir, '%s_scaffold_info_MEA.txt' % args.cohort_name)


if args.chunk_vcf_list is None:
    # 0. preprocess repbase file
    if args.merge_absent_me is True:
        import reshape_rep, blastn
        print()
        log.logger.info('Preprocess started.')
        reshape_rep.reshape(args, params, filenames)
        for ext in ['nhr', 'nin', 'nog', 'nsd', 'nsi', 'nsq']:
            f='%s.%s' % (filenames.repdb, ext)
            if os.path.exists(f) is True:
                os.remove('%s.%s' % (filenames.repdb, ext))


    # 1. make scaffold, fill-in
    import make_scaffold_vcf,fill_in_vcf
    log.logger.info('File check started.')
    make_scaffold_vcf.check_paths(args)
    if args.merge_mei is True:
        if args.input_scaffold is None:
            log.logger.info('Making scaffold started.')
            make_scaffold_vcf.merge_vcf_ins(args, params, filenames)
        else:
            log.logger.info('Making scaffold skipped.')
        if args.make_scaffold is True:
            make_scaffold_vcf.output_summary_for_fill_in(args, params, filenames)
            log.logger.info('Making scaffold VCF finished.')
        else:
            log.logger.info('Making final VCF started.')
            fill_in_vcf.fill_in_ins(args, params, filenames)
    elif args.merge_absent_me is True:
        if args.input_scaffold is None:
            log.logger.info('Making scaffold started.')
            make_scaffold_vcf.merge_vcf_abs(args, params, filenames)
            os.remove(filenames.reshaped_rep)
        else:
            log.logger.info('Making scaffold skipped.')
        if args.make_scaffold is True:
            make_scaffold_vcf.output_summary_for_fill_in(args, params, filenames)
            log.logger.info('Making scaffold VCF finished.')
        else:
            log.logger.info('Making final VCF started.')
            fill_in_vcf.fill_in_abs(args, params, filenames)
            if os.path.exists(filenames.reshaped_rep) is True:
                os.remove(filenames.reshaped_rep)

else:
    # merge chunks
    import merge_chunk_vcf
    log.logger.info('File check started.')
    merge_chunk_vcf.merge(args, params, filenames)


# all finish!
utils.output_finish_comment_merge_vcf(args, filenames)
