#!/usr/bin/perl -w
# $Id: killres.pl,v 1.4 2003/08/01 13:22:57 oliver Exp $

use strict;
use lib $ENV{HOME}."/lib/perl";
use Messages; 

use Getopt::Long;

my $Version_ID = q($Id: killres.pl,v 1.4 2003/08/01 13:22:57 oliver Exp $ );

# imported from Messages
# DEBUG level says how chatty we are:
#  <0  dead silent
#   0  no diagnostic messages, only output
#   1  a few additional important informations
#   3  verbose
#   5  extra verbose
#  10  start with real debugging stuff
#  15  many intermediate results
#  16  print input files line-by-line
#  20  everything   
$DEBUG = 1;

my $DataFn;
my ($RESNAME, $Z1, $Z2, $bSLAB) = ("SOL", -1000, 1000, 0);
my $OutputFn = "killed.gro";

my ($opt_verbose, $opt_version, $opt_help);

# how to indicate a comment line in the data file (regex)
my $commentchar = '[;#@]';




sub print_usage () {
    my $Prog = Progname;
    print <<EOF;
usage: $Prog [OPTIONS] [FILE]
       $Prog --slab [OPTIONS] [FILE]

This program reads a gromos file FILE (.gro) and does evil things to
molecules with name RESNAME. Normally, it kills all
residues with z above z1 and below z2 but in slab-mode it carves a slab
out of the input with identical x and y dimensions but a reduced z=z2-z1
and all residues kept are shifted by -z1.

Read data from all FILEs or from standard input.

The OPTIONS can be abbreviated; defaults are given in [].	 
--help              help
--verbose           (alias for DEBUGLEVEL=3)
--version           print version id
--debug=DEBUGLEVEL  0 (almost) quiet, <0 silent,
                    1..3 verbose, 
                    >5 really debugging stuff (20 max) [$DEBUG]

--output=FILE       output gro file [$OutputFn]

--resname=RESNAME   name or pattern for the residue to kill [$RESNAME]
--z1=Z1             kill RESNAME with  z1 < z < z2   [$Z1]
--z2=Z2                                              [$Z2]

--slab              keep all residues  z1 <= z <= z2   and change the 
                    boxsize from (bx, by, bz) to (bx, by, z2-z1)
		    (RESNAME is ignored)

BUGS:
=====
* If the output has more than 100,000 atoms water will be in random places
  No idea why--it is some very strange bug.
  At the moment I also write a tmp output file on the fly ( *_tmp.gro)
* The first residue in the gro file should have the number 1.
EOF
    exit;
}




######################################
#
# MAIN
#
######################################


# options
&GetOptions( "output=s" => \$OutputFn,
	     "resname=s"=> \$RESNAME,
	     "z1=f"     => \$Z1,
	     "z2=f"     => \$Z2,
	     "slab"     => \$bSLAB,
	     "debug=i" => \$DEBUG,
	     "verbose" => \$opt_verbose,
	     "version" => \$opt_version,
             "help" => \$opt_help) or &print_usage;
if ($opt_help) {&print_usage};
if ($opt_version) { msg (0, "%s\n", $Version_ID); exit 0 };
if ($opt_verbose) { $DEBUG = 3 };


open_log or die "Cannot open logfile for writing, stopped ";

if ($ARGV[0]) {
    $DataFn = $ARGV[0];
    -e $DataFn or 
	die "The first file '" . $DataFn . 
	    "' to read from does not exist, stopped ";
};

######### Initialize variables ############
my $TmpOut = $OutputFn . "_tmp.gro";
my ($line, $head, $head_nr_res, @unitcell);
my ($length, $nr_res) = (0, 0);

open FILE, "<$DataFn" or die "Cannot open '$DataFn', ";

# write on the fly, see if buffer gets corrupted
open TMPOUT, ">${TmpOut}" or die "Cannot open '$TmpOut', ";

# Header (comment)
$head        = <FILE>;
$head_nr_res = <FILE>;

# extract unitcell and count ATOM entries
while (<FILE>) {
    if (/^ +[0-9.]+ +[0-9.]+ +[0-9.]+ *$/) {@unitcell = split};
    # count ATOMS
    if (/^ *[0-9]+[a-zA-Z]+/) {$nr_res++};
    # all lines
    $length++;
};
msg (3, "Number of residues before (counted): $nr_res\n");
msg (3, "Number of residues before (header):  $head_nr_res\n");

msg (3, $bSLAB ? 
        "Making a slab out of the input with new bz = " . ($Z2 - $Z1) ."\n" : 
        "Going to remove all water molecules between $Z1 and $Z2!\n");
 

# This routine goes through a coordinate file and removes groups with a
# certain name (SOL in this case) that have one or more atoms with a z-
# coordinate of < Z1 or > Z2. One entire residue is analyzed and
# written to the output, with renumbered groups and atoms.

logger(1, "file: $DataFn\n");
logger(1, "original length: $length lines\n");
logger(1, "unitcell: %10.5f %10.5f %10.5f\n", @unitcell);

# write bits and pieces of online gro file
# (post process and subst NR_RES_KILLED witrh the real number
print  TMPOUT $head;
print  TMPOUT "NR_RES_FINAL\n";

###
# debugging... 
# 1. output garbled (water in wrong places) if output > 100,000 atoms (it seems)
# 2. BUT: on the fly output and buffer output identical
# ... no idea whats wrong!!


