############################################################## 
# PRS based on cognitive function GWAS
##############################################################
#===========================================================
#-----------------------------------------------------------
# Format GWAS for PRS-CS
#-----------------------------------------------------------
echo "SNP A1 A2 BETA P" > Cognition_savage_noUKB.EUR_maf005info3_finalqc.samplesizechecked.prscs.txt
zcat /home/cchen11/cog_gwas_sumstats/Cognition_savage_noUKB.1000G.EUR.ref.SumstatsQC.AF_0.005.INFO_0.3.AFB_0.15.results.finalqc.txt.gz | awk 'FNR>1{if($17 > 74210*0.8) print $5,$8,$9,$14,$16}' >> Cognition_savage_noUKB.EUR_maf005info3_finalqc.samplesizechecked.prscs.txt
# check sample size
# 74210

#-----------------------------------------------------------
# Run PRS-CS
#-----------------------------------------------------------
for i in Cognition_savage_noUKB,74210
do
for j in {1..22}
do set -- $i $j
a=$( echo $1 | cut -d, -f 1)
b=$( echo $1 | cut -d, -f 2)
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=24
#$ -l h_vmem=12G
#$ -N $a.chr$2.prscs_phi01
#$ -cwd
#$ -V
#$ -e /MYPATH/$a.chr$2.prscs_phi01.err
#$ -o /MYPATH/$a.chr$2.prscs_phi01.out
module load python/2.7.11
module load pymods/h5py/2.7.0
module load pymods/scipy/0.19.1
python /MYPATH/PRScs.py --ref_dir=/MYPATH/PRScs_ldr/ldblk_1kg/ --bim_prefix=/MYPATH/1kg_phase3/g1000_eur.chr$2 --sst_file=/MYPATH/$a.EUR_maf005info3_finalqc.samplesizechecked.prscs.txt --n_gwas=$b --out_dir=/MYPATH/$a --chrom=$2 --phi=1e-2
EOS
    done
done

#-----------------------------------------------------------
# calculate scores
#-----------------------------------------------------------
for i in {1..22}
do set -- $i
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=12
#$ -l h_vmem=16G
#$ -N Cognition_savage_noUKB.chr$1.calc_prs
#$ -cwd
#$ -V
#$ -e /MYPATH/Cognition_savage_noUKB.chr$1.calc_prs.err
#$ -o /MYPATH/Cognition_savage_noUKB.chr$1.calc_prs.out
plink2 --pfile /MYPATH/ukb_imp_chr$1_info08maf001 --score /MYPATH/Cognition_savage_noUKB_pst_eff_a1_b0.5_phi1e-02_chr$1.txt 2 4 6 ignore-dup-ids no-mean-imputation cols=nallele,dosagesum,scoreavgs,scoresums --out /MYPATH/Cognition_savage_noUKB.chr$1 --threads 12
EOS
done

cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=8
#$ -l h_vmem=8G
#$ -N collect_score
#$ -cwd
#$ -e /MYPATH/collect_score.err
#$ -o /MYPATH/collect_score.out
#$ -V
cd /MYPATH
Rscript --vanilla ./collect_score_chr.R Cognition_savage_noUKB
EOS

#===========================================================
# PRS analysis
#===========================================================
library(ggplot2)
library(data.table)
library(dplyr)

#-----------------------------------------------------------
# Get data
# R script below
#-----------------------------------------------------------
# PRS
prs <- read.table("/MYPATH/Cognition_savage_noUKB.totalscore",header=TRUE)
colnames(prs)[1] <- "f_eid"
# PTV
filelist <- c("ptv.gw", "ptv.pLI_intol", "ptv.pLI_tol")
burden_raw <- fread(paste0("/MUYPATH/",filelist[1],".txt"))
for (i in 2:length(filelist)){
tmpdat <- fread(paste0("/MYPATH/",filelist[i],".txt"))
burden_raw <- inner_join(burden_raw,tmpdat,by="eid_sample")
}

