diff --git a/src/gradients.R b/src/gradients.R
index 963e9c4699b60c9ecd30871053c1a79ede0eaf60..0dbb9a3686e1d154115fc2af9e8a346f029d9861 100644
--- a/src/gradients.R
+++ b/src/gradients.R
@@ -3,7 +3,7 @@ source('src/00_base_functions.R')
 source('src/gradients_fcts.R')
 
 rds.name <- 'seuObj_wt_preprocessed_filtered'
-output.dir <- 'results/filtered/230217_wt'
+output.dir <- 'results/filtered/230220_wt'
 dir.create(output.dir)
 output.name <- 'ryu2019'
 
@@ -23,24 +23,28 @@ seuObj[['background']] <- 'wt'
 # similarity score to final id
 # select terminal ref
 hist(seuObj@meta.data[,3])
-ref_ids <- c('Atrichoblast', 'Trichoblast')
+ref_ids <- c('Cortex', 'Endodermis')
+ref_colors <- paste0('#', colors[colors$V1 %in% ref_ids, 3])
+names(ref_colors) <- ref_ids
 ref.i1.cells <- seuObj@meta.data %>%
-  filter(zone == 'Mature' & id == ref_ids[1] & nFeature_RNA < quantile(nFeature_RNA, 0.1)) %>%
+  filter(zone == 'Mature' & id == ref_ids[1] & nFeature_RNA < quantile(nFeature_RNA, 0.99)) %>%
   rownames()
 length(ref.i1.cells)
 png(sprintf("%s/refcells_%s_%s.png", output.dir, output.name, ref_ids[1]), 
     width = 130, height = 130, res = 300, units = 'mm')
-DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = ref.i1.cells) +
+DimPlot(seuObj, reduction = "umap", 
+        label = FALSE, cells.highlight = ref.i1.cells, cols.highlight = ref_colors[1]) +
   theme_void() + theme(legend.position = 'none')
 dev.off()
 
 ref.i2.cells <- seuObj@meta.data %>%
-  filter(zone == 'Mature' & id == ref_ids[2] & nFeature_RNA < quantile(nFeature_RNA, 0.1)) %>%
+  filter(zone == 'Mature' & id == ref_ids[2] & nFeature_RNA < quantile(nFeature_RNA, 0.99)) %>%
   rownames()
 length(ref.i2.cells)
 png(sprintf("%s/refcells_%s_%s.png", output.dir, output.name, ref_ids[2]), 
     width = 130, height = 130, res = 300, units = 'mm')
-DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = ref.i2.cells) +
+DimPlot(seuObj, reduction = "umap", label = FALSE, 
+        cells.highlight = ref.i2.cells, cols.highlight = ref_colors[2]) +
   theme_void() + theme(legend.position = 'none')
 dev.off()
 
