require("SummarizedExperiment")
require("sva")
require("data.table")
require("qsmooth")
require("edgeR")
set.seed(1234)

#####################
### tcga data #######
#####################

# load in clinical data
load(file.path(data_dir, "rse_gene.Rdata"))
clin <- colData(rse_gene)

# get the expression data
exp_tcga <- assays(rse_gene)$counts
exp_tcga <- as.matrix(exp_tcga)
genes_tcga <- sapply(strsplit(rownames(exp_tcga), "\\."), function(x) x[1])
length(genes_tcga) #58073
length(unique(genes_tcga)) #57992

duplicates <- genes_tcga[duplicated(genes_tcga)] #45
exp_tcga <- exp_tcga[!genes_tcga %in% duplicates, ]
genes_tcga <- genes_tcga[!genes_tcga %in% duplicates]
rownames(exp_tcga) <- genes_tcga

# get feature annotations
anno_tcga <- rowData(rse_gene)
# remove all ensg ids that are not associated with unique gene symbols
ensg <- substr(anno_tcga[,1],1,15)
dupl <- which(duplicated(ensg))
ensgfull <- as.character(anno_tcga[,1])
pary <- substr(ensgfull, nchar(ensgfull)-4, nchar(ensgfull))
length(which(pary == "PAR_Y")) == length(dupl)
idx <- which(pary == "PAR_Y")
# all of those duplicated end with "PAR_Y" but have the same gene symbol
# therefore, remove those ending with "PAR_Y"
anno_tcga <- anno_tcga[-idx, ]
anno_tcga[,1] <- substr(anno_tcga[,1],1,15)
# remove ensg ids that are not associated with a gene symbol is.na 
anno_tcga <- as.data.frame(anno_tcga)
anno_tcga <- anno_tcga[!is.na(anno_tcga$symbol), ]
# remove ensg ids that are associated with multiple gene symbol 
anno_tcga <- anno_tcga[-grep("\\,|\\, ", anno_tcga$symbol), ]
dupl <- anno_tcga$symbol[duplicated(anno_tcga$symbol)]
anno_tcga_tab <- anno_tcga[!anno_tcga$symbol %in% dupl, ]
anno_tcga_tab$symbol <- as.character(anno_tcga_tab$symbol)
dim(anno_tcga_tab) 
head(anno_tcga_tab)

gene_symbols_tcga <- 
            anno_tcga_tab$symbol[match(genes_tcga, anno_tcga_tab$gene_id)]
anno_tcga <- data.table("ens_id" = genes_tcga,
            "gene_symbol" = gene_symbols_tcga)

anno_tcga <- anno_tcga[!is.na(anno_tcga$gene_symbol), ]
length(unique(anno_tcga$gene_symbol)) == length(anno_tcga$gene_symbol)
length(anno_tcga$ens_id) == length(unique(anno_tcga$ens_id))

#####################
### lms dkfz data ###
#####################
lms_dkfz <- fread(file.path(data_dir, "LMS_37_readCount.txt"))
dim(lms_dkfz)
head(lms_dkfz)

genes_dkfz <- sapply(strsplit(lms_dkfz$GeneID, "\\."), function(x) x[1])
lms_dkfz <- lms_dkfz[, -1]
length(genes_dkfz) == length(unique(genes_dkfz))
rownames(lms_dkfz) <- genes_dkfz

anno_dkfz <- fread(file.path(data_dir, "GN_ensemblID_symbol.txt")) ##57820
head(anno_dkfz); dim(anno_dkfz)
length(unique(anno_dkfz$gn)) ##55765
anno_dkfz$ens_short <- sapply(strsplit(anno_dkfz$V1, "\\."), function(x) x[1])
length(unique(anno_dkfz$ens_short)) ##57820
# remove gene names duplicates
duplicates_dkfz <- anno_dkfz$gn[duplicated(anno_dkfz$gn)]
anno_dkfz <- anno_dkfz[!anno_dkfz$gn %in% duplicates_dkfz, ]
gene_symbols_dkfz <- anno_dkfz$gn[match(genes_dkfz, anno_dkfz$ens_short)]
anno_dkfz <- data.table("ens_id" = genes_dkfz,
            "gene_symbol" = gene_symbols_dkfz )