# MPC1_0: missense with 2 >= MPC > 0
# MPC1_2: missense with 3 >= MPC > 2
# MPC1_3: missense with MPC > 3
filelist <- c("gw.maf00001.MPC1_0.missensesum", "gw.maf00001.MPC1_2.missensesum", "gw.maf00001.MPC1_3.missensesum", "pLI_intol.maf00001.MPC1_0.missensesum", "pLI_intol.maf00001.MPC1_2.missensesum", "pLI_intol.maf00001.MPC1_3.missensesum", "pLI_tol.maf00001.MPC1_0.missensesum", "pLI_tol.maf00001.MPC1_2.missensesum", "pLI_tol.maf00001.MPC1_3.missensesum")
missense_colnames_list <- c("missense_MPC1_0.gw", "missense_MPC1_2.gw", "missense_MPC1_3.gw", "missense_MPC1_0.pLI_intol", "missense_MPC1_2.pLI_intol", "missense_MPC1_3.pLI_intol", "missense_MPC1_0.pLI_tol", "missense_MPC1_2.pLI_tol", "missense_MPC1_3.pLI_tol")
for (i in 1:length(filelist)){
tmpdat <- fread(paste0("/MYPATH/ukbb.",filelist[i],".txt"))
burden_raw <- inner_join(burden_raw,tmpdat,by="eid_sample")
colnames(burden_raw)[length(burden_raw)] <- missense_colnames_list[i]
}
burden_raw$missense_ALL.gw <- burden_raw$missense_MPC1_0.gw + burden_raw$missense_MPC1_2.gw + burden_raw$missense_MPC1_3.gw
burden_raw$missense_ALL.pLI_intol <- burden_raw$missense_MPC1_0.pLI_intol + burden_raw$missense_MPC1_2.pLI_intol + burden_raw$missense_MPC1_3.pLI_intol
burden_raw$missense_ALL.pLI_tol <- burden_raw$missense_MPC1_0.pLI_tol + burden_raw$missense_MPC1_2.pLI_tol + burden_raw$missense_MPC1_3.pLI_tol

# Synonymous
filelist <- c("gw.maf00001.synonymoussum", "pLI_intol.maf00001.synonymoussum", "pLI_tol.maf00001.synonymoussum")
synonymous_colnames_list <- c("synonymous.gw", "synonymous.pLI_intol", "synonymous.pLI_tol")
for (i in 1:length(filelist)){
tmpdat <- fread(paste0("/MYPATH/ukbb.",filelist[i],".txt"))
burden_raw <- inner_join(burden_raw,tmpdat,by="eid_sample")
colnames(burden_raw)[length(burden_raw)] <- synonymous_colnames_list[i]
}

colnames(burden_raw)[1] <- "IID"

# phenotype file
pheno_unrel <- fread("/MYPATH/MYFILE.EUR.unrel.txt")


summary(pheno_unrel$edu_baseline,na.rm=TRUE)
sd(pheno_unrel$edu_baseline,na.rm=TRUE)

summary(pheno_unrel$fiall_online,na.rm=TRUE)
sd(pheno_unrel$fiall_online,na.rm=TRUE)

# merge PRS and burden
pheno_burden_raw <- merge(pheno_unrel, burden_raw, by="IID", all=FALSE)
dat_raw <- merge(pheno_burden_raw, prs, by="f_eid", all=FALSE)

# create PRS and burden categories
dat <- as.data.frame(dat_raw)
dat$prs_std <- scale(dat$TOTALSCORE_SUM)
burden_list <- c("ptv.gw", "ptv.pLI_intol", "ptv.pLI_tol", "missense_MPC1_0.gw", "missense_MPC1_2.gw", "missense_MPC1_3.gw", "missense_MPC1_0.pLI_intol", "missense_MPC1_2.pLI_intol", "missense_MPC1_3.pLI_intol", "missense_MPC1_0.pLI_tol", "missense_MPC1_2.pLI_tol", "missense_MPC1_3.pLI_tol", "missense_ALL.gw", "missense_ALL.pLI_intol", "missense_ALL.pLI_tol", "synonymous.gw", "synonymous.pLI_intol", "synonymous.pLI_tol")
for (i in burden_list){
 dat <- cbind.data.frame(dat, scale(dat[,i]))
 colnames(dat)[length(dat)] <- paste0(i,".std")
}

for (i in burden_list){
dat[,paste0(i,".bin")] <- NA
dat[which(dat[,paste0(i,".std")]>0),paste0(i,".bin")] <- 1
dat[which(dat[,paste0(i,".std")]<=0),paste0(i,".bin")] <- 0
dat[,paste0(i,".bin")] <- factor(dat[,paste0(i,".bin")])	
}

# top 10% vs rest 90%
for (i in burden_list[1:3]){
dat[,paste0(i,".decile")] <- cut(dat[,i], quantile(dat[,i],seq(0,1,0.1))[c(1,10,11)], labels = NULL, include.lowest = TRUE)
levels(dat[,paste0(i,".decile")]) <- c(0,1)
}

dat$TOTALSCORE_SUM_quan <- cut(dat$TOTALSCORE_SUM, quantile(dat$TOTALSCORE_SUM,seq(0,1,0.1)), labels = NULL, include.lowest = TRUE)
levels(dat$TOTALSCORE_SUM_quan) <- c(1,2,3,4,5,6,7,8,9,10)

# output data
fwrite(dat, "UKBB.PRS_codingburden_pheno.EUR.tsv.gz", sep=" ")