#!/usr/bin/R
# ------------------------------------------
# Plot basic stats for the CKP25 mouse data:
# ------------------------------------------
library(cbrbase)
set_proj('DEVTRAJ', 'mosaicism')

library(tidyr)
library(RColorBrewer)
library(viridis)
library(ggplot2)
library(ggpubr)
library(ComplexHeatmap)

source(paste0(sbindir, 'CKP25/load_counts.R'))

# Directories:
imgdir = paste0(img, 'CKP25/')
imgdir2 = paste0(imgdir, 'genes/')
imgpref = paste0(imgdir2, 'genes_')
cmd = paste('mkdir -p', imgdir, imgdir2)
system(cmd)


# Plot specific genes on the UMAP:
# --------------------------------
source(paste0(sbindir, 'CKP25/analysis/load_umap_basis.R'))

norm = norm[,smeta$cell]

gene.col_fun = function(x, pal=rev(viridis(100))){
    palette = c('grey', rev(pal))
    bin <- cut(x, seq(0, max(x), length.out=length(palette)), include.lowest=T) 
    palette[bin] }

plot.gene = function(gene, x=NULL, u1='U2', u2='U2'){
    fname = paste0(imgpref, 'UMAP_', gene,'.png')
    if (u1 != 'U1'){ fname = paste0(imgpref, 'UMAP_', u1, u2, '_', gene, '.png') }
    if (is.null(x)){ x = norm[gene,] }
    ind = order(x)
    png(fname, res=300, width=3,height=3, units='in')
    par(mar=c(0,0,0,0))
    plot(smeta[[u1]][ind], smeta[[u2]][ind], pch=19, col=gene.col_fun(x[ind]),
         axes=F, ylab='', xlab='', cex=.5)
    dev.off()
}

pltgenes <- c("Gria2", "Lig1", "Cdkn1a",
    'Lmnb1', 'Lmnb2',
    "Cadm2", "Hjurp", "Ubb",
    'mt-Nd1','Apoe','Cst3',
    'Lig1','Hjurp','Exoc4',
    'Nek3', 'Mbp', 'Csf1r', 
    'Hfm1', 'Alk','Axl', 'Vcan','Gfap')

coh.genes.mm = c('Smc1a', 'Smc3', 'Rad21', 
    'Stag1', 'Stag2', 'Nipbl', 'Wapl', 'Rif1')

for (gene in pltgenes){
    print(gene)
    plot.gene(gene)
}

plot.gene('Cdkn1a', u1='V1', u2='V2')
plot.gene('Lmnb1', u1='V1', u2='V2')

# Standard legend:
# ----------------
r.legend = Legend(at=seq(0, 1), 
                  labels_gp = gpar(fontsize=5),
                  title_gp = gpar(fontsize=5, fontface='bold'),
                  col_fun=gene.col_fun, title_position = "topleft", 
                  background='black', legend_height=unit(.25,'in'),
                  legend_width=unit(1.25,'in'),
                  title='Norm. expr.', direction = 'horizontal')
rlegend = packLegend(r.legend)

pdf(paste0(imgpref, 'UMAP_standard_gene_legend.pdf'), width=2, height=2)
draw(rlegend, x = unit(.90,'in'), y=unit(.6,'in'), just = "top")
dev.off()




# Tests for specific genes:
# -------------------------
mat = norm[coh.genes.mm,]
df = data.frame(mat, check.names=FALSE)
df$gene = rownames(mat)
df = gather(df, cell, value, -gene)
df = merge(df, meta)

subdf = df[df$label %in% c('Excitatory', 'Stage 2'),]
aggdf = aggregate(value ~ mouse_id + sublabel + label + short.sublabel + gene + genotype, subdf, mean)
aggdf$short.sublabel = factor(aggdf$short.sublabel, levels=c('In0','In1','Ex0','Ex1','Ex2','Ex3','St.2'))

runcomp <- list(c("Ex0", "Ex2"), c("Ex0", "Ex3"), c("Ex0", "St.2"), 
                c('Ex2', 'Ex3'), c('Ex2', 'St.2'), c('Ex3', 'St.2'))

gplot = ggplot(aggdf, aes(short.sublabel, value, fill=sublabel)) + 
    facet_wrap(~gene, nrow=1, scale='free_y') + 
    scale_fill_manual(values=sublabel.col) + 
    geom_boxplot(outlier.shape=NA, alpha=1, lwd=.5) + 
    geom_jitter(position=position_jitterdodge(jitter.width=.15, dodge.width=.75), cex=.8) + 
    stat_compare_means() + 
    theme_pubr() + theme(legend.position='none') + 
    labs(x=paste('Sub-celltype'), y='Norm. expression, log(X+1)') + 
    theme_pubr() + theme(legend.position='none')

pltpref = paste0(imgpref, 'cohesin_genes_boxplots')
ggsave(paste0(pltpref, '.png'), gplot, dpi=450, units='in', width=10, height=4)
ggsave(paste0(pltpref, '.pdf'), gplot, dpi=450, units='in', width=10, height=4)



# Plot overall mitochondrial and ribosomal content
# as well as number of genes,  number of counts:
# ------------------------------------------------
plot.gene('ncounts', smeta$n_counts)
plot.gene('ngenes', smeta$n_genes)
plot.gene('ngenes2', 1 * (smeta$n_genes > 500))

mtgenes = genes[grep("^mt[-nr]",genes)]
mt.pct = colSums(amat[mtgenes, smeta$cell]) / colSums(amat[, smeta$cell])
plot.gene('mito.pct', mt.pct)

rpgenes = genes[grep("^Rp[ls][0-9]",genes)]
rp.pct = colSums(amat[rpgenes, smeta$cell]) / colSums(amat[, smeta$cell])
plot.gene('ribo.pct', rp.pct)


# Plot violin plots by annotation:
# --------------------------------
m1 = c('Gria2','Mbp','Vcan','Gfap','Csf1r','Gad1')
df = smeta[,c('cell','sublabel2')]
df = cbind(df, data.frame(t(norm[m1,])))
df = gather(df, gene, expr, -cell, -sublabel2)

gplot = ggplot(df, aes(sublabel2, expr, fill=sublabel2)) + 
    facet_wrap(~gene) + 
    scale_y_continuous(expand=c(0,0)) + 
    scale_fill_manual(values=sublabel.col) +
    labs(x='scRNA sub-celltype', y='Normalized expression') +
    geom_boxplot() + 
    theme_pubr() + theme(legend.position = 'none') + 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))
ggsave(paste0(imgpref, 'boxplots_markers.png'), gplot, width=5.5, height=7, units='in', dpi=450)


m2 = c('Cdkn1a','Ubb','Apoe','Cst3','mt-Nd1')
df = smeta[,c('cell','sublabel2')]
df = cbind(df, data.frame(t(norm[m2,])))
df = gather(df, gene, expr, -cell, -sublabel2)

gplot = ggplot(df, aes(sublabel2, expr, fill=sublabel2)) + 
    facet_wrap(~gene) + 
    scale_y_continuous(expand=c(0,0)) + 
    scale_fill_manual(values=sublabel.col) +
    geom_boxplot() + 
    labs(x='scRNA sub-celltype', y='Normalized expression') +
    theme_pubr() + theme(legend.position = 'none') + 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))
ggsave(paste0(imgpref, 'boxplots_trajgenes.png'), gplot, width=5.5, height=7, units='in', dpi=450)