library(tidyverse)

cov_path<-"results/var_sets/EURs_unrel/B2/cov.cov"
prs_path<-"results/PRS/scores.profile"
cov_out_path<-"results/var_sets/EURs_unrel/B2/cov.cov"

cov_path<-snakemake@input[[1]]
prs_path<-snakemake@input[[2]]
cov_out_path<-snakemake@output[[1]]

COV<-read_tsv(file=cov_path, col_names=TRUE)
PRS<-read.table(file = prs_path, header = TRUE)

print("Percentage of individuals in the COV file that also got a PRS calculated:")
mean(COV$IID %in% PRS$IID)

PRS_red<-PRS %>%
  select(IID, SCORESUM)

COV<-COV %>% 
  left_join(PRS_red, by="IID")

write_tsv(x=COV, file=cov_out_path, col_names = TRUE)

library(tidyverse)
library(readxl)

args=c("config/post_df3_HGI_sample_QC_summary.xlsx", "data/hail_gather_data/for_sample_QC.tsv", "data/plotPhenoInfo/")

args = commandArgs(trailingOnly=TRUE)
print(args)
PhenoFileLocation=args[1]
qc_file<-args[2]
outFolder=args[3]

samplesheet<-read_excel(PhenoFileLocation, sheet = "qc_table")

qc_table<-read_tsv(qc_file)

dir.create(outFolder,recursive = TRUE, showWarnings = FALSE)
setwd(outFolder)

samplesheet<-samplesheet%>% 
  rename(sex_to_use=`Sex (f=fem, m=m)`)%>%
  mutate(sex_for_fam=ifelse(sex_to_use=="f", 2, ifelse(sex_to_use=="m",1,0)))

sample_sex<-samplesheet %>% select(sex_to_use, s)

qc_table<-qc_table %>%
  left_join(sample_sex, by=(c("s"="s")))

qc_table<-qc_table %>% 
  mutate(Age=as.integer(Age),
         A2=as.integer(A2),
         B2=as.integer(B2),
         C2=as.integer(C2),
         is_case=as.integer(is_case)
         )



# Check Age vs. Sex

SexPlot<-ggplot(qc_table)+
  geom_histogram(aes(x=Age, fill=sex_to_use), bins=20)+
  theme_bw()

print(SexPlot)
ggsave(filename="SexPlot.pdf",
       plot=SexPlot,
       width=7,
       height=3)

SexInfo<-qc_table %>% 
  summarise(mean(sex_to_use=="m"), num=sum(sex_to_use!="noSex"))
write_tsv(x=SexInfo, file="SexInfo.tsv")

# Check Age vs. Severity
qc_table <- qc_table %>% 
  mutate(is_case_F=as.factor(is_case))

AgePlot<-ggplot(qc_table)+
  geom_histogram(aes(x=Age, fill=is_case_F), bins=20)+
  geom_vline(data=qc_table %>% filter(is_case==0), aes(xintercept = mean(Age)), col="red")+
  geom_vline(data=qc_table %>% filter(is_case==1), aes(xintercept = mean(Age)), col="green")+
  geom_vline(data=qc_table %>% filter(is_case==2), aes(xintercept = mean(Age)), col="blue")+
  theme_bw()

print(AgePlot)
ggsave(filename="AgePlot.pdf",
       plot=AgePlot,
       width=7,
       height=3)

AgeInfo<-qc_table %>% 
  group_by(is_case)%>%
  summarise(MEAN=mean(Age), SD=sd(Age), N=length(Age))%>%
  mutate(SEM=SD/sqrt(N))


write_tsv(x=AgeInfo, file="AgeInfo.tsv")

library(tidyverse)
library(readxl)

# Params:
# 1. input pheno-excel
# 2. input PC file
# 3. output file

args=c("config/post_df3_HGI_sample_QC_summary.xlsx","data/PCAcovar/A2_EUR_PCA.eigenvec")

# first: excel_file, second: fam-file, third: Phenotype column in excel

args = commandArgs(trailingOnly=TRUE)
print(args)

# File formats:
#Line 1 : Header with FID, IID and C covariate names.
#Followed by lines of C+2 values. Space/tab separated.

samplesheet<-read_tsv(args[1])%>%
  mutate(sex_for_fam=ifelse(is_female==TRUE, 2, ifelse(is_female==FALSE,1,0)))




