# usage: main(<input CSV file>  ,  <output PNG filename>  ,  "neutral/charged")
# See last line of this file for typical execution under R environment
# the last term in the command is just a procedural keyword
# one of either "neutral" or "charged" can be used.
# 2DSE plots can be produced without these keywords
# but will require some re-writing of this code
# The input.csv file MUST possess the following columns:
# 'AtomID','x','y','z','potential'
# The GAG-binding site predictor code provided herein
# produces an appropriately formatted CSV file that
# can be used as input for this R script.
# The values of charged_max and neutral_max must be set.
# charged_max is basically -1*minimum potential observed.
# This functionality is provided to enable an apples-to-apples 
# comparison across multiple protein structures.
# You will require an appropriately functioning R installation,
# along with the GRID and the GGPLOT2 packages.
# Please contact the authors for support.
# Created by Aurijit Sarkar, May 11, 2015

library('ggplot2')
library('grid')

charged_max <- c(3000)
neutral_max <- c(550)

rescale <- function(dat,chargestat) {
	minval <- min(dat)
	if (minval < 0) {
		dat <- dat + abs(min(minval)) + 1
	}
	if (chargestat=='charged') {maxval <- charged_max} else {maxval <- neutral_max}
	newdat <- dat/maxval
	return(newdat)
}


main <- function(readf,writef,chargestat) {
dat <- read.table(readf,sep=',',col.names=c('ID','x','y','z','potential'))

if(chargestat == "charged") {xlim <- c(-55,charged_max[1])} else {xlim <- c(-55,neutral_max[1])}

plot1 <- ggplot(dat,aes(x=sqrt(x**2+y**2), y=z, col=-1*potential, fill=-1*potential )) + 
  geom_point(shape=21, size=15**(rescale(-1*dat$potential,chargestat))) +
  scale_color_gradient(limits=xlim,low='green',high='red',name='Free Energy\n(kcal/mol)') +
  scale_fill_gradient(limits=xlim,low='green',high='red',guide=FALSE) +
  theme(text=element_text(size=15),
  axis.text.x=element_blank(), axis.text.y=element_blank(),
  axis.title.x=element_blank(), axis.title.y=element_blank(), 
  panel.grid.major.x=element_blank(), panel.grid.minor.x=element_blank(), 
  panel.grid.major.y=element_blank(), panel.grid.minor.y=element_blank(),
  panel.background=element_rect(fill="white",color="black",size=2.0),
  legend.position="top",legend.key.height = unit(1,"cm"), legend.key.width = unit(2.5,"cm")) +
  scale_size_area()

ggsave(plot1, file=writef, height=7.2, width=7.2, dpi=600, units='in')

plot2 <- ggplot(dat,aes(x=sqrt(x**2+y**2), y=z, col=-1*potential, fill=-1*potential )) + 
  geom_point(shape=21, size=15**(rescale(-1*dat$potential,chargestat))) +
  scale_color_gradient(limits=xlim,low='green',high='red',name='Free Energy\n(kcal/mol)') +
  scale_fill_gradient(limits=xlim,low='green',high='red',guide=FALSE) +
  theme(text=element_text(size=15),
  axis.text.x=element_blank(), axis.text.y=element_blank(),
  axis.title.x=element_blank(), axis.title.y=element_blank(), 
  panel.grid.major.x=element_blank(), panel.grid.minor.x=element_blank(), 
  panel.grid.major.y=element_blank(), panel.grid.minor.y=element_blank(),
  panel.background=element_rect(fill="white",color="black",size=2.0),
  legend.position="top",legend.key.height = unit(1,"cm"), legend.key.width = unit(2.5,"cm")) +
  geom_text(aes(label=ID),size=1,color='black') +
  scale_size_area()

ggsave(plot2, file=paste(writef,'_withResNames.png',sep=''), height=7.2, width=7.2, dpi=600, units='in')
}

main('input.csv','output.png','neutral')