import pandas as pd
import numpy as np
import os
import gzip
from functools import reduce
import scipy.stats as st
import statsmodels.stats.multitest as multi
from sklearn.decomposition import PCA


cohort_nms= ['harvestm12', 'harvestm24','rotterdam1', 'rotterdam2', 'normentfeb', 'normentmay']
smpl_nms= ['maternal','paternal', 'fetal']
batch_nms= ['m12', 'm24']
CHR_nms= [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,12 ,13 ,14 ,15 ,16 ,17 ,18 ,19 ,20 ,21 ,22]

# Other arguments:

pruning_nms= ['none', 'soft', 'moderate']

dens_nms= [5]
SNP_nms= [15, 25, 50, 75, 100, 150, 200, 300, 400]
length_nms= [0.0000001]
het_nms= [0, 1]
GAP_nms= [5]

dens_bp= [5000]
SNP_bp= [15, 25, 50, 75, 100, 150, 200, 300, 400]
length_bp= [0.0000001]
het_bp= [0, 1]
GAP_bp= [5000]

# Functions

def isfloat(str):
    try:
        float(str)
        return True
    except ValueError:
        return False

rule surv_analysis:
        'AFT analysis analyis on bined ROH calls.'
        input:
                '/mnt/work/pol/ROH/genotypes/maps/{sample}/segments_maps_{sample}_chr{CHR}.txt.gz',
                '/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
        output:
                temp('/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}_chr{CHR}'),
		'/mnt/work/pol/ROH/results/{sample}/eff_ROH_{CHR}.txt'
        script:
                'surv_segments.R'

rule surv_reduced:
	'Loglik for the reduced model AFT'
	input:
		'/mnt/work/pol/ROH/genotypes/maps/{sample}/segments_maps_{sample}_chr1.txt.gz',
		'/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
	output:
		'/mnt/work/pol/ROH/results/{sample}/loglik_{sample}.txt'
	script:
		'loglik.R'

rule concat_surv:
	'Concat surv results from multiple chromosomes.'
	input:
		expand('/mnt/work/pol/ROH/results/{{sample}}/surv_spont_{{sample}}_chr{CHR}', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}'
	shell:
		'cat {input} > {output[0]}'

