#!/usr/bin/env python3
"""
NOTE: This script is a hard parallelization of the last step in the meta_split.py
script. It can distribute replicates across clusters.

More info:
By meta-splits we mean the splitting of the dataset of "n_chunks" number of chunks
in half and combining the summary statistics for all the chunks in a given half.

We combine sumstats using inverse-variance weighted meta-analysis.

To calculate the genetic correlation of the splits, we can use rg_split_parallel.py,
which requires the h2part tsv holding the file paths of all the meta-analyzed 
sumstats generated by this method. The h2part tsv can be generated by code in 
the meta_split_supp.py file.

@author: nbaya
"""

import hail as hl
import numpy as np
import datetime
#from multiprocessing import Pool

phen_dict = {
    '50':'height',
    '20160':'smoking',
    '2443':'diabetes',
    '23107':'legimp_r',
    '6138_100':'qual_none',
    '50_raw':'height_raw',
    '30100':'platelet_vol',
    '50_sim_inf':'height_sim_infinitesimal',
    '50_sim_inf_h2_0.1':'height_sim_infinitesimal, h2=0.1',
    '50_raw_res':'Residuals from 50_raw linreg'
}

import argparse
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--phen', type=str, required=True, help="phenotype code")
parser.add_argument('--n_chunks', type=int, required=True, help="number of subgroups (or 'chunks')")
parser.add_argument('--variant_set', type=str, required=True, help="set of variants to use")
parser.add_argument('--batch', type=str, required=True, help="batch number")
parser.add_argument('--reps', type=int, required=True, help="number of permutation replicates, default: 100")
parser.add_argument('--parsplit', type=int, required=True, help="number of parallel batches (should be factor of reps), default: 5")
parser.add_argument('--paridx', type=int, required=True, help="which of the parallel batches to run")

args = parser.parse_args()

#idx = range(args.iters) #non-parallelized
idx = range(args.paridx-1,args.reps, args.parsplit) #hard-parallelized

#Globals
phen = args.phen
desc = phen_dict[phen]
batch = args.batch
n_chunks = args.n_chunks #number of subgroups (or "chunks")
variant_set = args.variant_set

print('####################')
print('Phenotype: '+phen)
print('Description: '+desc)
print('Batch: '+batch)
print('n chunks: '+str(n_chunks))
print('variant set: '+variant_set)
print('####################')


"""
╔═════════════════════════════════════════════════════════╗
║ Part 4: Split and calculate meta-analyzed summary stats ║
╚═════════════════════════════════════════════════════════╝
Split n groups into two populations (A, B) and calculate meta summary statistics (via inverse-variance weighting meta-analysis)
"""
print('Starting Part 4: Splitting into populations A and B, calculatng summary statistics')
print('Time: {:%H:%M:%S (%Y-%b-%d)}'.format(datetime.datetime.now()))

gmt = hl.read_matrix_table('gs://nbaya/split/meta_split/ukb31063.'+variant_set+'_'+phen+'_gmt'+str(n_chunks)+'_batch_'+batch+'.mt') 

gmt = gmt.add_col_index()
gmt = gmt.rename({'rsid': 'SNP'})
    
def run_meta_split(i):
    print('####################')
    print('Starting split '+str(i))
    print('####################')
    starttime = datetime.datetime.now()
    pi = ['A']*int(n_chunks/2) + ['B']*int(n_chunks/2)
    seed_id = int(batch+str(i).zfill(4)) #create a seed_id unique to every split
    randstate = np.random.RandomState(seed_id) #seed with seed_id

    randstate.shuffle(pi)
    gmt_shuf = gmt.annotate_cols(label = hl.literal(pi)[hl.int32(gmt.col_idx)])
    
    mt = gmt_shuf.group_cols_by(gmt_shuf.label).aggregate(unnorm_meta_beta=hl.agg.sum(gmt_shuf.beta / gmt_shuf.standard_error ** 2),
                                inv_se2 = hl.agg.sum(1 / gmt_shuf.standard_error ** 2)).key_rows_by('SNP')

    ht = mt.make_table()
    
    ht = ht.annotate(A_Z = ht['A.unnorm_meta_beta'] / hl.sqrt(ht['A.inv_se2']),
                     B_Z = ht['B.unnorm_meta_beta'] / hl.sqrt(ht['B.inv_se2']))
    
    
    ht = ht.drop('A.unnorm_meta_beta','B.unnorm_meta_beta','A.inv_se2','B.inv_se2')
    
    variants = hl.import_table('gs://nbaya/rg_sex/50_snps_alleles_N.tsv.gz',types={'N': hl.tint64})
    variants = variants.key_by('SNP')
#    mt_all = hl.read_matrix_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_'+phen+'_grouped'+str(n_chunks)+'_batch_'+batch+'.mt') #matrix table containing individual samples. OUTDATED
    ht_all = hl.read_table('gs://nbaya/split/ukb31063.hm3_variants.gwas_samples_'+phen+'_grouped'+str(n_chunks)+'_batch_'+batch+'.ht') #hail table containing individual samples
    variants = variants.annotate(N = hl.int32(ht_all.count()/2))
    variants.show()
    
    metaA = variants.annotate(Z = ht[variants.SNP].A_Z)
    metaB = variants.annotate(Z = ht[variants.SNP].B_Z)
    
#    metaA_path = 'gs://nbaya/rg_sex/'+phen+'_meta_A_n'+str(n_chunks)+'_batch_'+batch+'_s'+str(i)+'.tsv.bgz' 
#    metaB_path = 'gs://nbaya/rg_sex/'+phen+'_meta_B_n'+str(n_chunks)+'_batch_'+batch+'_s'+str(i)+'.tsv.bgz' 
    metaA_path = 'gs://nbaya/rg_sex/'+variant_set+'_'+phen+'_meta_A_n'+str(n_chunks)+'_batch_'+batch+'_s'+str(i)+'.tsv.bgz'  #only used by qc_pos variant set and later hm3 phens
    metaB_path = 'gs://nbaya/rg_sex/'+variant_set+'_'+phen+'_meta_B_n'+str(n_chunks)+'_batch_'+batch+'_s'+str(i)+'.tsv.bgz' #only used by qc_pos variant set and later hm3 phens
    metaA.export(metaA_path)
    metaB.export(metaB_path)
    
    endtime = datetime.datetime.now()
    elapsed = endtime-starttime
    print('####################')
    print('Completed iteration '+str(i))
    print('Files written to:')
    print(metaA_path+'\t'+metaB_path)
    print('Time: {:%H:%M:%S (%Y-%b-%d)}'.format(datetime.datetime.now()))
    print('Iteration time: '+str(round(elapsed.seconds/60, 2))+' minutes')
    print('####################')    

for i in idx:
    run_meta_split(i)

print('####################')
print('Finished Part 5')
print('Meta-split complete')
print('Time: {:%H:%M:%S (%Y-%b-%d)}'.format(datetime.datetime.now()))
print('####################')










