source("install_dependencies.R")

# 加载stringr包，提供强大的字符串处理功能
# Load the stringr package, which provides powerful string processing functions
library(stringr)

# 设置系统环境，使R显示英文错误信息
# Set the system environment to display English error messages
Sys.setenv(LANGUAGE = "en") 

# 全局选项设置：禁止将字符型变量自动转换为因子类型
# Global option setting: Prevent automatic conversion of character variables to factors
options(stringsAsFactors = FALSE) 

# 功能：基于主成分分析(PCA)绘制批次效应图
# Function: Draw batch effect plot based on Principal Component Analysis (PCA)
# 参数：
# indata: 输入数据矩阵，行表示特征，列表示样本
# batch: 批次信息数据框
# batchvar: 批次变量名称，默认为批次数据框的列名
# fig.dir: 图片输出目录
# PCA.fig.title: PCA图的标题
# pos1: 第一个图例的位置，默认为"bottomright"
# pos2: 第二个图例的位置，默认为"topright"
# xy: 要绘制的主成分，默认为前两个主成分
# pch.orginal: 原始点的形状，默认为16
# cols: 自定义颜色，默认为NULL
# showID: 是否显示样本ID，默认为FALSE
# cex: 点的大小，默认为1
# showLegend: 是否显示图例，默认为TRUE
# batch1move: 第一个批次图例的移动距离，默认为0.3
# batch2move: 第二个批次图例的移动距离，默认为0.4
# width: 图片宽度，默认为5
# height: 图片高度，默认为5
# withoutgrid: 是否不带网格，默认为TRUE
# Parameters:
# indata: Input data matrix with features as rows and samples as columns
# batch: Data frame containing batch information
# batchvar: Names of batch variables, default to column names of batch data frame
# fig.dir: Directory to save figures
# PCA.fig.title: Title for the PCA plot
# pos1: Position of the first legend, default is "bottomright"
# pos2: Position of the second legend, default is "topright"
# xy: Principal components to plot, default is the first two components
# pch.orginal: Original point shape, default is 16
# cols: Custom colors, default is NULL
# showID: Whether to show sample IDs, default is FALSE
# cex: Point size, default is 1
# showLegend: Whether to show legends, default is TRUE
# batch1move: Movement distance for the first batch legend, default is 0.3
# batch2move: Movement distance for the second batch legend, default is 0.4
# width: Figure width, default is 5
# height: Figure height, default is 5
# withoutgrid: Whether to plot without grid, default is TRUE
pca2batch <- function(indata, batch, batchvar = colnames(batch), fig.dir, PCA.fig.title, pos1 = "bottomright", pos2 = "topright", xy=c(1,2), pch.orginal=16, cols=NULL, showID=FALSE, cex=1, showLegend=T, batch1move=0.3, batch2move=0.4, width=5, height=5, withoutgrid=TRUE) {
 
  # 加载ClassDiscovery包，用于PCA分析
  # Load the ClassDiscovery package for PCA analysis
  library(ClassDiscovery) 
  
  # 计算第一个批次变量的水平数
  # Calculate the number of levels for the first batch variable
  N.batch1 = length(unique(batch[,batchvar[1]]))    
  # 计算第二个批次变量的水平数
  # Calculate the number of levels for the second batch variable
  N.batch2 = length(unique(batch[,batchvar[2]]))    
  
  # 检查是否提供了自定义颜色
  # Check if custom colors are provided
  if (is.null(cols)) { 
    # 如果没有，使用彩虹色
    # If not, use rainbow colors
    cols <- rainbow(N.batch1) 
  }else{
    # 检查颜色数量是否与批次数量匹配
    # Check if the number of colors matches the number of batches
    if (length(cols) != N.batch1) {stop("cols length not equal to batch length")} 
  }           
  
  # 移除数据中的缺失值
  # Remove missing values from the data
  indata=na.omit(indata)
  # 执行PCA分析，不进行标准化
  # Perform PCA analysis without standardization
  pca <- SamplePCA(indata, usecor=F, center=F) 
  # 计算第一个主成分的方差解释比例
  # Calculate the variance explained by the first principal component
  pct1 <- round (pca@variances[xy[1]]/sum(pca@variances), digits=3)*100 
  # 计算第二个主成分的方差解释比例
  # Calculate the variance explained by the second principal component
  pct2 <- round (pca@variances[xy[2]]/sum(pca@variances), digits=3)*100 
  # 设置X轴标签文本
  # Set X-axis label text
  xlab.text = paste("PC", xy[1], ": ", as.character(pct1), "% variance", sep="")
  # 设置Y轴标签文本
  # Set Y-axis label text
  ylab.text = paste("PC", xy[2], ": ", as.character(pct2), "% variance", sep="")
  
  # 根据是否需要网格绘制不同的图形
  # Draw different plots based on whether grid is needed
  if(withoutgrid) {
    # 设置不带网格的输出文件名
    # Set output file name without grid
    outfile = file.path(fig.dir, paste(PCA.fig.title, ".withoutgrid.pdf",sep="")) 
    
    # 创建PDF设备
    # Create PDF device
    pdf(file=outfile, width = width, height = height)
    # 调整图形边距，为图例留出空间
    # Adjust plot margins to make space for legends
    par(mar = par()$mar + c(0,0,0,6)) 
    # 绘制PCA散点图
    # Plot PCA scatter plot
    plot(pca@scores[,xy[1]], pca@scores[,xy[2]], 
         cex=cex, xlab=xlab.text, ylab=ylab.text,
         col=cols[factor(batch[,batchvar[1]])], 
         pch=(pch.orginal:(pch.orginal-1+N.batch2))[factor(batch[,batchvar[2]])])
    # 添加水平和垂直参考线
    # Add horizontal and vertical reference lines
    abline(h=0, v=0, col="brown", lty=2)
  } else {
    # 设置带网格的输出文件名
    # Set output file name with grid
    outfile = file.path(fig.dir, paste(PCA.fig.title, ".withgrid.pdf",sep="")) 
    
    # 保存到PDF文件
    # Save to PDF file
    pdf(file=outfile, width = width, height = height)
    # 调整图形边距，为图例留出空间
    # Adjust plot margins to make space for legends
    par(mar = par()$mar + c(0,0,0,6)) 
    # 先绘制一次固定视窗
    # Plot once to fix the viewport
    plot(pca@scores[,xy[1]], pca@scores[,xy[2]], 
         cex=cex, xlab=xlab.text, ylab=ylab.text,
         col=cols[factor(batch[,batchvar[1]])], 
         pch=(pch.orginal:(pch.orginal-1+N.batch2))[factor(batch[,batchvar[2]])])
    # 添加灰色背景
    # Add gray background
    rect(par("usr")[1],par("usr")[3],par("usr")[2],par("usr")[4],col = "gray80") 
    
    # 再画一次覆盖灰色背景
    # Plot again to cover the gray background
    par(new=TRUE)
    plot(pca@scores[,xy[1]], pca@scores[,xy[2]], 
         cex=cex, xlab=xlab.text, ylab=ylab.text,
         col=cols[factor(batch[,batchvar[1]])], 
         pch=(pch.orginal:(pch.orginal-1+N.batch2))[factor(batch[,batchvar[2]])])
    # 添加白色网格线
    # Add white grid lines
    grid (lty = 1, col = "white",lwd=2)
  }
  
  # 如果showID为TRUE，则在图上显示样本ID
  # If showID is TRUE, display sample IDs on the plot
  if (showID) { 
    text(pca@scores[,xy[1]], pca@scores[,xy[2]], colnames(indata), lwd=1, cex=cex)
  }
  # 如果showLegend为TRUE，则添加图例
  # If showLegend is TRUE, add legends
  if(showLegend){ 
    # 允许图形绘制超出边界
    # Allow plotting outside the plot area
    par(xpd = TRUE) 
    # 添加第一个图例（批次1）
    # Add the first legend (batch 1)
    legend(pos1,fill = cols,
           legend=names(table(factor(batch[,batchvar[1]]))),
           inset=c(-batch1move,0), 
           border = NA,bty = "n")
    # 添加第二个图例（批次2）
    # Add the second legend (batch 2)
    legend(pos2,legend=names(table(factor(batch[,batchvar[2]]))),
           border = NA,bty = "n",
           inset=c(-batch2move,0), 
           pch = (pch.orginal:(pch.orginal-1+N.batch2)))
  }
  # 关闭图形设备
  # Close the graphics device
  invisible(dev.off())
}