library(tidyverse)
library(readxl)
library(gridExtra)
library(optparse)

option_list = list(
  make_option(c("-r", "--input_recalibrated"), type="character", default="C:/Users/Karo.PC-Karo/Master/Semester_3/Lab_rotation_Ludwig/for_Karola/SSH/gnomad_extracted_prepro_rec.csv.gz", 
              help="location of csv.gz file"),
  make_option(c("-e", "--excel_location"), type="character", default="C:/Users/Karo.PC-Karo/Master/Semester_3/Lab_rotation_Ludwig/for_Karola/local_files/available_colnames_regular.xlsx", 
              help="location of excel file"),
  make_option(c("-o", "--out_folder"), type="character", default="C:/Users/Karo.PC-Karo/Master/Semester_3/Lab_rotation_Ludwig/for_Karola/local_files", 
              help="name of folder to store output")

)

opt = parse_args(OptionParser(option_list=option_list))

variants<-read_csv(opt$input_recalibrated, na=c(".","NA"))
colnames_usage <- read_excel(opt$excel_location)

dir.create(opt$out_folder, recursive=TRUE)
setwd(opt$out_folder)

write_tsv(as_tibble(colnames(variants)), file="available_colnames_W_surr.tsv")

for_prediction<-(colnames_usage %>% filter(!is.na(for_prediction)))$value

mean_no_na_no_round<-function(x) {(mean(x, na.rm=TRUE))}


mean_properties<-variants %>% filter(b_factor>80) %>% 
  group_by(from_AS, outcome)%>% 
  dplyr::select(for_prediction) %>%
  summarise_all(mean_no_na_no_round)%>%
  ungroup()%>%
  dplyr::select(-starts_with("RESIDUE_NAME"))

mean_properties<-mean_properties %>% 
  gather(variable, value, -outcome, -from_AS) %>%
  group_by(variable,from_AS) %>%
  summarise(value = (value[outcome == 1] - value[outcome == 0]))

####Both Outcomes, total count####
mean_properties_freq<-variants %>%  
  group_by(from_AS, to_AS) %>%
  select(from_AS, to_AS) %>%
  filter(!to_AS=="%3D") %>%
  filter(!to_AS=="???") %>%
  count(from_AS,to_AS,sort=TRUE)

plot_tot_AA<-ggplot(mean_properties_freq, aes(x = from_AS, y = to_AS, fill = n)) +
  geom_tile()+
  theme_minimal()+  
  scale_fill_viridis_c()+
  theme(plot.title= element_text(size = 6),
        axis.text.x = element_text(size=5,angle = 90, vjust = 0.5, hjust=1),
        axis.text.y = element_text(size=5, vjust = 0.5, hjust=1),
        legend.key.height = unit(3, 'mm'),
        legend.key.width = unit(2.5, 'mm'),
        legend.title = element_text(size=5), 
        legend.text = element_text(size=5),
        legend.box.spacing = unit(0.1,"mm"),
        axis.title.y=element_text(size=6),
        axis.title.x=element_text(size=6),
        plot.margin=unit(c(1,1,1,1), "mm"))+
  ggtitle("AA exchanges, raw count")+
  geom_text(aes(from_AS, to_AS, label=n), colour = "white", check_overlap = TRUE, size=1.5)

#ggsave("AA_ex_tot.pdf", width=28.575, height=19.05, units="cm", dpi=300)

####Relative fraction outcome1/outcome0####
mean_properties_freq_outcome<-variants %>% 
  filter(!to_AS=="%3D") %>%
  filter(!to_AS=="???") %>%
  group_by(from_AS, to_AS) %>%
  select(from_AS, to_AS, outcome) %>%
  summarise(frac=mean(outcome))

frac_round<- round(mean_properties_freq_outcome$frac, digits = 2)

plot_AA_freq<-ggplot(mean_properties_freq_outcome, aes(x = from_AS, y = to_AS, fill = frac)) +
  geom_tile()+
  theme_minimal()+  
  scale_fill_viridis_c()+
  theme(plot.title= element_text(size = 6),
        axis.text.x = element_text(size=5,angle = 90, vjust = 0.5, hjust=1),
        axis.text.y = element_text(size=5, vjust = 0.5, hjust=1),
        legend.key.height = unit(3, 'mm'),
        legend.key.width = unit(2.5, 'mm'),
        legend.title = element_text(size=5), 
        legend.text = element_text(size=5),
        legend.box.spacing = unit(0.1,"mm"),
        axis.title.y=element_text(size=6),
        axis.title.x=element_text(size=6),
        plot.margin=unit(c(1,1,1,1), "mm"))+
  ggtitle("Fraction of AA exchanges with outcome = 1")+
  geom_text(aes(from_AS, to_AS, label=frac_round), colour = "white", check_overlap = TRUE, size=1.5)

