# Clear the current workspace
rm(list = ls())

# Load the required libraries
library(pheatmap)         

# Define input files
expFile <- "geneexpr.txt"          # Expression input file
clusterFile <- "Cluster.txt"      # Clustering results file
cliFile <- "clinicaldata.txt"     # Clinical data file

# Read input files
exp <- read.table(expFile, header = TRUE, sep = "\t", check.names = FALSE, row.names = 1)
exp <- t(exp)
cluster <- read.table(clusterFile, header = TRUE, sep = "\t", check.names = FALSE, row.names = 1)

# Merge expression and clustering data
sameSample <- intersect(row.names(exp), row.names(cluster))
exp <- exp[sameSample, , drop = FALSE]
cluster <- cluster[sameSample, , drop = FALSE]
expCluster <- cbind(exp, cluster)

# Extract project information
Project <- gsub("(.*?)\\_.*", "\\1", rownames(expCluster))
library(tidyverse)
rownames(expCluster) <- str_replace_all(rownames(expCluster), "TCGA_", "")
rownames(expCluster) <- str_replace_all(rownames(expCluster), "GSE17538_", "")
rownames(expCluster) <- str_replace_all(rownames(expCluster), "GSE39582_", "")
expCluster <- cbind(expCluster, Project)

# Merge clinical data
cli <- read.table(cliFile, header = TRUE, sep = "\t", check.names = FALSE, row.names = 1)
sameSample <- intersect(row.names(expCluster), row.names(cli))
expCluster <- expCluster[sameSample,, drop = FALSE]
cli <- cli[sameSample,, drop = FALSE]
data <- cbind(expCluster, cli)
data$fustat[data$fustat == 1] <- "Dead"
data$fustat[data$fustat == 0] <- "Alive"

# Extract heatmap data
data <- data[order(data$cluster), ]
Type <- data[, ((ncol(exp) + 1):ncol(data))]
data <- t(data[, 1:ncol(exp)])

# Cluster colors
bioCol <- c("#0066FF", "#FF9900", "#FF0000", "#6E568C", "#7CC767", "#223D6C", "#D20A13", "#FFD121", "#088247", "#11AA4D")
ann_colors <- list()
CluCol <- bioCol[1:length(levels(factor(Type$cluster)))]
names(CluCol) <- levels(factor(Type$cluster))
ann_colors[["cluster"]] <- CluCol

# Visualize the heatmap
pdf("heatmap.pdf", height = 6, width = 9)
pheatmap(data,
         annotation = Type,
         annotation_colors = ann_colors,
         color = colorRampPalette(c(rep("green", 5), "white", rep("red", 5)))(100),
         cluster_cols = FALSE,
         cluster_rows = FALSE,
         scale = "row",
         show_colnames = FALSE,
         fontsize = 6,
         fontsize_row = 6,
         fontsize_col = 6)
dev.off()