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

# Load required libraries
library(RColorBrewer)
library(circlize)
library(gplots)
library(viridis)
library(oompaBase)
library(tidyverse)
Sys.setenv(LANGUAGE = "en") # Display error messages in English
options(stringsAsFactors = FALSE) # Disable converting characters to factors

# Standardize function definition
standarize.fun <- function(indata=NULL, halfwidth=NULL, centerFlag=T, scaleFlag=T) {  
  outdata = t(scale(t(indata), center = centerFlag, scale = scaleFlag))
  if (!is.null(halfwidth)) {
    outdata[outdata > halfwidth] = halfwidth
    outdata[outdata < (-halfwidth)] = -halfwidth
  }
  return(outdata)
}

# Load data from merge.RDATA file
load("merge.RDATA")
dataa = outTab
colnames(dataa) = str_replace_all(colnames(dataa), "TCGA_", "")
colnames(dataa) = str_replace_all(colnames(dataa), "GSE17538_", "")
colnames(dataa) = str_replace_all(colnames(dataa), "GSE39582_", "")

# Read data from Cytokine.csv file
aaa = read.csv("Cytokine.csv", header = T, sep = ",")
genename1 = intersect(aaa$ID, rownames(dataa))
dataa = dataa[genename1, ]
aaa = dplyr::filter(aaa, ID %in% genename1)

# Read risk data from score.group.txt file
risk = read.table("score.group.txt", header = T, sep = "\t", row.names = 1)
sameid = intersect(colnames(dataa), rownames(risk))
dataa = dataa[, sameid]
risk = risk[sameid, ]
hmdat = as.data.frame(t(dataa))

# Define classification information
type = data.frame(Type = aaa$Type)
row.names(type) = aaa$ID
Typeid <- type$Type

# Create annotations
annCol <- data.frame(score = scale(risk$score),
                     RiskType = risk$group,
                     row.names = rownames(risk),
                     stringsAsFactors = F)

# Perform gene differential analysis
bbb = annCol[rownames(hmdat), ]
sampleType = bbb$RiskType
sigVec = c()
for(i in colnames(hmdat)) {
  test = wilcox.test(hmdat[, i] ~ sampleType)
  pvalue = test$p.value
  Sig = ifelse(pvalue < 0.001, "***", ifelse(pvalue < 0.01, "**", ifelse(pvalue < 0.05, "*", "")))
  sigVec = c(sigVec, paste0(i, Sig))
}
colnames(hmdat) = sigVec

# Standardize the data
indata = t(hmdat)
indata = indata[, colSums(indata) > 0]

# Define colors
Type.col <- brewer.pal(n = length(unique(Typeid)), name = "Paired")

# Row annotations on the left side of the heatmap
annRow <- data.frame(Type = factor(Typeid, levels = unique(Typeid)),
                     row.names = colnames(hmdat),
                     stringsAsFactors = F)

# Set colors for annotations
annColors <- list(Type = c("Chemokines and receptors" = Type.col[1],
                           "Interieukines and receptors" = Type.col[2],
                           "Interferons and receptors" = Type.col[3],
                           "Other cytokines" = Type.col[4]),
                  "score" = greenred(64),
                  "RiskType" = c("High" = "red", "Low" = "blue"))

# Sort samples by risk score
samorder <- rownames(risk[order(risk$score), ])

# Standardize the data
plotdata <- standarize.fun(indata, halfwidth = 2)

# Create a heatmap using pheatmap
pheatmap::pheatmap(mat = as.matrix(plotdata[, samorder]),
                   border_color = NA,
                   color = bluered(64),
                   cluster_rows = F,
                   cluster_cols = F,
                   show_rownames = T,
                   show_colnames = F,
                   annotation_col = annCol[samorder,, drop = F],
                   annotation_row = annRow,
                   annotation_colors = annColors,
                   gaps_col = table(annCol$RiskType)[2],
                   gaps_row = cumsum(table(annRow$Type)),
                   cellwidth = 0.8,
                   cellheight = 10,
                   filename = "immune heatmap by pheatmap_score.pdf")

# Calculate means for High and Low risk groups
ccc = hmdat[rownames(annCol)[annCol$RiskType == "High"], ]
ddd = hmdat[rownames(annCol)[annCol$RiskType == "Low"], ]
ccc1 = colMeans(ccc)
ddd1 = colMeans(ddd)
mewdata = data.frame(High = ccc1, Low = ddd1)

# Create a new heatmap for means
plotdata1 <- mewdata
annCol1 = data.frame(RiskType = c("High", "Low"))
row.names(annCol1) = c("High", "Low")

pheatmap::pheatmap(mat = as.matrix(mewdata[, row.names(annCol1)]),
                   border_color = "grey",
                   color = bluered(64),
                   cluster_rows = F,
                   cluster_cols = F,
                   show_rownames = T,
                   show_colnames = F,
                   annotation_col = annCol1,
                   annotation_row = annRow,
                   annotation_colors = annColors,
                   gaps_row = cumsum(table(annRow$Type)),
                   cellwidth = 20,
                   cellheight = 10,
                   filename = "immune heatmap by pheatmap_score_mean.pdf")