#!/usr/bin/env python
'''
(c) 2015-Present Brendan Bulik-Sullivan

Multi-trait genetic prediction.

'''

from __future__ import division
import argparse
import time
from core import bayes
from core import ld
import traceback
import sys
import pandas as pd

try:
    x = pd.DataFrame({'A': [1, 2, 3]})
    x.drop_duplicates(subset='A')
except TypeError:
    raise ImportError('mtpred requires pandas version > 0.15.2')

__version__ = '0.1'
MASTHEAD = "*************************************\n"
MASTHEAD += "* Multi-Trait Prediction (mtpred)\n"
MASTHEAD += "* Version {V}\n".format(V=__version__)
MASTHEAD += "* (C) 2015 Brendan Bulik-Sullivan\n"
MASTHEAD += "* Broad Institute of MIT and Harvard\n"
MASTHEAD += "* GNU General Public License v3\n"
MASTHEAD += "*************************************\n"


def sec_to_str(t):
    '''Convert seconds to days:hours:minutes:seconds'''
    [d, h, m, s, n] = reduce(
        lambda ll, b: divmod(ll[0], b) + ll[1:], [(t, 1), 60, 60, 24])
    f = ''
    if d > 0:
        f += '{D}d:'.format(D=int(d))
    if h > 0:
        f += '{H}h:'.format(H=int(h))
    if m > 0:
        f += '{M}m:'.format(M=int(m))
    f += '{S}s'.format(S=s)
    return f


class Logger(object):
    '''Print to log file and stdout.'''
    def __init__(self, fh):
        self.log_fh = open(fh, 'wb', 0)

    def log(self, msg, stdout=True):
        '''Print to log file and stdout.'''
        print >>self.log_fh, msg
        if stdout:
            x = str(msg).split('\n')
            if len(x) > 20:
                msg = '\n'.join(x[1:10])
                msg += '\nOutput truncated. See log file for full list.'
            sys.stdout.write(str(msg)+'\n')
            sys.stdout.flush()

    def close(self):
        self.log_fh.close()


class ThrowingArgumentParser(argparse.ArgumentParser):

    def error(self, message):
        raise ValueError(message)

    def print_help(self, masthead=True):
        if masthead:
            print MASTHEAD
        argparse.ArgumentParser.print_help(self)


parser = ThrowingArgumentParser()
subparsers = parser.add_subparsers()

# eigendecomposition of LD matrices
eigen = subparsers.add_parser('ld', help='Calculate eigen-decompositions'
                                         ' of LD matrices.')
eigen.add_argument('--bfile', default=None, type=str, required=True,
                   help='Filename prefix for plink .bed/.bim/.fam fileset.')
ld_wind = eigen.add_mutually_exclusive_group(required=True)
ld_wind.add_argument('--ld-wind-snp', default=None, type=int, nargs=2,
                     help='First arg window size, second arg buffer size.'
                          ' Units are # SNPs.',
                     metavar=('WIND', 'BUF'))
ld_wind.add_argument('--ld-wind-kb', default=None, type=int, nargs=2,
                     help='Not implemented yet.',
                     metavar=('WIND', 'BUF'))
ld_wind.add_argument('--ld-wind-cm', default=None, type=int, nargs=2,
                     help='Not implemented yet.',
                     metavar=('WIND', 'BUF'))
eigen.add_argument('--covar', default=None, type=str,
                   help='Covariates.')
eigen.add_argument('--out', default='mtpred', type=str,
                   help='Output filename prefix, default mtpred.')

# posterior mean beta
pmb = subparsers.add_parser('pmb', help='Convert summary statistics to '
                                        'prediction weights.')
pmb.add_argument('--sumstats', default=None, type=str, required=True,
                 nargs='+',
                 help='List of filenames in .sumstats format containing GWAS '
                      'summary statistics.')
pmb.add_argument('--mat', default=None, type=str, required=True, nargs=2,
                 help='Genetic covariance matrix '
                      'followed by a sample overlap matrix.',
                      metavar=('VG', 'P'))
pmb.add_argument('--eigen-ld', default=None, type=str, required=True,
                 help='Pre-computed eigen-decomposition of the LD matrix'
                      ' (from mypred.py ld).')
pmb.add_argument('--png', default=False, action='store_true',
                 help='Save per-normalized-genotype weights? ' +
                 '(default per-allele).')
pmb.add_argument('--out', default='mtpred', type=str,
                 help='Output filename prefix, default mtpred.')
pmb.add_argument('--block-buf', default=3, type=int,
                 help='Eigenblock buffer size.')

if __name__ == '__main__':
    try:
        args = parser.parse_args()
    except ValueError:  # override bad help msgs w/ subparsers
        print MASTHEAD
        ex_type, ex, tb = sys.exc_info()
        print traceback.format_exc(ex)
        parser.print_help(masthead=False)
        sys.exit(2)
    if args.out is None:
        raise ValueError('--out is required.')
    log = Logger(args.out+'.log')
    start_time = time.time()
    try:
        opts = vars(args)
        opts = {x: ' '.join([str(i) for i in opts[x]])
                if type(opts[x]) is list
                else opts[x] for x in
                filter(lambda y: opts[y] is not None, opts)}
        header = MASTHEAD
        header += "\nOptions: \n"
        options = [
            '  --' + x.replace('_', '-') + ' ' + str(opts[x]) for x in opts]
        header += '\n'.join(options) + '\n'
        log.log(header)
        log.log('Beginning analysis at {T}'.format(T=time.ctime()))
        try:
            args.bfile
            which = 'ld'
        except AttributeError:
            which = 'pmb'
        if which == 'ld':
            ld.eigenblocks(args, log)
        elif which == 'pmb':
            bayes.mtpred(args, log)
        else:
            log.log('Something went horribly wrong.')
    except Exception:
        ex_type, ex, tb = sys.exc_info()
        log.log(traceback.format_exc(ex), stdout=False)
        raise
    finally:
        log.log('Analysis finished at {T}'.format(T=time.ctime()))
        time_elapsed = round(time.time() - start_time, 2)
        log.log('Total time elapsed: {T}'.format(T=sec_to_str(time_elapsed)))
        log.close()
