import pandas as pd
import os
import numpy as np

# configs
result_paths = snakemake.input['enrichment_results']

results_all_path = snakemake.output['results_all']
results_sig_path = snakemake.output['results_sig']

group = snakemake.wildcards["group"] #"testgroup"
tool = snakemake.wildcards["tool"] #"ORA_GSEApy"
db = snakemake.wildcards["db"] #"WikiPathways_2019_Human"

term_col = snakemake.config["column_names"][tool]["term"] #'Term'
adjp_col = snakemake.config["column_names"][tool]["adj_pvalue"] #'Adjusted_P_value'
# effect_col = snakemake.config["column_names"][tool]["effect_size"] #'Odds_Ratio'

adjp_th = snakemake.config["adjp_th"][tool] #0.05

dir_results = os.path.dirname(results_all_path)
if not os.path.exists(dir_results):
    os.mkdir(dir_results)

# load results
results_list = list()
for result_path in result_paths:
    if os.stat(result_path).st_size != 0:
        tmp_name = os.path.basename(result_path).replace("_{}.csv".format(db),"")
        tmp_res = pd.read_csv(result_path, index_col=0)
        tmp_res['name'] = tmp_name
        results_list.append(tmp_res)


# move on if results are empty
if len(results_list)==0:
    open(results_all_path, mode='a').close()
    open(results_sig_path, mode='a').close()
    sys.exit(0)

# concatenate all results into one results dataframe
result_df = pd.concat(results_list, axis=0)

# save all enirchment results
result_df.to_csv(results_all_path)

# find union of statistically significant terms
sig_terms = result_df.loc[result_df[adjp_col]<adjp_th, term_col].unique()

# filter by significant terms
result_sig_df = result_df.loc[result_df[term_col].isin(sig_terms), :]

# save filtered enirchment results by significance
result_sig_df.to_csv(results_sig_path)

library(ggplot2)
library(svglite)

# source utility functions
# source("workflow/scripts/utils.R")
snakemake@source("./utils.R")

# configs
enrichment_result_path <- snakemake@input[["enrichment_result"]]
enrichment_plot_path <- snakemake@output[["enrichment_plot"]]

tool <- snakemake@params[["tool"]] #"ORA_GSEApy"
database <- snakemake@params[["database"]] #"KEGG_2021_Human"
feature_set <- snakemake@params[["feature_set"]] 
plot_cols <- snakemake@config[["column_names"]][[tool]]

top_n <- plot_cols[["top_n"]] #25
pval_col <- plot_cols[["p_value"]] #'P.value'
adjp_col <- plot_cols[["adj_pvalue"]] #'Adjusted.P.value'
effect_col <- plot_cols[["effect_size"]] #'Odds.Ratio'
overlap_col <- plot_cols[["overlap"]] #'Overlap'
term_col <- plot_cols[["term"]] #'Term'

# load enrichment result
if (file.size(enrichment_result_path) != 0L){
    enrichment_result <- read.csv(enrichment_result_path, row.names = NULL, header= TRUE)
}else{
    file.create(enrichment_plot_path)
    quit()
}

# evaluate overlap numerically if necessary
if(class(enrichment_result[[overlap_col]])=="character"){
    enrichment_result[[overlap_col]] <- as.numeric(lapply(enrichment_result[[overlap_col]], evaltext))
}

# calculate comparable effect size either NES or odds-ratio/fold based
if (tool!="preranked_GSEApy"){
    # calculate log2(effect-size) and put in new column
    effect_col_new <- paste0("log2_",effect_col)
    enrichment_result[[effect_col_new]] <- log2(enrichment_result[[effect_col]])
    effect_col <- effect_col_new
}

# determine ranks
enrichment_result$PValue_Rnk <- rank(enrichment_result[[pval_col]])
enrichment_result$Fold_Rnk <- rank(-abs(enrichment_result[[effect_col]]))
enrichment_result$Coverage_Rnk <- rank(-enrichment_result[[overlap_col]])
# calculate and sort by mean rank
enrichment_result$meanRnk <- rowMeans(enrichment_result[,c('PValue_Rnk', 'Fold_Rnk','Coverage_Rnk')])
enrichment_result <- enrichment_result[order(enrichment_result$meanRnk, decreasing=FALSE),]

# format term column that order is kept and values are unique
enrichment_result[[term_col]] <- make.names(enrichment_result[[term_col]], unique=TRUE)
enrichment_result[[term_col]] <- factor(enrichment_result[[term_col]], levels = enrichment_result[[term_col]])