#ggsave("AA_ex_frac.pdf",width=28.575, height=19.05, units="cm", dpi=300)

g<-arrangeGrob(plot_tot_AA, plot_AA_freq, ncol=1) #combine both plots in one
ggsave("AA_ex_compared.pdf", width=95, height=120,units="mm", dpi=300,g)

library(tidyverse)
library(pROC)
library(gridExtra)
library(optparse)
library(boot)
source("workflow/scripts/existing_scores_glm_functions.R")
BOOT_REPETITIONS=1000
N_CPU=3
COL_ORDER=c("Alph_null","AlphScore","CADD_raw","DEOGEN2_score_med","REVEL_score","glm_AlphCadd","glm_AlphDeogen",
  "glm_AlphRevel","glm_CaddDeogen","glm_DeogenRevel","glm_RevelCadd","glm_AlphDeogenRevel",
  "glm_AlphCaddDeogen","glm_AlphCaddRevel","glm_CaddDeogenRevel")

set.seed(1)


option_list = list(
  make_option(c("-t", "--variants"), type="character", default="results/prediction_final/pre_final_model_regular_variants.csv.gz", 
              help="csv.gz file, test dataset"),
  make_option(c("-v", "--validation_set"), type="character", default="results/validation_set/validation_set_w_AlphScore.csv.gz", 
              help="Validation set"),
  make_option(c("-o", "--out_folder"), type="character", default="results/analyse_score", 
              help="Output folder"))

opt = parse_args(OptionParser(option_list=option_list))

variants<-read_csv(opt$variants)
validation_dataset<-read_csv(opt$validation_set)

dir.create(opt$out_folder, recursive=TRUE)
setwd(opt$out_folder)

# get DEOGEN2 Scores:
variants$pos_in_VEP_and_Uniprot<-get_index_col(variants)
variants$DEOGEN2_score_med<-unlist_score(variants$DEOGEN2_score, variants$pos_in_VEP_and_Uniprot)

validation_dataset$pos_in_VEP_and_Uniprot<-get_index_col(validation_dataset)
validation_dataset$DEOGEN2_score_med<-unlist_score(validation_dataset$DEOGEN2_score, validation_dataset$pos_in_VEP_and_Uniprot)

variants$AlphScore<-variants$predicted_Alph

variants<-variants%>% 
  mutate(ID=paste(`#chr`, `pos(1-based)`,ref, alt, sep=":"))

### combine scores on interim dataset:
interim_dataset<-variants %>% 
  filter(clinvar_interim_test==TRUE)
gnomad_train_dataset<-variants %>% 
  filter(gnomad_train==TRUE)

validation_dataset<-validation_dataset %>%
  mutate(ID=paste(`#chr`, `pos(1-based)`,ref, alt, sep=":"))%>%
  filter(!ID %in% interim_dataset$ID)%>%
  filter(!ID %in% gnomad_train_dataset$ID)


set_of_models<-fit_set_of_models(interim_dataset)
validation_dataset<-predict_set_of_models(set_of_models, validation_dataset)


selected_scores<-c("HRAS_DMS_g12v",
                   "P53_DMS_WT_Nutlin",
                   "ADRB2_DMS_0.625",          
                   "BRCA1_DMS_E3_high_(a)",    
                   "BRCA1_DMS_(b)",            
                   "MSH2_DMS",      
                   "TPMT_DMS",                 
                   "PTEN_DMS_(a)",             
                   "PTEN_DMS_highqual_(b)",    
                   "SUMO1_DMS_flipped",        
                   "UBE2I_DMS_flipped",       
                   "VKOR1_DMS_abundance_score",
                   "TPK1_DMS_flipped")


validation_dataset<-validation_dataset%>%
  mutate(scoreID=paste(gene_dms, DMS, sep="_"))%>%
  filter(scoreID %in% selected_scores)



###### validation Dataset

