library(plyr)
library(dplyr)
library(seqinr)
library(stringr)

source("SOURCE/loadFunctions.r")

print("--------------------------------------------")
print("1) PARSE SEARCH RESULT FILES AND CREATE MSDB")
print("--------------------------------------------")

project_name = snakemake@params[["project_name"]]
include_scanNum = snakemake@params[["include_scanNum"]]

SEARCHRESULTS_PATH = "INPUT/search_results/"


### INPUT ###
sample_list = read.csv(file = snakemake@input[["sample_list"]],
                       sep = ",", header = T, stringsAsFactors = F)

# sample_list = read.table(file = "INPUT/sample_list.csv",
#                          sep = ",", header = T, stringsAsFactors = F)

# filter sample list
sample_list = sample_list[sample_list$project_name == project_name, ]

# add meta-information
if (!any(is.na(sample_list$metainformation)) & length(unique(sample_list$metainformation)) == 1) {
  metainfo = read.csv(sample_list$metainformation[1], stringsAsFactors = F)
  meta = T
} else {
  meta = F
}

### MAIN PART ###

if (!dir.exists(paste0("OUTPUT/", project_name))) {
  dir.create(paste0("OUTPUT/", project_name))
}

columns = c(            "runID" ,
                        "substrateID" ,
                        "substrateSeq" ,
                        "digestTime" ,
                        "pepSeq" ,
                        "productType",
                        "spliceType",
                        "positions",
                        "scanNum",
                        "rank",
                        "ionScore",
                        "qValue",
                        "charge",
                        "PTM")

MSDB = list()
pb = txtProgressBar(min = 0, max = nrow(sample_list), style = 3)

for (i in 1:nrow(sample_list)) {
  # setTxtProgressBar(pb, i)

  # load search result file
  con = file(paste0(SEARCHRESULTS_PATH, sample_list$filename[i]),"r")
  i1 = readLines(con)
  close(con)

  # remove header from search result file
  # and remove not assigned queries if there are any
  start = grep("prot_hit_num",i1)

  currentSearchFile = read.table(paste(SEARCHRESULTS_PATH, sample_list$filename[i],sep="/"),
                                 skip = start-1,
                                 sep = ",", header = TRUE, fill = TRUE)

  kk = which(currentSearchFile$prot_hit_num == "")
  if (length(kk) > 0) {
    kk = min(kk)
    currentSearchFile = currentSearchFile[c(1:(kk-2)),]
  }

  substrateID = sample_list$substrateID[i]
  digestTime = sample_list$digestTime[i]


  if(nrow(currentSearchFile) > 0){

    # empty DB for current search result file
    currentDB = data.frame(matrix(ncol = length(columns), nrow = nrow(currentSearchFile)))
    colnames(currentDB) = columns

    # fill data frame with info straight outta .csv
    currentDB$charge <- currentSearchFile$pep_exp_z
    currentDB$rank <- currentSearchFile$pep_rank
    currentDB$pepSeq <- currentSearchFile$pep_seq
    currentDB$ionScore <- currentSearchFile$pep_score
    currentDB$qValue <- currentSearchFile$pep_expect


    currentDB$PTM <- currentSearchFile$pep_var_mod
    currentDB$scanNum <- currentSearchFile$pep_query

    # currentDB[which(duplicated(currentDB) | duplicated(currentDB, fromLast = T)),] %>% nrow()
    # cntS = currentSearchFile[, c("pep_exp_z","pep_rank","pep_seq", "pep_score", "pep_expect", "pep_var_mod", "pep_query")]
    # cntS[which(duplicated(cntS) | duplicated(cntS, fromLast = T)),] %>% nrow()

    #get actual scan numbers, not query numbers
    # varies between Mascot and Mascot Distiller
    if(include_scanNum != "no") {

      warn = F

      scans = rep(NA,dim(currentSearchFile)[1])
      for(ii in 1:dim(currentSearchFile)[1]){
        tit = currentSearchFile$pep_scan_title[ii]

        if (str_detect(tit, pattern = "scan=")) {

          scans[ii] = as.numeric(strsplit(strsplit(tit,split="scan=")[[1]][2],split="~")[[1]][1])

        } else if (str_detect(tit, pattern = "Scan ")) {

          scans[ii] = as.numeric(unlist(strsplit(tit, " "))[3])

        } else if (str_detect(tit, pattern = coll(".")) | str_detect(tit, pattern = "TSN")) {
          scans[ii] = as.numeric(strsplit(tit,split="\\.")[[1]][2])

        } else {
          warn = T
        }

      }
      currentDB$scanNum <- scans

      if (warn) {
        print("could not extract scan numbers - check search result formatting!")
        currentDB$scanNum = currentSearchFile$pep_query
      }

    }


    #extract product type and position information
    pepName <- as.character(currentSearchFile$prot_acc)
    pepType <- c()
    pepPos <- c()
    for (z in 1:length(pepName)){
      splitPepName <- unlist(strsplit(pepName[z], "_"))
      pepType <- c(pepType, splitPepName[1])
      splitPepName <- splitPepName[-1]
      pepPos <- c(pepPos, paste(splitPepName, collapse = "_"))

    }

    currentDB$productType <- pepType
    currentDB$positions <- pepPos

    # substrate sequence
    if (grepl(pattern = ".fasta", sample_list$substrateSeq[i])) {

      cntSeq = read.fasta(file = paste0("INPUT/sequences/", sample_list$substrateSeq[i]),
                          seqtype = "AA", strip.desc = T)[[1]] %>%
        unlist() %>%
        paste(collapse = "")

    } else {
      cntSeq = sample_list$substrateSeq[i]
    }

    # add metainfo
    currentDB$substrateID = substrateID
    currentDB$substrateSeq <- cntSeq
    currentDB$digestTime <- digestTime
    currentDB$runID = paste(substrateID, digestTime,sample_list$filename[i], sample_list$replicate[i], sep = "-")

    # add further meta-information if provided
    if (meta) {

      cntMeta = metainfo[metainfo$filename == sample_list$filename[i], ]
      currentDB = cbind(currentDB, cntMeta[rep(1,nrow(currentDB)),]) %>%
        as.data.frame()

      dp = which(duplicated(names(currentDB)))
      if (length(dp) > 0) { currentDB = currentDB[,-dp]}

    }

    # assign substrate hits as PCP
    substratehits = which(currentDB$substrateSeq == currentDB$pepSeq)
    if (length(substratehits) > 0) {
      currentDB$productType[substratehits] = "PCP"
      currentDB$positions[substratehits] = paste(1,
                                                 nchar(currentDB$substrateSeq[substratehits]),
                                                 sep = "_")
    }

    print("")
    paste0("obtain correct annotations for positions and splice types for run: ",
           currentDB$runID[1]) %>%
      print()

    currentDB = currentDB %>% mapping()
    currentDB$spliceType[currentDB$productType == "PCP"] = NA

  } else {
    currentDB = data.frame(matrix(ncol = length(columns), nrow = 0))
    colnames(currentDB) = columns
    print("No peptides listed in this file! Skipping...")
  }

  # append data frame for current search file to master data frame
  MSDB[[i]] = currentDB

}

