## Rotate given PDB by x, y and z axes
## Usage: Rscript rotate_PDB.R [a,b,c] pdb_file
## a,b,c are the angles in degrees to rotate about x,y and z axes.



args = commandArgs(trailingOnly=TRUE)
if (length(args) != 2) {
  cat("Rotate the given PDB by x,y,z axes \n")
  cat("Usage:\n")
  cat("       Rscript rotate_PDB.R [a,b,c] pdb_file\n")
  cat("       [a,b,c]     a,b,c are the angles in degrees(+ve or -ve) to rotate about x,y and z axes resp.\n")
  cat("       pdb_file    the pdb file to rotate\n")
  cat("       The output pdb file is given extension of _rot \n")
  cat("       The order of rotation is X, Y and Z axes resp. \n")
  cat("       [0,0,0] denotes no rotation \n")
  quit()
}

## load dependency
library(bio3d)
library(dplyr)
library(stringr)

#setwd("/Users/bhandariyr/Documents/projects/afm/scripts/rotate")
#input.pdb = "/Users/bhandariyr/Documents/projects/afm/scripts/rotate/4gma_amber_offset.pdb"
angles.3d = as.numeric(c(unlist(strsplit(str_sub(args[1], 2, -2),","))))
input.pdb = args[2]
pdb = read.pdb2(input.pdb)

## make a matrix of x,y,z coordinates
coor.mat = cbind(pdb$atom$x, pdb$atom$y, pdb$atom$z)
#class(coor.mat)

## Angles to radians
a = angles.3d[1]
b= angles.3d[2]
c= angles.3d[3]
a.rad = a*pi/180
b.rad = b*pi/180
c.rad = c*pi/180

## Write the rotation matrices
rot.x = cbind(c(1,0,0), c(0, cos(a.rad), sin(a.rad)), c(0, -sin(a.rad), cos(a.rad)))
rot.y = cbind(c(cos(b.rad), 0, -sin(b.rad)), c(0,1,0), c(sin(b.rad), 0, cos(b.rad)))
rot.z = cbind(c(cos(c.rad), sin(c.rad), 0), c(-sin(c.rad), cos(c.rad), 0), c(0, 0, 1))
rot.3d = rot.z %*% (rot.y %*% rot.x)
#coor.new = t(rot.3d %*% coor.mat[1,])

## Function to transform the coordinates
rotate <- function(coor.xyz){
  new.coor = rot.3d %*% coor.xyz
  return(new.coor)
}
## Rotate the coordinates
#coor.transform = (apply(coor.mat, 1, rotate))
coor.transform = as.vector(apply(coor.mat, 1, rotate))
#class(as.vector(coor.transform))
#dim(coor.mat); length(coor.transform); 
#head(coor.mat); head(coor.transform)
## set the spacing with respect to resid
pdb$atom$resid = paste0(pdb$atom$resid, " ")

out.pdb = paste0(tools::file_path_sans_ext(basename(input.pdb)),"_rot.pdb")
cat(paste("The rotated pdb is:", out.pdb, "\n", sep = " "))
write.pdb(pdb = pdb, xyz = coor.transform, file = out.pdb)