prot_info_for_join<-validation_dataset %>% 
  distinct(Uniprot_acc_split, protein_length, protein_mean_b_factor)

SCORES<-c("Alph_null","AlphScore", "CADD_raw", "glm_AlphCadd", 
          "REVEL_score", "glm_AlphRevel", "glm_RevelCadd", 
          "DEOGEN2_score_med", "glm_AlphDeogen", "glm_CaddDeogen", 
          "glm_AlphCaddDeogen", "glm_DeogenRevel","glm_AlphDeogenRevel",
          "glm_AlphCaddRevel",    "glm_CaddDeogenRevel")

get_correlations<-function(val_dat){


  spearman_joined_private<-tibble()

for (un_ID in unique(val_dat$Uniprot_acc_split)){
  temp_values_joined<-val_dat %>% 
    filter(Uniprot_acc_split==un_ID)
  temp_values_joined <- temp_values_joined[complete.cases(temp_values_joined[, SCORES]),]
  for (dms in unique(temp_values_joined$DMS)){
    temp_inner_loop_values_joined<-temp_values_joined %>% 
      filter(DMS==dms)
    for (score in SCORES){
     score_DMS_tibble<-tibble(score_val=unlist(temp_inner_loop_values_joined %>% select(one_of(score))),
                              DMS_val=temp_inner_loop_values_joined$DMS_val)
     score_DMS_tibble<-score_DMS_tibble %>% filter(complete.cases(.))
      spearmans<-tibble(
      gene=unique(temp_inner_loop_values_joined$gene_dms), 
      UP_ID=un_ID, 
      DMS=dms, 
      spearm=ifelse(nrow(score_DMS_tibble)>2,
      cor.test(score_DMS_tibble$DMS_val, score_DMS_tibble$score_val,
               method = "spearman", na.action="na.ommit")$estimate,NA),
      method=score,
      number_of_vals=nrow(score_DMS_tibble))
      spearman_joined_private=rbind(spearman_joined_private,spearmans)
      }
  }
}

bad_UP_ids<-(spearman_joined_private %>% 
               filter(number_of_vals<21))$UP_ID

spearman_joined_private<-spearman_joined_private %>% 
  mutate(scoreID=paste(gene, DMS, sep="_"))%>%
  mutate(method=factor(method, levels= SCORES))%>%
  filter(! UP_ID %in% bad_UP_ids)%>%
  left_join(prot_info_for_join, by=c("UP_ID"="Uniprot_acc_split"))%>%
  mutate(abs_spear=abs(spearm)) %>%
  mutate(alphInclude=grepl("Alph", method))

return(spearman_joined_private)
}

spearmans_joined<-get_correlations(validation_dataset)

scatter_dms<-ggplot(data=validation_dataset)+
  geom_point(aes(x=DMS_val, y=AlphScore), color="darkblue", alpha=0.1)+
  geom_text(data=spearmans_joined %>% filter(method=="AlphScore"), aes(label=round(abs_spear, 3)),x = Inf, y = Inf, hjust = 1, vjust = 1)+
  labs(x="DMS score")+
  facet_wrap(~ scoreID, scales = "free", ncol = 3)+
  theme_minimal()

ggsave(filename = "scatter_dms_alphscore.png",
       plot = scatter_dms,
       height = 10,
       width=8, 
       dpi=300)

scatter_dms<-ggplot(data=validation_dataset)+
  geom_point(aes(x=DMS_val, y=glm_AlphDeogenRevel), color="darkblue", alpha=0.1)+
  geom_text(data=spearmans_joined %>% filter(method=="glm_AlphDeogenRevel"), aes(label=round(abs_spear, 3)),x = Inf, y = Inf, hjust = 1, vjust = 1)+
  labs(y="AlphScore + DEOGEN2 + REVEL", x="DMS score")+
  facet_wrap(~ scoreID, scales = "free", ncol = 3)+
  theme_minimal()

ggsave(filename = "scatter_dms_alph_doegen_revel.png",
       plot = scatter_dms,
       height = 10,
       width=8, 
       dpi=300)



pairwise.t.test(x = spearmans_joined$abs_spear, g = spearmans_joined$method, paired = TRUE, 
                p.adjust.method = "none", alternative = "greater")
# has been performed first; 
spearmans_joined %>% 
  group_by(gene, UP_ID, DMS) %>% 
  summarise(mean_abs_spear=mean(abs(spearm)))