# plot top_n terms by mean_rnk
do_enrichment_plot(plot_data=enrichment_result[1:top_n,], 
               title=paste0(tool, ' results of \n',feature_set,' in ',database), 
               x=effect_col, 
               y=term_col, 
               size=overlap_col, 
               colorBy=adjp_col, 
               font.size=10, 
               path=file.path(dirname(enrichment_plot_path)), 
               filename=paste0(feature_set,"_",database),
                   top_n = top_n
              )

import pandas as pd
import json
# import pickle
import os
import numpy as np
import gseapy as gp
import sys


# utils for manual odds ratio calculation -> not used anymore
def overlap_converter(overlap_str, bg_n, gene_list_n):
    overlap_n, gene_set_n = str(overlap_str).split('/')
    return odds_ratio_calc(bg_n, gene_list_n, int(gene_set_n), int(overlap_n))

def odds_ratio_calc(bg_n, gene_list_n, gene_set_n, overlap_n): 
    import scipy.stats as stats
    table=np.array([[gene_set_n, bg_n-gene_set_n],[overlap_n, gene_list_n-overlap_n]])
    oddsratio, pvalue = stats.fisher_exact(table)
    return (1/oddsratio)


# configs

# input
query_genes_path = snakemake.input['query_genes']
background_genes_path = snakemake.input['background_genes']
database_path = snakemake.input['database']

# output
result_path = snakemake.output['result_file']

# parameters
db = snakemake.params["database"]

dir_results = os.path.dirname(result_path)

if not os.path.exists(dir_results):
    os.mkdir(dir_results)

# check if genes file exists & load or handle exception
if os.path.exists(query_genes_path):
    genes = open(query_genes_path, "r")
    gene_list = genes.read()
    gene_list = gene_list.split('\n')
    gene_list.remove('')
    genes.close()
else:
    with open(os.path.join(dir_results,"no_genes_found.txt"), 'w') as f:
        f.write('no genes found')
    quit()

# load background genes
bg_file = open(background_genes_path, "r")
background = bg_file.read()
background = background.split('\n')
background.remove('')
bg_file.close()

# move on if query-genes are empty
if len(gene_list)==0:
    open(result_path, mode='a').close()
    sys.exit(0)

# load database .pkl file
# with open(enrichr_databases, 'rb') as f:
#     db_dict = pickle.load(f)
with open(database_path) as json_file:
    db_dict = json.load(json_file)

# convert gene lists and database to upper case
gene_list=[str(x).upper() for x in list(gene_list)]
background=[str(x).upper() for x in list(background)]
db_dict = {key: [ele.upper() for ele in db_dict[key] ] for key in db_dict}

# count number of background genes for odds-ratio calculation
bg_n = len(background)

# if background-genes are empty provide number of genes as 20,000 as heuristic
# this excpetion only occurs if background region set exceeds 500,000 regions
if bg_n==0:
    bg_n = 20000
    background = bg_n

# perform ORA (hypergeometric test) in database using GSEApy (barplots are generated automatically)
try:
    res = gp.enrich(gene_list=gene_list,
                     gene_sets=db_dict,
                     background=background,
                     outdir=None,#os.path.join(dir_results),
                     top_term=25,
                     cutoff=0.05,
                     format='png',
                     verbose=True,
                    ).res2d
except ValueError:
    print("Result is empty")
    res = pd.DataFrame()

# move on if result is empty
if res.shape[0]==0:
    open(result_path, mode='a').close()
    sys.exit(0)

# annotate used gene set
res['Gene_set'] = db

# odds ratio calculation -> results still differ, but correlation is >0.99 and with large OR values the difference shows up at third decimal place
# gene_list_n=len(gene_list)
# res['Odds Ratio'] = res['Overlap'].apply(overlap_converter, args=(bg_n, gene_list_n))

# make column names language agnostic (i.e., R compatible)
column_names = list(res.columns.values)
column_names = [col.replace(" ","_") for col in column_names]
column_names = [col.replace("-","_") for col in column_names]
res.columns = column_names

# separate export
res.to_csv(result_path)

import pandas as pd
import json
# import pickle
import os
import numpy as np
import gseapy as gp
import sys


# configs

# input
query_genes_path = snakemake.input['query_genes']
database_path = snakemake.input['database']

# output
result_path = snakemake.output['result_file']

# parameters
db = snakemake.params["database"]

dir_results = os.path.dirname(result_path)

