## Modify the given cafemol input file to manipulate pixel size, and max dimension 
## and recenter the PDB file in new image field
## Rscript pixel_size max_dim cafe.inputfile

library(dplyr)

args = commandArgs(trailingOnly=TRUE)
if (length(args) != 3) {
  stop("Please provide pixel size, maximum dimension and cafemol inputfile", call.=FALSE)
}
pixel_size = as.numeric(args[1])
max_dim = as.numeric(args[2])
cafe.input = args[3]

## units in Angstroms
# pixel_size = 1
# max_dim = 100
xy_pair = data.frame(x=numeric(), y=numeric())
## Create the set of x, y pixels
pix1 = pixel_size
pix2 = pixel_size
for (i in 1:(round(max_dim/pixel_size))){
  pix2 = pixel_size
  if (pix1 <= max_dim){
  for (j in 1:(round(max_dim/pixel_size))){
    if (pix2 <= max_dim){
      # print(paste0(pix2, " ", pix1))
      #print(j)
      xy = data.frame(x = pix2, y = pix1)
      xy_pair = rbind(xy_pair, xy)
      pix2 = pix2 + pixel_size
    }
  } 
    pix1 = pix1 + pixel_size
}
}
#xy_pair
# format the data frame with xy_pair
heights = as.numeric(rep(1.0, nrow(xy_pair)))
df = data.frame(x = xy_pair$x, y = xy_pair$y, z = heights)
df$x = format(df$x, digits = 1, nsmall = 1)
df$y = format(df$y, digits = 1, nsmall = 1)
df$z = format(df$z, digits = 3, nsmall = 3)
## add the id column
df$id = rep("AFM_IMAGE", nrow(df))
#move last column to first
df <- df %>%   select(id, everything())

num.pixels.x = round(max_dim/pixel_size)
#num.pixels.y = round(max_dim/pixel_size)

## Read the cafemol input file and edit afm fitting section accordingly
# cafe.input = "/Users/bhandariyr/Documents/projects/afm/scripts/image_setup/format.inp"
lines_vec = vector(mode="character") ## initialize a character vector

flag.fit = 0 ## flag to read the afm_fitting section
for (lines in readLines(cafe.input)){
  if (startsWith(lines, "<<<< afm_fitting")){
    lines_vec = append(lines_vec, lines)
    flag.fit = 1
  }
  else if (flag.fit == 1 & (startsWith(lines, "n_x") )){
    lines = paste0("n_x = ", num.pixels.x)
    lines_vec = append(lines_vec, lines)
  }
  else if (flag.fit == 1 & (startsWith(lines, "n_y") )){
    lines = paste0("n_y = ", num.pixels.x)
    lines_vec = append(lines_vec, lines)
  }
  else if (flag.fit == 1 & (startsWith(lines, "pixel_size_x") )){
    lines = paste0("pixel_size_x = ", format(pixel_size, digits = 1, nsmall = 1))
    lines_vec = append(lines_vec, lines)
  }
  else if (flag.fit == 1 & (startsWith(lines, "pixel_size_y") )){
    lines = paste0("pixel_size_y = ", format(pixel_size, digits = 1, nsmall = 1))
    lines_vec = append(lines_vec, lines)
  }
  else 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)),"_customField.inp")
#out.file = "cafe_input_2.inp"
if (file.exists(out.file))
  #Delete file if it exists
  file.remove(out.file)

#setwd("/Users/bhandariyr/Documents/projects/afm/scripts/image_setup")
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)