## selected the scores with the highest overall correlation in mean_corr_per_score

spearmans_joined<-spearmans_joined%>%
  mutate(method=factor(method, levels=COL_ORDER))

plot_spearmans<-ggplot(spearmans_joined, aes(x=method, y=abs(spearm)))+
  stat_summary(fun.y = mean, geom = "bar", width=0.8, color="black", fill="grey") + 
  stat_summary(fun.data = mean_se, geom = "errorbar", width=0.4)+
  geom_jitter(width=0.15, color="blue4", alpha=1, size=1)+
  theme_minimal()+ 
  theme(axis.text.x = element_text(angle = 45, vjust = 1.0, hjust=1, color="black", size=10),
                          axis.text.y = element_text(color="black", size=10))+
  labs(x = "")+
  labs(y = "absolute Spearman correlation", size=12)

plot_spearmans
ggsave(filename= "spearman_plot.pdf", plot=plot_spearmans, height=5, width=4)

getmeansSD_spearman<-function(spearmans_joined_private){
  table_spearm_private<-spearmans_joined_private %>% 
    mutate(abs_spearman=abs(spearm))%>%
    group_by(method)%>%
    summarise(mean_spearm=mean(abs_spearman), sd=sd(abs_spearman), num=n(), sem=sd(abs_spearman)/sqrt(n())) %>%
    arrange(-mean_spearm)
  return(table_spearm_private)
}

table_spearm<-getmeansSD_spearman(spearmans_joined)
print(table_spearm)

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

getMeanAbsCorr<-function(data_priv, i){
  spearman_joined_priv<-get_correlations(data_priv[i,])
  table_spearm_private<-getmeansSD_spearman(spearman_joined_priv)
  return(unlist(table_spearm_private[,"mean_spearm"]))
}

getMeanAbsCorr_scoreids<-function(data_priv, i){
  filtered_data<-tibble()

  for (score_id in unlist(data_priv[i,]) ){
  filtered_data<-rbind(filtered_data,
    validation_dataset %>%
      filter(scoreID %in% score_id) )
  }

  spearman_joined_priv<-get_correlations(filtered_data)
  table_spearm_private<-getmeansSD_spearman(spearman_joined_priv)
  table_spearm_private$used_scores<-as.numeric(paste(i, collapse = ""))
  # for check print
  #print(data_priv[i,])
  #print(nrow(filtered_data))

  return(table_spearm_private)
}

getMeanAbsCorr_scoreids_return_vector<-function(data_priv, i){
  table_spearm_private<-getMeanAbsCorr_scoreids(data_priv, i)
  correlations_used_scores<-c(unlist(table_spearm_private[,"mean_spearm"]),
                              unlist(table_spearm_private[1,"used_scores"]))

  return(correlations_used_scores)
}

compare_scores<-function(booted_values_private){
  booted_values_private_summarised<-booted_values_private %>%
    rowwise%>%
    select(-starts_with("used_cols_vector"))%>%
    mutate(rowmax=max(across()))%>%
    mutate(glm_AlphDeogenRevel_biggest=glm_AlphDeogenRevel==rowmax)%>%
    mutate(glm_AlphCadd_biggerThanCadd=glm_AlphCadd>CADD_raw)%>%
    mutate(glm_AlphDeogen_biggerThanDeogen=glm_AlphDeogen>DEOGEN2_score_med)%>%
    mutate(glm_AlphRevel_biggerThanRevel=glm_AlphRevel>REVEL_score)
  return(booted_values_private_summarised)
}

get_p_value_table<-function(booted_values_summarised_private){
  relevant_scores<-c("glm_AlphDeogenRevel_biggest","glm_AlphDeogen_biggerThanDeogen","glm_AlphRevel_biggerThanRevel","glm_AlphCadd_biggerThanCadd")
  p_val_tabl<-tibble()
  for (columname in relevant_scores){
    p_val_tabl<- rbind(p_val_tabl, 
                       tibble(name=columname, num_true=sum(unlist(booted_values_summarised_private[,columname])), total_num=nrow(booted_values_summarised_private)))
  }
  return(p_val_tabl)
}



# do bootstrapping for means
mean_original_mat<-getmeansSD_spearman(get_correlations(validation_dataset))
set.seed(1)

