#!/usr/bin/env python

import argparse
import bz2
import gzip
import logging
import os
import re
import sys

import pyfizi
import numpy as np
import pandas as pd

from scipy.stats import chi2

np.seterr(invalid='ignore')


def parse_annots(alist):
    values = re.split(",", alist)
    if not values:
        raise argparse.ArgumentTypeError("annot-list should be comma-separated list with no spaces")

    return values


def parse_pos(pos, option):
    """
    Parse a specified genomic position.
    Should be digits followed optionally by case-insensitive Mb or Kb modifiers.

    :param pos: string to be parsed
    :param option: parameter name being parsed

    :return: float position
    """

    match = re.match("^(([0-9]*[.])?[0-9]+)(mb|kb)?$", pos, flags=re.IGNORECASE)
    if match:
        pos_tmp = float(match.group(1))  # position
        pos_mod = match.group(3)  # modifier
        if pos_mod:
            pos_mod = pos_mod.upper()
            if pos_mod == "MB":
                pos_tmp *= 1000000
            elif pos_mod == "KB":
                pos_tmp *= 1000

        position = pos_tmp
    else:
        raise ValueError("Option {} {} is an invalid genomic position".format(option, pos))

    return position


def parse_locations(locations, chrom=None, start_bp=None, stop_bp=None):
    """
    Parse user-specified BED file with [CHR, START, STOP] windows defining where to perform
    imputation.

    If user also specified chr, start-bp, or stop-bp arguments filter on those as well.

    :param locations: enumerable object containing locations
    :param chrom: string chromosome to filter by (default None)
    :param start_bp: float starting base pair position
    :param stop_bp: float ending base pair position

    :return: enumerable of (chrom, start, stop) tuples
    """

    for idx, line in enumerate(locations):
        # skip comments
        if "#" in line:
            continue

        row = line.split()

        if len(row) < 3:
            raise ValueError("Line {} in locations file does not contain [CHR, START, STOP]".format(idx))

        chrom_arg = row[0]
        start_arg = parse_pos(row[1], "start argument in locations file")
        stop_arg = parse_pos(row[2], "stop argument in locations file")

        if chrom is not None and chrom_arg != chrom:
            continue
        elif start_bp is not None and start_arg < start_bp:
            continue
        elif stop_bp is not None and stop_arg > stop_bp:
            continue

        yield [chrom_arg, start_arg, stop_arg]

    return


def get_command_string(args):
    """
    Format fizi call and options into a string for logging/printing

    :param args: argparse arguments to fizi

    :return: string containing formatted arguments to fizi
    """

    base = "fizi {}{}".format(args[0], os.linesep)
    rest = args[1:]
    rest_strs = []
    needs_tab = True
    for cmd in rest:
        if "-" == cmd[0]:
            if cmd in ["--quiet", "-q", "--verbose", "-v"]:
                rest_strs.append("\t{}{}".format(cmd, os.linesep))
                needs_tab = True
            else:
                rest_strs.append("\t{}".format(cmd))
                needs_tab = False
        else:
            if needs_tab:
                rest_strs.append("\t{}{}".format(cmd, os.linesep))
                needs_tab = True
            else:
                rest_strs.append(" {}{}".format(cmd, os.linesep))
                needs_tab = True

    return base + "".join(rest_strs) + os.linesep


def read_header(fh):
    """Read the first line of a file and returns a list with the column names."""
    (openfunc, compression) = get_compression(fh)
    return [x.rstrip('\n') for x in openfunc(fh).readline().split()]


def get_cname_map(flag, default, ignore):
    """
    Figure out which column names to use.

    Priority is
    (1) ignore everything in ignore
    (2) use everything in flags that is not in ignore
    (3) use everything in default that is not in ignore or in flags

    The keys of flag are cleaned. The entries of ignore are not cleaned. The keys of defualt
    are cleaned. But all equality is modulo clean_header().

    """
    clean_ignore = {clean_header(x) for x in ignore}
    both = clean_ignore | set(flag)
    cname_map = {x: flag[x] for x in flag if x not in clean_ignore}
    cname_map.update({x: default[x] for x in default if x not in both})

    return cname_map


def get_compression(fh):
    """
    Read filename suffixes and figure out whether it is gzipped,bzip2'ed or not compressed
    """
    if fh.endswith('gz'):
        compression = 'gzip'
        openfunc = lambda x: gzip.open(x, "rt")
    elif fh.endswith('bz2'):
        compression = 'bz2'
        openfunc = lambda x: bz2.open(x, "rt")
    else:
        openfunc = open
        compression = None

    return openfunc, compression


def clean_header(header):
    """
    For cleaning file headers.
    - convert to uppercase
    - replace dashes '-' with underscores '_'
    - replace dots '.' (as in R) with underscores '_'
    - remove newlines ('\n')
    """
    return header.upper().replace('-', '_').replace('.', '_').replace('\n', '')


def filter_pvals(pvals, log):
    """Remove out-of-bounds P-values"""

    ii = (pvals >= 0) & (pvals <= 1)
    bad_p = (~ii).sum()
    if bad_p > 0:
        msg = '{N} SNPs had P outside of [0,1]. The P column may be mislabeled'
        log.warning(msg.format(N=bad_p))

    return ii


