#!/usr/bin/perl -w

# rotatepdb.pl -- rotate & translate the atoms from a PDB file.
#                 
# Copyright © 2003 Russell Hanson
#
# Permission to use, copy, modify, distribute, and sell this software
# and its documentation for any purpose is hereby granted without fee,
# provided that the above copyright notice appear in all copies and
# that both that copyright notice and this permission notice appear in
# supporting documentation.  No representations are made about the
# suitability of this software for any purpose.  It is provided "as
# is" without express or implied warranty.

# Usage: rotatepdb -i <pdb input> -o <pdb output> -p <anglephi (degree)> -d <distance (Ang)> -r <rotation axis {x,y,z}> -t <translation axis {x,y,z}>

use strict;
use warnings;
require 5;
use Getopt::Std;
use Switch;
use POSIX qw(ceil floor);
use Math::Trig;

# --------------- NOTES ---------------
# rot mats
# x axis
#   1 0 0  | x
#   0 c -s | y
#   0 s c  | z

# y axis
#   c  0  s
#   0  1  0
#  -s  0  c

# z axis
#   c -s 0
#   s  c 0
#   0  0 1

# Sample PDB line, the coords X,Y,Z are fields 5,6,7 on each line.
# ATOM      1  N   ARG     1      29.292  13.212 -12.751  1.00 33.78      1BPT 108

# use sprintf %8.3f for coords
# printf PDB2 ("ATOM  %5d %4s %3s A%4d    %8.3f%8.3f%8.3f%6.2f%6.2f      %4s%2s\n", $index,$atname[$i],$resname[$i],$resnum[$i],$x[$i],$y[$i],$z[$i],$occ[$i],$bfac[$i]),$segid[$i],$element[$i];

my $M_PI = 3.14159265358979323846;

sub read_atoms ($) {
    my ($inpdb) = @_;
    my (@atomdata, 		# a list of 3-tuples
		$pdbline, @fields, $x, $y, $z);
    open(ED, "<$inpdb") || die("can't open $inpdb: $!");
    while (<ED>) {
		next unless (/^ATOM/);
		$pdbline = $_;
		$x = substr($pdbline,30,8);
		$y = substr($pdbline,38,8);
		$z = substr($pdbline,46,8);
		push (@atomdata, [ $x,$y,$z ]);
   	}

# 	warn "read_atoms()";	# DEBUG
#     for my $i ( 0 .. $#atomdata ) {
#  		print "\t [ @{$atomdata[$i]} ],\n";
#     }

    close(ED);
    return @atomdata;
}

sub transatoms ($$@) {
    my ($dist, $transaxis, @atomdata) = @_;
    my (@transatoms, $x,$y,$z, $xp,$yp,$zp, @atom);
    warn "transatoms dist = $dist, transaxis = $transaxis\n";
    for my $tup ( @atomdata ) {
		# print "tup \t [ @$tup ],\n";
		$x = @$tup[0]; $y = @$tup[1]; $z = @$tup[2]; 
#		warn "atoms pre-trans ", $x, " ", $y, " ", $z;
		# if $transaxis uninitialized use x-axis
		if (length($transaxis) == 0) { 
			$transaxis = "x";
		}
		if ($transaxis eq "x") {
			$xp = $x + $dist;
			@atom = [ $xp,$y,$z ];
		} elsif ($transaxis eq "y") {
			$yp = $y + $dist;
			@atom = [ $x,$yp,$z ];
		} elsif ($transaxis eq "z") {
			$zp = $z + $dist;
			@atom = [ $x,$y,$zp ];
		}
#		warn "atoms post-trans ", $xp, " ", $yp, " ", $zp;
		push (@transatoms, @atom);
    }
#  	warn "transatoms()";	# DEBUG
#  	for my $tup ( @transatoms ) {
#  		print " \t [ @$tup ],\n";
#  	}
    return @transatoms;
}

sub rotatoms ($$@) {
    my ($phi, $rotaxis, @atomdata) = @_;
    my (@newatoms, $x,$y,$z, $xp,$yp,$zp, @atom);
    $phi = $phi/360 * 2 * $M_PI;
    my $rcos = cos($phi);
    my $rsin = sin($phi);
    warn "rotatoms phi = $phi, rotaxis = $rotaxis, rcos = $rcos, rsin = $rsin\n";
    for my $tup ( @atomdata ) {
		# print "tup \t [ @$tup ],\n";
		$x = @$tup[0]; $y = @$tup[1]; $z = @$tup[2]; 
#		warn "atoms pre-rotation ", $x, " ", $y, " ", $z;
		
		# if $rotaxis uninitialized use z-axis
		if (length($rotaxis) == 0) { 
			$rotaxis = "z";
		}
		if ($rotaxis eq "x") {
#			warn "x axis \n";
			$xp = $x;
			$yp = $y * $rcos - $z * $rsin;
			$zp = $y * $rsin + $z * $rcos;
		} elsif ($rotaxis eq "y") {
#			warn "y axis \n";
			$xp = $x * $rcos + $z * $rsin;
			$yp = $y;
			$zp = $z * $rcos - $x * $rsin;
		} elsif ($rotaxis eq "z") {
#			warn "z axis \n";
			$xp = $x * $rcos - $y * $rsin;
			$yp = $x * $rsin + $y * $rcos;
			$zp = $z;
		}
#		warn "atoms post-rotation ", $xp, " ", $yp, " ", $zp;
		@atom = [ $xp,$yp,$zp ];
		push (@newatoms, @atom);
    }
#     	warn "rotatoms()";			# DEBUG
#     	for my $tup ( @newatoms ) {
#     		print "\t [ @$tup ],\n";
#     	}
    return @newatoms;
}

