library(Seurat)
library(data.table)
maher_smooth <- function(seur_obj,nn_graph,n_samples=NULL,n_nbrs=30,assay="Spatial",
                         layer="counts"){
    print("Converting data to data.table")
    #Grab expression data
    mat = GetAssayData(seur_obj,assay=assay,slot=layer)
    #Dense matrix has better column access...
    mat = as.data.table(mat)
    print("Conversion complete.")
    if(is.null(n_samples)){
        n_samples = round(n_nbrs/3,0)
    }
    #Sample neighbors for each cell
    print("Sampling neighbors")
    sampled = apply(nn_graph,1,function(x) sample(x,n_samples))
    #Turn into list
    sampled = as.list(as.data.frame(sampled))
    #Calculate new representation. Can take a while, using DT for it now. <1 min for 55k spots
    new_mat = lapply(sampled, function(x) rowMeans(mat[,..x]))
    new_mat = do.call(cbind, new_mat)
    return(new_mat)
}
#Not implementing the integration part for now. Suppose we only have 1 sample
maher_get_neighbors <- function(seur_obj,n_nbrs){
    #Grab coordinates. If your coordinates are not there, just add them to 
    #@images$slice_1@coordinates as a named df (cells x c(x,y))
    coords = GetTissueCoordinates(seur_obj)[,c(1,2)]
    #Incudes self for now. Use seurat to get the kNN, grab indices of top k
    nns = FindNeighbors(as.matrix(coords),k.param = n_nbrs,compute.SNN=F,return.neighbor = TRUE)
    nns = [email protected][,1:n_nbrs]
    return(nns)
}
#Wrapper
maher_spin <- function(seur_obj,n_nbrs = 30, n_samples=NULL,n_pcs=30,
                       random_state= 0,assay="Spatial",layer="counts",
                      resolution=0.5){
    set.seed(random_state)
    print("Computing nearest neighbors")
    nns = maher_get_neighbors(seur_obj,n_nbrs = n_nbrs)
    print("Computing Smoothed representations")
    new_repr = maher_smooth(seur_obj,nns,n_nbrs=n_nbrs,assay = assay,layer=layer,n_samples=n_samples)
    #Add new representation as Assay
    rownames(new_repr) = rownames(seur_obj)
    print("Normalizing")
    #Sparsify
    new_repr = as.sparse(new_repr)
    seur_obj_spin = CreateSeuratObject(counts = new_repr,meta.data = [email protected],verbose=F)
    seur_obj_spin = NormalizeData(seur_obj_spin,verbose=F)
    seur_obj_spin = FindVariableFeatures(seur_obj_spin,verbose=F)
    seur_obj_spin = ScaleData(seur_obj_spin,assay="RNA",verbose=F)
    print("PCA, SNN, UMAP, Louvain.")
    seur_obj_spin <- RunPCA(seur_obj_spin,verbose=F) 
    seur_obj_spin <- FindNeighbors(seur_obj_spin,dims=1:n_pcs,verbose=F)
    seur_obj_spin <- RunUMAP(seur_obj_spin,
                        dims=1:n_pcs,verbose=F)
    seur_obj_spin <- FindClusters(seur_obj_spin,resolution=resolution,verbose=F)
    domains = seur_obj_spin$seurat_clusters
    names(domains) = NULL
    seur_obj$SPINDomain = domains
    return(list(seur_obj,seur_obj_spin))
}