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

library(ggplot2)
library(ggpubr)
library(ggbeeswarm)
library(tidyr)
library(scales)
library(RColorBrewer)
library(viridis)

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

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

source(paste0(sbindir, 'CKP25/analysis/load_umap_basis.R'))


plot.gene = function(gene, x=NULL, u1='U1', 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)
    mtext(gene, 1, line=-1.1)
    dev.off()
}


plot.umap = function(var, col, w=3, cex=.5, u1='U1', u2='U2'){
    fname = paste0(imgpref, 'UMAP_', var, '_notext.png')
    if (u1 != 'U1'){ fname = paste0(imgpref, 'UMAP_', u1, u2, '_', var, '_notext.png') }
    png(fname,res=300, width=w, height=w, units='in')
    par(mar=rep(0, 4))
    plot(smeta[[u1]], smeta[[u2]], pch=19, col=col[smeta[[var]]], 
        axes=F, ylab='', xlab='', cex=cex)
    dev.off()
}


# Create plots of the UMAPs:
# --------------------------
mids = sort(unique(smeta$mouse_id))
j = 13; mcols = snap.cols[(1:length(mids)) + j]
names(mcols) = mids

red.ramp = colorRampPalette(c('white', 'red'))(20)

sublabel.vd.col = sublabel.col
sublabel.vd.col['Stage 2'] = '#8b0000'
sublabel.vd.col['Ex0'] = 'black'
sublabel.vd.col['Oligodendrocyte'] = sublabel.col['Ex1']
sublabel.vd.col[c('Ex1', 'Ex2', 'Ex3')] = c(red.ramp[5], red.ramp[10], red.ramp[20])
smeta$sublabel.vd = smeta$sublabel2
smeta$label.vd = smeta$sublabel2
smeta$label.vd[smeta$sublabel2 %in% c('Ex1', 'Ex2', 'Ex3')] = 'Stage 1'
label.vd.col = sublabel.vd.col
label.vd.col['Stage 1'] = 'red'


varmap = list(
    'label' = label.col,
    'sublabel' = sublabel.col,
    'sublabel2' = sublabel.col,
    'sublabel.vd' = sublabel.vd.col,
    'label.vd' = label.vd.col,
    'leiden' = snap.cols[3:40]
    # The following need legends, automate separately
    # 'genotype' = c('CK'='lightblue', 'CK-p25'='indianred'),
    # 'timepoint' = c('1wk'=tcols[2], 'p25'=tcols[8]),
    # 'has_p25' = c(TRUE='slateblue3', FALSE='grey80'),
    # 'mouse_id'=mcols
    
)

for (var in names(varmap)){
    print(var)
    plot.umap(var, varmap[[var]])
    plot.umap(var, varmap[[var]], u1='V1', u2='V2')
}


png(paste0(imgpref, 'UMAP_genotype_notext.png'),res=300, width=3,height=3, units='in')
par(mar=c(0,0,0,0))
plot(smeta[[u1]], smeta[[u2]], pch=19, col=ifelse(smeta$genotype=='CK','lightblue','indianred'),
     axes=F, ylab='', xlab='', cex=.5)
legend('bottomleft', title = 'Genotype', legend=c('CK','CK-p25'), 
       col=c('lightblue','indianred'), pch=19, bty='n', ncol=1, cex=.75, pt.cex=1)
dev.off()


png(paste0(imgpref, 'UMAP_timepoint_notext.png'),res=300, width=3,height=3, units='in')
par(mar=c(0,0,0,0))
plot(smeta[[u1]], smeta[[u2]], pch=19, col=ifelse(smeta$timepoint == '1wk', tcols[2], tcols[8]),
     axes=F, ylab='', xlab='', cex=.5)
legend('bottomleft', title = 'Timepoint', legend=c('1 Week','2 Weeks'), 
       col=tcols[c(2,8)], pch=19, bty='n', ncol=1, cex=.75, pt.cex=1)
dev.off()


png(paste0(imgpref, 'UMAP_hasp25_notext.png'),res=300, width=3,height=3, units='in')
par(mar=c(0,0,0,0))
plot(smeta[[u1]], smeta[[u2]], pch=19, col=ifelse(smeta$has_p25, 'slateblue3', 'grey80'),
     axes=F, ylab='', xlab='', cex=.5)
legend('bottomleft', title = 'p25 detected', legend=c('No','Yes'), 
       col=c('grey80','slateblue3'), pch=19, bty='n', ncol=1, cex=.75, pt.cex=1)
dev.off()


png(paste0(imgpref, 'UMAP_mouse_notext.png'),res=300, width=3,height=3, units='in')
par(mar=c(0,0,0,0))
plot(smeta[[u1]], smeta[[u2]], pch=19, col=mcols[smeta$mouse_id],
     axes=F, ylab='', xlab='', cex=.5)
legend('bottomleft', legend=names(mcols), col=mcols, pch=19, bty='n', ncol=1, cex=.75, pt.cex=1)
dev.off()


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

meta = merge(meta, updf2, all.x=TRUE)

# Subset matrix to match metadata subset:
if (reclustered){
    smeta = meta[!is.na(meta$U1),]
} else {
    smeta = meta[meta$label != 'Batch',]
}

ind = sample(1:nrow(smeta), nrow(smeta), replace=FALSE)
smeta = smeta[ind,]
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 = norm[gene,]
    ind = order(x)
    png(paste0(imgpref, 'UMAP_gene_', gene,'.png'),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)
    mtext(gene, 1, line=-1.1)
    dev.off()
}


plot.gene('Cdkn1a')