def filter_info(info, log, info_min):
    """Remove INFO < info_min (default 0.9) and complain about out-of-bounds INFO."""

    if type(info) is pd.Series:  # one INFO column
        jj = ((info > 2.0) | (info < 0)) & info.notnull()
        ii = info >= info_min
    elif type(info) is pd.DataFrame:  # several INFO columns
        jj = (((info > 2.0) & info.notnull()).any(axis=1) | (
            (info < 0) & info.notnull()).any(axis=1))
        ii = (info.sum(axis=1) >= info_min * (len(info.columns)))
    else:
        raise ValueError('Expected pd.DataFrame or pd.Series.')

    bad_info = jj.sum()
    if bad_info > 0:
        msg = '{N} SNPs had INFO outside of [0,1.5]. The INFO column may be mislabeled'
        log.warning(msg.format(N=bad_info))

    return ii


def filter_frq(frq, log, maf_min):
    """
    Filter on MAF. Remove MAF < maf_min and out-of-bounds MAF.
    """
    jj = (frq < 0) | (frq > 1)
    bad_frq = jj.sum()
    if bad_frq > 0:
        msg = '{N} SNPs had FRQ outside of [0,1]. The FRQ column may be mislabeled'
        log.warning(msg.format(N=bad_frq))

    frq = np.minimum(frq, 1 - frq)
    ii = frq > maf_min
    return ii & ~jj


def parse_dat(dat_gen, convert_colname, log, args):
    """Parse and filter a sumstats file chunk-wise"""

    tot_snps = 0
    dat_list = []
    msg = 'Reading sumstats from {F} into memory {N} SNPs at a time'
    log.info(msg.format(F=args.sumstats, N=int(args.chunksize)))
    drops = {'NA': 0, 'P': 0, 'INFO': 0, 'FRQ': 0, 'A': 0, 'SNP': 0}
    numeric_cols = ['P', 'N', 'N_CAS', 'N_CON', 'Z', 'OR', 'BETA', 'LOG_ODDS', 'INFO', 'FRQ', 'SIGNED_SUMSTAT', 'NSTUDY', 'SE']

    for block_num, dat in enumerate(dat_gen):
        log.info('Reading SNP chunk {}'.format(block_num + 1))
        tot_snps += len(dat)
        old = len(dat)

        for c in dat.columns:
            # sometimes column types change when streaming the data
            if c in numeric_cols and not np.issubdtype(dat[c].dtype, np.number):
                log.warning('Column {} expected to be numeric. Attempting to convert'.format(c))
                dat[c] = pd.to_numeric(dat[c], errors="coerce")

        dat = dat.dropna(axis=0, how="any", subset=filter(lambda x: x != 'INFO', dat.columns)).reset_index(drop=True)
        drops['NA'] += old - len(dat)
        dat.columns = map(lambda x: convert_colname[x], dat.columns)

        ii = np.ones(len(dat), dtype=bool)

        if 'INFO' in dat.columns:
            old = ii.sum()
            ii &= filter_info(dat['INFO'], log, args.info_min)
            new = ii.sum()
            drops['INFO'] += old - new

        if 'FRQ' in dat.columns:
            old = ii.sum()
            ii &= filter_frq(dat['FRQ'], log, args.maf_min)
            new = ii.sum()
            drops['FRQ'] += old - new

        old = ii.sum()
        if args.keep_maf:
            dat.drop(
                [x for x in ['INFO'] if x in dat.columns], inplace=True, axis=1)
        else:
            dat.drop(
                [x for x in ['INFO', 'FRQ'] if x in dat.columns], inplace=True, axis=1)

        ii &= filter_pvals(dat.P, log)
        new = ii.sum()
        drops['P'] += old - new

        if ii.sum() == 0:
            continue

        dat_list.append(dat[ii].reset_index(drop=True))

    log.info('Done reading SNP chunks')
    dat = pd.concat(dat_list, axis=0).reset_index(drop=True)
    log.info('Read {N} SNPs from --sumstats file'.format(N=tot_snps))
    log.info('Removed {N} SNPs with missing values'.format(N=drops['NA']))
    log.info('Removed {N} SNPs with INFO <= {I}'.format(N=drops['INFO'], I=args.info_min))
    log.info('Removed {N} SNPs with MAF <= {M}'.format(N=drops['FRQ'], M=args.maf_min))
    log.info('Removed {N} SNPs with out-of-bounds p-values'.format(N=drops['P']))
    log.info('{N} SNPs remain'.format(N=len(dat)))

    return dat