# copy inpdb to outpdb.  All ATOMs are replaced by rotated atoms.
sub write_pdb ($$@) {
    my ($inpdb, $outpdb, @newatoms) = @_;
    my ($n, $pdbline, $tmp0, $tmp1, $tmp2);
    $n=0;
    open(IN, "<$inpdb") || die("can't open $inpdb: $!");
    open(OUT, ">$outpdb") || die("can't open $outpdb: $!");
    while (<IN>) {
		if (/^ATOM/) {
			$pdbline = $_;
			# warn "old ", $pdbline, "\n";
			# this has to be the stupidest way to unpack a list 
			$tmp0 = sprintf "%8.3f",$newatoms[$n][0];
			$tmp1 = sprintf "%8.3f",$newatoms[$n][1]; 	    
			$tmp2 = sprintf "%8.3f",$newatoms[$n][2]; 	    

			substr($pdbline,30,8) = $tmp0;
			substr($pdbline,38,8) = $tmp1;
			substr($pdbline,46,8) = $tmp2;
			# warn "new ", $pdbline, "\n";
			$n++;
			print OUT $pdbline;
		} else {
			print OUT $_;
		}
	}
    close(IN);
    close(OUT);
}

sub max ($$) {
	my ($a, $b) = @_;
	return ($a < $b) ? ($b) : ($a);
}

# normalize to spatial centroid and zero degrees theta and phi
sub normatoms (@) {
    my (@atomdata) = @_;
    my (@normatoms, $x,$y,$z, $sumx,$sumy,$sumz, $n, $theta, $phi, @atom, @cntratom);
    warn "normatoms()\n";
	$n = 0;

	# calculate (spatial) centroid
	for my $i ( 0 .. $#atomdata ) {
		$x = $atomdata[$i][0]; $y = $atomdata[$i][1]; $z = $atomdata[$i][2];
		$sumx += $x; $sumy += $y; $sumz += $z;
# 		print "\t [ @{$atomdata[$i]} ],\n";
		$n++;
    }
	$atom[0] = $sumx/$n; $atom[1] = $sumy/$n; $atom[2] = $sumz/$n;

	# subtract centroid coord off all atomdata
	for my $i ( 0 .. $#atomdata ) {
		$normatoms[$i][0] = $atomdata[$i][0] - $atom[0];
		$normatoms[$i][1] = $atomdata[$i][1] - $atom[1];
		$normatoms[$i][2] = $atomdata[$i][2] - $atom[2];
	}

	# find spherical coord angles theta and phi
	#    y    p
	#    |   /|
	#    |  / |
	#    | /  |
	#    |----|---x
	#   /-----|
	#  /
	# /z
	#   p = (x,y,z);
	#   theta = arctan(y/x)
	#   phi = arctan(z/x)

	# find center of one "average" helical turn, num of atoms
	# >>> (31 + 32 + 29 + 32)/4  # Voet & Voet
	# 31
	# >>> (44 + 41 + 39 + 43)/4  # PDB
	# 41
	# if it's really short, use the whole thing
	my $i; my $AGCTatoms = 41; $sumx = $sumy = $sumz = 0;
	if ( $n < $AGCTatoms ) {
		$i = 0;
	} else {
		$i = $n - $AGCTatoms;
	}
	for ( ; $i < $n; $i++ ) {
		$x = $atomdata[$i][0]; $y = $atomdata[$i][1]; $z = $atomdata[$i][2];
		$sumx += $x; $sumy += $y; $sumz += $z;
    }
	$cntratom[0] = $sumx/$n; $cntratom[1] = $sumy/$n; $cntratom[2] = $sumz/$n;
	
	$theta = atan ($cntratom[1]/$cntratom[0]); # (x,y,0)
	$phi   = atan ($cntratom[2]/$cntratom[0]); # (x,0,z)
	@normatoms = rotatoms ($theta, "y", @normatoms); ## ???
	@normatoms = rotatoms ($phi, "x", @normatoms);
	
	warn "normatoms()";	# DEBUG
 	for my $tup ( @normatoms ) {
 		print "\t [ @$tup ],\n";
 	}
 	return @normatoms;
}

sub do_file ($$$$$$) {
    my ($inpdb, $outpdb, $phi, $dist, $rotaxis, $transaxis) = @_;
    my (@atomdata, @newatomdata);
    warn "$inpdb, $outpdb, $phi, $dist, $rotaxis, $transaxis";

    @atomdata    = read_atoms ($inpdb);
#    @newatomdata = normatoms (@atomdata);
	@newatomdata = rotatoms ($phi, $rotaxis, @atomdata);
	@newatomdata = transatoms ($dist, $transaxis, @newatomdata);
    write_pdb ($inpdb, $outpdb, @newatomdata);

    print "have a nice day!\n";
}

# --- main ---
$0 =~ s,.*/,,;

my (%args);
@ARGV && getopts("i:o:p:d:r:t:", \%args) || die "Usage: rotatepdb -i <pdb input> -o <pdb output> -p <anglephi (degree)> -d <distance (Ang)> -r <rotation axis {x,y,z}> -t <translation axis {x,y,z}>";

print "-i $args{i}\n";
print "-o $args{o}\n";
print "-p $args{p}\n";
print "-d $args{d}\n";
print "-r $args{r}\n";
print "-t $args{t}\n";

do_file($args{i}, $args{o}, $args{p}, $args{d}, $args{r}, $args{t});

__END__