my $keep = 2;            # thats why the first res must have resnr 1
my $current_resnum = 1;  #  -"-
my $current_atom = 0;
my $still_valid = 1;

my $outbuffer = "";
my $buffer = "";
my $reject_count = 0;
my $last_atom = 0;
my $count_res = 0;
my ($bKeepRes);

my $template = "A5 A5 A5 A5 A8 A8 A8 A8 A8 A8 A1";
my ($resnr,$resname,$atomname,$atomnr,$x,$y,$z,$vx,$vy,$vz);

my $debug_count = 0;

# reset file pointer to beginning
seek FILE, 0, 0;
LINE: while ($line = <FILE>)
{
    chomp $line;
    # only process atom entries
    #    1234ATOM ....
    next LINE if ($line !~ /^ *[0-9]+[a-zA-Z]+/);

    # get data from fixed format gro file
    # and remove all spaces
    { 
	my @entry = unpack($template,$line);
	foreach (@entry) { tr/ //d; };
	($resnr,$resname,$atomname,$atomnr,$x,$y,$z,$vx,$vy,$vz) = @entry;
    };

    if ($resnr == $keep) # we're in the next residue now
    {
	if ($buffer)	# if there was something in the buffer,
	{		# print it.
	    $outbuffer .= $buffer;
            
            # for debugging: write atom entries on the fly
            print TMPOUT $buffer;

	    logger (10, "++ KEEP: " . 
               "outbuffer resnum [%g], was orig resnr [%g]:\n%s",
		    $current_resnum,  $resnr-1, $buffer);
	    $current_resnum++;          # update residue number
	    $last_atom = $current_atom; # update total number of atoms
	}			

	logger (10, 
		"##### new residue[%4d] ##### last kept [%4d] ######\n",
		$resnr, $current_resnum - 1);

	# this group is now current group
	$keep = $resnr+1;       
	$buffer = "";
	$still_valid = 1;	

	# counts atoms in each  residue again
	$count_res = 0;
	$reject_count = 0;
    };

    if ($DEBUG > 10) {
	$debug_count++;
	logger (10, "# line $debug_count in gro:  $line\n");
    };


    if ($bSLAB) {
	$bKeepRes = (($z >= $Z1 && $z <= $Z2) && $still_valid);
    } else {
	my $hasRESNAME;
	($hasRESNAME) = $resname =~ /$RESNAME/; # check if group contains SOL
	$bKeepRes = (($z < $Z1 || $z > $Z2) && $still_valid)  || !$hasRESNAME;
    };

    if ($bKeepRes)  # if in range, or not a RESNAME, add this atom to buffer
	            # (if it turns out later to be invalid the buffer is 
	            #  discarded)
    {
       my $buf;
       $count_res++;
       $current_atom = $last_atom + $count_res;

       # shift residues to the bottom of the new unitcell
       if ($bSLAB) { $z -= $Z1 };

       if ( $vx && $vy && $vz ) {
	   $buf = sprintf("%5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f\n",
		      $current_resnum,$resname,$atomname,$current_atom,
		      $x,$y,$z,$vx,$vy,$vz);
       } else {
	   $buf = sprintf("%5d%-5s%5s%5d%8.3f%8.3f%8.3f\n",
		      $current_resnum,$resname,$atomname,$current_atom,
		      $x,$y,$z);
       };
       $buffer .= $buf;
   }
   else
   {
       # reset $current_atom by subtracting all atoms of this residue which 
       # will not be counted because $count_res+1 (the current_atom) does
       # not fullfill the  $bKeepRes criterion
       #
       # does this correct the pains with current_current_atom ?
       $current_atom -= $count_res;

       $buffer = "";	        # delete value of $buffer;
       $still_valid = 0;	# don't print to buffer anymore for this group
       $reject_count++;
       logger (10, 
          "-- REJ: atomnr [orig:%g, now:%g] in orig RES [%s], resnr [%g]; ".
	  "rej_cnt %g\n", 
	       $atomnr, ($current_atom+$reject_count), 
	       $resname, $resnr, $reject_count);
   }
}

if ($buffer)	# if there was something in the buffer,
{		# print it.
    $outbuffer .= $buffer;
    logger (11, "Printing last buffer:\n$buffer\n");
    # for debugging: write atom entries on the fly
    print TMPOUT $buffer;
}			


# header of new gro file
# NEW number of atoms (after kill)
$outbuffer  = $head .  "$current_atom\n" . $outbuffer;

if ($bSLAB) {
    $unitcell[2] = abs ($Z2 - $Z1);
    logger(1, "Slab mode: all z-coordinates were shifted by %g.\n" . 
	      "New unitcell: %10.5f %10.5f %10.5f\n",
	   -$Z1, @unitcell);
};

$outbuffer .= sprintf "%10.5f %10.5f %10.5f\n", @unitcell;

open (OUT, "> $OutputFn") or die "Cannot open output file '$OutputFn': $!";
print OUT $outbuffer;
close (OUT);

logger(1, "Final number of residues: $current_atom\n");

# tmpout postproc
# todo:
# s/NR_RES_FINAL/$current_atom/

printf TMPOUT "%10.5f %10.5f %10.5f\n", @unitcell;

close(TMPOUT);

exit(0);	