# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_Main.ipynb (unless otherwise specified).

__all__ = ['Args', 'checkParams', 'main']

# Cell
from argparse import ArgumentParser, ArgumentTypeError, RawDescriptionHelpFormatter, SUPPRESS
import os, glob, platform
from multiprocessing import cpu_count, Queue
from .Utils import *
from .Runner import *
from .Core import *
from multiprocessing import Process, Queue
from collections import OrderedDict
import itertools
from copy import deepcopy
import sys, faulthandler, platform
import numpy as np
import os
if sys.version_info.major == 2:
    from cstatgen import cstatgen_py2 as cstatgen
    from cstatgen.egglib import Align
else:
    from cstatgen import cstatgen_py3 as cstatgen
    import egglib
    from egglib import Align

class Args:
    def __init__(self):
        self.parser = ArgumentParser(description = '''\t{}, linkage analysis using sequence data\n\t'''.format("SEQLinkage V"+VERSION))
        self.getEncoderArguments(self.parser)
        self.getLinkageArguments(self.parser)


    def get(self):
        return self.parser.parse_args()

    def getEncoderArguments(self, parser):
        vargs = parser.add_argument_group('Collapsed haplotype pattern method arguments')
        vargs.add_argument('--single-marker', action='store_true', dest = "single_marker",help='''Use single variant as the marker. Otherwise, use CHP markers.''')
        vargs.add_argument('--fam', metavar='FILE', required=True, dest = "tfam", help='''Input pedigree and phenotype information in FAM format.''')
        vargs.add_argument('--vcf', metavar='FILE', required=True, help='''Input VCF file, bgzipped.''')
        vargs.add_argument('--anno', metavar='FILE', required=False, help='''Input annotation file from annovar.''')
        vargs.add_argument('--pop', metavar='FILE', required=False, help='''Input two columns file, first column is family ID, second column population information.''')
        vargs.add_argument('--included-vars', metavar='FILE', dest='included_vars', help='''Variants to be included for linkage analysis, if None, the analysis won't filter any variants. But you can still set AF cutoff by -c argment.''')
        vargs.add_argument('-b', '--blueprint', metavar = 'FILE',
                           help='''Blueprint file that defines regional marker (format: "chr startpos endpos name avg.distance male.distance female.distance").''')
        vargs.add_argument('-c', '--maf-cutoff', metavar='P', default=None, type=float, dest = "maf_cutoff",
                           help='''MAF cutoff to define variants to be excluded from analyses. this is useful, if you need to analyse multiple population together.''')
        vargs.add_argument('-o', '--output', metavar='Name', default='data/linkage_analysis/', help='''Output name prefix.''')

        vargs.add_argument('--build', metavar='STRING', default='hg38', choices = ["hg19", "hg38"], help='''Reference genome version for VCF file.''')
        vargs.add_argument('--window', metavar='INT', type=int, default=1000,help='If no blueprint, seprate chromosome to pseudogenes with 1000 (as default) variants.')
        vargs.add_argument('--freq', metavar='INFO', default = 'AF',help='''Info field name for allele frequency in VCF file.''')
        vargs.add_argument('--chrom-prefix', metavar='STRING', dest = 'chr_prefix',
                           help='''Prefix to chromosome name in VCF file if applicable, e.g. "chr".''')

    def getLinkageArguments(self, parser):
        vargs = parser.add_argument_group('LINKAGE options')
        vargs.add_argument('--run-linkage', action='store_true', dest = "run_linkage",
                           help='''Perform Linkage analysis.''')
        vargs.add_argument('-K', '--prevalence', metavar='FLOAT', type=float,default = 0.001,
                           help='Disease prevalence. Default to 0.001.')
        vargs.add_argument('--moi', metavar='STRING', dest = "inherit_mode",default = 'AD',
                           # choices=['AD', 'AR', 'Xlinked', 'Y'],
                           choices=['AD', 'AR'],
                           help='Mode of inheritance, AD/AR: autosomal dominant/recessive. Default to AD.')
        vargs.add_argument('-W', '--wt-pen', metavar='FLOAT', type=float, dest = "wild_pen",default=0.01,
                           help='Penetrance for wild type. Default to 0.01.')
        vargs.add_argument('-M', '--mut-pen', metavar='FLOAT', type=float, dest = "muta_pen",default=0.9,
                           help='Penetrance for mutation. Default to 0.9.')
        vargs.add_argument('--theta-max', metavar='FLOAT', type=float, dest = "theta_max", default = 0.5,
                           help='Theta upper bound. Default to 0.5.')
        vargs.add_argument('--theta-inc', metavar='FLOAT', type=float, dest = "theta_inc", default = 0.05,
                           help='Theta increment. Default to 0.05.')