if not os.path.exists(dir_results):
    os.mkdir(dir_results)


# load gene-score file
genes = pd.read_csv(query_genes_path, index_col=0)

# load database JSON file
with open(database_path) as json_file:
    db_dict = json.load(json_file)

# convert all genes to upper case
genes.index = [str(x).upper() for x in list(genes.index)]
db_dict = {key: [ele.upper() for ele in db_dict[key] ] for key in db_dict}

# remove duplicates: only retain largest absolute value
# sort by absolute value of "score" column (i.e., first column)
genes = genes.iloc[abs(genes.iloc[:, 0]).argsort()[::-1]]
# drop duplicates based on index
genes = genes[~genes.index.duplicated(keep='first')]

# run prerank GSEA of database with GSEApy
res = gp.prerank(rnk=genes, 
                 gene_sets=db_dict,
                 #threads=4,
                 min_size=1, # Minimum allowed number of genes from gene set also the data set. Default: 15.
                 max_size=100000, # Maximum allowed number of genes from gene set also the data set. Defaults: 500.
                 permutation_num=1000, # Number of permutations. Reduce number to speed up testing;  Default: 1000. Minimial possible nominal p-value is about 1/nperm.
                 outdir=os.path.join(dir_results),
                 graph_num = 25, # Plot graphs for top sets of each phenotype.
                 format='png',
                 seed=42,
                 verbose=True,
                ).res2d

# move on if result is empty
if res.shape[0]==0:
    open(result_path, mode='a').close()
    sys.exit(0)

# annotate used gene set
res['Gene_set'] = db

# make column names language agnostic (i.e., R compatible)
column_names = list(res.columns.values)
column_names = [col.replace(" %","") for col in column_names]
column_names = [col.replace(" ","_") for col in column_names]
column_names = [col.replace("-","_") for col in column_names]
res.columns = column_names

# separate export
res.to_csv(result_path)

library(ggplot2)
library(svglite)
library(reshape2)
library(pheatmap)

# source utility functions
# source("workflow/scripts/utils.R")
snakemake@source("./utils.R")

# configs
results_all_path <- snakemake@input[["results_all"]]

plot_path <- snakemake@output[["summary_plot"]]
adjp_hm_path <- snakemake@output[["adjp_hm"]]
effect_hm_path <- snakemake@output[["effect_hm"]]

tool <- snakemake@wildcards[["tool"]] #"ORA_GSEApy"
database <- snakemake@wildcards[["db"]] #"GO_Biological_Process_2021"
group <- snakemake@wildcards[["group"]] #"testgroup"

term_col <- snakemake@config[["column_names"]][[tool]][["term"]] #'Term'
adjp_col <- snakemake@config[["column_names"]][[tool]][["adj_pvalue"]] #'Adjusted_P_value'
effect_col <- snakemake@config[["column_names"]][[tool]][["effect_size"]] #'Odds_Ratio'

top_n <- snakemake@config[["top_terms_n"]] #5
adjp_cap <- snakemake@config[["adjp_cap"]] #4
effect_cap <- if (tool=="preranked_GSEApy") snakemake@config[["nes_cap"]] else snakemake@config[["or_cap"]] #5

adjp_th <- as.numeric(snakemake@config[["adjp_th"]][tool]) #0.05

# stop early if results are empty
if(file.size(results_all_path) == 0L){
    file.create(plot_path)
    file.create(adjp_hm_path)
    file.create(effect_hm_path)
    quit()
}

# load aggregated result dataframe
results_all <- read.csv(results_all_path, header= TRUE)

# stop early if results consist of only one query
if(length(unique(results_all$name))==1){
    file.create(plot_path)
    file.create(adjp_hm_path)
    file.create(effect_hm_path)
    quit()
}

# determine top_n most significant terms (not necessarily statistically significant!)
top_terms <- c()
for (query in unique(results_all$name)){
    tmp_result <- results_all[results_all$name==query,]
    tmp_terms <- tmp_result[order(tmp_result[[adjp_col]]), term_col][1:top_n]
    top_terms <- unique(c(top_terms,tmp_terms))
}

# make adjusted p-vale and effect-size (odds-ratio or normalized enrichment scores) dataframes
adjp_df <- dcast(results_all, as.formula(paste(term_col, "~ name")), value.var = adjp_col)
rownames(adjp_df) <- adjp_df[[term_col]]
adjp_df[[term_col]] <- NULL

