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

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

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

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

# Load after defining imgpref
source(paste0(sbindir, 'CKP25/analysis/load_basic_fusions.R'))
source(paste0(sbindir, 'CKP25/analysis/load_umap_basis.R'))


# Perform analysis at the intra-chromosomal level:
# ------------------------------------------------
idf = aggregate(close.fusion ~ cell, intradf, sum)
idf = merge(smeta, idf, all.x=TRUE)
idf$close.fusion[is.na(idf$close.fusion)] = 0
idf$has.fusion = idf$close.fusion > 0
aidf = aggregate(cell ~ has.fusion + short.sublabel + genotype, idf, length)

gplot = ggplot(aidf, aes(short.sublabel, cell, fill=has.fusion)) + 
    facet_wrap(~genotype) + 
    scale_fill_manual(values=c('grey85','slateblue4'), name='Has Intra-chromosomal Fusion < 1Mb:') + 
    scale_y_continuous(expand=c(0,0)) + 
    theme_pubr() + 
    labs(x='Sub-celltype', y='Percent of cells')+ 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))

g1 = gplot + geom_bar(stat='identity', position='stack') + labs(y='Number of cells', x='')
g2 = gplot + geom_bar(stat='identity', position='fill') + labs(y='Fraction of cells')
garr = ggarrange(g1, g2, ncol=1, nrow=2, align='hv',
                 common.legend = TRUE, legend = "top")
ggsave(paste0(imgpref, 'intrafusion_subcelltype.png'), garr, width=5, height=6, units='in', dpi=450)
ggsave(paste0(imgpref, 'intrafusion_subcelltype.pdf'), garr, width=5, height=6)

# Test for enrichment of Stage 2:
mat = matrix(c(nrow(idf), sum(idf$has.fusion), 
  sum(idf$label == 'Stage 2'), sum(idf$label == 'Stage 2' & idf$has.fusion)), 2,2)
fisher.test(mat)

# Just in ckp25 mice
idf2 = idf[idf$genotype == 'CKp25' & idf$sublabel %in% c('Ex0','Ex1','Ex2','Ex3','In0','In1','Stage 2'),]
mat = matrix(c(nrow(idf2), sum(idf2$has.fusion), 
  sum(idf2$label == 'Stage 2'), sum(idf2$label == 'Stage 2' & idf2$has.fusion)), 2,2)
fisher.test(mat)

# Test for differences between time-points:
mat = matrix( c(sum(idf$label == 'Stage 2'), sum(idf$label == 'Stage 2' & idf$has.fusion), sum(idf$label == 'Stage 2' & idf$timepoint == '2wk'), sum(idf$label == 'Stage 2' & idf$has.fusion & idf$timepoint == '2wk')) , 2,2)
fisher.test(mat)


# Compare to analysis with the Deid celltype instead:
# ---------------------------------------------------
aidf = aggregate(cell ~ has.fusion + sublabel2 + genotype, idf, length)

gplot = ggplot(aidf, aes(sublabel2, cell, fill=has.fusion)) + 
    facet_wrap(~genotype) + 
    scale_fill_manual(values=c('grey85','slateblue4'), name='Has Intra-chromosomal Fusion < 1Mb:') + 
    scale_y_continuous(expand=c(0,0)) + 
    theme_pubr() + 
    labs(x='Sub-celltype', y='Percent of cells')+ 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))

g1 = gplot + geom_bar(stat='identity', position='stack') + labs(y='Number of cells', x='')
g2 = gplot + geom_bar(stat='identity', position='fill') + labs(y='Fraction of cells')
garr = ggarrange(g1, g2, ncol=1, nrow=2, align='hv',
                 common.legend = TRUE, legend = "top")

ggsave(paste0(imgpref, 'intrafusion_subcelltype2.png'), garr, width=5, height=6, units='in', dpi=450)
ggsave(paste0(imgpref, 'intrafusion_subcelltype2.pdf'), garr, width=5, height=6)

# Test for enrichment of Stage 2:
mat = matrix(c(nrow(idf), sum(idf$has.fusion), 
  sum(idf$sublabel2 == 'Stage 2'), sum(idf$sublabel2 == 'Stage 2' & idf$has.fusion)), 2,2)
fisher.test(mat)

# Just in ckp25 mice (up to p=4.6e-6 when only considering "proper" Stage 2 cells)
idf2 = idf[idf$genotype == 'CKp25' & idf$sublabel %in% c('Ex0','Ex1','Ex2','Ex3','In0','In1','Stage 2'),]
mat = matrix(c(nrow(idf2), sum(idf2$has.fusion), 
  sum(idf2$sublabel2 == 'Stage 2'), sum(idf2$sublabel2 == 'Stage 2' & idf2$has.fusion)), 2,2)