boot_meanAbsCor<-boot(data=validation_dataset, statistic=getMeanAbsCorr, R=BOOT_REPETITIONS, parallel="multicore", ncpus=N_CPU)
booted_meanAbsCor_values<-as_tibble(boot_meanAbsCor$t)
colnames(booted_meanAbsCor_values)<-mean_original_mat$method
booted_meanAbsCor_values<-compare_scores(booted_meanAbsCor_values)
pValTable_AbsCor<-get_p_value_table(booted_meanAbsCor_values)
pValTable_AbsCor<-pValTable_AbsCor%>% 
  mutate(pval=1-num_true/total_num)
write_tsv(x=pValTable_AbsCor, "pValTable_AbsCor.tsv")


# bootstrap the choice of scores
score_ids<-as_tibble(unique(validation_dataset$scoreID))
mean_original_mat<-getMeanAbsCorr_scoreids(score_ids,1:13)
set.seed(1)
boot_meanAbsCor<-boot(data=score_ids, statistic=getMeanAbsCorr_scoreids_return_vector, R=BOOT_REPETITIONS, parallel="multicore", ncpus=N_CPU)
booted_meanAbsCor_values<-as_tibble(boot_meanAbsCor$t)
colnames(booted_meanAbsCor_values)<-c(as.character(mean_original_mat$method), "used_cols_vector")
booted_meanAbsCor_values<-compare_scores(booted_meanAbsCor_values)
pValTable_AbsCor<-get_p_value_table(booted_meanAbsCor_values)
pValTable_AbsCor<-pValTable_AbsCor%>% 
  mutate(pval=1-num_true/total_num)
write_tsv(x=pValTable_AbsCor, "pValTable_AbsCor_choice_of_experiments.tsv")
pValTable_AbsCor

rank_score_spearm<-function(spearm_joined_private){
  rank_scored_priv<-spearm_joined_private%>%group_by(scoreID)%>%
    mutate(minCor=min(abs_spear), maxCor=max(abs_spear))%>%
    ungroup()%>%
    mutate(norm_cor=(abs_spear-minCor)/(maxCor-minCor))%>%
    group_by(method,gene)%>%
    summarise(gene_level_cor=mean(norm_cor))%>%
    group_by(method)%>%
    summarise(rank_score=sum(gene_level_cor)/n() , sd_score=sd(gene_level_cor)) %>%
    mutate(sd_score_norm=sd_score/rank_score)%>%
    arrange(-method)
  return(rank_scored_priv)
}

getSpearmanRank<-function(data_priv, i){
  spearman_joined_priv<-get_correlations(data_priv[i,])# %>%
    #filter(method!="Alph_null")

  rank_scored<-rank_score_spearm(spearman_joined_priv)
  return(unlist(rank_scored[,2]))
}


# do bootstrapping for rankscores
rank_scored_original_mat<-rank_score_spearm(get_correlations(validation_dataset))
set.seed(1)
boot_rank_rankscore<-boot(data=validation_dataset, statistic=getSpearmanRank, R=BOOT_REPETITIONS, parallel="multicore", ncpus=N_CPU) # parallel="multicore", ncpus=3
booted_rankscores_values<-as_tibble(boot_rank_rankscore$t)
colnames(booted_rankscores_values)<-rank_scored_original_mat$method
booted_rankscores_values<-compare_scores(booted_rankscores_values)
vals<-booted_rankscores_values %>% pivot_longer(cols=1:length(unique(rank_scored_original_mat$method)))

vals<-vals%>%
  mutate(name=factor(name, levels=COL_ORDER))



rankscore_plot<-ggplot(vals)+
  geom_boxplot(aes(x=name,y=value))+
  theme_minimal()+
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1, color="black", size=10))

ggsave(filename = "rankscore_plot.pdf", plot = rankscore_plot, width=5, height=5.5)

pValTable_RankScore<-get_p_value_table(booted_rankscores_values)
pValTable_RankScore<-pValTable_RankScore%>% 
  mutate(pval=1-num_true/total_num)
pValTable_RankScore
write_tsv(x=pValTable_RankScore, "pValTable_RankScore.tsv")


ggplot(spearmans_joined, aes(x=scoreID, y=abs_spear, color=method, label=method))+
  geom_jitter(position = position_jitter(seed = 1))+
  #geom_text(position = position_jitter(seed = 1))+
  theme_minimal()+
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1, color="black", size=10),
        axis.text.y = element_text(color="black", size=10))