effect_df <- dcast(results_all, as.formula(paste(term_col, "~ name")), value.var = effect_col)
rownames(effect_df) <- effect_df[[term_col]]
effect_df[[term_col]] <- NULL

# filter by top terms
adjp_df <- adjp_df[top_terms,]
effect_df <- effect_df[top_terms,]

# fill NA for effect_df with 1 or 0 (i.e., neutral enrichment) and for adjp_df with 1 (i.e., no significance)
effect_df[is.na(effect_df)] <- if (tool=="preranked_GSEApy") 0 else 1
adjp_df[is.na(adjp_df)] <- 1

# make stat. sign. annotation for effect-size plot later
adjp_annot <- adjp_df
adjp_annot[adjp_df < adjp_th] <- "*"
adjp_annot[adjp_df >= adjp_th] <- ""

# log2 transform odds ratios
if (tool!="preranked_GSEApy"){
    effect_df <- log2(effect_df)
}

# cap effect_df for plotting depending on tool  abs(log2(or)) < or_cap OR abs(NES) < nes_cap
effect_df[effect_df > effect_cap] <- effect_cap
effect_df[effect_df < -effect_cap] <- -effect_cap

# log10 transform adjp & cap -log10(adjpvalue) < adjp_cap
adjp_df <- -log10(adjp_df)
adjp_df[adjp_df > adjp_cap] <- adjp_cap

# plot hierarchically clustered heatmap for adjp and effect
width_hm <- 0.2 * dim(adjp_df)[2] + 5
height_hm <- 0.2 * dim(adjp_df)[1] + 3

pheatmap(adjp_df,
         display_numbers=adjp_annot,
         main="-log10(adj. p-values)",
         treeheight_row = 10,
         treeheight_col = 10,
         fontsize = 6,
         fontsize_number = 10,
         silent=TRUE,
         width=width_hm,
         height=height_hm,
         angle_col=45, 
         cellwidth=10,
         cellheight=10,
         filename=adjp_hm_path,
         breaks=seq(0, max(adjp_df), length.out=200),
         color=colorRampPalette(c("white", "red"))(200)
        )

pheatmap(effect_df,
         display_numbers=adjp_annot,
         main = if (tool=="preranked_GSEApy") effect_col else paste0("log2(",effect_col,")"),
         treeheight_row = 10,
         treeheight_col = 10,
         fontsize = 6,
         fontsize_number = 10,
         silent=TRUE,
         width=width_hm,
         height=height_hm,
         angle_col=45, 
         cellwidth=10,
         cellheight=10,
         filename=effect_hm_path,
         breaks=seq(-max(abs(effect_df)), max(abs(effect_df)), length.out=200),
         color=colorRampPalette(c("blue", "white", "red"))(200)
        )


# perform hierarchical clustering on the effect-sizes (NES or log2 odds ratios) of the terms and reorder dataframe
hc_rows <- hclust(dist(effect_df))
hc_row_names <- rownames(effect_df)[hc_rows$order]
hc_cols <- hclust(dist(t(effect_df)))
hc_col_names <- colnames(effect_df)[hc_cols$order]
effect_df <- effect_df[hc_rows$order, hc_cols$order]

# add a column for the terms
effect_df$terms <- rownames(effect_df)
# melt data frame for plotting
plot_df <- melt(data=effect_df,
                id.vars="terms",
                measure.vars=colnames(adjp_df),
                variable.name = "feature_set",
               value.name = "effect")

# add adjusted p-values to plot dataframe
plot_df$adjp <- apply(plot_df, 1, function(x) adjp_df[x['terms'], x['feature_set']])

# set effect-size and adjusted p-value conditional to NA (odds ratios == 0 and NES == 0) for plotting
plot_df$effect[plot_df$effect==0] <- NA
plot_df$adjp[plot_df$adjp==0] <- NA
# if (tool=="preranked_GSEApy"){
#     plot_df$effect[plot_df$effect==0] <- NA
# }else{
#     plot_df$effect[plot_df$effect<=0] <- NA
# }

# ensure that the order of terms and feature sets is kept
plot_df$terms <- factor(plot_df$terms,levels=hc_row_names)
plot_df$feature_set <- factor(plot_df$feature_set, levels=hc_col_names)

