#!/usr/bin/perl -w

use strict;

((@ARGV==1) or (@ARGV==2)) or die "".
"Error in syntax!\n".
"  ./gen_rasmol_script.pl <Tab_deliminated.predictions> [R,G,B]\n".
"Note:\n".
"  The default color for the predicted PTM sites is 115,254,2, you can also\n".
"set the color in the format of \"R,G,B\" with integer values.\n";


my $egTag = "";
my $egChain = "";
my $egStart = 0;
my $egEnd = 0;

my $tline = "";
my @tlist = ();

my $tempi = 0;
my $tempj = 0;
my $tempk = 0;

my $egFileIn = $ARGV[0];
my $egColor = "115,254,2";
if( @ARGV==2 )
{
	@tlist = split(/,/, $ARGV[1]);
	if( (@tlist==3) and ($tlist[0]=~/^\d+$/) and ($tlist[1]=~/^\d+$/) and ($tlist[2]=~/^\d+$/) )
	{
		$egColor = $ARGV[1];
	}
	else
	{
		die "Error! Wrong color format!\n";
	}
}

my $tpos = 0;

# Load the prediction results

my %egIdx = ();

open(efIn, "$egFileIn") or die "Error when loading prediction results!\n";
while(<efIn>)
{
	if( $_=~/^>(.*)$/ )
	{
		#if( (length($egTag)>0) and (length($egChain)>0) )
		#{
		#}
		$tline = $1; $tline=~s/[\r\n]//g;
		@tlist = split(/;;/, $tline);
		if( @tlist == 4 )
		{
			if( length($egTag)==0 )
			{
				$egTag = $tlist[0];
			}

			if( $egTag eq $tlist[0] )
			{
				$egChain = $tlist[1];
				$egStart = $tlist[2];
				$egEnd   = $tlist[3];
			}
			else
			{
				$egChain = "";
				$egStart = 0;
				$egEnd   = 0;
			}
		}
	}
	elsif( $_=~/^(\d+)	(.*)$/ )
	{
		$tpos = $1;
		if( (length($egTag)>0) and (length($egChain)>0) and ($egStart>0) and ($egEnd>0) )
		{
			if( $tpos <= ($egEnd-$egStart+1) )
			{
				$tempi = $egStart + $tpos - 1;
				$egIdx{"$egChain;;$tempi"} = "";
			}
		}
	}
}
close(efIn);

print "Number of predicted PTM sites: ".(scalar keys %egIdx)."\n";

# Scan PDB file for ATOM IDs

$egFileIn = "$egTag.pdb";

my $tAtomID = 0;
my $tChain = "";
my $tAAID = 0;

my $tS = 0;
my $tE = 0;

my $tl = "";

my $tcount = 0;

open(efIn, "$egFileIn") or die "Error! Cannot open PDB file from $egFileIn\n";
while(<efIn>)
{
	if( $_=~/^ATOM\s*\d+/ )
	{
		$tline = $_; $tline=~s/[\r\n]//g;
		@tlist = split(//, $tline);
		
		$tl = $tlist[4].$tlist[5].$tlist[6].$tlist[7].$tlist[8].$tlist[9].$tlist[10];
		$tl=~s/^\s+//g;
		$tAtomID = int($tl);

		$tl = $tlist[20].$tlist[21];
		$tl=~s/^\s+//g;
		if( $tl eq "" )
		{
			$tl = "DefaultChain";
		}
		$tChain = $tl;

		$tl = $tlist[22].$tlist[23].$tlist[24].$tlist[25];
		$tl=~s/^\s+//g;
		$tAAID = int($tl);

		$tl = $egIdx{"$tChain;;$tAAID"};
		if( defined $tl )
		{
			if( $tl eq "" )
			{
				$tl = "$tAtomID";
			}
			else
			{
				$tl = $tl."	$tAtomID";
			}
			$egIdx{"$tChain;;$tAAID"} = $tl;
			$tcount++;
		}
	}
}
close(efIn);
print "Number of atoms selected: $tcount\n";

print "Generating RasMol script ... ";
my $egFileOut = "$egTag.rasmol.script";
open(efOut, ">$egFileOut") or die "Error when saving RasMol script for $egTag.pdb\n";
print efOut "#!rasmol -script\n";
print efOut "# File: $egFileOut\n";
print efOut "# PDB: $egTag.pdb\n";
print efOut "# Creator: gen_rasmol_script.pl version 1.0 (c) Fengfeng Zhou 2006-08-01\n";
print efOut "# Date: ".`date`."\n";
print efOut "\n";

print efOut "zap\n";
print efOut "load pdb \"$egTag.pdb\"\n";
print efOut "background [0,0,0]\n";
print efOut "set ambient 40\n";
print efOut "set specular off\n";
print efOut "\n";

print efOut "reset\n";
print efOut "slab off\n";
print efOut "\n";

print efOut "set axes off\n";
print efOut "set boundingbox off\n";
print efOut "set unitcell off\n";
print efOut "set bondmode and\n";
print efOut "dots off\n";
print efOut "\n";

print efOut "# To avoid color problems\n";
print efOut "select all\n";
print efOut "colour bonds none\n";
print efOut "colour backbone none\n";
print efOut "colour hbonds none\n";
print efOut "colour ssbonds none\n";
print efOut "colour ribbons none\n";
#print efOut "colour white\n";
print efOut "colour cpk\n";
print efOut "\n";

print efOut "# Select those atoms in the predicted amino acids, and color them specifically\n";

$tS = 0; $tE = 0;
$tChain = ""; $tAAID = 0;
$tline = ""; $tl = "";

my @tIdxList = keys %egIdx;
for($tempi=0;$tempi<@tIdxList;$tempi++)
{
	$tS = 0; $tE = 0;
	$tline = $egIdx{"$tIdxList[$tempi]"};
	@tlist = split(/	/, $tline);

	if( ($tline ne "") and (@tlist>=1) )
	{
		$tS = int($tlist[0]); $tE = int($tlist[@tlist-1]);
		if( $tS > $tE ) { $tempj = $tE; $tE = $tS; $tS = $tempj; }
		print efOut "select (atomno>=$tS) and (atomno<=$tE)\n";
		print efOut "colour atoms [$egColor]\n";
	}
}

print efOut "select all\n";
print efOut "spacefill on\n";
print efOut "set shadow off\n";

print efOut "wireframe off\n";
print efOut "ribbons off\n";
print efOut "backbone off\n";
print efOut "labels off\n";
print efOut "monitors off\n";
print efOut "hbonds off\n";
print efOut "ssbonds off\n";
print efOut "\n";


close(efOut);
print "[done]\n\n";
print "Please copy the two files $egTag.pdb and $egFileOut to the same directory, and load the script $egFileOut from RasMol.\n";