col_names_eigenv<-c("FID","IID", paste0("PC",1:10))

# sex coding: 1=male, 2=female
# pheno coding: 1=control, 2=case
eigenv<-read_delim(args[2], col_names=col_names_eigenv, delim=" ")

cov_info<-eigenv %>% left_join(samplesheet, by=c("IID"="s"))%>%
  mutate(Age=as.integer(Age)) %>%
  mutate(age_sex=Age*sex_for_fam)%>%
  mutate(Age2=Age*Age)%>%
  select(all_of(col_names_eigenv), Age, Age2, age_sex, sex_for_fam)

write_tsv(x = cov_info, file = args[3], col_names = TRUE) # pheno file for regenie

library(tidyverse)
library(readxl)
# Params:
# 1. input pheno-excel
# 2. input fam file
# 3. pheno column name in excel
# 4. Output: fam file (Plink)
# 5. Output: pheno file (Regenie)

#testing
#setwd("/media/axel/Dateien/Arbeit_Gen/COVID_WGS/WGS_GWAS_pipeline")
args=c("../../resources/EURs_unrel.tsv","data/anc_vcf/EUR_vcf_not_rel.vcf.bgz.fam", "A1", 
       "data/regenie_pheno/fam_file.fam", "data/regenie_pheno/phenox")
# first: excel_file, second: fam-file, third: Phenotype column in excel

args = commandArgs(trailingOnly=TRUE)
print(args)

# File formats:
#FAM file: 
# A text file with no header line, and one line per sample with the following six fields:
#Family ID ('FID')
#Within-family ID ('IID'; cannot be '0')
#Within-family ID of father ('0' if father isn't in dataset)
#    Within-family ID of mother ('0' if mother isn't in dataset)
#Sex code ('1' = male, '2' = female, '0' = unknown)
#Phenotype value ('1' = control, '2' = case, '-9'/'0'/non-numeric = missing data if case/control)

# Regnie pheno file
# FID IID Y1 Y2
#  0=control, 1=case, missing values must be coded as NA

samplesheet<-read_tsv(args[1])%>%
  mutate(sex_for_fam=ifelse(is_female==TRUE, 2, ifelse(is_female==FALSE,1,0)))

pheno_col_num<-which(colnames(samplesheet) == args[3])
pheno_PLINK<-as.integer(unlist(samplesheet[,pheno_col_num]))+1
samplesheet$pheno_P=replace_na(pheno_PLINK, -9)
pheno_REGENIE<-as.integer(unlist(samplesheet[,pheno_col_num]))
samplesheet$pheno_R=pheno_REGENIE

# add sex stratification / age strat
samplesheet <- samplesheet %>%
  mutate(pheno_R_female=ifelse(is_female==TRUE, pheno_R, NA), # just keep females
         pheno_R_male=ifelse(is_female==FALSE, pheno_R, NA), # just keep males
         pheno_R_GE60=ifelse(pheno_R==1 & Age < 60, NA, pheno_R),  # set cases 59 or younger to missing
         pheno_R_LT60=ifelse(pheno_R==1 & Age >= 60, NA, pheno_R),
         pheno_R_GE60CC=ifelse(Age < 60, NA, pheno_R),  # set persons 59 or younger to missing
         pheno_R_LT60CC=ifelse(Age >= 60, NA, pheno_R)) # set persons 60 or older to missing 

col_names_fam<-c("FID","IID","father","mother","sex","pheno")

# sex coding: 1=male, 2=female
# pheno coding: 1=control, 2=case
fam<-read_delim(args[2],
delim=" ", 
col_names=col_names_fam,
col_types=c(
  FID = col_character(),
  IID = col_character(),
  father = col_character(),
  mother = col_character(),
  sex = col_integer(),
  pheno = col_integer()
)
)

print("fam:")
head(fam)
print("samplesheet:")
head(samplesheet)

all_info_joined<-fam %>% left_join(samplesheet, by=c("IID"="s"))%>%
  mutate(sex=sex_for_fam)

fam_new<-all_info_joined %>% 
  select(all_of(col_names_fam[1:4]), sex_for_fam, pheno_P)
write_tsv(x = fam_new, file = args[4], col_names = FALSE) # fam_file for plink