def process_n(dat, args, log):
    """Determine sample size from --N* flags or N* columns. Filter out low N SNPs.s"""

    if all(i in dat.columns for i in ['N_CAS', 'N_CON']):
        N = dat.N_CAS + dat.N_CON
        P = dat.N_CAS / N
        dat['N'] = N * P / P[N == N.max()].mean()
        dat.drop(['N_CAS', 'N_CON'], inplace=True, axis=1)
        # NB no filtering on N done here -- that is done in the next code block

    if 'N' in dat.columns:
        n_min = args.n_min if args.n_min else dat.N.quantile(0.9) / 1.5
        old = len(dat)
        dat = dat[dat.N >= n_min].reset_index(drop=True)
        new = len(dat)
        log.info('Removed {M} SNPs with N < {MIN} ({N} SNPs remain)'.format(M=old - new, N=new, MIN=n_min))

    elif 'NSTUDY' in dat.columns and 'N' not in dat.columns:
        nstudy_min = args.nstudy_min if args.nstudy_min else dat.NSTUDY.max()
        old = len(dat)
        dat = dat[dat.NSTUDY >= nstudy_min].drop(
            ['NSTUDY'], axis=1).reset_index(drop=True)
        new = len(dat)
        log.info('Removed {M} SNPs with NSTUDY < {MIN} ({N} SNPs remain)'.format(M=old - new, N=new, MIN=nstudy_min))

    if 'N' not in dat.columns:
        if args.N:
            dat['N'] = args.N
            log.info('Using N = {N}'.format(N=args.N))
        elif args.N_cas and args.N_con:
            dat['N'] = args.N_cas + args.N_con
            if args.daner is None:
                msg = 'Using N_cas = {N1}; N_con = {N2}'
                log.info(msg.format(N1=args.N_cas, N2=args.N_con))
        else:
            raise ValueError(
                'Cannot determine N. This message indicates a bug. N should have been checked earlier in the program.')

    return dat


def p_to_z(pvals):
    """Convert P-value to Z-score"""
    return np.sqrt(chi2.isf(pvals, 1))


def pass_median_check(m, expected_median, tolerance):
    """Check that median(x) is within tolerance of expected_median."""
    return np.abs(m - expected_median) <= tolerance


def parse_flag_cnames(args):
    """
    Parse flags that specify how to interpret nonstandard column names.

    flag_cnames is a dict that maps (cleaned) arguments to internal column names
    """

    cname_options = [
        [args.nstudy, 'NSTUDY', '--nstudy'],
        [args.snp, 'SNP', '--snp'],
        [args.N_col, 'N', '--N'],
        [args.N_cas_col, 'N_CAS', '--N-cas-col'],
        [args.N_con_col, 'N_CON', '--N-con-col'],
        [args.a1, 'A1', '--a1'],
        [args.a2, 'A2', '--a2'],
        [args.p, 'P', '--P'],
        [args.frq, 'FRQ', '--nstudy'],
        [args.info, 'INFO', '--info']
    ]
    flag_cnames = {clean_header(x[0]): x[1]
                   for x in cname_options if x[0] is not None}
    if args.info_list:
        try:
            flag_cnames.update(
                {clean_header(x): 'INFO' for x in args.info_list.split(',')})
        except ValueError:
            raise ValueError('The argument to --info-list should be a comma-separated list of column names')

    null_value = None
    if args.signed_sumstats:
        try:
            cname, null_value = args.signed_sumstats.split(',')
            null_value = float(null_value)
            flag_cnames[clean_header(cname)] = 'SIGNED_SUMSTAT'
        except ValueError:
            raise ValueError('The argument to --signed-sumstats should be column header comma number')

    return [flag_cnames, null_value]