anno_dkfz  <- anno_dkfz[!is.na(anno_dkfz$gene_symbol), ]
length(anno_dkfz$gene_symbol) == length(unique(anno_dkfz$gene_symbol))

####################################
### combining tcga and dkfz data ###
####################################

olap_tcga <- anno_tcga[anno_tcga$gene_symbol %in% anno_dkfz$gene_symbol, ]
olap_dkfz <- anno_dkfz[anno_dkfz$gene_symbol %in% anno_tcga$gene_symbol, ]

rn_dkfz <- rownames(lms_dkfz)[rownames(lms_dkfz) %in% olap_dkfz$ens_id]
rn_dkfz <- olap_dkfz$gene_symbol[match(rn_dkfz, olap_dkfz$ens_id)]
lms_dkfz <- lms_dkfz[rownames(lms_dkfz) %in% olap_dkfz$ens_id, ]
rownames(lms_dkfz) <- rn_dkfz

rn_tcga <- rownames(exp_tcga)[rownames(exp_tcga) %in% olap_tcga$ens_id]
rn_tcga <- olap_tcga$gene_symbol[match(rn_tcga, olap_tcga$ens_id)]
exp_tcga <- exp_tcga[rownames(exp_tcga) %in% olap_tcga$ens_id, ]
rownames(exp_tcga) <- rn_tcga


exp_tcga <- exp_tcga[rownames(lms_dkfz), ]
all(rownames(lms_dkfz) == rownames(exp_tcga)) # TRUE

# merge the data
data_merged  <- cbind(exp_tcga, lms_dkfz)
rownames(data_merged) <- rownames(exp_tcga)

rn_data_merged <- rownames(data_merged)
#################################################
### perform batch correction with combat-seq ###
#################################################

groups_tcga <- colData(rse_gene)$gdc_cases.project.project_id
ids <- rownames(colData(rse_gene))
all(ids == colnames(exp_tcga))
id_cancer_tab <- data.table("id" = colnames(exp_tcga),
                "tcga_cancer" = groups_tcga)

groups_lms <- rep("TCGA-SARC", 37)
groups_all <- c(groups_tcga, groups_lms)

mdata <- cbind.data.frame("sample" = colnames(data_merged), 
                "batch" = ifelse(grepl('LMS', colnames(data_merged)),
                 "batch_1", "batch_2"), "cancer" = groups_all)

idx_sarc <- grep("SARC|LMS", groups_all)
groups_sarc <- groups_all[idx_sarc]



data_sarc <- data_merged[, idx_sarc, with = FALSE]
mdata_sarc <- mdata[idx_sarc, ]

adjusted_counts <- ComBat_seq(as.matrix(data_sarc), batch = mdata_sarc$batch)

#################################################
### perform qsmooth normalization  ##############
#################################################

remain_idx <- !colnames(data_merged) %in% colnames(adjusted_counts)
remain_tcga <- data_merged[, remain_idx, with = FALSE]
dim(remain_tcga)

data_merged_corr <- cbind(remain_tcga, adjusted_counts)
group <- mdata$cancer[match(colnames(data_merged_corr), mdata$sample)]
norm_counts <- qsmooth(object = as.matrix(data_merged_corr),
                group_factor = group)

norm_counts <- qsmoothData(norm_counts)
log_norm_counts <- log(norm_counts + 1)
rownames(log_norm_counts) <- rn_data_merged

#################################################
### preparing exp for PANDA + LIONESS  ##########
#################################################

# read in PPI table
ppi <- fread(file.path(data_dir, "ppi.txt"))
# read in motif prior
prior <- fread(file.path(data_dir, "prior.txt"))

# check for overlaps
table(ppi$V1 %in% rownames(log_norm_counts))
table(ppi$V2 %in% rownames(log_norm_counts))
table(prior$V1 %in% rownames(log_norm_counts))
table(prior$V2 %in% rownames(log_norm_counts))
table(rownames(log_norm_counts) %in% prior$V2)

# clean exp to contain target genes in prior
exp <- log_norm_counts[rownames(log_norm_counts) %in% prior$V2, ]
write.table(exp, file.path(data_dir, "exp.txt"),
            col.names = T, row.names = T, sep = "\t", quote = F)
