# Laura Aarnos 190213
# meaRtools: write burst summary
# original functions from meaRtools 2018, modification made by LAa and marked with _mod

# code for write_plate_summary_for_bursts

write_plate_summary_for_bursts_mod <- function(s, outputdir) {
  master_sum <- .get_mean_burst_info_per_well_mod(s)
  csvwell <- paste(outputdir, "/", get_project_plate_name(s[[1]]$file),
                   "_well_bursts.csv", sep = "")
  
  for (i in 1:length(s)) {
    div <- .get_div(s[[i]])
    basename <- get_file_basename(s[[i]]$file)
    csvfile <- paste(outputdir, "/", basename, "_bursts.csv", sep = "")
    
    ########## data frame summarized over well
    # get number of object in master_sum[[1]] list
    tempdf <- c(); tempcolnames <- c()
    for (j in 2:length(master_sum[[i]])) {
      tempc <- unlist(master_sum[[i]][j])
      tempdf <- cbind(tempdf, tempc)
      tempcolnames <- c(tempcolnames, names(master_sum[[i]][j]))
    } # end of loop through master_sum list objects
    
    # need to switch around columns so first columns come first
    if (dim(tempdf)[2] > 20) {
      if (dim(tempdf)[1] == 1) {
        df <- cbind(t(tempdf[, 21:25]), t(tempdf[, 1:20]))
      } else {
        df <- cbind(tempdf[, 21:25], tempdf[, 1:20])
      }
      colnames <- c(tempcolnames[21:25], tempcolnames[1:20])
      colnames(df) <- colnames
    }
    
    ################## channel by channel burst summary
    
    # meta data and misc
    # get vector of which wells are active
    wellindex <- which(is.element(names(s[[i]]$treatment),
                                  unique(s[[i]]$cw)))
    well <- c(); file <- c();
    
    file <- rep(strsplit(basename(s[[i]]$file), ".RData")[[1]][1],
                length(s[[i]]$cw))
    well <- s[[i]]$cw
    
    # channel data frame
    df2 <- cbind(file, well, as.data.frame(s[[i]]$bs[1:length(s[[i]]$bs)]))
    
    
    # write a title
    write.table("Burst Analysis Averaged Over Each Well",
                csvfile, sep = ",", append = FALSE, row.names = FALSE, col.names = FALSE)
    
    
    write.table(paste("file= ", strsplit(basename(s[[i]]$file),
                                         ".RData")[[1]][1], sep = ""),
                csvfile, sep = ",", append = TRUE, row.names = FALSE, col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    # recording time
    write.table(paste("recording time (s): [", paste(s[[i]]$rec_time[1],
                                                     round(s[[i]]$rec_time[2]), sep = " ,"),
                      "]", sep = ""), csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    
    # summary write data, add back genotype (column 1)
    if (dim(df)[1] == 1) {
      suppressWarnings(write.table(t(df[, - c(2:3)]),
                                   csvfile, sep = ",", append = TRUE, row.names = FALSE,
                                   col.names = TRUE))
      if (i == 1) {
        suppressWarnings(write.table(cbind(div, t(df[, - c(2:3)])),
                                     csvwell, sep = ",", append = TRUE, row.names = FALSE,
                                     col.names = TRUE))
      } else {
        suppressWarnings(write.table(cbind(div, t(df[, - c(2:3)])),
                                     csvwell, sep = ",", append = TRUE, row.names = FALSE,
                                     col.names = FALSE))
      }
      
    } else {
      suppressWarnings(write.table(df[, - c(2:3)],
                                   csvfile, sep = ",", append = TRUE, row.names = FALSE,
                                   col.names = TRUE))
      if (i == 1) {
        suppressWarnings(write.table(cbind(div, df[, - c(2:3)]),
                                     csvwell, sep = ",", append = TRUE, row.names = FALSE,
                                     col.names = TRUE))
      } else {
        suppressWarnings(write.table(cbind(div, df[, - c(2:3)]),
                                     csvwell, sep = ",", append = TRUE, row.names = FALSE,
                                     col.names = FALSE))
      }
      
    }
    
    # new lines
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    
    # title
    write.table("Channel Burst Summary", csvfile, sep = ",", append = TRUE,
                row.names = FALSE, col.names = FALSE)
    write.table(paste("file= ", strsplit(basename(s[[i]]$file),
                                         ".RData")[[1]][1], sep = ""),
                csvfile, sep = ",", append = TRUE, row.names = FALSE, col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    
    # channel data
    suppressWarnings(write.table(df2,
                                 csvfile, sep = ",", append = TRUE, row.names = FALSE, col.names = TRUE))
    # new lines
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
    write.table(" ", csvfile, sep = ",", append = TRUE, row.names = FALSE,
                col.names = FALSE)
  } # end of loop through writting tables
}


.calc_all_isi <- function(s, allb) {
  ## Compute ISI within bursts for all spike trains.
  
  calc_isi <- function(spikes, b) {
    ## for one spike train, get all isis within bursts in that train.
    if (.num_bursts(b) == 0) {
      return(NA)
    }
    
    ## as.vector is needed below in case each burst is of the same
    ## length (in which case an array is returned by apply).  In
    ## normal cases, bursts are of different lengths, so "apply"
    ## returns a list.
    
    isis <- as.vector(
      unlist(apply(b, 1,
                   function(x) {
                     diff(spikes[x[1]:x[2]])
                   })))
  }
  
  nchannels <- s$NCells
  isis <- list()
  for (i in 1:nchannels) {
    isis[[i]] <- calc_isi(s$spikes[[i]], allb[[i]])
  }
  
  isis
}

.calc_ibi <- function(spikes, b) {
  ## Compute the interburst intervals (IBI) for one spike train.
  ## Only valid if more than one burst.
  
  nburst <- .num_bursts(b)
  if (nburst == 0) {
    res <- NA # no bursts
  } else {
    if (nburst == 1) {
      res <- NA # cannot compute  IBI w/only 1 burst.
    } else {
      ## find end spike in each burst.
      end <- b[, "beg"] + b[, "len"] - 1
      
      ## for n bursts, there will be n bursts minus 1 IBIs.
      start_spikes <- b[2:nburst, "beg"]
      end_spikes <- end[1:(nburst - 1)]
      ## NEX uses a strange definition of IBI: it counts the time between
      ## the first spike of n burst and the first spike of n burst
      ## plus one as the IBI.
      ## If we want to use that definition, use the following line:
      ## end.spikes equal b where 1 to nburst minus 1
      res <- spikes[start_spikes] - spikes[end_spikes]
    }
  }
  res
}

.num_bursts <- function(b) {
  ## Return the number of bursts found for a spike train.
  if (is.na(b[1]))
    0
  else
    nrow(b)
}

.calc_all_ibi <- function(s, allb) {
  ## Compute IBI for all spike trains.
  nchannels <- s$NCells
  ibis <- list()
  for (i in 1:nchannels) {
    ibis[[i]] <- .calc_ibi(s$spikes[[i]], allb[[i]])
  }
  
  ibis
}

.get_mean_burst_info_per_well_mod <- function(s) {
  master_sum <- list() # summary over all files
  for (i in 1:length(s)) {
    sum <- list() # summary for each timepoint
    # calculate bursting variables for current data File
    nbursts <- sapply(s[[i]]$allb, nrow)
    allb <- s[[i]]$allb
    tempsum <- calc_burst_summary(s[[i]])
    
    ###### NEW PART 1/3 ######
    
    for (t in 1:nrow(tempsum)) {
      if (tempsum[t,'nbursts'] == 0) {
        tempsum[t,'bursts_per_sec'] <- NA
        tempsum[t,'bursts_per_min'] <- NA
        tempsum[t,'mean_dur'] <- NA
        tempsum[t,'mean_spikes'] <- NA
      }
    }
    
    ###### NEW PART 1/3 ENDS ######
    
    # isis: gets the ISI for each channel of s[[i]]
    isis <- .calc_all_isi(s[[i]], allb)
    
    # IBIs get IBI's across all inter burst intervals across all data
    temp_ibis <- .calc_all_ibi(s[[i]], allb)
    
    # loop through goodwells
    for (j in 1:length(s[[i]]$goodwells)) {
      # indicator of current well
      icurrentwell <- (s[[i]]$goodwells[j] == s[[i]]$cw)
      
      # index of current well
      incurrentwell <- which(s[[i]]$goodwells[j] == s[[i]]$cw)
      
      if (sum(icurrentwell) != 0){
        ##### variables that need summing and averaging
        # total spikes across all AE in current well
        sum$nspikes[j] <- sum(tempsum$spikes[icurrentwell], na.rm = TRUE)
        sum$nAB[j] <- length(which(nbursts[incurrentwell] > 0))
        # Total recorded time on current well= recording time * nae
        sum$duration[j] <-
          length(incurrentwell) * (s[[i]]$rec_time[2] - s[[i]]$rec_time[1])
        
        # mean duration
        sum$mean_dur[j] <- mean(tempsum$mean_dur[incurrentwell], na.rm = TRUE)
        
        # mean spikes per second
        sum$mean_freq[j] <- mean(tempsum$mean_freq[incurrentwell], na.rm = TRUE)
        # total number of bursts
        sum$nbursts[j] <- sum(tempsum$nbursts[icurrentwell], na.rm = TRUE)
        # mean burst per second
        sum$bursts_per_sec[j] <- mean(tempsum$bursts_per_sec[incurrentwell], na.rm = TRUE)
        # mean burst per minute
        sum$bursts_per_min[j] <- sum$bursts_per_sec[j] * 60
        
        # finds the mean of duration for a particular well (across all channels)
        # get_burst_info(allb[icurrentwell],"durn")
        # takes out the column "durn" of all
        # matricies allb among the indicator set icurrentwell
        # get duration data across all channels of current well
        #sum$mean_dur[j] <-
        #  mean(unlist(get_burst_info(allb[icurrentwell], "durn")), na.rm = TRUE)
        
        ######## NEW PART 2/3 #########
        
        my.dur <- unlist(get_burst_info(allb[icurrentwell], "durn"))
        my.len <- unlist(get_burst_info(allb[incurrentwell], "len"))
        for (m in 1:length(my.dur)) {
          if (my.dur[m] == 0) {
            my.dur[m] <- NA
          }
        }
        for (n in 1:length(my.len)) {
          if (my.len[n] == 0) {
            my.len[n] <- NA
          }
        }
        
        # get duration data across all channels of current well
        sum$mean_dur[j] <-
          mean(my.dur, na.rm = TRUE)
        # sd of current well burst durations
        sum$sd_dur[j] <- sd(my.dur, na.rm = TRUE)
        
        # mean frequency within a burst
        sum$mean_freq_in_burst[j] <-
          mean(my.len /
                 my.dur, na.rm = TRUE)
        # sd frequency within a burst
        sum$sd_freq_in_burst[j] <-
          sd(my.len /
               my.dur, na.rm = TRUE)
        
        ######## NEW PART 2/3 ENDS ########
        
        # sd of current well burst durations
        #sum$sd_dur[j] <- sd(unlist(get_burst_info(allb[icurrentwell], "durn")))
        
        # mean frequency within a burst
        #sum$mean_freq_in_burst[j] <-
        #  mean(unlist(get_burst_info(allb[incurrentwell], "len")) /
        #         unlist(get_burst_info(allb[incurrentwell], "durn")), na.rm = TRUE)
        
        # sd frequency within a burst
        #sum$sd_freq_in_burst[j] <-
        #  sd(unlist(get_burst_info(allb[incurrentwell], "len")) /
        #       unlist(get_burst_info(allb[incurrentwell], "durn")), na.rm = TRUE)
        
        # mean of ISI across all channels in current well
        sum$mean_isis[j] <- mean(unlist(isis[incurrentwell]), na.rm = TRUE)
        
        # finds sd of ISI across all channels in current well
        sum$sd_isis[j] <- sd(unlist(isis[incurrentwell]), na.rm = TRUE)
        
        ######## NEW PART 3/3 ########
        
        # mean spikes in burst
        sum$mean_spikes_in_burst[j] <- round(mean(my.len, na.rm = TRUE), 3)
        # sd of spikes in burst
        sum$sd_spikes_in_burst[j] <- round(sd(my.len, na.rm = TRUE), 3)
        
        # total number of spikes arcross all bursts
        sum$total_spikes_in_burst[j] <- sum(unlist(get_burst_info(allb[incurrentwell], "len")), na.rm = TRUE)
        
        ######## NEW PART 3/3 ENDS ########        
        
        # len=#spikes in burst (length of burst in bursts)
        # mean_spikes_in_burst
        #ns <- unlist(get_burst_info(allb[icurrentwell], "len"))
        #sum$mean_spikes_in_burst[j] <- round(mean(ns, na.rm = TRUE), 3)
        
        # sd of spikes in burst
        #sum$sd_spikes_in_burst[j] <- round(sd(ns, na.rm = TRUE), 3)
        
        # total number of spikes arcross all bursts
        #sum$total_spikes_in_burst[j] <- sum(ns, na.rm = TRUE)
        
        # percent of spikes in bursts
        sum$per_spikes_in_burst[j] <-
          round(100 * (sum$total_spikes_in_burst[j] / sum$nspikes[j]), 3)
        
        # mean IBI
        sum$mean_ibis[j] <-
          round(mean(unlist(temp_ibis[incurrentwell]), na.rm = TRUE), 3)
        # sd IBI
        sum$sd_ibis[j] <-
          round(sd(unlist(temp_ibis[incurrentwell]), na.rm = TRUE), 3)
        # cv IBI
        sum$cv_ibis[j] <- round(sum$mean_ibis[j] / sum$sd_ibis[j], 3)
        
      } else {
        sum$nspikes[j] <- NA
        sum$nAB[j] <- NA
        sum$duration[j] <- NA
        sum$mean_dur[j] <- NA
        sum$mean_freq[j] <- NA
        sum$nbursts[j] <- NA
        sum$bursts_per_sec[j] <- NA
        sum$bursts_per_min[j] <- NA
        sum$mean_dur[j] <- NA
        sum$sd_dur[j] <- NA
        sum$mean_freq_in_burst[j] <- NA
        sum$sd_freq_in_burst[j] <- NA
        sum$mean_isis[j] <- NA
        sum$sd_isis[j] <- NA
        sum$mean_spikes_in_burst[j] <- NA
        sum$sd_spikes_in_burst[j] <- NA
        sum$total_spikes_in_burst[j] <- NA
        sum$per_spikes_in_burst[j] <- NA
        sum$mean_ibis[j] <- NA
        sum$sd_ibis[j] <- NA
        sum$cv_ibis[j] <- NA
        
      }
      
    }
    
    ### Set all names
    for (k in 1:length(names(sum))) {
      names(sum[[k]]) <- s[[i]]$goodwells
    }
    
    # make a master_sum, that is a list of all the summaries
    goodwellindex <- which(is.element(s[[i]]$well, s[[i]]$goodwells))
    
    master_sum[[i]] <- sum
    master_sum[[i]]$file <- strsplit(basename(s[[i]]$file), ".RData")[[1]][1]
    master_sum[[i]]$treatment <- s[[i]]$treatment[goodwellindex]
    master_sum[[i]]$size <- s[[i]]$size[goodwellindex]
    master_sum[[i]]$dose <- s[[i]]$dose[goodwellindex]
    master_sum[[i]]$well <- s[[i]]$well[goodwellindex]
    master_sum[[i]]$nae <- s[[i]]$nae[goodwellindex]
    master_sum[[i]]$timepoint <-
      rep(s[[i]]$timepoint[1], length(s[[i]]$goodwells))
    master_sum[[i]]$start_rec_time <-
      rep(s[[i]]$rec_time[1], length(s[[i]]$goodwells))
    master_sum[[i]]$end_rec_time <-
      rep(s[[i]]$rec_time[2], length(s[[i]]$goodwells))
    master_sum[[i]]$goodwells <- s[[i]]$goodwells
    
  }
  
  master_sum
}


##### needed additional functions

.get_div <- function(s) {
  div <- NA
  if(length(s$div>0)){
    div<-s$div
  } else {
    t1 <- strsplit(s$file, split = "_", fixed = TRUE)
    for (i in t1[[1]]) {
      i <- toupper(i)
      if (nchar(i) > 2 && substr(i, 1, 3) == "DIV") {
        if (nchar(i) > 5) {
          i <- unlist(strsplit(i, split = ".", fixed = T))[1]
        }
        div <- as.numeric(substr(i, 4, nchar(i)))
      }
    }
  }
  # set default div as 1
  if (is.na(div)){ div <- 1}
  div
}