library(dyno)

library(tidyverse)

library(Matrix)

library(hdf5r)

library(Seurat)

method_names <- c('paga','slingshot','monocle')

container_names <- c('jaydu/ti_paga_tree:v2.0.0','dynverse/ti_slingshot:v2.0.1','jaydu/ti_monocle3:v3.0.0')

path <- "../../data/*.h5"#"/Users/dujinhong/Downloads/new dataset/data_h5/*.h5"#

save_path <- "/Users/dujinhong/Downloads/result/"#"result/other methods/"

for(load_file in Sys.glob(path)){  

    filename <- gsub(".h5", "", basename(load_file))

    if(grepl('mouse_brain',filename)){

        next

    }

    dir.create(file.path(save_path, filename))

    file.h5 <- H5File$new(load_file, mode="r")

    counts <- t(as.matrix(file.h5[['count']][,]))

    grouping <- data.frame(cell_id = paste0('c-',as.character(c(1:length(file.h5[['grouping']][])))), 

                           group_id = as.character(as.vector(file.h5[['grouping']][])))    

    milestone_net <- data.frame(file.h5[['milestone_network']][])

    root_milestone_id <- file.h5[['root_milestone_id']][]    

    if('w' %in% colnames(milestone_net)){

        cell_ids <- which(milestone_net$from==root_milestone_id)

        start_id <- paste0('c-',as.character(cell_ids[which.max(milestone_net[cell_ids,'w'])]))

    }else{

        start_id <- paste0('c-',sample(which(grouping$group_id==root_milestone_id), 1))

    }    

    file.h5$close_all()

    rownames(counts) <- paste0('c-',as.character(c(1:dim(counts)[1])))

    colnames(counts) <- paste0('g-',as.character(c(1:dim(counts)[2])))

    data <- CreateSeuratObject(counts = t(counts), project = "miller")

    data <- NormalizeData(data, normalization.method = "LogNormalize", scale.factor = 10000)

    data <- FindVariableFeatures(data, selection.method = "vst", nfeatures = 2000)

    data <- ScaleData(data, features = rownames(data))

    counts <- Matrix(counts[,VariableFeatures(data)], sparse = TRUE)

    expression <- Matrix(t(data[["RNA"]]@scale.data)[,VariableFeatures(data)], sparse = TRUE)

    dataset <- wrap_expression(

      expression = expression,

      counts = counts

    )

    for(i in c(1:3)){

        method_name <- method_names[i]                

        method <- create_ti_method_container(container_names[i])        

        if(i<3){

            K <- length(unique(grouping$group_id))

            if(i==1){

                K <- K - 1

            }

            else

                ndim <- 20

            pca <- irlba::prcomp_irlba(expression, n = ndim)

            rd <- pca$x

            cluster_labels <- cluster::pam(rd, K)$clustering

            grouping <- data.frame(cell_id = grouping$cell_id, 

                                   group_id = cluster_labels)   

        }

        if(i==2){

            dataset <- dataset %>% add_prior_information(start_id = c(start_id), groups_id = grouping)

        }else{

            dataset <- dataset %>% add_prior_information(start_id = start_id, groups_id = grouping)

        }

        model <- infer_trajectory(dataset, method(), give_priors = c("start_id"), verbose = TRUE)

        ix <- wrapr::match_order(model$progressions$cell_id, grouping$cell_id)

        model$progressions <- model$progressions[ix,]

        rownames(model$progressions) <- NULL

        write.csv(x=model$progressions, file=file.path(save_path, filename, sprintf("%s_progressions.csv", method_name)))

        write.csv(x=model$milestone_network, file=file.path(save_path, filename, sprintf("%s_milestone_network.csv", method_name)))

    }

}