#!/usr/bin/R
# ------------------------------------
# Load metadata for the CKP25 project:
# ------------------------------------
library(cbrbase)
set_proj('DEVTRAJ', 'mosaicism')

library(RColorBrewer)

# ---------------------
# Directories/metadata:
# ---------------------
projdir = paste0('CKP25/')

# -------------------
# Genome Annotations:
# -------------------
# TODO: Get the mouse annotation

# --------------------
# Metadata Annotation:
# --------------------
meta = read.delim('Annotation/CKP25_UMAP_clusters_dataframe.tsv', header=T, stringsAsFactors=F)

# Short labels:
meta$celltype[meta$celltype == 'Undamaged'] = 'Baseline'
meta$short.label = paste0(substr(meta$label,1,3), '.')
meta$short.sublabel = substr(meta$sublabel,1,3)
ind = meta$short.sublabel != meta$sublabel
meta$short.sublabel[ind] = paste0(meta$short.sublabel[ind],'.')
meta$short.sublabel[meta$short.sublabel == 'Sta.'] = 'St.2'
meta$short.label[meta$short.label == 'Sta.'] = 'St.2'

# Add deidentified - weaker stage 2 signature, 
# separate from Stage 2 to some extent
meta$sublabel2 = meta$sublabel
meta$sublabel2[meta$leiden %in% c(7, 15)] = 'Deid'

# --------------------
# Colors for plotting:
# --------------------
# To look at these colors:
snap.cols = scan('Annotation/snap_colors.tsv', 'c')
# x = 1:length(snap.cols)
# nx = 8
# plot(x %% nx, x %/% nx, pch=15, cex=8, col=snap.cols)
# text(x %% nx, x %/% nx, x, pch=19, cex=1, col='white')

# Cleaner UMAP colors:
tcols = brewer.pal(12,'Paired')
label.col = c(tcols[c(4,2,6,8)], snap.cols[48], tcols[10], 'grey')
# sublabel.col = c(tcols[c(6,4)], exramp,tcols[c(2,1,10,8)], snap.cols[48], 'grey')
names(label.col) = c('Excitatory','Inhibitory','Stage 2', 'OPC', 'Oligodendrocyte','Microglia','Batch')

exramp = colorRampPalette(c(tcols[3],tcols[5]))(3)
sublabel.col = c(tcols[c(6,4)], exramp,tcols[c(2,1,10,8)], snap.cols[c(48)], 'grey', 'grey50')
names(sublabel.col) = c('Stage 2','Ex0','Ex1','Ex2','Ex3','In0', 'In1',
                        'Microglia', 'OPC','Oligodendrocyte','Batch', 'Deid')

celltype.col = c('Glia' = tcols[9], 'Baseline'=tcols[4], 'Stage 1'=exramp[3],'Stage 2' = '#B22222')
time.col = c('1wk' = tcols[9], '2wk'=tcols[10])


# Covariate colors:
colvals = list()
colvals[['sex']] = c("female" = "pink", "male" = "lightblue")
colvals[['amyloid.level']] =  c("low" = "royalblue", "high" = "indianred")
colvals[['tangles.level']] =  c("low" = "royalblue", "high" = "indianred")
colvals[['apoe_genotype']] =  c("23" = "lightgrey", "33" = "lightgrey", "34" = "darkgrey", "44" = "darkgrey", "0" = "black")
colvals[['sex.and.amyloid']] =  c("female, high amyloid" = "indianred", "female, low amyloid" = "pink", "male, high amyloid" = "royalblue", "male, low amyloid" = "lightblue")


# Genes for cohesin:
# ------------------
coh.genes.hg = c('SMC1A', 'SMC3', 'RAD21', 'STAG1', 'STAG2', 'NIPBL', 'WAPL', 'RIF1')
coh.genes.mm = c('Smc1a', 'Smc3', 'Rad21', 'Stag1', 'Stag2', 'Nipbl', 'Wapl', 'Rif1')