pheno_file_regenie <- all_info_joined %>%
  select(FID, IID, pheno_R, pheno_R_female, pheno_R_male, pheno_R_GE60, pheno_R_LT60, pheno_R_GE60CC, pheno_R_LT60CC)

colnames(pheno_file_regenie)<-c("FID","IID",
                                args[3],
                                paste0(args[3],"_female"),
                                paste0(args[3],"_male"),
                                paste0(args[3],"_GE60"),
                                paste0(args[3],"_LT60"),
                                paste0(args[3],"_GE60CC"),
                                paste0(args[3],"_LT60CC"))

print("out:")
head(pheno_file_regenie)

write_tsv(x = pheno_file_regenie, file = args[5], col_names = TRUE) # pheno file for regenie

if (!require("GenomicRanges", quietly = TRUE) & (!require("rtracklayer", quietly = TRUE)) ){
  install.packages("BiocManager", repos='http://cran.us.r-project.org')
  BiocManager::install("GenomicRanges")
  BiocManager::install("rtracklayer")
}

library(tidyverse)
library(data.table)
library(GenomicRanges)
library(rtracklayer)


gene_set_path<-"resources/gene_sets.tsv"
gene_set_path<-snakemake@input[[1]]
bed_path<-"resources/genomic_ranges.bed"
bed_path<-snakemake@input[[2]]
aaf_path<-"results/all_vars_for_RVAS/data.anno.aaf.file.txt"
aaf_path<-snakemake@input[[3]]
anno_path<-"/home/aschmidt/Arbeit_Gen/COVID_WGS/COVID_annotation/results/for_RVAS/all_contigs_anno.file.txt"
anno_path<-snakemake@input[[4]]
bim_path<-"results/all_vars_for_RVAS/eurs.bim"
bim_path<-snakemake@input[[5]]
#order_csq<-"resources/genomic_ranges.bed"
#order_csq<-snakemake@input[[6]]

new_anno_path<-"set.annos.tsv"
new_anno_path<-snakemake@output[[1]]
new_aafs_path<-"aafs.tsv"
new_aafs_path<-snakemake@output[[2]]
new_sets_path<-"sets.tsv"
new_sets_path<-snakemake@output[[3]]
relevant_variants_path<-"relevant_variants.tsv"
relevant_variants_path<-snakemake@output[[4]]
new_bim_path<-"new_bim.bim"
new_bim_path<-snakemake@output[[5]]

aafs<-read_delim(aaf_path, delim=" ", col_names = c("ID", "AAF"))
print("aafs:")
head(aafs)
annos<-read_delim(anno_path, delim=" ", col_names = c("ID", "gene_id", "annot"))
print("annos")
head(annos)

bim<-fread(bim_path, sep ="\t", col.names = c("CHR", "ID", "cM","BP","A1","A2"))

print("bim")
head(bim)

#bim_test<-head(bim, n=5000000)

bim_new_id<-bim %>%
  mutate(IDnew=paste("1", 1:n(), A1, A2, sep=":"))%>%
  mutate(BPnew=1:n())%>%
  mutate(CHRnew=1)

comb_sources<-bim_new_id %>%
  left_join(aafs, by="ID")%>%
  left_join(annos, by="ID")

print("combsource")
head(comb_sources)
head(comb_sources %>% filter(!is.na(annot)))

# make sets based on gene definitions
gene_sets<-read_tsv(gene_set_path)
new_sets<-tibble()
new_annos<-tibble()
new_aaf<-tibble()

score_conseq<-function(conseq){
  conseq_tbl<-c("pLoF",
                "missense.revel07",
                "missense.revel05",
                "missense.revel03",
                "moderate",
                "synonymous",
                "UTR5_CADD",
                "UTR5",
		            "UTR3_CADD",
                "UTR3",
                "promoter_CADD",
                "promoter",
                "enhancer_CADD",
                "enhancer",
                "regionBased")


conseq_rank<-c()
for (single_conseq in conseq){
  single_rank<-which(conseq_tbl %in% single_conseq)
  single_rank<-ifelse(is.na(single_rank), 99, single_rank)
  conseq_rank<-c(conseq_rank, single_rank)
}
conseq_rank<-
return(conseq_rank)
}


