#ref: Single cell transcriptional and chromatin accessibility profiling redefine cellular heterogeneity in the adult human kidney

#'@param ref.npcs = 30, dims = 30, are the parameters in the FindTransferAnchors function, used to control the number of dimensions used by reference data and anchor search space respectively
#'@param scATAC.assay = "ACTIVITY", scATAC.normalize = T, PC = 30, is used to process scATAC data.
# scATAC.assay indicates the name of the assay where the gene activity is stored in the scATAC number
# scATAC.normalize Whether the gene activity has been logNormalized
# PC indicates the dimension of scATAC dimensionality reduction, which comes from the lsi of scATAC.object itself. Note that the default dimensionality reduction assay of scATAC.object is lsi

Integrate.scRNA.scATAC <- function(scATAC.object, scRNA.object, class.label = "cellType", scRNA.assay = "RNA", ref.npcs = 30, dims = 30, scATAC.assay = "ACTIVITY", scATAC.normalize = T, PC = 30, nfeatures = 3000, SCT = F){

    ##1. load data
    #scRNA-seq data
    DefaultAssay(scRNA.object) <- scRNA.assay
    scRNA.object <- NormalizeData(scRNA.object)
    if(SCT){
        #Need to use sctransform to normalize data in advance
        var.genes <- scRNA.object@assays$SCT@var.features[1:nfeatures]
    }else{
        var.genes <- VariableFeatures(scRNA.object)[1:nfeatures]
    }
    #scATAC-seq data
    DefaultAssay(scATAC.object) <- scATAC.assay
    if(!scATAC.normalize){
        scATAC.object <- NormalizeData(scATAC.object)
    }
    scATAC.object <- ScaleData(scATAC.object, features = rownames(scATAC.object))

    ##2. Identify anchors
    #npcs = 30; dim = 1:30
    transfer.anchors <- FindTransferAnchors(reference = scRNA.object, 
                                            query = scATAC.object, 
                                            features = var.genes, #SCT model：scRNA.merge@assays$SCT@var.features
                                            reference.assay = scRNA.assay,
                                            query.assay = scATAC.assay, 
                                            reduction = "cca",
                                            npcs = ref.npcs, # Number of PCs to compute on reference if reference
                                            dims = 1:dims) # Which dimensions to use from the reduction to specify the neighbor search space
    #saveRDS(transfer.anchors, file = file.path(saveDir, "transfer.anchors.rds"))

    #The default dims = NULL, it is consistent with FindTransferAnchors
    celltype.predictions <- TransferData(anchorset = transfer.anchors, 
                                refdata = scRNA.object@meta.data[,class.label], 
                                weight.reduction = scATAC.object[["lsi"]], 
                                dims = 2:PC)    


    scATAC.object <- AddMetaData(scATAC.object, metadata = celltype.predictions)
    scATAC.object <- AddMetaData(scATAC.object, metadata = celltype.predictions$predicted.id, class.label)

    p <- gghistogram(scATAC.object@meta.data, x = "prediction.score.max", y = "..density..", main = "Prediction Score for scATAC", add_density = TRUE)
    print(p)    
    p <- DimPlot(scATAC.object, group.by = class.label, label = TRUE, repel = TRUE, reduction = "umap") +
        ggtitle("scATAC-seq Predicted Celltypes Before Harmony") + NoLegend()
    print(p)
    p <- DimPlot(scATAC.object, group.by = class.label, label = TRUE, repel = TRUE, reduction = "umap") +
        ggtitle("scATAC-seq Predicted Celltypes Before Harmony")
    print(p)    
    p <- DimPlot(scATAC.object, group.by = class.label, label = TRUE, repel = TRUE, reduction = "tsne") +
        ggtitle("scATAC-seq Predicted Celltypes Before Harmony") + NoLegend()
    print(p)
    p <- DimPlot(scATAC.object, group.by = class.label, label = TRUE, repel = TRUE, reduction = "umap") +
        ggtitle("scATAC-seq Predicted Celltypes Before Harmony")
    print(p)
    return(list(transfer.anchors = transfer.anchors, scATAC.object = scATAC.object))
}