def munge(args):
    # This is a modified munge_sumstat.py from LDSC project
    # FIZI needed a similar tool but with a few extra output columns
    # Credit to Brendan Bulik-Sullivan and Hilary Finucane

    log = logging.getLogger(pyfizi.LOG)

    null_values = {
        'LOG_ODDS': 0,
        'BETA': 0,
        'OR': 1,
        'Z': 0
    }

    default_cnames = {
        # Chromosome
        'CHR': 'CHR',
        'CHROM': 'CHR',
        # BP
        'BP': 'BP',
        'POS': 'BP',
        # RS NUMBER
        'SNP': 'SNP',
        'MARKERNAME': 'SNP',
        'SNPID': 'SNP',
        'RS': 'SNP',
        'RSID': 'SNP',
        'RS_NUMBER': 'SNP',
        'RS_NUMBERS': 'SNP',
        # NUMBER OF STUDIES
        'NSTUDY': 'NSTUDY',
        'N_STUDY': 'NSTUDY',
        'NSTUDIES': 'NSTUDY',
        'N_STUDIES': 'NSTUDY',
        # P-VALUE
        'P': 'P',
        'PVALUE': 'P',
        'P_VALUE':  'P',
        'PVAL': 'P',
        'P_VAL': 'P',
        'GC_PVALUE': 'P',
        # ALLELE 1
        'A1': 'A1',
        'ALLELE1': 'A1',
        'ALLELE_1': 'A1',
        'EFFECT_ALLELE': 'A1',
        'REFERENCE_ALLELE': 'A1',
        'INC_ALLELE': 'A1',
        'EA': 'A1',
        # ALLELE 2
        'A2': 'A2',
        'ALLELE2': 'A2',
        'ALLELE_2': 'A2',
        'OTHER_ALLELE': 'A2',
        'NON_EFFECT_ALLELE': 'A2',
        'DEC_ALLELE': 'A2',
        'NEA': 'A2',
        # N
        'N': 'N',
        'NCASE': 'N_CAS',
        'CASES_N': 'N_CAS',
        'N_CASE': 'N_CAS',
        'N_CASES': 'N_CAS',
        'N_CONTROLS': 'N_CON',
        'N_CAS': 'N_CAS',
        'N_CON': 'N_CON',
        'NCONTROL': 'N_CON',
        'CONTROLS_N': 'N_CON',
        'N_CONTROL': 'N_CON',
        'WEIGHT': 'N',  # metal does this. possibly risky.
        # SIGNED STATISTICS
        'ZSCORE': 'Z',
        'Z-SCORE': 'Z',
        'GC_ZSCORE': 'Z',
        'Z': 'Z',
        'OR': 'OR',
        'B': 'BETA',
        'BETA': 'BETA',
        'LOG_ODDS': 'LOG_ODDS',
        'EFFECTS': 'BETA',
        'EFFECT': 'BETA',
        'SIGNED_SUMSTAT': 'SIGNED_SUMSTAT',
        # Standard Error
        'SE': 'SE',
        'std_error': 'SE',
        # INFO
        'INFO': 'INFO',
        # MAF
        'EAF': 'FRQ',
        'FRQ': 'FRQ',
        'MAF': 'FRQ',
        'FRQ_U': 'FRQ',
        'F_U': 'FRQ',
    }

    describe_cname = {
        'CHR': 'Chromsome',
        'BP': 'Base position',
        'SNP': 'Variant ID (e.g., rs number)',
        'P': 'p-Value',
        'A1': 'Allele 1, interpreted as ref allele for signed sumstat',
        'A2': 'Allele 2, interpreted as non-ref allele for signed sumstat',
        'N': 'Sample size',
        'N_CAS': 'Number of cases',
        'N_CON': 'Number of controls',
        'Z': 'Z-score (0 --> no effect; above 0 --> A1 is trait/risk increasing)',
        'OR': 'Odds ratio (1 --> no effect; above 1 --> A1 is risk increasing)',
        'BETA': '[linear/logistic] regression coefficient (0 --> no effect; above 0 --> A1 is trait/risk increasing)',
        'SE': 'Standard error of the beta',
        'LOG_ODDS': 'Log odds ratio (0 --> no effect; above 0 --> A1 is risk increasing)',
        'INFO': 'INFO score (imputation quality; higher --> better imputation)',
        'FRQ': 'Allele frequency',
        'SIGNED_SUMSTAT': 'Directional summary statistic as specified by --signed-sumstats',
        'NSTUDY': 'Number of studies in which the SNP was genotyped'
    }

    try:
        if args.daner and args.daner_n:
            raise ValueError('--daner and --daner-n are not compatible. Use --daner for sample ' +
                             'size from FRQ_A/FRQ_U headers, use --daner-n for values from Nca/Nco columns')

        file_cnames = read_header(args.sumstats)  # note keys not cleaned
        flag_cnames, signed_sumstat_null = parse_flag_cnames(args)
        if args.ignore:
            ignore_cnames = [clean_header(x) for x in args.ignore.split(',')]
        else:
            ignore_cnames = []

        # remove LOG_ODDS, BETA, Z, OR from the default list
        if args.signed_sumstats is not None or args.a1_inc:
            mod_default_cnames = {x: default_cnames[x] for x in default_cnames if default_cnames[x] not in null_values}
        else:
            mod_default_cnames = default_cnames

        cname_map = get_cname_map(flag_cnames, mod_default_cnames, ignore_cnames)
        if args.daner:
            frq_u = filter(lambda x: x.startswith('FRQ_U_'), file_cnames)[0]
            frq_a = filter(lambda x: x.startswith('FRQ_A_'), file_cnames)[0]
            N_cas = float(frq_a[6:])
            N_con = float(frq_u[6:])
            log.info(
                'Inferred that N_cas = {N1}, N_con = {N2} from the FRQ_[A/U] columns.'.format(N1=N_cas, N2=N_con))
            args.N_cas = N_cas
            args.N_con = N_con
            # drop any N, N_cas, N_con or FRQ columns
            for c in ['N', 'N_CAS', 'N_CON', 'FRQ']:
                for d in [x for x in cname_map if cname_map[x] == c]:
                    del cname_map[d]

            cname_map[frq_u] = 'FRQ'

        if args.daner_n:
            frq_u = filter(lambda x: x.startswith('FRQ_U_'), file_cnames)[0]
            cname_map[frq_u] = 'FRQ'
            try:
                dan_cas = clean_header(file_cnames[file_cnames.index('Nca')])
            except ValueError:
                raise ValueError('Could not find Nca column expected for daner-n format')
        
            try:
                dan_con = clean_header(file_cnames[file_cnames.index('Nco')])
            except ValueError:
                raise ValueError('Could not find Nco column expected for daner-n format')
    
            cname_map[dan_cas] = 'N_CAS'
            cname_map[dan_con] = 'N_CON'

        cname_translation = {x: cname_map[clean_header(x)] for x in file_cnames if
                             clean_header(x) in cname_map}  # note keys not cleaned
        cname_description = {
            x: describe_cname[cname_translation[x]] for x in cname_translation}
        if args.signed_sumstats is None and not args.a1_inc:
            sign_cnames = [
                x for x in cname_translation if cname_translation[x] in null_values]
            if len(sign_cnames) > 1:
                raise ValueError(
                    'Too many signed sumstat columns. Specify which to ignore with the --ignore flag.')
            if len(sign_cnames) == 0:
                raise ValueError(
                    'Could not find a signed summary statistic column.')

            sign_cname = sign_cnames[0]
            signed_sumstat_null = null_values[cname_translation[sign_cname]]
            cname_translation[sign_cname] = 'SIGNED_SUMSTAT'
        else:
            sign_cname = 'SIGNED_SUMSTATS'

        # check that we have all the columns we need
        if not args.a1_inc:
            req_cols = ['CHR', 'BP', 'SNP', 'A1', 'A2', 'P', 'SIGNED_SUMSTAT', 'SE']
        else:
            req_cols = ['CHR', 'BP', 'SNP', 'A1', 'A2', 'P', 'SE']

        for c in req_cols:
            if c not in cname_translation.values():
                raise ValueError('Could not find {C} column.'.format(C=c))

        # check aren't any duplicated column names in mapping
        for field in cname_translation:
            numk = file_cnames.count(field)
            if numk > 1:
                raise ValueError('Found {num} columns named {C}'.format(C=field, num=str(numk)))
    
        # check multiple different column names don't map to same data field
        check = set([])
        for head in cname_translation.values():
            if head in check:
                raise ValueError('Found multiple {C} columns'.format(C=head))
            else:
                check.add(head)

        if (not args.N) and (not (args.N_cas and args.N_con)) and ('N' not in cname_translation.values()) and\
                (any(x not in cname_translation.values() for x in ['N_CAS', 'N_CON'])):
            raise ValueError('Could not determine N.')
        if ('N' in cname_translation.values() or all(x in cname_translation.values() for x in ['N_CAS', 'N_CON']))\
                and 'NSTUDY' in cname_translation.values():
            nstudy = [
                x for x in cname_translation if cname_translation[x] == 'NSTUDY']
            for x in nstudy:
                del cname_translation[x]
        if not all(x in cname_translation.values() for x in ['A1', 'A2']):
            raise ValueError('Could not find A1/A2 columns.')

        log.info('Interpreting column names as follows:')
        for x in cname_description:
            log.info(x + ': ' + cname_description[x])

        _, compression = get_compression(args.sumstats)

        # figure out which columns are going to involve sign information, so we can ensure
        # they're read as floats
        signed_sumstat_cols = [k for k, v in cname_translation.items() if v == 'SIGNED_SUMSTAT']
        dtypes = {c: np.float64 for c in signed_sumstat_cols}
        dtypes['CHR'] = "category"
        dat_gen = pd.read_csv(args.sumstats, delim_whitespace=True, header=0,
                              compression=compression, usecols=cname_translation.keys(),
                              na_values=['.', 'NA'], iterator=True, chunksize=args.chunksize,
                              dtype=dtypes)

        dat = parse_dat(dat_gen, cname_translation, log, args)
        if len(dat) == 0:
            raise ValueError('After applying filters, no SNPs remain.')

        old = len(dat)
        dat = dat.drop_duplicates(subset='SNP').reset_index(drop=True)
        new = len(dat)
        log.info('Removed {M} SNPs with duplicated rs numbers ({N} SNPs remain).'.format(M=old - new, N=new))
        # filtering on N cannot be done chunkwise
        dat = process_n(dat, args, log)
        dat.P = p_to_z(dat.P)
        dat.rename(columns={'P': 'Z'}, inplace=True)
        if not args.a1_inc:
            m = np.median(dat.SIGNED_SUMSTAT)
            if not pass_median_check(m, signed_sumstat_null, 0.1):
                msg = 'Median value of {F} is {V} (should be close to {M}). This column may be mislabeled.'
                raise ValueError(msg.format(F=sign_cname, M=signed_sumstat_null, V=round(m, 2)))
            else:
                msg = 'Median value of {F} was {C}, which seems sensible.'.format(C=m, F=sign_cname)
                log.info(msg)

            dat.Z *= (-1) ** (dat.SIGNED_SUMSTAT < signed_sumstat_null)
            dat.drop('SIGNED_SUMSTAT', inplace=True, axis=1)

        out_fname = args.output + '.sumstats.gz'
        print_colnames = []
        for c in ['CHR', 'SNP', 'BP', 'A1', 'A2', 'FRQ', 'Z', 'SE', 'N']:
            if c in dat.columns:
                if c == 'FRQ' and not args.keep_maf:
                    continue

                print_colnames.append(c)

        msg = 'Writing summary statistics for {M} SNPs ({N} with nonmissing beta) to {F}.'
        log.info(
            msg.format(M=len(dat), F=out_fname, N=dat.N.notnull().sum()))

        dat.to_csv(
            out_fname, sep="\t", index=False, columns=print_colnames, float_format='%.3f', compression="gzip")

        CHISQ = (dat.Z ** 2)
        mean_chisq = CHISQ.mean()

        log.info('METADATA - Mean chi^2 = ' + str(round(mean_chisq, 3)))
        if mean_chisq < 1.02:
            log.warning("METADATA: Mean chi^2 may be too small")

        log.info('METADATA - Lambda GC = ' + str(round(CHISQ.median() / 0.4549, 3)))
        log.info('METADATA - Max chi^2 = ' + str(round(CHISQ.max(), 3)))
        log.info('METADATA - {N} Genome-wide significant SNPs (some may have been removed by filtering)'.format(N=(CHISQ > 29).sum()))

    except Exception as err:
        log.error(err)
    finally:
        log.info('Conversion finished')

    return


