## Scale x,y,z values from the cafemol input file
## Usage: Rscript scale_afem_CafeInput.R cafe_input [a,b,c] Zcutoff 
### a,b,c are scaling factors for x,y,z
## Zcutoff denotes z values below which are not scaled

suppressPackageStartupMessages(library("dplyr"))
suppressPackageStartupMessages(library("stringr"))
#library(dplyr, quietly=T)
#library(stringr, quietly=T)

args = commandArgs(trailingOnly=TRUE)
if (length(args) != 3) {
  stop("Please provide cafemol inputfile, x,y,z scale factor and Zcutoff ", call.=FALSE)
}

#setwd("/Users/bhandariyr/Documents/projects/afm/scripts/scale_cafeinput")
#string1 = "[2,4,6]"
scale_factor <- as.numeric(c(unlist(strsplit(str_sub(args[2], 2, -2),","))))

cafe.input = args[1]
#cafe.input = "/Users/bhandariyr/Documents/projects/afm/scripts/scale_cafeinput/input.inp"
## grep the "AFM_IMAGE" identifier
data.lines = readLines(cafe.input)[grep("AFM_IMAGE", readLines(cafe.input))]
df =  data.lines %>%
  strsplit("\\s+") %>%
  do.call(rbind, .) %>%
  data.frame()
# dim(df)
colnames(df) = c("header","x","y","z")
#head(df)

df$x = format(as.numeric(df$x)*scale_factor[1], digits = 1, nsmall = 1)
df$y = format(as.numeric(df$y)*scale_factor[2], digits = 1, nsmall = 1)
Zcutoff = as.numeric(args[3])
df$z = unlist(lapply(as.numeric(df$z), function(x){
  if (x > Zcutoff){
    x = x * scale_factor[3]
  }
  return(x)
}))
df$z = format(as.numeric(df$z), digits = 3, nsmall = 3)

## Read the cafemol input file and edit afm fitting section accordingly
lines_vec = vector(mode="character") ## initialize a character vector

for (lines in readLines(cafe.input)){

  if (startsWith(lines, "AFM_IMAGE") ){
    ## Do nothing, skip
    }
  else {
    lines_vec = append(lines_vec, lines)
  }
}

lines_vec = head(lines_vec, -1) ## remove the last line ">>>>"

out.file = paste0(tools::file_path_sans_ext(basename(cafe.input)),"_scaled.inp")

if (file.exists(out.file)){
  #Delete file if it exists
  file.remove(out.file)
}

writeLines(lines_vec, out.file)
write.table(df, out.file,
            sep = " ", quote = FALSE, row.names = FALSE, col.names = FALSE, append = TRUE)
write(">>>>", out.file, append = TRUE)