director_combined_files=$1
out_dir=$2
first_file="$director_combined_files/"$(ls $director_combined_files | head -n1)

echo "using header of: $first_file"

mkdir -p $out_dir
mkdir -p $out_dir"/tmp"
zcat $first_file | head -n1 > $out_dir"/header.csv"
zcat_sting=""

# loop over directory content and extract collect non empty files with TRUE in filename 
for file in $(ls $director_combined_files)
do
if [[ $file == *"TRUE"* ]]; then
f_size=$(stat --printf="%s" $director_combined_files"/"$file)
if [ $f_size -gt 100 ]
then
zcat_string+=" ${director_combined_files}/${file}"
fi
fi
done

#relevant_cols=$(zcat results/validation_set/validation_set_w_AlphScore.csv.gz | awk -v RS=',' '/^b_factor$/{print NR} /^SOLVENT_ACCESSIBILITY_core$/{print NR} /^Uniprot_acc_split$/{print NR} /^gnomad_train$/{print NR} /AlphScore/{print NR; exit}' | tr "\n" ",")"2,3,4,5,6,7,8,9,10"


echo $zcat_string

cat $out_dir"/header.csv" > $out_dir"/header_short.csv"

zcat $zcat_string  | grep -v "chr,pos(1-based)," | sort -V -t, -T $out_dir"/tmp/" -k1 -k2 - | \
cat $out_dir"/header_short.csv" - | tr "," "\t" | bgzip -c > $out_dir"/all_possible_values_concat.csv.gz"

tabix -s 1 -b 2 -e 2 $out_dir"/all_possible_values_concat.csv.gz"

rm -rf $out_dir"/tmp"

import sys
import pandas as pd
import numpy as np
from operator import itemgetter


### for debuging
class testclass:
	def __init__(self):
		self.wildcards=["AF-A0A087WUL8-F11-model_v1"]
		self.output=["results/AF-A0A087WUL8-F11-model_v1__combined_features.csv.gz"]
		self.params=["results/pdb_features/"]

#snakemake=testclass()
offset=0
offset=( int(snakemake.wildcards[0].split("-")[2][1:] ) - 1) * 200

# generate a list with the number of interactions
possible_interactions=["aroaro", "arosul", "cationpi", "disulphide", "hbond_main_side", "hbond_side_side", "hydrophobic", "ionic", "within_radius"]
interaction_matrix=np.zeros((40000,len(possible_interactions)),np.int) # 40000 - longer than Titin
i=0
edges=[]
for interaction in possible_interactions:
	interactions_pd=pd.read_csv(snakemake.params[0] + "/result_" + interaction + ".csv", low_memory=False)
	interactions_np=interactions_pd.to_numpy()
	interactions_np[:,1] = interactions_np[:,1] + offset
	interactions_np[:,3] = interactions_np[:,3] + offset
	interactions_joined=np.append(interactions_np[:,1], interactions_np[:,3])
	interactions_np[:,0]=np.array([s.strip() for s in interactions_np[:,0]])
	interactions_np[:,2]=np.array([s.strip() for s in interactions_np[:,2]])
	for ele in interactions_joined:
		interaction_matrix[ele,i]=interaction_matrix[ele,i]+1
	i=i+1


# Read features from FEATURIZE Framework
header=pd.read_csv(snakemake.params[1],skipinitialspace=True).columns.to_list()

to_remove=["RAUTE","coord1", "coord2", "coord3", "RAUTE2", "SECONDARY_STRUCTURE2_IS_UNKNOWN", "SECONDARY_STRUCTURE1_IS_UNKNOWN", "RESIDUE_CLASS2_IS_UNKNOWN","RESIDUE_CLASS1_IS_UNKNOWN", "RESIDUE_NAME_IS_OTHER", "ATOM_TYPE_IS_OTHER", "ELEMENT_IS_OTHER", "Atom_name"]