fisher.test(mat)
intra.mat  = mat

# Test for differences between time-points:
mat = matrix( c(sum(idf$sublabel2 == 'Stage 2'), sum(idf$sublabel2 == 'Stage 2' & idf$has.fusion),
                sum(idf$sublabel2 == 'Stage 2' & idf$timepoint == '2wk'), sum(idf$sublabel2 == 'Stage 2' & idf$has.fusion & idf$timepoint == '2wk')) , 2,2)
fisher.test(mat)



# ------------------------------------
# Analysis at inter-chromosomal level:
# ------------------------------------
idf = aggregate(inter.fusion ~ cell, interdf, sum)
idf = merge(meta, idf, all.x=TRUE)
idf$inter.fusion[is.na(idf$inter.fusion)] = 0
idf$has.fusion = idf$inter.fusion > 0
aidf = aggregate(cell ~ has.fusion + short.sublabel + genotype, idf, length)

gplot = ggplot(aidf, aes(short.sublabel, cell, fill=has.fusion)) + 
    facet_wrap(~genotype) + 
    scale_fill_manual(values=c('grey85','slateblue4'), name='Has Inter-chromosomal Fusion:') + 
    scale_y_continuous(expand=c(0,0)) + 
    theme_pubr() + 
    labs(x='Sub-celltype', y='Percent of cells')+ 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))

g1 = gplot + geom_bar(stat='identity', position='stack') + labs(y='Number of cells', x='')
g2 = gplot + geom_bar(stat='identity', position='fill') + labs(y='Fraction of cells')
garr = ggarrange(g1, g2, ncol=1, nrow=2, align='hv',
                 common.legend = TRUE, legend = "top")
ggsave(paste0(imgpref, 'interfusion_subcelltype.png'), garr, width=5, height=6, units='in', dpi=450)
ggsave(paste0(imgpref, 'interfusion_subcelltype.pdf'), garr, width=5, height=6)

# Test for enrichment of Stage 2:
mat = matrix(c(nrow(idf), sum(idf$has.fusion), 
  sum(idf$label == 'Stage 2'), sum(idf$label == 'Stage 2' & idf$has.fusion)), 2,2)
fisher.test(mat)

mat[2,1] = mat[2,1] + 65
mat[2,2] = mat[2,2] + 36
fisher.test(mat)

# Just in ckp25 mice
idf2 = idf[idf$genotype == 'CKp25' & idf$sublabel %in% c('Ex0','Ex1','Ex2','Ex3','In0','In1','Stage 2'),]
mat = matrix(c(nrow(idf2), sum(idf2$has.fusion), 
  sum(idf2$label == 'Stage 2'), sum(idf2$label == 'Stage 2' & idf2$has.fusion)), 2,2)
fisher.test(mat)

mat[2,1] = mat[2,1] + 54
mat[2,2] = mat[2,2] + 36
fisher.test(mat)


idf2 = idf[idf$genotype == 'CKp25' & idf$sublabel %in% c('Ex0','Ex1','Ex2','Ex3','In0','In1','Stage 2'),]
mat = matrix(c(nrow(idf2), sum(idf2$has.fusion), 
  sum(idf2$sublabel2 == 'Stage 2'), 
  sum(idf2$sublabel2 == 'Stage 2' & idf2$has.fusion)), 2,2)
fisher.test(mat)


fisher.test(mat + intra.mat)


# Plot for subct2
aidf = aggregate(cell ~ has.fusion + sublabel2 + genotype, idf, length)

gplot = ggplot(aidf, aes(sublabel2, cell, fill=has.fusion)) + 
    facet_wrap(~genotype) + 
    scale_fill_manual(values=c('grey85','slateblue4'), name='Has Inter-chromosomal Fusion:') + 
    scale_y_continuous(expand=c(0,0)) + 
    theme_pubr() + 
    labs(x='Sub-celltype', y='Percent of cells')+ 
    theme(axis.text.x=element_text(angle=90,vjust=.5, hjust=1))

g1 = gplot + geom_bar(stat='identity', position='stack') + labs(y='Number of cells', x='')
g2 = gplot + geom_bar(stat='identity', position='fill') + labs(y='Fraction of cells')
garr = ggarrange(g1, g2, ncol=1, nrow=2, align='hv',
                 common.legend = TRUE, legend = "top")
ggsave(paste0(imgpref, 'interfusion_subcelltype2.png'), garr, width=5, height=6, units='in', dpi=450)
ggsave(paste0(imgpref, 'interfusion_subcelltype2.pdf'), garr, width=5, height=6)