#!/usr/bin/R
# ------------------------------------------------
# Plots on aggregate statistics for fusions in 10x
# At the celltype / subtype level in this case
# --follows P301S, updated for human
# Updated: 01/13/23
# ------------------------------------------------
library(cbrbase)
library(ggplot2)
library(ggpubr)

set_proj('DEVTRAJ', 'mosaicism')
source(paste0(bindir, 'mosaicism/fusions_AD430/load_metadata.R'))
source(paste0(bindir, 'mosaicism/genefusions/auxiliary_fusion_plotting.R'))
source(paste0(bindir, 'multiRegion/auxiliary_plotting_settings.R'))


# Directories:
plotdir = paste0(imgdir, 'fusions/')
imgpref = paste0(plotdir, 'AD430_fusions_')
cmd = paste('mkdir -p', plotdir)
system(cmd)


# Load in the gene fusion data:
# -----------------------------
source(paste0(bindir, 'mosaicism/fusions_AD430/load_fusioncalls.R'))
source(paste0(bindir, 'mosaicism/genefusions/auxiliary_fusion_plotting.R'))


# Compute rates for the PFC alone
ctdf$ncell = 1
ctdf$n_counts[is.na(ctdf$n_counts)] = 1
aggdf = aggregate(cbind(count, n_counts, ncell) ~ dataset + projid + region + bd + nrad + celltype, ctdf, sum)
aggdf$rate = with(aggdf, count / n_counts)
aggdf$rate.cell = with(aggdf, count / ncell)
aggdf = aggdf[order(aggdf$rate),]
aggdf = aggdf[aggdf$region == 'PFC',]

aggdf = aggdf[order(aggdf$rate, decreasing=T),]
aggdf = merge(aggdf, unique(metadata[,c('projid','braaksc','niareagansc')]))


gp = ggplot(aggdf[aggdf$ncell > 25,], aes(celltype, rate, fill=bd)) + 
    scale_fill_manual(values=colvals$nrad, name='Late-AD?') + 
    geom_boxplot() + coord_flip() +
    labs(x='Celltype', y='Rate (counts / # reads (PFC) or # cells (HC))') +
    stat_compare_means(label='p.format') + 
    theme_pubr()
pltprefix = paste0(imgpref, 'boxplot_bd_rate_PFC_bycelltype')
saveGGplot(gp, pltprefix, w=5, h=5)


gp = ggplot(aggdf[aggdf$ncell > 25,], aes(celltype, rate, fill=nrad)) + 
    scale_fill_manual(values=colvals$nrad, name='NIA-Reagan diagnosis:') + 
    geom_boxplot() + coord_flip() +
    labs(x='Only excitatory neurons?', y='Rate (counts / # reads (PFC) or # cells (HC))') +
    stat_compare_means(label='p.format') + 
    theme_pubr()
pltprefix = paste0(imgpref, 'boxplot_nrad_rate_PFC_bycelltype')
saveGGplot(gp, pltprefix, w=5, h=5)


gp = ggplot(aggdf, aes(only.exc, rate, fill=factor(niareagansc))) + 
    scale_fill_manual(values=colvals$niareagansc, name='NIA-Reagan diagnosis:') + 
    geom_boxplot() + coord_flip() + 
    labs(x='Only excitatory neurons?', y='Rate (counts / # reads (PFC) or # cells (HC))') +
    stat_compare_means(label='p.format') + 
    theme_pubr()
pltprefix = paste0(imgpref, 'boxplot_niareagan_rate_PFC_bycelltype')
saveGGplot(gp, pltprefix, w=5, h=5)


gp = ggplot(aggdf, aes(only.exc, rate, fill=factor(braaksc))) + 
    scale_fill_manual(values=colvals$braaksc, name='Braak Stage:') + 
    geom_boxplot() + coord_flip() + 
    labs(x='Only excitatory neurons?', y='Rate (counts / # reads (PFC) or # cells (HC))') +
    stat_compare_means(label='p.format') + 
    theme_pubr()
pltprefix = paste0(imgpref, 'boxplot_braak_rate_PFC_bycelltype')
saveGGplot(gp, pltprefix, w=7, h=5)




# Simple binomial fit on these aggregated results:
# ------------------------------------------------
df = aggdf[aggdf$ncell > 25,]
fit = glm(cbind(count, n_counts) ~ bd * celltype, df, family='binomial')
cfit = processLmFit(fit)
cfit$col = ifelse(cfit$p < 0.05, ifelse(cfit$Est > 0, 'Up','Down'), 'NS')

# Plot the errorbars:
gp = coefErrorbarPlot(cfit, title='binomial, indlevel')

resdf = resdf[order(resdf$p),]
resdf$padj = p.adjust(resdf$p, method='fdr')


library(MASS)
fit = glm.nb(count ~ offset(log(n_counts)) + bd * celltype, df)
cfit = processLmFit(fit)
cfit$col = ifelse(cfit$p < 0.05, ifelse(cfit$Est > 0, 'Up','Down'), 'NS')
gp = coefErrorbarPlot(cfit, title='NB, indlevel')


library(lme4)
fit = glmer.nb(count ~ offset(log(n_counts)) + bd * celltype + (1|dataset), ctdf[ctdf$region == 'PFC',])
cfit = processLmFit(fit)
cfit$col = ifelse(cfit$p < 0.05, ifelse(cfit$Est > 0, 'Up','Down'), 'NS')
print(cfit)
gp = coefErrorbarPlot(cfit, title='NB, LMM')





aggdf = merge(aggdf, aggregate(cbind(nft_mf, plaq_n_mf, gpath) ~ projid + braaksc, metadata, mean))
gp = ggplot(aggdf, aes(plaq_n_mf, rate, color=bd)) + 
    facet_wrap(~ region, scales='free') + 
    scale_color_manual(values=colvals$nrad) + 
    geom_point(cex=.75) + 
    stat_cor(color='black') + 
    geom_smooth(method='lm', color='black') + 
    theme_pubr()
pltprefix = paste0(imgpref, 'scatter_plaq_n_mf_bd_rate_PFC_bycelltype')
saveGGplot(gp, pltprefix, w=7, h=5)