i=1
for (single_gene_set in unique(gene_sets$set)){
  rel_genes<-gene_sets %>%
    filter(set==single_gene_set)

  tmp_annos<-comb_sources %>%
    filter(gene_id %in% rel_genes$gene_id)%>%
    filter(annot!="not_relevant")%>%
    mutate(gene_id=single_gene_set)
  if (nrow(tmp_annos)==0){
    next()	
  }

  # check for variants that are added to the gene set more than once, take the most severe conseq
  mltpl_occuring_tmp<-tmp_annos %>% 
    add_count(IDnew)%>%
    filter(n>1)%>%
    select(-n)

  if (nrow(mltpl_occuring_tmp)>0){
    tmp_annos<-tmp_annos %>%
      filter(!IDnew %in% mltpl_occuring_tmp$IDnew)

    for (IDnew_tmp in unique(mltpl_occuring_tmp$IDnew)){
      mltpl_occuring_single<-mltpl_occuring_tmp %>%
        filter(IDnew==IDnew_tmp)%>%
        mutate(eff_score=score_conseq(annot))%>%
        arrange(eff_score) 

      print(IDnew_tmp)
      print(mltpl_occuring_single)

      mltpl_occuring_single<-mltpl_occuring_single %>%
        select(-eff_score)

      tmp_annos<-rbind(tmp_annos, mltpl_occuring_single[1,])
    }
  }

  new_annos<-rbind(new_annos, tmp_annos) 
  tmp_IDs<-paste(tmp_annos$IDnew, collapse=",")
  tmp_set<-tibble(gene_id=single_gene_set, chr="1", pos=i, IDs=tmp_IDs)
  new_sets<-rbind(new_sets, tmp_set)
  i=i+1 
}



#### 

# make sets based on bed file

region_sets<-import(bed_path)
head(region_sets)
region_sets_names<-str_split(region_sets$name, ";", simplify=TRUE)
region_sets_names<-as_tibble(region_sets_names)
colnames(region_sets_names)<- c("set", "gene")
set_names<-unique(region_sets_names$set)

sources_w_AF<-comb_sources %>%
  filter(!is.na(AAF))%>%
  mutate(pos_start=BP,
         pos_end=BP+nchar(A1))

head(sources_w_AF)
chroms<-paste0("chr",sources_w_AF$CHR)
chroms<-str_replace(chroms, "chr23", "chrX")

annos_granges <- GRanges(
  seqnames = chroms,
  ranges = IRanges(start=sources_w_AF$pos_start, end = sources_w_AF$pos_end) )
head(annos_granges)
for (set_name in set_names){

  tmp_region<-region_sets[region_sets_names$set==set_name]
  tmp_intersect<-GenomicRanges::intersect(annos_granges, tmp_region)

  tmp_chr_pos_inReg <-sources_w_AF[annos_granges %in% tmp_intersect,]%>%
    distinct(ID, .keep_all=TRUE)%>%
    mutate(gene_id=set_name,
           annot="regionBased")
  if (nrow(tmp_chr_pos_inReg)==0){
    next()
  }

  new_annos<-rbind(new_annos, tmp_chr_pos_inReg, fill=TRUE)

  tmp_IDs<-paste(tmp_chr_pos_inReg$IDnew, collapse=",")
  tmp_set<-tibble(gene_id=set_name, chr="1", pos=i, IDs=tmp_IDs)
  new_sets<-rbind(new_sets, tmp_set)
  i=i+1 
}


####

new_annos_export<-new_annos %>%
 distinct(IDnew, gene_id, annot)

new_aaf<-comb_sources %>% distinct(IDnew, AAF) %>%
  filter(!is.na(AAF))

relevant_variants<-unique(new_annos_export$ID)

write_delim(x=new_annos_export,
            file = new_anno_path,
            col_names = FALSE )


write_delim(x=new_aaf,
            file = new_aafs_path,
            col_names = FALSE )

write_delim(x=new_sets,
            file = new_sets_path,
            col_names = FALSE )

write(x = relevant_variants,
      file=relevant_variants_path)

bim_new_id<-bim_new_id%>%
  select(CHRnew,IDnew,cM,BPnew,A1,A2)%>%
  relocate(CHRnew,IDnew,cM,BPnew,A1,A2)

write_tsv(x=bim_new_id,
            file = new_bim_path,
            col_names = FALSE )

import pandas as pd
import pandas_plink
from xarray import DataArray
import shutil
import dask.array