def impute(args):
    log = logging.getLogger(pyfizi.LOG)
    try:
        # perform sanity arguments checking before continuing
        chrom = None
        start_bp = None
        stop_bp = None
        if any(x is not None for x in [args.chr, args.start, args.stop]):
            if args.start is not None and args.chr is None:
                raise ValueError("Option --start cannot be set unless --chr is specified")
            if args.stop is not None and args.chr is None:
                raise ValueError("Option --stop cannot be set unless --chr is specified")

            chrom = args.chr

            # parse start/stop positions and make sure they're ordered (if exist)
            if args.start is not None:
                start_bp = parse_pos(args.start, "--start")
            else:
                start_bp = None

            if args.stop is not None:
                stop_bp = parse_pos(args.stop, "--stop")
            else:
                stop_bp = None

            if args.start is not None and args.stop is not None:
                if start_bp >= stop_bp:
                    raise ValueError("Specified --start position must be before --stop position")

        window_size = parse_pos(args.window_size, "--window-size")
        buffer_size = parse_pos(args.buffer_size, "--buffer-size")

        min_prop = args.min_prop
        if min_prop <= 0 or min_prop >= 1:
            raise ValueError("--min-prop must be strictly bounded in (0, 1)")

        # load GWAS summary data
        log.info("Preparing GWAS summary file")
        gwas = pyfizi.GWAS.parse_gwas(args.gwas)

        # filter GWAS down to chromosome if it is specified here to save later on
        if chrom is not None:
            gwas = gwas.subset_by_pos(chrom, start_bp, stop_bp)

        # load reference genotype data
        log.info("Preparing reference SNP data")
        ref = pyfizi.RefPanel.parse_plink(args.ref)

        # load functional annotations and taus
        # parity check for functional data
        if args.annot is None and args.taus is not None:
            raise ValueError("LDSC estimates require corresponding annotation (--annot) file")

        run_fizi = args.annot is not None
        if run_fizi:
            log.info("Preparing annotation file")
            annot = pyfizi.Annot.parse_annot(args.annot, args.annot_names)
            if len(annot.names) == 0:
                raise ValueError("No remaining annotations for inference")

            n_annot = len(annot)
            annot = annot.drop_duplicates(pyfizi.Annot.SNPCOL)
            n_annot_dd = len(annot)
            if n_annot != n_annot_dd:
                diff = n_annot - n_annot_dd
                log.info("Dropped {} duplicated SNP annotations".format(diff))

            if args.taus is not None:
                log.info("Preparing SNP effect-size variance file")
                taus = pyfizi.Taus.parse_taus(args.taus, args.annot_names)
                n_t = len(taus)
                taus = taus.subset_by_enrich_pvalue(args.alpha)
                n_tt = len(taus)
                log.info("Dropped {} annotations with p-value less than {}".format(n_t - n_tt, args.alpha))
                for cn in taus.names:
                    if cn not in annot.names:
                        raise KeyError("Prior variance for {} not found in annotation file".format(cn))
                if len(taus) == 0:
                    raise ValueError("No tau estimates remaining for inference")
            else:
                taus = None
                log.info("Using {} annotation(s) for tau estimate(s)".format(len(annot.names)))

            # we should check columns in taus are contained in annot here...
        log.info("Starting summary statistics imputation with window size {} and buffer size {}".format(window_size, buffer_size))
        written = False
        with open("{}.sumstat".format(args.output), "w") as output:

            if args.locations is not None:
                log.info("Preparing user-defined locations")
                partitions = parse_locations(args.locations, chrom, start_bp, stop_bp)
            else:
                partitions = ref.get_partitions(window_size, chrom, start_bp, stop_bp)

            for chrom, start, stop in partitions:
                pstart = max(1, start - buffer_size)
                pstop = stop + buffer_size

                log.debug("Subsetting GWAS data by {}:{} - {}:{}".format(chrom, int(pstart), chrom, int(pstop)))
                part_gwas = gwas.subset_by_pos(chrom, pstart, pstop)
                if len(part_gwas) == 0:
                    log.warning("No GWAS SNPs found at {}:{} - {}:{}. Skipping".format(chrom, int(pstart), chrom, int(pstop)))
                    continue

                log.debug("Subsetting reference SNP data by {}:{} - {}:{}".format(chrom, int(pstart), chrom, int(pstop)))
                part_ref = ref.subset_by_pos(chrom, pstart, pstop)
                if len(part_ref) == 0:
                    log.warning("No reference SNPs found at {}:{} - {}:{}. Skipping".format(chrom, int(pstart), chrom, int(pstop)))
                    imputed_gwas = pyfizi.create_output(part_gwas, start=start, stop=stop)
                    pyfizi.write_output(imputed_gwas, output, append=written)
                    written = True
                    continue

                if run_fizi:
                    log.debug("Subsetting annotation data by {}:{} - {}:{}".format(chrom, int(pstart), chrom, int(pstop)))
                    part_annot = annot.subset_by_pos(chrom, pstart, pstop)
                    if len(part_annot) == 0:
                        log.warning("No annotations found at {}:{} - {}:{}. Defaulting to ImpG".format(chrom, int(pstart), chrom, int(pstop)))
                        imputed_gwas = pyfizi.impute_gwas(part_gwas, part_ref, prop=min_prop, start=start, stop=stop,
                                                        ridge=args.ridge_term)
                    else:
                        imputed_gwas = pyfizi.impute_gwas(part_gwas, part_ref, annot=part_annot, taus=taus,
                                                          prop=min_prop, start=start, stop=stop, ridge=args.ridge_term)

                else:
                    imputed_gwas = pyfizi.impute_gwas(part_gwas, part_ref, prop=min_prop, start=start, stop=stop,
                                                      ridge=args.ridge_term)

                # format floats to only use 3 decimal places
                pyfizi.write_output(imputed_gwas, output, append=written)
                written = True

    except Exception as err:
        log.error(err)
    finally:
        log.info("Finished summary statistic imputation")

    return 0


