# Module 1: from 8 kallisto outputs to gene-level count and TPM tables.
# Run with the working directory set to the "workshop" folder:  setwd("path/to/workshop")
dir.create("results", showWarnings = FALSE)

# ---- 1. read every abundance.tsv into transcript-level tables -------------------------------
result_dir <- "data/kallisto"          # one sub-folder per sample (unzipped KallistoResults.zip)
samples <- list.dirs(result_dir, recursive = FALSE, full.names = FALSE)

count_tr <- NULL
tpm_tr <- NULL
for (s in samples) {
  ab <- read.delim(file.path(result_dir, s, "abundance.tsv"))
  if (is.null(count_tr)) {
    count_tr <- data.frame(row.names = ab$target_id)
    tpm_tr <- data.frame(row.names = ab$target_id)
    len_tr <- data.frame(row.names = ab$target_id)
  }
  stopifnot(identical(rownames(count_tr), ab$target_id))   # same transcripts, same order
  count_tr[[s]] <- ab$est_counts
  tpm_tr[[s]] <- ab$tpm
  len_tr[[s]] <- ab$eff_length
}
dim(count_tr)
head(count_tr)

# ---- 2. transcripts -> genes, keeping protein coding transcripts only ------------------------
mapping <- read.delim("data/mart_export.txt")          # from Ensembl BioMart (see handout)
gene_of <- mapping$Gene.name[match(rownames(count_tr), mapping$Transcript.stable.ID.version)]
keep <- !is.na(gene_of)
cat("transcripts in the index:", nrow(count_tr), "\n")
cat("protein coding transcripts kept:", sum(keep), "\n")

count_gn <- rowsum(as.matrix(count_tr[keep, ]), group = gene_of[keep])   # sum transcripts per gene
tpm_gn <- rowsum(as.matrix(tpm_tr[keep, ]), group = gene_of[keep])
cat("genes:", nrow(count_gn), "\n")

# ---- 3. save ----------------------------------------------------------------------------------
write.table(data.frame(Gene.name = rownames(count_gn), count_gn, check.names = FALSE),
            "results/count_gene.txt", sep = "\t", quote = FALSE, row.names = FALSE)
write.table(data.frame(Gene.name = rownames(tpm_gn), tpm_gn, check.names = FALSE),
            "results/tpm_gene.txt", sep = "\t", quote = FALSE, row.names = FALSE)
saveRDS(list(count = count_tr, tpm = tpm_tr, eff_length = len_tr), "results/transcript_level.rds")

# ---- 4. check against the provided tables -----------------------------------------------------
ref <- read.delim("data/count_gene.txt", row.names = 1, check.names = FALSE)
common <- intersect(rownames(ref), rownames(count_gn))
cat("genes in the provided table:", nrow(ref), " shared with ours:", length(common), "\n")
cat("largest absolute difference:", max(abs(as.matrix(ref[common, samples]) - count_gn[common, samples])), "\n")

# ---- 5. a gene with several transcripts: Pkm ---------------------------------------------------
pkm <- rownames(count_tr)[keep & gene_of == "Pkm"]
round(count_tr[pkm, ], 1)
colSums(count_tr[pkm, ])