### Scale the given image file and add Gaussian noise to z column
### Usage: Rscript scale.R [2,3,1] [-2,2] afm_image.txt 
### first vector is scale factor for (x,y,z) and second is (min,max) range for signal noise


args = commandArgs(trailingOnly=TRUE)
if (length(args) != 3) {
  cat("Scale the given image file and add Gaussian noise to z column\n")
  cat("Usage:\n")
  cat("       Rscript scale.R [a,b,c] [p,q] afm_image.txt\n")
  cat("       image.txt    afm image file with headers x, y, z\n") 
  cat("       [a,b,c]      scale factors that multiply x, y, z values resp., scaling is applied before adding noise\n")
  cat("       [p,q]        min and max limits of gaussian noise, p & q have to be same numbers with opposite signs\n")
  cat("                    eg. [-10,10] will add gaussian distributed noise around mean 0 and limits -10 to +10, sd ~ 10/3 \n") 
  cat("                    noise is added to all z values (signal as well as background) \n")
  cat("                    Two final files ending in A.txt (angstoms units) and nm.txt (nanometer units) are produced\n")
  quit()
}

library(dplyr)
library(stringr)

# image="/Users/bhandariyr/Documents/projects/afm/scripts/addNoise/CG-f5511199c_3132.pdb_calc_data.txt"
image = args[3]
#setwd("/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/scale")

data = read.table(image, header = TRUE)
data.orig = data

# scale_factor = c(1,1,1)
# min.noise = -10
# max.noise = 10
scale_factor <- as.numeric(c(unlist(strsplit(str_sub(args[1], 2, -2),","))))
noise_factor <- as.numeric(c(unlist(strsplit(str_sub(args[2], 2, -2),","))))
min.noise = noise_factor[1]
max.noise = noise_factor[2]
# colnames(data)

## Scale the data
data$x = (data$x)*scale_factor[1]
data$y = (data$y)*scale_factor[2]
data$z = (data$z)*scale_factor[3]

## Generate normal distribution using min/max values
## 99.7% of values fall within 3 standard deviations of the mean
mn <- (max.noise - min.noise)/2
rtruncnorm <- function(N, mean = 0, sd = 1, a = -Inf, b = Inf) {
   if (a > b) stop('Error: Truncation range is empty');
   U <- runif(N, pnorm(a, mean, sd), pnorm(b, mean, sd));
   qnorm(U, mean, sd); }
noise.gauss = rtruncnorm(nrow(data), mean = 0, sd = mn/3, a = min.noise, b = max.noise)
# hist(data$noise)

## Iterate through the rows of data and add the noise
# dim(data);
# print(noise.gauss)
data$z = data$z + noise.gauss
# hist(data$z); boxplot(data$z)
# boxplot(data.orig$z)

# View(data); View(data.orig)
# dim(data_bg); dim(data)
## Plot to see the distribution
# hist(data[data$z < 10, "z"])
# hist(data[, "z"])


#data in nanometers
data.nm = data.frame( x = formatC(data$x/10, format="e", digits = 6),
                      y = formatC(data$y/10, format="e", digits = 6),
                      z = formatC(data$z/10, format="e", digits = 6))

#data in angstroms
data$x = format(data$x, digits = 1, nsmall = 1)
data$y = format(data$y, digits = 1, nsmall = 1)
data$z = format(round(data$z, digits = 3))

out.file1 = paste0(tools::file_path_sans_ext(basename(image)),"_scaled_A.txt")
out.file2 = paste0(tools::file_path_sans_ext(basename(image)),"_scaled_nm.txt")

write.table(data, out.file1,
            sep = " ", quote = FALSE, row.names = FALSE, col.names = TRUE, append = FALSE)
write.table(data.nm, out.file2,
            sep = " ", quote = FALSE, row.names = FALSE, col.names = TRUE, append = FALSE)