def main(argsv):
    # setup main parser
    argp = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # add subparser for the  munge summary statistics program
    subp = argp.add_subparsers(title="Subcommands",
                               help="'munge' to clean up summary statistics (see fizi munge --help)." +
                               "'impute' to perform imputation (see fizi impute --help).")

    munp = subp.add_parser("munge",
                           description="Munge summary statistics input to conform to FIZI requirements")

    munp.add_argument('sumstats',
                      help="Input filename.")
    munp.add_argument('--N', default=None, type=float,
                      help="Sample size If this option is not set, will try to infer the sample "
                           "size from the input file. If the input file contains a sample size "
                           "column, and this flag is set, the argument to this flag has priority.")
    munp.add_argument('--N-cas', default=None, type=float,
                      help="Number of cases. If this option is not set, will try to infer the number "
                           "of cases from the input file. If the input file contains a number of cases "
                           "column, and this flag is set, the argument to this flag has priority.")
    munp.add_argument('--N-con', default=None, type=float,
                      help="Number of controls. If this option is not set, will try to infer the number "
                           "of controls from the input file. If the input file contains a number of controls "
                           "column, and this flag is set, the argument to this flag has priority.")
    munp.add_argument('--info-min', default=0.9, type=float,
                      help="Minimum INFO score.")
    munp.add_argument('--maf-min', default=0.01, type=float,
                      help="Minimum MAF.")
    munp.add_argument('--daner', default=False, action='store_true',
                      help="Use this flag to parse Stephan Ripke's daner* file format.")
    munp.add_argument('--daner-n', default=False, action='store_true',
                      help="Use this flag to parse more recent daner* formatted files, which "
                           "include sample size column 'Nca' and 'Nco'.")
    munp.add_argument('--n-min', default=None, type=float,
                      help='Minimum N (sample size). Default is (90th percentile N) / 2.')
    munp.add_argument('--chunksize', default=5e6, type=int,
                      help='Chunksize.')

    # optional args to specify column names
    munp.add_argument('--snp', default=None, type=str,
                      help='Name of SNP column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--N-col', default=None, type=str,
                      help='Name of N column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--N-cas-col', default=None, type=str,
                      help='Name of N column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--N-con-col', default=None, type=str,
                      help='Name of N column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--a1', default=None, type=str,
                      help='Name of A1 column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--a2', default=None, type=str,
                      help='Name of A2 column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--p', default=None, type=str,
                      help='Name of p-value column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--frq', default=None, type=str,
                      help='Name of FRQ or MAF column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--signed-sumstats', default=None, type=str,
                      help='Name of signed sumstat column, comma null value (e.g., Z,0 or OR,1). NB: case insensitive.')
    munp.add_argument('--info', default=None, type=str,
                      help='Name of INFO column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--info-list', default=None, type=str,
                      help='Comma-separated list of INFO columns. Will filter on the mean. NB: case insensitive.')
    munp.add_argument('--nstudy', default=None, type=str,
                      help='Name of NSTUDY column (if not a name that FIZI understands). NB: case insensitive.')
    munp.add_argument('--nstudy-min', default=None, type=float,
                      help='Minimum # of studies. Default is to remove everything below the max, unless there is an N column,'
                           ' in which case do nothing.')
    munp.add_argument('--ignore', default=None, type=str,
                      help='Comma-separated list of column names to ignore.')
    munp.add_argument('--a1-inc', default=False, action='store_true',
                      help='A1 is the increasing allele.')
    munp.add_argument('--keep-maf', default=False, action='store_true',
                      help='Keep the MAF column (if one exists).')

    # misc options
    munp.add_argument("-q", "--quiet", default=False, action="store_true",
                      help="Do not print anything to stdout.")
    munp.add_argument("-v", "--verbose", default=False, action="store_true",
                      help="Verbose logging. Includes debug info.")
    munp.add_argument("-o", "--output", default="FIZI",
                      help="Prefix for output data.")

    munp.set_defaults(func=munge)

    # add imputation parser
    impp = subp.add_parser("impute", description="GWAS summary statistics imputation with functional data.")

    # main arguments
    impp.add_argument("gwas",
                      help="GWAS summary data. Supports gzip and bz2 compression.")
    impp.add_argument("ref",
                      help="Path to reference panel PLINK data.")

    # functional arguments
    impp.add_argument("--annot", default=None,
                      help="Path to SNP functional annotation data. Should be in LDScore regression-style format. Supports gzip and bz2 compression.")
    impp.add_argument("--annot-names", type=parse_annots,
                      help="Comma-separated list of annotations to use in the `annot` file (e.g., `--annot-names Coding_UCSC.bed,DHS_Trynka.bed`). Baseline is always included.")
    impp.add_argument("--taus", default=None,
                      help="Path to LDScore regression output. Must contain coefficient estimates. Supports gzip and bz2 compression.")
    impp.add_argument("--alpha", default=1.00, type=float,
                      help="Significance threshold to determine which functional categories to keep.")

    # GWAS options
    impp.add_argument("--gwas-n", default=None, type=int,
                      help="GWAS sample size.")

    # imputation location options
    impp.add_argument("--chr", default=None,
                      help="Perform imputation for specific chromosome.")
    impp.add_argument("--start", default=None,
                      help="Perform imputation starting at specific location (in base pairs). Accepts kb/mb modifiers. Requires --chr to be specified.")
    impp.add_argument("--stop", default=None,
                      help="Perform imputation until at specific location (in base pairs). Accepts kb/mb modifiers. Requires --chr to be specified.")
    impp.add_argument("--locations", default=None, type=argparse.FileType("r"),
                      help="Path to a BED file containing windows (e.g., CHR START STOP) to impute. Start and stop values may contain kb/mb modifiers.")

    # imputation options
    impp.add_argument("--window-size", default="250kb",
                      help="Size of imputation window (in base pairs). Accepts kb/mb modifiers.")
    impp.add_argument("--buffer-size", default="250kb",
                      help="Size of buffer window (in base pairs). Accepts kb/mb modifiers.")
    impp.add_argument("--min-prop", default=0.5, type=float,
                      help="Minimum required proportion of gwas/reference panel overlap to perform imputation.")
    impp.add_argument("--ridge-term", default=0.1, type=float,
                      help="Diagonal adjustment for linkage-disequilibrium (LD) estimate.")

    # misc options
    impp.add_argument("-q", "--quiet", default=False, action="store_true",
                      help="Do not print anything to stdout.")
    impp.add_argument("-v", "--verbose", default=False, action="store_true",
                      help="Verbose logging. Includes debug info.")
    impp.add_argument("-o", "--output", default="FIZI",
                      help="Prefix for output data.")

    # set defaults
    impp.set_defaults(func=impute)

    # parse arguments
    args = argp.parse_args(argsv)

    # hack to check that at least one sub-command was selected in 3.6
    # 3.7 -might- have fixed this bug
    if not hasattr(args, "func"):
        argp.print_help()
        return 2  # command-line error

    cmd_str = get_command_string(argsv)

    masthead =  "====================================" + os.linesep
    masthead += "               FIZI v{}             ".format(pyfizi.VERSION) + os.linesep
    masthead += "====================================" + os.linesep

    # setup logging
    FORMAT = "[%(asctime)s - %(levelname)s] %(message)s"
    DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
    log = logging.getLogger(pyfizi.LOG)
    log.setLevel(logging.INFO if not args.verbose else logging.DEBUG)
    fmt = logging.Formatter(fmt=FORMAT, datefmt=DATE_FORMAT)

    # write to stdout unless quiet is set
    if not args.quiet:
        sys.stdout.write(masthead)
        sys.stdout.write(cmd_str)
        sys.stdout.write("Starting log..." + os.linesep)
        stdout_handler = logging.StreamHandler(sys.stdout)
        stdout_handler.setFormatter(fmt)
        log.addHandler(stdout_handler)

    # setup log file, but write PLINK-style command first
    disk_log_stream = open("{}.log".format(args.output), "w")
    disk_log_stream.write(masthead)
    disk_log_stream.write(cmd_str)
    disk_log_stream.write("Starting log..." + os.linesep)

    disk_handler = logging.StreamHandler(disk_log_stream)
    disk_handler.setFormatter(fmt)
    log.addHandler(disk_handler)

    # launch either munge or impute
    args.func(args)

    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
