<gh_stars>0
##
# Implement idea 1
##

import argparse
parser = argparse.ArgumentParser(prog='impute_parent_of_origin.py', description='''
    Impute parent of origin from PRS scores and observed phenotypes.
''')

parser.add_argument('--prs-matrix', help='''
    PRS matrix. 
    (generated by ../prs/prs_matrix.py)
    Dimension: hap x trait x cutoff x sample
''')
parser.add_argument('--map-prs-yaml', help='''
    Specify which PRS should be used for
    which phenotype.
    trait1: 
        name: 'name in prs matrix'
        pval: 'pval cutoff to use'
''')
parser.add_argument('--father-phenotype-yaml', help='''
    Observed phenotype of father
    with all loading specified in YAML:
    path: 'path_to_file'
    param:   # kwargs in pandas.read_csv
        header: 0
        sep: ','
    col:  # column name of phenotype of interest
        - 'alzheimer_disease'
        - 'heart_disease'
    indiv_col: 'eid'  # column name of individual id
''')
parser.add_argument('--mother-phenotype-yaml', help='''
    Observed phenotype  of mother 
    with all loading specified in YAML 
    (follow the same structure as father's)
''')
parser.add_argument('--impute-mode', type=str, default='basic_em', help='''
    Imputation approach.
''')
parser.add_argument('--output', help='''
    Output in TSV.GZ format.
''')
args = parser.parse_args()


import logging, sys
# configing util
logging.basicConfig(
    level = logging.INFO, 
    stream = sys.stderr, 
    format = '%(asctime)s  %(message)s',
    datefmt = '%Y-%m-%d %I:%M:%S %p'
)

import sys
sys.path.insert(0, '../logistic_gpu')
import table_reader
import prs_reader
import haplotype_imputer

logging.info('Loading observed phenotypes')
df_father = table_reader.load_table_from_yaml(
    args.father_phenotype_yaml,
    rename_cols=True
)
df_mother = table_reader.load_table_from_yaml(
    args.mother_phenotype_yaml,
    rename_cols=True
)
# print(df_father)
# print(df_mother)

logging.info('Loading PRS')
df_h1, df_h2 = prs_reader.load_prs(args.prs_matrix, args.map_prs_yaml)
# print(df_h1)
# print(df_h2)


logging.info('Run imputation: mode = {}'.format(args.impute_mode))
imputer = haplotype_imputer.HaploImputer()
out = imputer.impute(df_father, df_mother, df_h1, df_h2, mode=args.impute_mode)

logging.info('Output')
out.to_csv(args.output, compression='gzip', sep='\t', index=False)