rule concat_eff:
	'Concat effective number of ROHs.'
	input:
		'/mnt/work/pol/ROH/results/maternal/eff_ROHX.txt',
		expand('/mnt/work/pol/ROH/results/{{sample}}/eff_ROH_{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/eff_ROH.txt'
	run:
		infiles= input[1:]
		if wildcards.sample== 'maternal':
			shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal':
			shell('cat {infiles} > {output[0]}')

rule surv_cM_to_bp:
	'Convert cM to bp in surv results.'
	input:
		'/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}',
		expand('/mnt/work/pol/ROH/{cohort}/runs/{{sample}}_input_ROH_geno.txt', cohort= cohort_nms),
		expand('/mnt/work/pol/ROH/{cohort}/genotypes/{pruning}/pruned{cohort}_{{sample}}.bim', cohort= cohort_nms, pruning= pruning_nms)
	output:
		'/mnt/work/pol/ROH/results/autoz_surv_spont_{sample}'
	run:
		if os.stat(input[0]).st_size> 0:
			d= pd.read_csv(input[0], header= None, sep= '\t', names= ['segment', 'n', 'freq', 'beta', 'sd', 'pvalue', 'loglik'])
			d[['chr', 'cM1', 'cM2']]= d['segment'].str.split(':', expand= True)
			d[['chr', 'cM1', 'cM2']]= d[['chr', 'cM1', 'cM2']].apply(lambda x: x.astype('float'))
			flist= [infile for infile in input if 'geno.txt' in infile]
			bim_list= list()
			for infile in flist:
				x= [line.strip() for line in open("".join(infile), 'r')]
				bim= [i for i in x if 'bim' in i]
				bim= pd.read_csv("".join(bim), delim_whitespace= True, header= None, names= ['chr', 'id', 'pos', 'cM', 'A1', 'A2'])
				bim_list.append(bim)
			bim= pd.concat(bim_list)
			bim.drop_duplicates(subset= ['chr', 'pos'], keep= 'first', inplace= True)
			bim.drop_duplicates(subset= ['chr', 'cM'], keep= 'first', inplace= True)
			bim= bim[['chr', 'pos', 'cM']]	
			d= pd.merge(d, bim, left_on= ['chr', 'cM1'], right_on= ['chr', 'cM'])
			d.rename(columns= {'pos': 'pos1'}, inplace= True)
			d= pd.merge(d, bim, left_on= ['chr', 'cM2'], right_on= ['chr', 'cM'])
			d.rename(columns= {'pos': 'pos2'}, inplace= True)
			d.drop(['cM_x', 'cM_y'], axis= 1, inplace= True)
			d.to_csv(output[0], sep= '\t', header= True, index= False)
		else:
			open(output[0], 'a').close()

rule concat_surv_chrX:
	''
	input:
		'/mnt/work/pol/ROH/results/autoz_surv_spont_{sample}',
		'/mnt/work/pol/ROH/results/chrX/chrX_surv_spont'
	output:
		'/mnt/work/pol/ROH/results/surv_spont_{sample}'
	run:
		if wildcards.sample== 'maternal':
			shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal':
			shell('mv {input[0]} {output[0]}')

rule surv_imputed:
        'Cox regression for imputed variants within HC segments.'
        input:
                '/mnt/work/pol/ROH/genotypes/GT/{sample}_HC_temp',
                '/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt',
                '/mnt/work/pol/ROH/genotypes/imputed/{sample}_ids_toextract'
        output:
                '/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}'
        script:
                'surv_imputed.R'

rule filter_variants:
	'Select variants with INFO score> 0.4 in all cohorts.'
	input:
		'/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}',
		'/mnt/archive/MOBAGENETICS/genotypes-base/aux/markerinfo/all-markerinfo.gz'
	output:
		temp('/mnt/work/pol/ROH/results/variants_info_score_{sample}.txt')
	run:
		d= pd.read_csv(input[0], header= None, sep= '\t', names= ['snp', 'ref', 'eff', 'model', 'n', 'het', 'hom', 'coef', 'sd', 'pvalue'], usecols= ['snp', 'eff'])
		df= pd.read_csv(input[1], header= 0, sep= '\t', compression= 'gzip', usecols= ['CHROM', 'POS', 'INFO', 'REF', 'ALT'])
		df= df.loc[df.INFO> 0.4, :]
		df.columns= ['chr', 'pos', 'REF', 'EFF', 'INFO']
		df["REF"], df["EFF"] = np.where(df['REF'] > df['EFF'], [df["EFF"], df["REF"]], [df["REF"], df["EFF"]])
		df['snp']= df['chr'].map(str) + ':' + df.pos.map(str) + ':' + df.REF + ':' + df.EFF
		df= df.loc[df.snp.isin(d.snp), :]
		d.to_csv(output[0], sep= '\t', header= False, index= False, columns= ['snp'])

rule filter_surv:
	'Select top genetic variants associated with spontaneous delivery risk (filter INFO score >0.4)'
	input:
		'/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}',
		'/mnt/work/pol/ROH/results/variants_info_score_{sample}.txt'
	output:
		'/mnt/work/pol/ROH/results/imputed/filtered_surv_spont_{sample}'
	run:
		d= pd.read_csv(input[0], sep= '\t', header= None, names= ['variant', 'ref', 'eff', 'model', 'n', 'het', 'hom', 'beta', 'sd', 'pvalue'])
		with open(input[1]) as f:
			variants_filter= f.read().splitlines()
		d= d.loc[d.variant.isin(variants_filter), :]
		d.to_csv(output[0], sep= '\t', header= False, index= False)

rule concat_imputed_surv_chrX:
	''
	input:
		'/mnt/work/pol/ROH/results/chrX/imputed/chrX_imputed_surv_spont',
		'/mnt/work/pol/ROH/results/imputed/filtered_surv_spont_{sample}'
	output:
		'/mnt/work/pol/ROH/results/imputed/surv_imputed_{sample}.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal': shell('mv {input[1]} {output[0]}')

rule gene_burden:
	''
	input:
		'/mnt/work/pol/ROH/genotypes/maps/gene/{sample}_genemap_CHR{CHR}.txt',
		'/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
	output:
		temp('/mnt/work/pol/ROH/results/temp_burden_survival_{sample}_CHR{CHR}.txt'),
		temp('/mnt/work/pol/ROH/results/{sample}/gene_burden_eff_ROH_CHR{CHR}.txt')
	script:
		'aft_genes.R'

rule concat_gene:
	''
	input:
		'/mnt/work/pol/ROH/results/chrX/maps_surv/temp_burden_surv_spont_temp',
		expand('/mnt/work/pol/ROH/results/temp_burden_survival_{{sample}}_CHR{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/burden_survival_{sample}.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
		input_non_maternal= input[1:]
		if wildcards.sample!= 'maternal': shell('cat {input_non_maternal} > {output[0]}')

rule concat_eff_gene:
	''
	input:
		'/mnt/work/pol/ROH/results/maternal/gene_burden_eff_ROHX.txt',
		expand('/mnt/work/pol/ROH/results/{{sample}}/gene_burden_eff_ROH_CHR{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/gene_burden_eff_ROH.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
                input_non_maternal= input[1:]
                if wildcards.sample!= 'maternal': shell('cat {input_non_maternal} > {output[0]}')

959b4fa118f6d30500f529ed390d65053e8dc6e7
import pandas as pd
import numpy as np
import os
import gzip
from functools import reduce
import scipy.stats as st
import statsmodels.stats.multitest as multi
from sklearn.decomposition import PCA


cohort_nms= ['harvestm12', 'harvestm24','rotterdam1', 'rotterdam2', 'normentfeb', 'normentmay']
smpl_nms= ['maternal','paternal', 'fetal']
batch_nms= ['m12', 'm24']
CHR_nms= [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,12 ,13 ,14 ,15 ,16 ,17 ,18 ,19 ,20 ,21 ,22]

# Other arguments:

pruning_nms= ['none', 'soft', 'moderate']

dens_nms= [5]
SNP_nms= [15, 25, 50, 75, 100, 150, 200, 300, 400]
length_nms= [0.0000001]
het_nms= [0, 1]
GAP_nms= [5]

dens_bp= [5000]
SNP_bp= [15, 25, 50, 75, 100, 150, 200, 300, 400]
length_bp= [0.0000001]
het_bp= [0, 1]
GAP_bp= [5000]

# Functions

def isfloat(str):
    try:
        float(str)
        return True
    except ValueError:
        return False

rule surv_analysis:
        'AFT analysis analyis on bined ROH calls.'
        input:
                '/mnt/work/pol/ROH/genotypes/maps/{sample}/segments_maps_{sample}_chr{CHR}.txt.gz',
                '/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
        output:
                temp('/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}_chr{CHR}'),
		'/mnt/work/pol/ROH/results/{sample}/eff_ROH_{CHR}.txt'
        script:
                'surv_segments.R'

rule surv_reduced:
	'Loglik for the reduced model AFT'
	input:
		'/mnt/work/pol/ROH/genotypes/maps/{sample}/segments_maps_{sample}_chr1.txt.gz',
		'/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
	output:
		'/mnt/work/pol/ROH/results/{sample}/loglik_{sample}.txt'
	script:
		'loglik.R'

rule concat_surv:
	'Concat surv results from multiple chromosomes.'
	input:
		expand('/mnt/work/pol/ROH/results/{{sample}}/surv_spont_{{sample}}_chr{CHR}', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}'
	shell:
		'cat {input} > {output[0]}'

rule concat_eff:
	'Concat effective number of ROHs.'
	input:
		'/mnt/work/pol/ROH/results/maternal/eff_ROHX.txt',
		expand('/mnt/work/pol/ROH/results/{{sample}}/eff_ROH_{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/eff_ROH.txt'
	run:
		infiles= input[1:]
		if wildcards.sample== 'maternal':
			shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal':
			shell('cat {infiles} > {output[0]}')

rule surv_cM_to_bp:
	'Convert cM to bp in surv results.'
	input:
		'/mnt/work/pol/ROH/results/{sample}/surv_spont_{sample}',
		expand('/mnt/work/pol/ROH/{cohort}/runs/{{sample}}_input_ROH_geno.txt', cohort= cohort_nms),
		expand('/mnt/work/pol/ROH/{cohort}/genotypes/{pruning}/pruned{cohort}_{{sample}}.bim', cohort= cohort_nms, pruning= pruning_nms)
	output:
		'/mnt/work/pol/ROH/results/autoz_surv_spont_{sample}'
	run:
		if os.stat(input[0]).st_size> 0:
			d= pd.read_csv(input[0], header= None, sep= '\t', names= ['segment', 'n', 'freq', 'beta', 'sd', 'pvalue', 'loglik'])
			d[['chr', 'cM1', 'cM2']]= d['segment'].str.split(':', expand= True)
			d[['chr', 'cM1', 'cM2']]= d[['chr', 'cM1', 'cM2']].apply(lambda x: x.astype('float'))
			flist= [infile for infile in input if 'geno.txt' in infile]
			bim_list= list()
			for infile in flist:
				x= [line.strip() for line in open("".join(infile), 'r')]
				bim= [i for i in x if 'bim' in i]
				bim= pd.read_csv("".join(bim), delim_whitespace= True, header= None, names= ['chr', 'id', 'pos', 'cM', 'A1', 'A2'])
				bim_list.append(bim)
			bim= pd.concat(bim_list)
			bim.drop_duplicates(subset= ['chr', 'pos'], keep= 'first', inplace= True)
			bim.drop_duplicates(subset= ['chr', 'cM'], keep= 'first', inplace= True)
			bim= bim[['chr', 'pos', 'cM']]	
			d= pd.merge(d, bim, left_on= ['chr', 'cM1'], right_on= ['chr', 'cM'])
			d.rename(columns= {'pos': 'pos1'}, inplace= True)
			d= pd.merge(d, bim, left_on= ['chr', 'cM2'], right_on= ['chr', 'cM'])
			d.rename(columns= {'pos': 'pos2'}, inplace= True)
			d.drop(['cM_x', 'cM_y'], axis= 1, inplace= True)
			d.to_csv(output[0], sep= '\t', header= True, index= False)
		else:
			open(output[0], 'a').close()

rule concat_surv_chrX:
	''
	input:
		'/mnt/work/pol/ROH/results/autoz_surv_spont_{sample}',
		'/mnt/work/pol/ROH/results/chrX/chrX_surv_spont'
	output:
		'/mnt/work/pol/ROH/results/surv_spont_{sample}'
	run:
		if wildcards.sample== 'maternal':
			shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal':
			shell('mv {input[0]} {output[0]}')

rule surv_imputed:
        'Cox regression for imputed variants within HC segments.'
        input:
                '/mnt/work/pol/ROH/genotypes/GT/{sample}_HC_temp',
                '/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt',
                '/mnt/work/pol/ROH/genotypes/imputed/{sample}_ids_toextract'
        output:
                '/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}'
        script:
                'surv_imputed.R'

rule filter_variants:
	'Select variants with INFO score> 0.4 in all cohorts.'
	input:
		'/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}',
		'/mnt/archive/MOBAGENETICS/genotypes-base/aux/markerinfo/all-markerinfo.gz'
	output:
		temp('/mnt/work/pol/ROH/results/variants_info_score_{sample}.txt')
	run:
		d= pd.read_csv(input[0], header= None, sep= '\t', names= ['snp', 'ref', 'eff', 'model', 'n', 'het', 'hom', 'coef', 'sd', 'pvalue'], usecols= ['snp', 'eff'])
		df= pd.read_csv(input[1], header= 0, sep= '\t', compression= 'gzip', usecols= ['CHROM', 'POS', 'INFO', 'REF', 'ALT'])
		df= df.loc[df.INFO> 0.4, :]
		df.columns= ['chr', 'pos', 'REF', 'EFF', 'INFO']
		df["REF"], df["EFF"] = np.where(df['REF'] > df['EFF'], [df["EFF"], df["REF"]], [df["REF"], df["EFF"]])
		df['snp']= df['chr'].map(str) + ':' + df.pos.map(str) + ':' + df.REF + ':' + df.EFF
		df= df.loc[df.snp.isin(d.snp), :]
		d.to_csv(output[0], sep= '\t', header= False, index= False, columns= ['snp'])

rule filter_surv:
	'Select top genetic variants associated with spontaneous delivery risk (filter INFO score >0.4)'
	input:
		'/mnt/work/pol/ROH/results/imputed/nofilter_surv_spont_{sample}',
		'/mnt/work/pol/ROH/results/variants_info_score_{sample}.txt'
	output:
		'/mnt/work/pol/ROH/results/imputed/filtered_surv_spont_{sample}'
	run:
		d= pd.read_csv(input[0], sep= '\t', header= None, names= ['variant', 'ref', 'eff', 'model', 'n', 'het', 'hom', 'beta', 'sd', 'pvalue'])
		with open(input[1]) as f:
			variants_filter= f.read().splitlines()
		d= d.loc[d.variant.isin(variants_filter), :]
		d.to_csv(output[0], sep= '\t', header= False, index= False)

rule concat_imputed_surv_chrX:
	''
	input:
		'/mnt/work/pol/ROH/results/chrX/imputed/chrX_imputed_surv_spont',
		'/mnt/work/pol/ROH/results/imputed/filtered_surv_spont_{sample}'
	output:
		'/mnt/work/pol/ROH/results/imputed/surv_imputed_{sample}.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
		if wildcards.sample!= 'maternal': shell('mv {input[1]} {output[0]}')

rule gene_burden:
	''
	input:
		'/mnt/work/pol/ROH/genotypes/maps/gene/{sample}_genemap_CHR{CHR}.txt',
		'/mnt/work/pol/ROH/pheno/runs_mfr_{sample}.txt'
	output:
		temp('/mnt/work/pol/ROH/results/temp_burden_survival_{sample}_CHR{CHR}.txt'),
		temp('/mnt/work/pol/ROH/results/{sample}/gene_burden_eff_ROH_CHR{CHR}.txt')
	script:
		'aft_genes.R'

rule concat_gene:
	''
	input:
		'/mnt/work/pol/ROH/results/chrX/maps_surv/temp_burden_surv_spont_temp',
		expand('/mnt/work/pol/ROH/results/temp_burden_survival_{{sample}}_CHR{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/burden_survival_{sample}.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
		input_non_maternal= input[1:]
		if wildcards.sample!= 'maternal': shell('cat {input_non_maternal} > {output[0]}')

rule concat_eff_gene:
	''
	input:
		'/mnt/work/pol/ROH/results/maternal/gene_burden_eff_ROHX.txt',
		expand('/mnt/work/pol/ROH/results/{{sample}}/gene_burden_eff_ROH_CHR{CHR}.txt', CHR= CHR_nms)
	output:
		'/mnt/work/pol/ROH/results/{sample}/gene_burden_eff_ROH.txt'
	run:
		if wildcards.sample== 'maternal': shell('cat {input} > {output[0]}')
                input_non_maternal= input[1:]
                if wildcards.sample!= 'maternal': shell('cat {input_non_maternal} > {output[0]}')