#' @param dims Indicates the dimensionality reduction used when coembedding
Coembedding.scATAC.scRNA <- function(transfer.anchors, scATAC.object, scRNA.assay = "RNA", scRNA.object, class.label = "cellType", PC = 30, dims = 30){
    
    ####Co-embedding scRNA-seq and scATAC-seq datasets
    # note that we restrict the imputation to variable genes from scRNA-seq, but could impute the
    # full transcriptome if we wanted to
    genes.use <- transfer.anchors@anchor.features
    DefaultAssay(scRNA.object) <- scRNA.assay
    scRNA.object <- NormalizeData(scRNA.object)
    refdata <- GetAssayData(scRNA.object, assay = scRNA.assay, slot = "data")[genes.use, ]
    # refdata (input) contains a scRNA-seq expression matrix for the scRNA-seq cells.  imputation
    # (output) will contain an imputed scRNA-seq matrix for each of the ATAC cells
    imputation <- TransferData(anchorset = transfer.anchors, refdata = refdata, weight.reduction = scATAC.object[["lsi"]], dims = 2:PC)
    scATAC.object[["RNA"]] <- imputation
    DefaultAssay(scATAC.object) <- "RNA"     

    scRNA.object <- AddMetaData(scRNA.object, metadata = rep("scRNA", nrow(scRNA.object@meta.data)), col.name = "type")
    scATAC.object <- AddMetaData(scATAC.object, metadata = rep("scATAC", nrow(scATAC.object@meta.data)), col.name = "type")

    coembed <- merge(x = scRNA.object, y = scATAC.object)
    DefaultAssay(coembed) <- "RNA"
    # Finally, we run PCA and UMAP on this combined object, to visualize the co-embedding of both
    # datasets
    coembed <- ScaleData(coembed, features = genes.use, do.scale = FALSE)
    coembed <- RunPCA(coembed, features = genes.use, verbose = FALSE)
    coembed <- RunHarmony(object = coembed, group.by.vars = "orig.ident", verbose = FALSE)
    coembed <- RunUMAP(coembed, dims = 1:dims, reduction = 'harmony', verbose = FALSE)
    coembed <- RunTSNE(coembed, dims = 1:dims, reduction = 'harmony', check_duplicates = FALSE, verbose = FALSE)
    #saveRDS(coembed, file = file.path(saveDir, "coembed.rds"))

    p <- ElbowPlot(object = coembed, ndims = dims)
    print(p)
    p <- DimPlot(coembed, group.by = c("orig.ident"), reduction = "umap", label = T) + NoLegend()
    print(p)
    p <- DimPlot(coembed, group.by = c("orig.ident"), reduction = "tsne", label = T) + NoLegend()
    print(p)
    p <- DimPlot(coembed, group.by = c("type"), reduction = "umap", label = T) + NoLegend()
    print(p)
    p <- DimPlot(coembed, group.by = class.label, reduction = "umap", label = T) + NoLegend()
    print(p)
    p <- DimPlot(coembed, group.by = c("type"), reduction = "tsne", label = T) + NoLegend()
    print(p)
    p <- DimPlot(coembed, group.by = class.label, reduction = "tsne", label = T) + NoLegend()
    print(p)

    p <- DimPlot(coembed, group.by = c("orig.ident"), reduction = "umap", label = T)
    print(p)
    p <- DimPlot(coembed, group.by = c("orig.ident"), reduction = "tsne", label = T)
    print(p)
    p <- DimPlot(coembed, group.by = c("type"), reduction = "umap", label = T)
    print(p)
    p <- DimPlot(coembed, group.by = class.label, reduction = "umap", label = T)
    print(p)
    p <- DimPlot(coembed, group.by = c("type"), reduction = "tsne", label = T)
    print(p)
    p <- DimPlot(coembed, group.by = class.label, reduction = "tsne", label = T)
    print(p)
    return(list(coembed = coembed, scATAC.object = scATAC.object))
}