@@ -95,13 +99,13 @@ saveRDS(tzs, file = sprintf('%s/tzs_%s.rds', output.dir, paste0(ref_ids, collaps
 
 # Plot selected cells on the similarity space
 factor <- seuObj@meta.data[filter.cells,]$zone
-subset.cells <- rownames(seuObj@meta.data[which(seuObj@meta.data$id %in% c('Atrichoblast') &
+subset.cells <- rownames(seuObj@meta.data[which(seuObj@meta.data$id %in% c('Cortex') &
                                                              seuObj@meta.data$background == 'wt'), ])
 png(sprintf("%s/gradient_%s_%s.png", output.dir, output.name, paste0(ref_ids, collapse = '_')), 
     width = 130, height = 130, res = 300, units = 'mm')
 plot(cor.ref[filter.cells,1], cor.ref[filter.cells,2], pch = 19, cex = 0.2, col = 'black',#scales::alpha(cols[factor], 1),
      xlab = ref_ids[1], ylab = ref_ids[2])
-points(cor.ref[subset.cells,1], cor.ref[subset.cells,2], pch = 19, cex = 0.4, col = 'red')
+points(cor.ref[subset.cells,1], cor.ref[subset.cells,2], pch = 19, cex = 0.4, col = ref_colors[1])
 dev.off()
 
 png(sprintf("%s/gradient_zones_%s_%s_1.png", output.dir, output.name, paste0(ref_ids, collapse = '_')), 
@@ -124,11 +128,13 @@ dev.off()
 # Selection of cells
 i1_wt <- rownames(seuObj@meta.data[which(seuObj@meta.data$id %in% ref_ids[1] &
                                                         seuObj@meta.data$background == 'wt'), ])
-DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = i1_wt)
+DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = i1_wt,
+        cols.highlight = ref_colors[1])
 
 i2_wt <- rownames(seuObj@meta.data[which(seuObj@meta.data$id %in% ref_ids[2] &
                                                        seuObj@meta.data$background == 'wt'), ])
-DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = i2_wt)
+DimPlot(seuObj, reduction = "umap", label = FALSE, cells.highlight = i2_wt,
+        cols.highlight = ref_colors[2])
 
 
 # Compute gradients
@@ -232,7 +238,11 @@ dev.off()
 #                   cluster_rows = F, show_rownames = F, show_colnames = T)
 
 
-# waves of down/upregulation
+
+
+# ------------------------------------------------------------------------------
+# Global waves of down/upregulation
+# ------------------------------------------------------------------------------
 png(sprintf("%s/gr_mean_%s.png", output.dir, ref_ids[2]), 
     width = 130, height = 100, res = 300, units = 'mm')
 y <- colMeans(abs(gr_i2_wt[[2]][bottom_i2$genes,]))
@@ -259,7 +269,10 @@ legend("topleft", legend = c('Down', 'Up'), pch = c(21, 19), col = c('blue', 're
 abline(v = tzs[,1], lty = 'dashed')
 dev.off()
 
-# Pcg
+
+# ------------------------------------------------------------------------------
+# PcG
+# ------------------------------------------------------------------------------
 head(labels)
 cols <- c('lightblue', 'darkblue', 'white', 'grey')
 names(cols) <- levels(labels)
@@ -274,8 +287,9 @@ proportions(table(labels[bottom_i1$genes]))
 proportions(table(labels[bottom_i2$genes]))
 proportions(table(labels))
 
+# proportions PcG
 png(sprintf("%s/props_%s_%s.png", output.dir, ref_ids[2], output.name), 
-    width = 100, height = 100, res = 300, units = 'mm')
+    width = 80, height = 100, res = 300, units = 'mm')
 rbind('top' = proportions(table(labels[top_i2$genes])),
   'bottom' = proportions(table(labels[bottom_i2$genes])),
   'all' = proportions(table(labels))) %>%
@@ -283,10 +297,31 @@ rbind('top' = proportions(table(labels[top_i2$genes])),
   rownames_to_column(var = 'set') %>%
   pivot_longer(cols = c(2, 3, 4, 5), names_to = 'type', values_to = 'p') %>%
   mutate(type = gsub('\\.o', '-o', type)) %>%
-  ggplot(aes(x = set, y = p, fill = type)) + geom_col() +
-  scale_fill_manual(values = cols)
+  ggplot(aes(x = set, y = p, fill = type)) + geom_bar(color = 'darkgrey', stat = 'identity') +
+  scale_fill_manual(values = cols) +
+  labs(x = '', y = 'Proportion', fill = '', title = ref_ids[2]) +
+  theme_pubclean() + theme(legend.position = 'right') +
+  scale_x_discrete(labels = c('All', 'Down', 'Up'))
 dev.off()
 
+png(sprintf("%s/props_%s_%s.png", output.dir, ref_ids[1], output.name), 
+    width = 80, height = 100, res = 300, units = 'mm')
+rbind('top' = proportions(table(labels[top_i1$genes])),
+      'bottom' = proportions(table(labels[bottom_i1$genes])),
+      'all' = proportions(table(labels))) %>%
+  data.frame() %>%
+  rownames_to_column(var = 'set') %>%
+  pivot_longer(cols = c(2, 3, 4, 5), names_to = 'type', values_to = 'p') %>%
+  mutate(type = gsub('\\.o', '-o', type)) %>%
+  ggplot(aes(x = set, y = p, fill = type)) + geom_bar(color = 'darkgrey', stat = 'identity') +
+  scale_fill_manual(values = cols) +
+  labs(x = '', y = 'Proportion', fill = '', title = ref_ids[1]) +
+  theme_pubclean() + theme(legend.position = 'right') +
+  scale_x_discrete(labels = c('All', 'Down', 'Up'))
+dev.off()
+
+
+# Heatmaps
 png(sprintf("%s/heatmap_bottom_%s_%s.png", output.dir, ref_ids[2], output.name), 
     width = 100, height = 150, res = 300, units = 'mm')
 bottom_i2$heatmap
@@ -307,18 +342,52 @@ png(sprintf("%s/heatmap_top_%s_%s.png", output.dir, ref_ids[1], output.name),
 top_i1$heatmap
 dev.off()
 
-write.table(names(which(lab[top_i2$genes] == 'PcG')), 
-            file = sprintf('%s/gr_PcG_genes_%s_%s.txt', output.dir, ref_ids[2], n_features), 
+# 
+labels_to_keep <- c('Bivalent', 'K27-only')
+length(top_i2$genes[which(labels[top_i2$genes] %in% labels_to_keep)])
+length(top_i1$genes[which(labels[top_i1$genes] %in% labels_to_keep)])
+write.table(top_i2$genes[which(labels[top_i2$genes] %in% labels_to_keep)], 
+            file = sprintf('%s/gr_PcG_genes_up_%s_%s.txt', output.dir, ref_ids[2], n_features), 
             quote = FALSE, col.names = F, row.names = F)
-write.table(names(which(lab[top_i1$genes] == 'PcG')), 
-            file = sprintf('%s/gr_PcG_genes_%s_%s.txt', output.dir, ref_ids[1], n_features), 
+write.table(top_i1$genes[which(labels[top_i1$genes] %in% labels_to_keep)], 
+            file = sprintf('%s/gr_PcG_genes_up_%s_%s.txt', output.dir, ref_ids[1], n_features), 
             quote = FALSE, col.names = F, row.names = F)
 
-# TFs
-lab <- ifelse(!features %in% tf_agris[,2], 'TF', 'PcG')
+length(bottom_i2$genes[which(labels[bottom_i2$genes] %in% labels_to_keep)])
+length(bottom_i1$genes[which(labels[bottom_i1$genes] %in% labels_to_keep)])
+write.table(bottom_i2$genes[which(labels[bottom_i2$genes] %in% labels_to_keep)], 
+            file = sprintf('%s/gr_PcG_genes_down_%s_%s.txt', output.dir, ref_ids[2], n_features), 
+            quote = FALSE, col.names = F, row.names = F)
+write.table(bottom_i1$genes[which(labels[bottom_i1$genes] %in% labels_to_keep)], 
+            file = sprintf('%s/gr_PcG_genes_down_%s_%s.txt', output.dir, ref_ids[1], n_features), 
+            quote = FALSE, col.names = F, row.names = F)
+
+# amplitude gradient
+boxplot(top_i2$max[top_i2$genes]~labels[top_i2$genes])
+boxplot(top_i1$max[top_i1$genes]~labels[top_i1$genes])
+
+# ------------------------------------------------------------------------------
+# Transcription factors
+# ------------------------------------------------------------------------------
+lab <- factor(ifelse(features %in% tf_agris[,2], 'TF', '-'))
 names(lab) <- features
-top_i1 <- plot_gr(gr_i1_wt[[1]][features,], n_features, which_genes = 'top', label = lab)
-top_i2 <- plot_gr(gr_i2_wt[[2]][features,], n_features, which_genes = 'top', label = lab)
+cols <- c('white', 'red')
+names(cols) <- levels(lab)
+
+top_i1 <- plot_gr(gr_i1_wt[[1]][features,], n_features, which_genes = 'top', label = lab, cols = cols)
+top_i2 <- plot_gr(gr_i2_wt[[2]][features,], n_features, which_genes = 'top', label = lab, cols = cols)
+bottom_i1 <- plot_gr(gr_i1_wt[[1]][features,], n_features, which_genes = 'bottom', label = lab, cols = cols)
+bottom_i2 <- plot_gr(gr_i2_wt[[2]][features,], n_features, which_genes = 'bottom', label = lab, cols = cols)
+
+table(lab[top_i1$genes], labels[top_i1$genes])
+table(lab[bottom_i1$genes], labels[bottom_i1$genes])
+table(lab[top_i2$genes], labels[top_i2$genes])
+table(lab[bottom_i2$genes], labels[bottom_i2$genes])
+
+bottom_i2$genes[lab[bottom_i2$genes] == 'TF' & labels[bottom_i2$genes] %in% labels_to_keep]
+bottom_i1$genes[lab[bottom_i1$genes] == 'TF' & labels[bottom_i1$genes] %in% labels_to_keep]
+top_i1$genes[lab[top_i1$genes] == 'TF' & labels[top_i1$genes] %in% labels_to_keep]
+top_i2$genes[lab[top_i2$genes] == 'TF' & labels[top_i2$genes] %in% labels_to_keep]
 
 ## first id
 tfs_gr <- names(which(lab[top_i2$genes] == 'TF'))
diff --git a/src/gradients_fcts.R b/src/gradients_fcts.R
index 27bafee24f57817f48602ba10ace4f490be75a1a..7901372f8cbef9ec8eeb833ce7e1174bc6e8b442 100644
--- a/src/gradients_fcts.R
+++ b/src/gradients_fcts.R
@@ -67,8 +67,8 @@ plot_gr <- function(matrix, n_features, which_genes, label, cols = NULL){
     } 
     hm <- pheatmap::pheatmap(matrix[g,], cluster_cols = FALSE, 
                              cluster_rows = F, show_rownames = F, show_colnames = T,
-                             annotation_row = data.frame('PcG' = label[g]),
-                             annotation_colors = list('PcG' = cols),
+                             annotation_row = data.frame('Grp' = label[g]),
+                             annotation_colors = list('Grp' = cols),
                              color = myColor,
                              breaks = myBreaks)
   }