#!/usr/bin/python

##  shrink_traj.py:
##   Python script for shrinking long and big molecular dynamics
##   "xyz" trajectory files into more easy-to-handle ones
##
## By Julen Larrucea
## www.larrucea.eu

import os,sys
version=0.3

nat=0
linec=0 #line count
frame=[]
struc=0 #counter for structures in trajectory
step=3 #every how many frames do we print to output
trajfile="TRAJEC.xyz"
outfile="smallTRAJEC.xyz"

if len(sys.argv)>1:
 for i in sys.argv:
  if i.startswith('-'):
   option=i.split('-')[1]
   if option=="f":
    trajfile= sys.argv[sys.argv.index('-f')+1]
   if option=="s":
    s= sys.argv[sys.argv.index('-s')+1]
    try:
     step=int(s)
    except:
     print "d must be an integer"
   if option=="o":
    outfile= sys.argv[sys.argv.index('-o')+1]
   if option=="v":
    print "shrink_traj.py v. " + float(version)
    sys.exit()
   if option=="h":
    print '''
  Now help:
   -h: print this help
   -v: print version
   -f: name of trajectory file (def. TRAJEC.xyz)
   -s: step (every how many steps does it store a frame. Def. 3)
   -o: output name (def. smallTRAJEC.xyz)
'''
    sys.exit()


out=open(outfile,'w')
for line in open(trajfile,'r'):
  if line.strip().isdigit():
    nat=int(line.strip())
    linec=0
    struc+=1
    if len(frame)>0:
      print>>out,  "  ",nat
      print>>out,  "  stucture No.: ", struc-1
      for i in frame:
       print>>out, i[0],i[1],i[2],i[3]  
        
    frame=[]
  if float(struc)/float(step) == struc/step or struc==1:
    if linec>1:
     frame.append([line.split()[0],line.split()[1],line.split()[2],line.split()[3]])
    linec+=1