timestart <- Sys.time()

library(Seurat)
library(SeuratDisk)
library(Matrix)
library(dplyr)

args<-commandArgs(trailingOnly = TRUE)
input=args[1]

setwd(dirname(input))

# 参数设置
## 样本批次
CONTROL = 'CTRL'
# QC
##Cells with fewer than 200 expressed gene, classified as blank control(gene = 'None'), or containing a large fraction of mitochondrial genes (over 10%) were filtered.
##Genes expressed in less than 3 cells were filtered.
##Perturbations except negative control with perturbed cells lower than 30 were filtered.

QC <- function(seurat_raw){
    seurat_raw[["percent.mt"]] <- PercentageFeatureSet(seurat_raw, pattern = "^Mt\\.|^MT\\.|^mt\\.|^Mt-|^MT-|^mt-")
    seurat_raw <- subset(seurat_raw, nFeature_RNA > 200 & percent.mt < 10 & gene != 'None')
    selected_f <- rownames(seurat_raw)[Matrix::rowSums(seurat_raw[['RNA']]@counts > 0 ) > 3]
    keep_perturb <- names(which(table(seurat_raw[[]]$gene)>30))
    keep_perturb <- if(CONTROL %in% keep_perturb) keep_perturb else c(keep_perturb, CONTROL)
    Idents(seurat_raw) <- 'gene'
    selected_c <- WhichCells(seurat_raw, idents = keep_perturb)
    seurat_raw <- subset(seurat_raw, features = selected_f, cells = selected_c)
    return(seurat_raw)
}

# 去除NP细胞
remove_np <- function(seurat_raw){
  seurat_raw <- subset(seurat_raw, subset = `mixscape_class.global` != 'NP')
  seurat_raw@meta.data$gene <- factor(seurat_raw@meta.data$gene)
  return(seurat_raw)
}

seurat_raw <- LoadH5Seurat(input)
DefaultAssay(object = seurat_raw) <- 'RNA'
CalcPerturbSig.split.by <- if ( length(unique(seurat_raw@meta.data$batch)) > 1) 'batch' else NULL

seurat_raw <- QC(seurat_raw)
print("QC complete")

Idents(seurat_raw) <- 'orig.ident'
seurat_raw <- NormalizeData(seurat_raw) %>%
              FindVariableFeatures() %>% 
              ScaleData() %>%
              RunPCA()
print("Dimensionality reduction of RNA assay complete")

# Mixscape
seurat_raw <- CalcPerturbSig(
    object = seurat_raw,
    assay = "RNA",
    slot = "data",
    gd.class ="gene", 
    nt.cell.class = CONTROL, 
    reduction = "pca", 
    ndims = 15, 
    num.neighbors = 20, 
    split.by = CalcPerturbSig.split.by, 
    new.assay.name = "PRTB")
print("Calculate perturbation signature (PRTB) complete")

DefaultAssay(object = seurat_raw) <- 'PRTB'
VariableFeatures(object = seurat_raw) <- VariableFeatures(object = seurat_raw[["RNA"]])
seurat_raw <- ScaleData(object = seurat_raw, do.scale = F, do.center = T)
seurat_raw <- RunPCA(object = seurat_raw, reduction.key = 'prtbpca', reduction.name = 'prtbpca')
print("Dimensionality reduction of PRTB assay complete")

# 识别NP细胞（运行时间较长）
seurat_raw <- RunMixscape(
  object = seurat_raw, 
  assay = "PRTB", 
  labels = "gene", 
  nt.class.name = CONTROL, 
  de.assay = "RNA", 
  prtb.type = "KO")
print("Run mixscape complete")

seurat_raw <- remove_np(seurat_raw)
print("Remove NP cells complete")

retain = names(which(table(seurat_raw$gene) > 30))
# Ensure 'CTRL' is included in the retained genes
retain <- if (CONTROL %in% retain) {
  retain
} else {
  c(retain, CONTROL)
}

bcs = rownames(seurat_raw@meta.data)[seurat_raw@meta.data$gene %in% retain]
seurat_raw = subset(seurat_raw, cells = bcs)
print("Remove perturbation which less than 30 cells complete")

seurat_raw <- FindVariableFeatures(seurat_raw, assay = 'RNA') %>% ScaleData(assay = 'RNA')
DefaultAssay(object = seurat_raw) <- 'PRTB'
VariableFeatures(object = seurat_raw) <- VariableFeatures(object = seurat_raw[["RNA"]])
seurat_raw <- ScaleData(object = seurat_raw, do.scale = F, do.center = T)
seurat_raw <- RunPCA(object = seurat_raw, reduction.key = 'prtbpca', reduction.name = 'prtbpca')
seurat_raw <- FindNeighbors(seurat_raw, dims = 1:10, reduction = "prtbpca")
seurat_raw <- FindClusters(seurat_raw, algorithm = 4, resolution = 1, method = 'igraph') 
seurat_raw <- RunUMAP(seurat_raw, dims = 1:10, reduction = 'prtbpca', reduction.key = 'prtbumap', reduction.name = 'prtbumap')
seurat_raw@meta.data <- seurat_raw@meta.data %>% rename(leiden = seurat_clusters)

print("Dimension reduction and clustering analysis of PRTB assay complete")

DefaultAssay(object = seurat_raw) <- 'RNA'
seurat_fil <- seurat_raw

save(seurat_fil, file = 'filtered_data.RData')
SaveH5Seurat(seurat_fil, filename = 'filtered_qupici.h5Seurat', overwrite = T)

timeend <- Sys.time()
runningtime <- timeend-timestart
print(runningtime)