for suff in ["_core", "_3A", "_6A", "_9A", "_12A"]:
	header_surround=[ el + suff  if el not in ["RESN_RESI_AT"] + to_remove else el for el in header] 
	features=pd.read_csv(snakemake.params[0] + "/" + snakemake.wildcards[0] + suff + ".txt", delimiter="\t", names=header_surround,skipinitialspace=True , low_memory=False)
	features["RESN_RESI"]=features.RESN_RESI_AT.str.split(":",expand=True).iloc[:,0]
	features["RESN"]=features.RESN_RESI.str[:3]
	features["RESI"]=features.RESN_RESI.str[3:]
	features=features.astype({"RESI": int})
	features.RESI=features.RESI + offset
	features["RESN_RESI"]=features["RESN"]+features["RESI"].astype(str)
	features.index=features.RESI
	features=features.drop(to_remove, axis=1)
	print(suff)
	if suff!="_core":
		features_joined=features_joined.join(features.drop(['RESN_RESI_AT', 'RESN_RESI', 'RESI','RESN'], axis=1))
	else:
		features_joined=features

# join interactions and FEATURE based values
interactions_for_join=pd.DataFrame(interaction_matrix, columns=possible_interactions) # offset already inserted
interactions_for_join=interactions_for_join.astype(int)
interactions_for_join.index.name="RESI" 
features_joined2=features_joined.join(interactions_for_join)

# join the bfactors
b_factors=pd.read_csv(snakemake.params[0] + "/b_factors.csv", delimiter=" ", names=["residue_number","atom_name","b_factor"], skipinitialspace=True, low_memory=False)
b_factors["residue_number"] = pd.to_numeric(b_factors["residue_number"]) + offset
b_factors.index=b_factors["residue_number"] ########### +
features_joined3=features_joined2.join(b_factors)

# join the HSEs
HSE=pd.read_csv(snakemake.params[0] + "/HSEs.csv", delimiter=",", skipinitialspace=True, low_memory=False)
print(HSE.columns)
HSE["RESI"] = pd.to_numeric(HSE["RESI"]) + offset
HSE.index=HSE["RESI"]
HSE=HSE.drop("RESI", axis=1) ########### +
features_joined4=features_joined3.join(HSE)


features_joined4[["pdb_file"]]=snakemake.wildcards[0]
features_joined4.to_csv(snakemake.output[0])

library(tidyverse)
library(readxl)
library(data.table)
library(optparse)

option_list = list(
  make_option(c("-m", "--mapping"), type="character", default="resources/mapping_genename_uniprot.tsv", 
              help="location of file that maps gene name in assay to UniProt ID"),
  make_option(c("-p", "--prot_folder"), type="character", default="results/predicted_prots_eval", 
              help="Folder with protein files with AlphScore predictions"),
  make_option(c("-a", "--suffix"), type="character", default="_w_AlphScore_red_FALSE.csv.gz", 
              help="suffix that protein files with AlphScore predictions have")
)
opt = parse_args(OptionParser(option_list=option_list))

list_uniprot_ids<-read_tsv(opt$mapping)
PREDICTED_PROT_SUFFIX=opt$suffix
folder_of_prots=opt$prot_folder

# MAPK1	/ Uniprot P28482 not in mapping file, as this was not mapped between Uniprot and dbNSFP

sheets<-c("ADRB2","BRCA1","CALM1", "HRAS", "MAPK1", "P53",  "PTEN", "SUMO1", "TPK1", "TPMT", "UBE2I")
scores_conc<-tibble()

for (gene in sheets){
  scores <- read_excel("resources/Source.xlsx", sheet=gene) %>% 
    select(variant, starts_with("DMS"))%>%
    mutate(gene=gene)%>%
    gather(key="DMS", value="DMS_val", -variant, -gene)
  scores_conc=rbind(scores_conc, scores)
}

#MSH2
scores <- read_excel("resources/MSH2_Jia_2020/1-s2.0-S0002929720304390-mmc2.xlsx", sheet="TableS5_AllVariants")%>%
  rename(variant=Variant, DMS=`LOF score`)%>%
  select(variant, starts_with("DMS"))%>%
  mutate(gene="MSH2")%>%
  gather(key="DMS", value="DMS_val", -variant, -gene)
scores_conc=rbind(scores_conc, scores)

#Abeta
scores <- read_excel("resources/Abeta_Seuma_2020/elife-63364-supp4-v2.xlsx", sheet="1 aa change")%>%
  mutate(variant=paste0(WT_AA, Pos, Mut))%>%
    rename( DMS_nscore=`nscore`)%>%
    select(variant, starts_with("DMS"))%>%
    mutate(gene="Abeta")%>%
    gather(key="DMS", value="DMS_val", -variant, -gene)
scores_conc=rbind(scores_conc, scores)