MSDB = plyr::ldply(MSDB)

### OUTPUT ###
save(MSDB, file = unlist(snakemake@output[["MSDB"]]))

library(dplyr)
library(stringr)
library(grid)
library(ggplot2)
library(gridExtra)

theme_set(theme_bw())

source("SOURCE/invitroSPI_utils.R")
source("SOURCE/plottingFunctions.R")

print("-------------------------------")
print("5) GENERATE DATABASE STATISTICS")
print("-------------------------------")


### INPUT ###
ProteasomeDB = read.csv(snakemake@input[["ProteasomeDB"]],
                        stringsAsFactors = F)
# ProteasomeDB = read.csv("OUTPUT/test_data/ProteasomeDB.csv", stringsAsFactors = F)

S = ProteasomeDB$substrateID %>% unique()
tps = ProteasomeDB$digestTime %>% unique() %>% as.numeric() %>% sort()


### MAIN PART ###

# ----- 1) pre-processing -----

Peps = ProteasomeDB %>%
  remSynthErrors() %>%
  ILredundancy() %>%
  filterPepLength() %>%
  disentangleMultimappers.Type() %>%
  disentangleMultimappers.SRlen() %>%
  disentangleMultimappers.IVSeqLen() %>%
  disentangleMultimappers.AA() %>%
  DBcosmetics()

ProteasomeDB = ProteasomeDB %>%
  remSynthErrors()


# ----- 2) general stats + peptides over time -----

uniquePeps = Peps %>% uniquePeptides()

# plotNumberofPeptides(ProteasomeDB,
#                      outname = "OUTPUT/test_data/number_of_products.pdf")
plotNumberofPeptides(ProteasomeDB,
                     outname = unlist(snakemake@output[["number_of_products"]]))

