#!/usr/bin/perl
use Getopt::Long;
Getopt::Long::Configure ('bundling');

#################################################################
#                                                               #
#  This script is designed to take the bond information from a  #
#  Lammps/reaxc *.trj file and output a molecular fraction file #
#  for each frame.                                              #
#                                                               #
#  written by Mike Russo, PSU                                   #
#  modified by Aidan Thompson, 5/12/2011                        #
#                                                               #
#  The required .trj file is generated by running LAMMPS        #
#  with pair_style reax/c and the following settings in         #
#  the corresponding reax/c control file                        #
#     write_freq              25   ! write trajectory after so many steps
#     traj_compress           0    ! 0: no compression  1: uses zlib to compress trajectory output
#     traj_title              TATB ! (no white spaces)
#     atom_info               0    ! 0: no atom info, 1: print basic atom info in the trajectory file
#     atom_forces             0    ! 0: basic atom format, 1: print force on each atom in the trajectory file
#     atom_velocities         0    ! 0: basic atom format, 1: print the velocity of each atom in the trajectory file
#     bond_info               1    ! 0: do not print bonds, 1: print bonds in the trajectory file
#     angle_info              0    ! 0: do not print angles, 1: print angles in the trajectory file 
#                                                               #
#################################################################

#################################################################
#  Setting up some default variables, and options for the user  #
#  to input.                                                    #
#################################################################
$in_file = "bonds.trj";
@test = qw(C H O N);
GetOptions ('f|file=s' => \$in_file, 'a|atoms=s' => \@atoms, 'h|help' => \$help);
if($help) {
    print "Options for this program:\n-f --file for input file default= bonds.trj\n-a --atoms atom types (in correct order and input separately) default= @test\n";
    exit;
}
open INPUT, "<$in_file" or die "Cannot open $in_file: $!";
open OUTPUT, ">frac.dat" or die "Cannot open output file: $!";

if(@atoms) {
    @test = @atoms;
}

print "Input for this run:\n  Input file = $in_file\n Atom types = @test\n";
#################################################################

#################################################################
#  Main loop of the script. Goes through each frames bond list. #
#################################################################
$i = 0;
$section = 0;
$at_count = -1;
while(<INPUT>) {
    if(/chars_to_skip_section/) {
        $section++;
        &bonds if($section > 2);
        next;
    }
    next if($section <= 0);  #skipping the header section
    next if(/\s*[A-Za-z]/); #skip text lines

    split;
    if ($section == 1) {
	$q = $_[0];
	$temp = $_[1];
	$q--;
	$at_type[$q] = $temp;
	next;
    }

    $_[0]--;
    $_[1]--;
# Add i-j and j-i entries
    push @{$list[$_[0]]}, $_[1];
    push @{$list[$_[1]]}, $_[0];
    $at_count++ if($section == 2);
}
$section++;
&bonds;
close(INPUT);
#################################################################

#################################################################
#  Subroutine bonds: Uses the bond information to generate a    #
#  count for each species, put them into molecules, and then    #
#  count the number of each molecule type.                      #
#################################################################
sub bonds {
  $flag = ();
  $k = 0;
  for(0..$#list) {
      if($flag[$_] == 0) {
          push @{$full_list[$k]}, $_;
          foreach $atom (@{$full_list[$k]}) {
              for($o = 0; $o <= $#{$list[$atom]}; $o++) {
                  unless(grep /^$list[$atom][$o]$/, @{$full_list[$k]}) {
                      push @{$full_list[$k]}, $list[$atom][$o];
                      $flag[$list[$atom][$o]] = 1;
                  }
              }
          }
      } else {
          next;
      }
      $k++;
  }

### Output section ###
  $frame = $section - 2;
  open OUTPUT2, ">temp_$frame.dat" or die "Cannot open temp file: $!";
  print OUTPUT2 "Frame # $frame\n";
  for($m = 0; $m < $k; $m++) {     #cycle through each molecule
      foreach $atom ((@{$full_list[$m]})) { #for each atom in this molecule
          ${"$test[$at_type[$atom]]"} += 1;        #Create variable named C,H,O, etc and set it to count
      }
      print OUTPUT2 "Mol $m = ";
      foreach $atom (@test) {
          print OUTPUT2 "$atom${$atom}" if(${$atom} > 0);
      }
      print OUTPUT2 "\n";
      for($r = 0; $r <= $#test; $r++) {
          ${"$test[$r]"} = 0;
      }
  }

  close (OUTPUT2); #close the temp file as output
  open INPUT3, "<temp_$frame.dat" or die "Cannot open temp file: $!"; # open it as input
  while(<INPUT3>) {
      next if(/Frame/);
      split;
      push @mol_list, $_[3] unless(grep /^$_[3]$/, @mol_list);
      ${"$_[3]"}++;
  }
  print OUTPUT "Frame # $frame\n";
  foreach $mol (@mol_list) {
      printf OUTPUT "%4d  %s\n", ${"$mol"}, $mol;
      ${"$mol"} = 0;
  }
  @mol_list = ();
###

### Cleanup between frames ###
  for(0..$at_count) {
      @{$full_list[$_]} = ();
      @{$list[$_]} = ();
  }
###
}
#################################################################