# plot
enr_plot <- ggplot(plot_df, aes(x=feature_set, y=terms, fill=effect, size=adjp))+ 
geom_point(shape=21, stroke=0.25) +
geom_point(data = plot_df[plot_df$adjp > -log10(adjp_th),], aes(x=feature_set, y=terms), shape=8, size=0.5, color = "black", alpha = 0.5) + # stars for statistical significance
# scale_fill_gradient(low="grey", high="red", breaks = c(1, 2, 3, 4), limits = c(0, 4), name="-log10(adjp)") +
scale_fill_gradient2(midpoint=0, low="royalblue4", mid="white", high="firebrick2", space ="Lab", name = if (tool=="preranked_GSEApy") effect_col else paste0("log2(",effect_col,")")) +
scale_y_discrete(label=addline_format) + 
scale_size_continuous(range = c(1,5), , name = "-log10(adjp)") + #not needed, because data already capped? limits = c(0, adjp_cap)
ggtitle(paste(tool, database, group, sep='\n')) +
clean_theme() +
theme(axis.text.x = element_text(angle = 45, hjust = 1, vjust=1),
      axis.title.x=element_blank(),
      axis.title.y=element_blank()) + guides(size = guide_legend(reverse=TRUE))

width <- 0.15 * dim(adjp_df)[2] + 3
height <- 0.2 * dim(adjp_df)[1] + 2
# options(repr.plot.width=width, repr.plot.height=height)
# enr_plot

# save plot
ggsave_new(filename=paste0(group,'_',database,'_summary'),
               results_path=dirname(plot_path),
               plot=enr_plot,
               width=width,
               height=height
              )

library(LOLA)
library(GenomicRanges)
library(rGREAT)

# configs
regions_file <- snakemake@input[["regions"]]
background_file <- snakemake@input[["background"]]
gene_file <- snakemake@output[["genes"]]
result_paths <- snakemake@output[["results"]]

genome <- snakemake@config[["genome"]] #"hg38"
databases <- snakemake@config[["great_dbs"]]
region_set <- snakemake@params[["region_set"]]

dir_result <- dirname(gene_file)

# make result directory, if not exist
dir.create(file.path(dir_result), showWarnings = FALSE, recursive = TRUE)

# load query and background/universe region sets (eg consensus region set)
regionSet_query <- unname(readBed(regions_file))

# create empty results if query region set is too large for GREAT and quit i.e. >500,000 https://great-help.atlassian.net/wiki/spaces/GREAT/pages/655402/File+Size
if (length(regionSet_query)>500000){
    print("query regions set is too large (>500,000)")
    for (path in result_paths){
        file.create(path)
    }
    file.create(file.path(gene_file))
    quit()
}

# if query and background the same, do not use a background (used to map background regions to genes for downstream gene-based analyses)
if(regions_file!=background_file){
    regionSet_background <- unname(readBed(background_file))
    # check if background region set is not too large https://great-help.atlassian.net/wiki/spaces/GREAT/pages/655402/File+Size
    if (length(regionSet_background)>1000000){
        print("background regions set is too large (>1,000,000) and set to NULL i.e. whole genome")
        regionSet_background <- NULL
    }
}else{
    regionSet_background <- NULL
}


###### GREAT
job = submitGreatJob(regionSet_query, species = genome, bg=regionSet_background)

# save job description
capture.output(job, file=file.path(dir_result, 'job_description.txt'), append=TRUE)

# plot, get and save gene-region association
pdf(file=file.path(dir_result, paste0('region_gene_associations.pdf')), width=20, height=15)
res = plotRegionGeneAssociationGraphs(job, request_interval=600)
write.csv(res, file.path(dir_result, 'region_gene_associations.csv'), row.names=TRUE)
dev.off()

# save unique associated genes by using mcols(), which returns a DataFrame object containing the metadata columns.
genes <- unique(mcols(res)$gene)
write(genes, file.path(gene_file))
# write.table(genes, file.path(gene_file), sep='\t', row.names=FALSE, col.names=FALSE, quote = FALSE)

# get & save enrichment results
if(regions_file!=background_file){
    for (db in databases){
        tb <- getEnrichmentTables(job,  ontology = db, category=NULL, download_by = 'tsv')
        tb$Desc <- paste(tb$Desc, tb$ID)
    #     write.table(tb[[db]], file.path(dir_result, db, paste0(region_set,'_',db,'.tsv')), sep='\t', row.names=FALSE)
        write.csv(tb[[db]], file.path(dir_result, gsub(" ", "_", db), paste0(region_set,'_',gsub(" ", "_", db),'.csv')), row.names=FALSE)
    }
}else{
    for (path in result_paths){
        file.create(path)
    }
}