# pdf("OUTPUT/test_data/DB_stats.pdf", height = 4, width = 6)
pdf(file = unlist(snakemake@output[["DB_stats"]]), height = 4, width = 6)
generalStats(uniquePeps, tp = "all") %>% grid.table()

for (t in 1:length(tps)) {
  cntDB = Peps[Peps$digestTime == tps[t], ] %>%
    uniquePeptides()
  grid.newpage()
  generalStats(DB = cntDB, tp = tps[t]) %>% grid.table()

}
dev.off()



# ----- 3) coverage maps -----

# pdf("OUTPUT/test_data/coverage_map.pdf", height = 12, width = 16)
# pdf(unlist(snakemake@output[["coverage_map"]]), height = 12, width = 16)
# for (s in 1:length(S)) {
#   
#   cntDB = uniquePeps[uniquePeps$substrateID == S[s], ]
#   out = plotCoverage(cntDB, name = S[s], tp = "all")
#   replayPlot(out)
#   
# }
# 
# dev.off()

# ----- 4) length distributions + product type frequencies -----

# pdf("OUTPUT/Delta03/length_distributions.pdf", height = 10, width = 14)
pdf(file = unlist(snakemake@output[["length_distributions"]]), height = 12, width = 16)

freq = TypeFrequency(uniquePeps)
pl = PepLength(uniquePeps, tp="all")
srlen = SRLength(uniquePeps, tp="all")
ivlen = IVSeqLength(uniquePeps, tp="all")

cntG = grid.arrange(freq,pl,srlen,ivlen,
                    nrow = 2, ncol = 2)

print(cntG)

for (t in 1:length(tps)) {

  cntDB = uniquePeps[uniquePeps$digestTime == tps[t], ]

  freq = TypeFrequency(cntDB)
  pl = PepLength(cntDB, tp=tps[t])
  srlen = SRLength(cntDB, tp=tps[t])
  ivlen = IVSeqLength(cntDB, tp=tps[t])

  cntG = grid.arrange(freq,pl,srlen,ivlen,
                      nrow = 2, ncol = 2)

  print(cntG)

}
dev.off()

library(data.table)
library(parallel)
library(dplyr)
library(seqinr)
library(stringr)
library(beepr)


# numCPU = if (detectCores() > 8) detectCores()-1 else if (detectCores() >= 4) 4 else 1
# cl <- makeForkCluster(numCPU)


# ----- user parameters -----
nmers = c(5, 40)
generateTrans = T  # if false, generate only cis peptides
outdir = "../DATA/polypeptideDBs/"
suppressWarnings(dir.create(outdir, recursive = T))

# change here (you can read in a dataframe with the name and the sequence)
sample_list = read.csv("~/Documents/SCIENCE/_tools/aSPIre+invitroSPI/data/sample_list_aSPIre.csv", stringsAsFactors = F)
substrates = sample_list %>%
  select(substrateID, substrateSeq) %>%
  unique()


# ----- functions -----

getCis = function(subSeq, N, Lext,L) {

  if (L-1 >= N) {

    allCis = sapply(seq(1,L-N), function(i){

      sapply(seq(i+Lext-1,N-Lext+i-1), function(j){

        sr2 = N-j+i-1
        sapply(seq(j+2, L-sr2+1), function(k){

          n = k+sr2-1
          pepSeq = paste(str_sub(subSeq, start=i, end=j),
                         str_sub(subSeq, start=k, end=n), sep = "")
          positions = paste(i,j,k,n, sep = "_")

          return(c(pepSeq, positions))

        })
      })
    })
    CIS = unlist(allCis)
    cisPep = data.frame(pepSeq = CIS[seq(1,length(CIS),2)],
                        PositionSubstrate = CIS[seq(2,length(CIS),2)])
  } else {
    cisPep=NA
  }

  return(cisPep)
}



getRevCis = function(subSeq, L, N, Lext=1) {


  if (L >= N) {

    allRevCis = sapply(seq(1,L-N+1), function(k){

      sapply(seq(k+Lext-1,N-Lext+k-1), function(n){

        sr1 = N-n+k-1
        sapply(seq(n+1, L-sr1+1), function(i){

          j = i+sr1-1

          pepSeq = paste(str_sub(subSeq, start=i, end=j),
                         str_sub(subSeq, start=k, end=n), sep = "")
          positions = paste(i,j,k,n, sep = "_")
          return(c(pepSeq, positions))
        })
      })
    })

    REVCIS = unlist(allRevCis)
    revcisPep = data.frame(pepSeq = REVCIS[seq(1,length(REVCIS),2)],
                           PositionSubstrate = REVCIS[seq(2,length(REVCIS),2)])

  } else {
    revcisPep = NA
  }


  return(revcisPep)
}