# Cell
def checkParams(args):
    '''set default arguments or make warnings'''
    env.setoutput(args.output)
    args.vcf = os.path.abspath(os.path.expanduser(args.vcf))
    args.tfam = os.path.abspath(os.path.expanduser(args.tfam))
    for item in [args.vcf, args.tfam]:
        if not os.path.exists(item):
            env.error("Cannot find file [{}]!".format(item), exit = True)
    if len([x for x in set(getColumn(args.tfam, 6)) if x.lower() not in env.ped_missing]) > 2:
        env.trait = 'quantitative'
    env.log('{} trait detected in [{}]'.format(env.trait.capitalize(), args.tfam))
    if not args.blueprint:
        if not args.anno:
            args.blueprint = os.path.join(env.resource_dir, 'genemap.{}.txt'.format(args.build))
        else:
            env.log('Generate regions by annotation')
    return True

# Cell
def main():
    '''the main encoder function'''
    args = Args().get()
    checkParams(args)
    env.log(args)
    # load data
    data = RData(args.vcf, args.tfam,args.anno,args.pop,allele_freq_info=args.freq,included_variant_file=args.included_vars)
    samples_vcf = data.samples_vcf

    if len(samples_vcf) == 0:
        env.error("Fail to extract samples from [{}]".format(args.vcf), exit = True)
    env.log('{:,d} samples found in [{}]'.format(len(samples_vcf), args.vcf))
    samples_not_vcf = data.samples_not_vcf

    if len(data.families) == 0:
        env.error('No valid family to process. ' \
                  'Families have to be at least trio with at least one member in VCF file.', exit = True)
    if len(data.samples) == 0:
        env.error('No valid sample to process. ' \
                  'Samples have to be in families, and present in both TFAM and VCF files.', exit = True)

    if args.blueprint:
        # load blueprint
        try:
            env.log('Loading marker map from [{}] ...'.format(args.blueprint))
            with open(args.blueprint, 'r') as f:
                regions = [x.strip().split() for x in f.readlines()]
        except IOError:
            env.error("Cannot load regional marker blueprint [{}]. ".format(args.blueprint), exit = True)
    else:
        env.log('separate chromosome to regions')
        regions=data.get_regions(step=args.window)  #whole-genome linkage analysis
    env.log('{:,d} families with a total of {:,d} samples will be scanned for {:,d} pre-defined units'.\
            format(len(data.families), len(data.samples), len(regions)))

    run_each_region(regions,data,RegionExtractor(args.vcf, build = args.build, chr_prefix = args.chr_prefix),MarkerMaker(maf_cutoff = args.maf_cutoff),LinkageWriter(len(samples_not_vcf)),
                    runlinkage=args.run_linkage,cutoff=args.maf_cutoff,chp=args.single_marker==False,rho=np.arange(0,args.theta_max,args.theta_inc),
                    model = args.inherit_mode,penetrances = [args.wild_pen,args.muta_pen,args.muta_pen],dfreq=args.prevalence)
    env.log('============= Finish analysis ==============')



if __name__ == '__main__':
    main()