#VKOR1
scores <- read_excel("resources/VKOR1_Chiasson_2020_activity/elife-58026-fig1-data1-v1.xlsx")%>%
  rename( DMS_abundance_score=`abundance_score`)%>%
  rename( DMS_activity_score=`activity_score`)%>%
  select(variant, starts_with("DMS"))%>%
  mutate(gene="VKOR1")%>%
  gather(key="DMS", value="DMS_val", -variant, -gene)
scores_conc=rbind(scores_conc, scores)

scores_conc<-scores_conc %>% 
  filter(!is.na(DMS_val))%>%
  left_join(list_uniprot_ids, by=c("gene"="gene_name"))%>%
  rename(gene_dms=gene)%>% 
  filter(!is.na(uniprot_id))


files<-paste0(folder_of_prots, list_uniprot_ids$uniprot_id, PREDICTED_PROT_SUFFIX) 

values_joined<-lapply(files, function(x) { #  mc.cores = 1,
  return(fread(x, na.strings = c("NA","."))  ) } )

values_joined<-rbindlist(values_joined)

values_joined<-values_joined%>%
  mutate(variant=paste0(
str_replace_all(RESN_RESI, c( 
 "ALA"="A", "ARG"="R", "ASN"="N", "ASP"="D",
 "CYS"="C", "GLU"="E", "GLN"="Q", "GLY"="G",
 "HIS"="H", "ILE"="I", "LEU"="L", "LYS"="K",
 "MET"="M", "PHE"="F", "PRO"="P", "SER"="S",
 "THR"="T", "TRP"="W", "TYR"="Y", "VAL"="V")),
str_replace_all(to_AS, c( 
  "ALA"="A", "ARG"="R", "ASN"="N", "ASP"="D",
  "CYS"="C", "GLU"="E", "GLN"="Q", "GLY"="G",
  "HIS"="H", "ILE"="I", "LEU"="L", "LYS"="K",
  "MET"="M", "PHE"="F", "PRO"="P", "SER"="S",
  "THR"="T", "TRP"="W", "TYR"="Y", "VAL"="V"))
))

values_joined<-values_joined %>% 
  left_join(scores_conc, by=c("variant"="variant", "Uniprot_acc_split"="uniprot_id"))%>%
  filter(!is.na(DMS_val))%>%
  mutate(DMS_val=as.double(DMS_val))

dir.create("results/validation_set")
setwd("results/validation_set")

fwrite(x=values_joined, file="validation_set_w_AlphScore.csv.gz")

library(readr)
library(dplyr)
library(optparse)
library(stringr)
source("workflow/scripts/existing_scores_glm_functions.R")
set.seed(1)

option_list = list(
  make_option(c("-t", "--training_dataset"), type="character", default="results/prediction_final/pre_final_model_regular_variants.csv.gz", 
              help="File containing the training dataset"),
  make_option(c("-c", "--combined_model"), type="character", default="results/prediction_final/combined_model.RData", 
              help="File where the final combined logistic regression model should be saved to"),
  make_option(c("-i", "--training_var_ids"), type="character", default="results/prediction_final/training_var_ids.tsv", 
              help="IDs of variants that were in the training / test set")
)

opt = parse_args(OptionParser(option_list=option_list))

variants<-read_csv(opt$training_dataset, na=c(".","NA", NA))
variants$AlphScore<-variants$predicted_Alph

index_col<-get_index_col(variants)

variants$DEOGEN2_score_med<- unlist_score(variants$DEOGEN2_score, index_col)

gnomad_vars<-variants %>% 
  filter(gnomadSet==1)

clinvar_vars_not_gnomad <-variants%>% 
  filter(gnomadSet==0) %>%
  filter(! var_id_genomic %in% gnomad_vars$var_id_genomic )

set_of_models<-fit_set_of_models(clinvar_vars_not_gnomad)

gnomad_set_position<-gnomad_vars$var_id_genomic
clinvar_set_position<-clinvar_vars_not_gnomad$var_id_genomic

training_vars<-rbind(
tibble(var_id_genomic=gnomad_set_position,
       gnomadSet=1),
tibble(var_id_genomic=clinvar_set_position,
       gnomadSet=0)
)
# write to disk
write_tsv(x=training_vars,
          file=opt$training_var_ids)

write_rds(x=set_of_models,
          file=opt$combined_model)