# %%
mask_def_file = snakemake.input.regenie_mask_def
aaf_file = snakemake.input.data_anno_aaf_file
anno_file = snakemake.input.regenie_anno_file

af_cutoffs = list(map(float, snakemake.params.af_cutoffs))
plink_in_prefix = snakemake.params.plink_in_prefix

out_prefix = snakemake.params.plink_out_folder
set_file = snakemake.output.setFile_rvtest

debug_joined_out = snakemake.output.debug_joined_out
#%%

# read plink with pandas_plink
(bim, fam, bed) = pandas_plink.read_plink(plink_in_prefix, verbose=False)
bim.cm = bim.cm.astype(int)
#%%
# Read regenie aaf_file, anno_file, mask_def
aafs = pd.read_csv(aaf_file, sep=' ', names=['ID', 'aaf'])

mask_def = pd.read_csv(mask_def_file, sep=' ', names=['mask', "annots"])
mask_def.annots = mask_def.annots.map(lambda a: a.split(','))
mask_def = mask_def.set_index("mask")

anno = pd.read_csv(anno_file, sep=' ', names=['ID', 'gene', 'annot'])
anno = anno.dropna()
# %%
# Create M0-M4 True/False columns in table
for mask, annots in mask_def.annots.items():
    anno[mask] = anno.annot.isin(annots)

# Join AF info into table
anno_aaf = anno.join(aafs.set_index("ID"), on="ID", how="outer")

# %%

# join bim (variant table) with regenie anno
joined = bim.join(anno_aaf.set_index("ID"), on="snp", how="inner")
joined = joined.sort_values(by=['gene', 'pos'])

joined.to_csv(debug_joined_out, sep="\t")

#%%

for mask in mask_def.index.unique():
    for af_cutoff in af_cutoffs:
        mask_out_prefix = f"{out_prefix}{mask}_{af_cutoff}"
        # subset variant table
        subset = joined[joined[mask] & (joined["aaf"] < af_cutoff)]
        # export genotype data (subset["i"] contains positions of variants in table)
        pandas_plink.write_plink1_bin(
            DataArray(bed[subset["i"]], dims=["variant", "sample"]),
            bed=mask_out_prefix+".bed"
            )
        # write variant subset to bim file
        subset[["gene","snp","cm","pos","a0","a1"]].to_csv(mask_out_prefix+".bim", sep=" ", header=False, index=False)
        # reuse original fam file
        shutil.copy(plink_in_prefix + ".fam", mask_out_prefix +".fam")

#%%

# create setfile, each gene has its own "chromosome"/contig
set_str = ""
for gene in joined.gene.unique():
    set_str += f"{gene}\t{gene}:0-{joined.pos.max() + 100}\n"

with open(set_file, "wt") as f:
    f.write(set_str)

shell:
	"""
	#wget -nc ftp.ncbi.nlm.nih.gov/genomes/all/GCA/000/001/405/GCA_000001405.15_GRCh38/seqs_for_alignment_pipelines.ucsc_ids/GCA_000001405.15_GRCh38_no_alt_analysis_set.fna.gz
	#gzip -d GCA_000001405.15_GRCh38_no_alt_analysis_set.fna.gz
	bcftools filter -r chr1,chr2,chr3,chr4,chr5,chr6,chr7,chr8,chr9,chr10,chr11,chr12,chr13,chr14,chr15,chr16,chr17,chr18,chr19,chr20,chr21,chr22,chrX \
	  {input.bcf} -Ou | \
	  bcftools norm --check-ref w -f {input.fasta} -Ou | \
	  bcftools annotate --set-id '%CHROM:%POS:%REF:%FIRST_ALT' -Oz > {output.vcf_gz}
	tabix -p vcf {output.vcf_gz}
	"""

shell:
	"""
	plink2 \
	--bfile {params.in_plink} \
	--maf 0.01 \
	--exclude bed0 {input.bed} \
	--make-bed \
	--out {params.tmp_plink}		

	plink \
	--bfile {params.tmp_plink} \
	--indep-pairwise 1000kb 50 0.2 \
	--out {params.out_plink}

	plink \
	--bfile  {params.in_plink} \
	--extract {params.out_plink}.prune.in \
	--split-x hg38 no-fail \
	--make-bed \
	--out {params.out_plink}

	rm {params.tmp_plink}*
	"""