#!/usr/bin/env python

from pyPheWAS.pyPhewasCorev2 import *
import os
import argparse
import time
import math

def parse_args():
    parser = argparse.ArgumentParser(description="pyPheWAS ICD-Phecode Lookup Tool")

    parser.add_argument('--feature_matrix', required=True, type=str, help='Name of the feature matrix file (e.g. feature_matrix_group.csv)')
    parser.add_argument('--group', required=True, type=str, help ='Name of the group file (e.g. groups.csv)')
    parser.add_argument('--reg_type', required=True, type=str, help='Type of regression that you would like to use (log, lin, or dur)')
    parser.add_argument('--path', required=False, default='.', type=str,help='Path to all input files and destination of output files')
    parser.add_argument('--outfile', required=False, default=None, type=str,help='Name of the output file for the feature matrix')
    parser.add_argument('--covariates', required=False, default='', type=str, help='Variables to be used as covariates')
    parser.add_argument('--response', required=False, default='', type=str, help='Variable to predict instead of genotype')
    parser.add_argument('--phewas_cov', required=False, default=None, type=float, help='PheCodes to use as covariates in regression')

    args = parser.parse_args()
    return args

"""
Retrieve and validate all arguments.
"""
start = time.time()

args = parse_args()
kwargs = {'path': os.path.join(os.path.abspath(args.path),''),
		  'feature_matrix': args.feature_matrix,
		  'groupfile': args.group,
		  'outfile':args.outfile,
		  'phewas_cov':args.phewas_cov,
          'covariates':args.covariates,
          'response':args.response
}
str_reg_type = args.reg_type

# Assert that a valid regression type was used
assert str_reg_type in regression_map.keys(), "%s is not a valid regression type" %str_reg_type
kwargs['reg_type'] = regression_map[str_reg_type]

# Assert that valid files were given
assert kwargs['feature_matrix'].endswith('.csv'), "%s is not a valid phenotype file, must be a .csv file" % (kwargs['feature_matrix'])
assert kwargs['groupfile'].endswith('.csv'), "%s is not a valid group file, must be a .csv file" % (kwargs['groupfile'])

# Assign the output file if none was assigned
if kwargs['outfile'] is None:
    kwargs['outfile'] = "regressions_" + kwargs['groupfile']

# Assert that the output file is valid
assert kwargs['outfile'].endswith('.csv'), "%s is not a valid outputfile, must be a .csv file" % (kwargs['outfile'])

# Check phewas_cov
if kwargs['phewas_cov']:
    kwargs['phewas_cov'] = float(kwargs['phewas_cov'])

if kwargs['response'] is None:
    kwargs['response'] = ''

# Print Arguments
display_kwargs(kwargs)

"""
Run Regressions
"""

# Make all arguments local variables
locals().update(kwargs)

print("Retrieving group data...")
genotypes = get_group_file(path, groupfile)

a = np.loadtxt(path + 'agg_measures_' + feature_matrix, delimiter=',')
b = np.loadtxt(path + 'icd_age_' + feature_matrix, delimiter=',')
c = np.loadtxt(path + 'phewas_cov_' + feature_matrix, delimiter=',')
fm = np.array([a,b,c])

print("Running PheWAS regressions...")
regressions = run_phewas(fm, genotypes,covariates, reg_type, response, phewas_cov)

print("Saving regression data to %s" % (path + outfile))
header = ','.join(['str_reg_type', str_reg_type, 'group', groupfile]) + '\n'
f = open(os.sep.join([path, outfile]), 'w')
f.write(header)
regressions.to_csv(f,index=False)
f.close()

interval = time.time() - start
hour = math.floor(interval/3600.0)
minute = math.floor((interval - hour*3600)/60)
second = math.floor(interval - hour*3600 - minute*60)

if hour > 0:
    time_str = '%dh:%dm:%ds' %(hour,minute,second)
elif minute > 0:
    time_str = '%dm:%ds' % (minute, second)
else:
    time_str = '%ds' % second

print('pyPhewasModel Complete\nRuntime: %s' %time_str)