getAllTrans = function(L, N, subSeq, Lext=1) {

  allSRs = lapply(seq(Lext,N-Lext), function(sr1){

    # SR1
    Is = seq(1, L-sr1+1)
    Js = Is+sr1-1

    # SR2
    Ks = seq(1,L-N+sr1+1)
    Ns = Ks+N-sr1-1

    SR1_pos = paste(Is,Js,sep = "_")
    SR2_pos = paste(Ks,Ns,sep = "_")

    SR1 = str_sub(subSeq,Is,Js)
    SR2 = str_sub(subSeq,Ks,Ns)

    PEP = outer(SR1, SR2, paste, sep="") %>% as.vector()
    positions = outer(SR1_pos,SR2_pos,paste, sep = "_") %>% as.vector()

    return(data.frame(pepSeq = PEP,
                      PositionSubstrate = positions
    ))

  })

  transPep = plyr::ldply(allSRs)


  return(transPep)
}


getPCP = function(L,subSeq,N) {

    if (L > N) {

    allPCP = sapply(seq(1,L-N+1), function(i){
      j = N+i-1
      positions = paste(i,j, sep = "_")
      pepSeq=str_sub(subSeq,i,j)
      return(c(pepSeq,positions))
    })

    PCP = unlist(allPCP)
    PCP_Pep = data.frame(pepSeq = PCP[1,],
                         PositionSubstrate = PCP[2,])
  } else {
    PCP = NA
  }

  return(PCP_Pep)
}


# ----- compute DB -----

for (i in 1:nrow(substrates)) {


  Name = substrates$substrateID[i]
  print("...............")
  print(Name)
  print("...............")

  # get substrate sequence and length
  subSeq = substrates$substrateSeq[i]
  if (grepl(".fasta", subSeq)) {
    subSeq = read.fasta(file = subSeq, seqtype = "AA", strip.desc = T, as.string = T) %>%
      unlist() %>%
      suppressWarnings()
  }

  L = nchar(subSeq)

  # initiate empty fasta
  write.fasta(sequences = subSeq, names = Name,
              file.out=paste0(outdir, Name, "_allPSP_0.fasta"),
              open="w")

  start_time = Sys.time()
  # generate peptides
  if (generateTrans) {
    print("generating cis and trans-spliced peptides.....")

    # all cis and trans
    allPSP <- lapply(seq(nmers[1], nmers[2]), function(N){
      paste0("generating ",N,"-mers") %>% print()

      DB = getAllTrans(subSeq = subSeq, L = L, N = N, Lext=1)
      write.fasta(as.list(DB$pepSeq), names = paste0("PSP_", DB$PositionSubstrate),
                  file.out = paste0(outdir, Name, "_allPSP_0.fasta"),
                  open="a")
    })


    print("generating non-spliced peptides.....")
    Nmax = if(nmers[2] >= L-1) L-1 else nmers[2]

    allPCP <- lapply(seq(nmers[1], Nmax), function(N){
      paste0("generating ",N,"-mers") %>% print()

      DB = getPCP(subSeq = subSeq, L = L, N = N)
      write.fasta(as.list(DB$pepSeq), names = paste0("PSP_", DB$PositionSubstrate),
                  file.out = paste0(outdir, Name, "_allPSP_0.fasta"),
                  open="a")
    })

  } else {
    print("generating cis-spliced peptides.....")

    # forward cis
    allCis <- lapply(seq(nmers[1], nmers[2]), function(N){
      paste0("generating fwd cis ",N,"-mers") %>% print()

      DB = getCis(subSeq = subSeq, L = L, N = N, Lext=1)
      write.fasta(as.list(DB$pepSeq), names = paste0("PSP_", DB$PositionSubstrate),
                  file.out = paste0(outdir, Name, "_allPSP_0.fasta"),
                  open="a")
    })

    # reverse cis
    allRevCis <- lapply(seq(nmers[1], nmers[2]), function(N){
      paste0("generating rev cis ",N,"-mers") %>% print()

      DB = getRevCis(subSeq = subSeq, L = L, N = N, Lext=1)
      write.fasta(as.list(DB$pepSeq), names = paste0("PSP_", DB$PositionSubstrate),
                  file.out = paste0(outdir, Name, "_allPSP_0.fasta"),
                  open="a")
    })


  }

  end_time = Sys.time()
  paste0("..... finished after: ", difftime(end_time, start_time, units = "s")) %>% print()

}


print("FINISHED!")
beep("mario")