obj <- readRDS ("/DATA/31/molteni/vexas_mouse_3/results/Full/01-seurat/Full_final_annot_filt.rds")
assay = "RNA"
reduction <- toupper(paste0("umap.harmony.orig.ident"))

clu_res <- "celltypes_annot"
obj <- SetIdent(obj, value = clu_res)

# convert to scExperiment object
obj_sc <- as.SingleCellExperiment(obj)

out_dir <- "/DATA/31/molteni/vexas_mouse_3/results/pseudotime_filt2_FINAL/"

cell_pal <- function(cell_vars, pal_fun,...) {
  if (is.numeric(cell_vars)) {
    pal <- pal_fun(100, ...)
    return(pal[cut(cell_vars, breaks = 100)])
  } else {
    categories <- sort(unique(cell_vars))
    pal <- setNames(pal_fun(length(categories), ...), categories)
    return(pal[cell_vars])
  }
}

#my_cols <- hue_pal()(length(levels(obj_sc[[clu_res]])))

#my_cols  <- colorRampPalette(pal_npg("nrc")(10))(16) # nature publishing group palette
my_cols <- colorRampPalette(brewer.pal(12, "Paired"))(length(levels(obj_sc[[clu_res]])))
my_cols <- brewer.pal(12, "Paired")[1:14] #inserire numero di cluster

names(my_cols) <- levels(obj_sc[[clu_res]])

#cell_colors_clust <- "black"
cell_colors_clust <- my_cols[obj_sc[[clu_res]]]

# Settings: resolution 1.2 harmony - starting point - cluster 3

namex <- toupper("umap.harmony.orig.ident")
scs <- slingshot(data = obj_sc,
                 reducedDim = reduction, 
                 start.clus = "HSC_MPP", 
                 extend = 'n', 
                 clusterLabels = obj_sc[[clu_res]])

pal <- viridis::inferno(100,end = 0.95)
sl_name <- 'Slingshot'
pt <- slingPseudotime(scs)
head(pt)



#se clu_res non è un numero#####
# Assuming clu_res is a character vector with cluster assignments like "A", "B", "C", etc.
# Convert it to a factor
factor_clu_res <- as.factor(obj_sc[[clu_res]])

# Then convert the factor to numeric
# The levels will be assigned numbers in the order they appear in the factor
numeric_clu_res <- as.numeric(factor_clu_res)

# Now you can use numeric_clu_res to index into your colors
cell_colors_clust <- my_cols[numeric_clu_res]
#grafici

for (num_curve in seq(1, length(colnames(pt)))) {
  lineagename <- paste('Curve',as.character(num_curve),sep='')
  colors <- pal[cut(pt[,num_curve], breaks = 100)]
  
  # Set the filename and create a PNG device
  pdf(paste(out_dir, paste0(sl_name, "_branch_", num_curve, ".pdf"), sep = "/"), width = 1500/90, height = 750/90) 
  
  # Start a new plot
  plot.new()
  #grid.newpage()
  pushViewport(viewport(layout = grid.layout(1, 2))) # Adjusted layout grid to 1 row and 2 columns
  
  # Pseudotime plot with adjusted margins
  pushViewport(viewport(layout.pos.row = 1, layout.pos.col = 1))
  par(mar = c(5, 5, 5, 2) + 0.1) # Adjust margins here; default is c(5, 4, 4, 2) + 0.1
  par(fig = gridFIG(), new = TRUE)
  plot(reducedDim(scs, type = namex), col = colors, 
       pch = 16, cex = 1, main = "Pseudotime", cex.axis=2.0, cex.lab=2.0, cex.main=2.2) # Adjusted cex.lab for label size
  lines(SlingshotDataSet(scs), type = 'lineages', show.constraints = TRUE, col = "black")
  lines(slingCurves(scs)[[num_curve]], lwd = 2, col = 'red')
  
  popViewport()
  
  # Sample density plot with adjusted margins
  pushViewport(viewport(layout.pos.row = 1, layout.pos.col = 2))
  par(mar = c(5, 5, 5, 2) + 0.1) # Adjust margins here; default is c(5, 4, 4, 2) + 0.1
  par(fig = gridFIG(), new = TRUE)
  ds <- list()
  for (samp in unique(colData(scs)$sample_label)) {
    ds[[samp]] <- density(pt[colData(scs)$sample_label == samp, num_curve], na.rm = TRUE)
  }
  # Setting the plot limits
  x_min <- min(sapply(ds, function(x) min(x$x, na.rm = TRUE)))
  x_max <- max(sapply(ds, function(x) max(x$x, na.rm = TRUE)))
  y_min <- min(sapply(ds, function(x) min(x$y, na.rm = TRUE)))
  y_max <- max(sapply(ds, function(x) max(x$y, na.rm = TRUE)))
  
  xlim <- range(c(x_min, x_max))
  ylim <- range(c(y_min, y_max))
  plot(xlim, ylim, type = "n", xlab = "Pseudotime", ylab = "", main = paste0("Sample density on branch ", num_curve),
       cex.axis=2.0, cex.lab=2.0, cex.main=2.2)
  
  for (samp in names(ds)) {
    polygon(c(min(ds[[samp]]$x), ds[[samp]]$x, max(ds[[samp]]$x)), c(0, ds[[samp]]$y, 0),
            col = alpha(brewer.pal(length(names(ds)), "Set1")[which(names(ds) == samp)], alpha = .5))
  }
  
  legend("topright", legend = names(ds), fill = alpha(brewer.pal(length(names(ds)), "Set1"), alpha = .5), bty = "n", cex = 2.0)
  popViewport()
  
  # Close the graphical device
  dev.off()
}