#!/usr/bin/R
# ----------------------------------------------------
# Stratifying cells by specific GO terms - diff by AD?
# Updated: 09/04/22
# ----------------------------------------------------
library(cbrbase)
set_proj('DEVTRAJ')

# Libraries:
library(tidyr)
library(ggplot2)
library(ggpubr)
library(ggrepel)
library(ggbeeswarm)
library(scales)
library(RColorBrewer)
library(Matrix)
options(width=170)

# Load in the two cohorts of candidate gene fusion counts:
# --------------------------------------------------------
gbdir = paste0(bindir, 'mosaicism/genefusions/')
source(paste0(gbdir, 'load_fusion_data.R'))
source(paste0(gbdir, 'auxiliary_fusion_plotting.R'))
source(paste0(gbdir, 'load_go_signatures.R'))
source(paste0(gbdir, 'load_joint_expression.R'))

source(paste0(bindir, 'multiRegion/auxiliary_plotting_settings.R'))

# Make directories for analysis:
# ------------------------------
analysis = 'fusion'
project = 'ALZ'
imgdir = paste0(img, project, '/', analysis , '/')
imgpref = paste0(imgdir, analysis, "_mspred_strat_", project, '_')
cmd = paste0('mkdir -p ', imgdir)
system(cmd)


# Load pre-computed relevant annotations:
# ---------------------------------------
filepref = paste0(dbdir, 'joint_cohort_annotations_scores')
enr.rds = paste0(filepref, '_broad.Rds')
fullenr.rds = paste0(filepref, '_specific.Rds')
enrdf = readRDS(enr.rds)
full.enrdf = readRDS(fullenr.rds)
kwds = names(siglist)
enrdf = merge(enrdf, unique(ext.meta[,c('projid', 'niareagansc')]))

enrdf$braaksc = as.numeric(as.character(enrdf$braaksc))
enrdf$bd = ifelse(enrdf$braaksc >= 5, 'AD', 'CTRL')
enrdf$nrad = ifelse(enrdf$niareagansc <= 2, 'AD', 'CTRL')
enrdf$bd = factor(enrdf$bd, levels=c('CTRL','AD'))
enrdf$nrad= factor(enrdf$nrad, levels=c('CTRL','AD'))

sigkwds = subkwds[5:length(subkwds)]  # Keep only GO sigs
subdf = enrdf[enrdf$tag %in% sigkwds,]
subdf = subdf[subdf$cohort == 'ALZ',]

# Are signatures different in original cohort:
advals = c('AD'='indianred', 'CTRL'='royalblue')
gp = ggplot(subdf, aes(relabel, enr, fill=nrad)) + 
    facet_wrap(~tag, nrow=2, scale='free_x') + 
    scale_fill_manual(values=advals) + 
    geom_boxplot(outlier.size=0.5) + 
    stat_compare_means(label='p.signif') + 
    labs(x=paste('Sub-celltype'), y='Norm. signature expression, log(X+1)') + 
    theme_pubr() + coord_flip() + 
    theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust=.5))
pltpref = paste0(imgpref, 'signatures_braak')
saveGGplot(gp, pltpref, w=11, h=6)