#!/usr/bin/env python
# coding: utf-8


import pandas as pd
import numpy as np
import subprocess
import os
import hashlib
import sys
import gzip
from argparse import ArgumentParser
import logging

'''
generates a MAF report (by calling plink)
'''

def get_args():
    p = ArgumentParser()
    # meaning of the different parameters explained below
    p.add_argument('-b','--bed', required=True)
    p.add_argument('-i','--iid', help='file containing the individual-IDs that we want to include in the calculation', default=None)
    p.add_argument('-o','--out_prefix', required=True)
    p.add_argument('--mac_report', required=True, help='MAC report generated by mac_report.py')
    p.add_argument('--anno', required=True, help='annotation generated by basic_annotation.py')
    p.add_argument('--max_missing', required=True, default=0.1, type=float)
    p.add_argument('--plink_path', required=False, default='bin/plink', help='path to plink 1.9 binary')
    p.add_argument('--plink2_path', required=False, default='bin/plink2', help='path to plink 2 binary')
    p.add_argument('--log', required=False, default='log.txt')
    p.add_argument('--threads', default=4, type=int)
    args = p.parse_args()
    return(args)

def read_ids(path):
    if path.endswith('.gz'):
        with gzip.open(path,'rt') as f:
            ids = [ l.rstrip() for l in f ]
    else:
        with open(path, 'r') as f:
            ids = [ l.rstrip() for l in f ]
    return ids


def write_ids(ids, path):
    with open(path, 'x') as outfile:
        for i in ids:
            outfile.write('{} {}\n'.format(str(i), str(i))) 

def write_ids_single(ids, path):
    with open(path, 'x') as outfile:
        for i in ids:
            outfile.write('{}\n'.format(i))

                          
def main():

    # arguments
    args = get_args()
    
    if args.log is not None:
        logging.basicConfig(filename=args.log)
    else:
        logging.basicConfig()
    
    
    if args.iid is not None:
        # read the individual files
        iids = np.unique(read_ids(args.iid))
        # temporary files for plink input
        tmpfilename_iid = hashlib.sha256((''.join(iids) + args.out_prefix).encode('utf-8')).hexdigest()
        tmpfilename_iid = '{}_{}_iid.txt'.format(args.out_prefix, tmpfilename_iid)
            
        write_ids(iids, tmpfilename_iid)
        
    
    logging.info(subprocess.check_output([args.plink_path, '--version']).decode('utf-8').strip())
    
    # load the MAC report, keep only variants that are observed and below missingness threshold
    mac_report = pd.read_csv(args.mac_report, sep='\t', usecols=['SNP','Minor','MAF','MISSING','alt_greater_ref'])
    mac_report = mac_report[(mac_report.Minor > 0) & (mac_report.MISSING < args.max_missing)]
    
    vids = mac_report.SNP.values
                          
    mac_report.set_index('SNP', inplace=True)
                          
    # load the variant annotation, we will later merge it with the resulting phenotype-specific table
    anno = pd.read_csv(args.anno, sep='\t', usecols = ['Name', 'hiconf_reg', 'exomeseq_target']).rename(columns={'Name':'SNP'}).set_index('SNP')
    mac_report = mac_report.join(anno, how='left')
    
    mac_report.loc[mac_report.hiconf_reg.isna(),'hiconf_reg'] = 0
    mac_report.loc[mac_report.exomeseq_target.isna(),'exomeseq_target'] = 0
    
    # vids_highconf = anno.Name[anno.hiconf_reg.astype(bool).values]
    # vids = np.intersect1d(vids, vids_highconf)
                                              
    tmpfilename_vid = hashlib.sha256(''.join(vids[0:100000] if len(vids) > 100000 else vids).encode('utf-8')).hexdigest()
    tmpfilename_vid = '{}_{}_vid.txt'.format(args.out_prefix, tmpfilename_vid)
    
    write_ids_single(vids, tmpfilename_vid)

    # plink command to create (temporary) count reports
    plink_command = [args.plink_path, '--bfile', args.bed]
    if args.iid is not None:
        plink_command += ['--keep', tmpfilename_iid, '--freq', 'counts', 'gz', '--threads', str(args.threads), '--out', args.out_prefix]
    else:
        plink_command += ['--freq', 'counts', 'gz', '--threads', str(args.threads), '--out', args.out_prefix ]
    
    # remove variants with missingness rate above threshold or that are never observed
    plink_command += ['--extract', tmpfilename_vid]
    logging.info('plink command: {}'.format(' '.join(plink_command)))
    
    # run plink command
    _ = subprocess.run(plink_command, check=True)

    # remove temporary files
    if args.iid is not None:
        os.remove(tmpfilename_iid)

    os.remove(tmpfilename_vid)

    # print the plink log
    with open(args.out_prefix + '.log', 'r') as infile:
        lines = infile.readlines()
        logging.info(lines)
    
    outfilepath = args.out_prefix + '.frq.counts.gz'
    counts = pd.read_csv(outfilepath, delim_whitespace=True, header=0)
    
    counts.loc[:,'CHR'] = counts['CHR'].astype(str)

    # get major / minor allele counts as they might be "flipped"
    counts['Minor'] = np.where(counts['C2'] < counts['C1'], counts['C2'], counts['C1'])
    
    # not needed, information is already in mac_report
    # counts['alt_greater_ref'] = (counts['C2'] < counts['C1']).astype(int)

    # set filters
    missing = (counts.Minor == 0).values
    print('{} unobserved variants.'.format(np.sum(missing)))
                          
    counts = counts[~missing]
    mac_report = mac_report.loc[counts.SNP]
    mac_report.loc[:,'Minor'] = counts['Minor'].values
    
    # not needed, information is already in mac_report
    # mac_report.loc[:,'alt_greater_ref'] = counts['alt_greater_ref'].values # store this information so we can remove them for missense / pLOF analysis
                          
    
    # export filtered MAC report (MAF calculated on all individuals, Minor calculated on subset)
    mac_report[['MAF','Minor','alt_greater_ref','exomeseq_target']].to_csv(args.out_prefix+'.tsv.gz', sep='\t', index=True)
                          
    os.remove(outfilepath)
    
    
    
if __name__ == '__main__':
    main()