# Pizza.py toolkit, www.cs.sandia.gov/~sjplimp/pizza.html
# Steve Plimpton, sjplimp@sandia.gov, Sandia National Laboratories
#
# Copyright (2005) Sandia Corporation.  Under the terms of Contract
# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
# certain rights in this software.  This software is distributed under 
# the GNU General Public License.

# data tool

oneline = "Read, write, manipulate LAMMPS data files"

docstr = """
d = data("data.poly")            read a LAMMPS data file, can be gzipped
d = data()			 create an empty data file

d.map(1,"id",3,"x")              assign names to atom columns (1-N)

coeffs = d.get("Pair Coeffs")    extract info from data file section
q = d.get("Atoms",4)

  1 arg = all columns returned as 2d array of floats
  2 args = Nth column returned as vector of floats

d.reorder("Atoms",1,3,2,4,5)     reorder columns (1-N) in a data file section

  1,3,2,4,5 = new order of previous columns, can delete columns this way

d.title = "My LAMMPS data file"	 set title of the data file
d.headers["atoms"] = 1500        set a header value
d.sections["Bonds"] = lines      set a section to list of lines (with newlines)
d.delete("bonds")		 delete a keyword or section of data file
d.delete("Bonds")
d.replace("Atoms",5,vec)      	 replace Nth column of section with vector
d.newxyz(dmp,1000)		 replace xyz in Atoms with xyz of snapshot N

  newxyz assumes id,x,y,z are defined in both data and dump files
    also replaces ix,iy,iz if they are defined
  
index,time,flag = d.iterator(0/1)          loop over single data file snapshot
time,box,atoms,bonds,tris,lines = d.viz(index)   return list of viz objects

  iterator() and viz() are compatible with equivalent dump calls
  iterator() called with arg = 0 first time, with arg = 1 on subsequent calls
    index = timestep index within dump object (only 0 for data file)
    time = timestep value (only 0 for data file)
    flag = -1 when iteration is done, 1 otherwise
  viz() returns info for specified timestep index (must be 0)
    time = 0
    box = [xlo,ylo,zlo,xhi,yhi,zhi]
    atoms = id,type,x,y,z for each atom as 2d array
    bonds = id,type,x1,y1,z1,x2,y2,z2,t1,t2 for each bond as 2d array
      NULL if bonds do not exist
    tris = NULL
    lines = NULL
    
d.write("data.new")             write a LAMMPS data file
"""

# History
#   8/05, Steve Plimpton (SNL): original version
#   11/07, added triclinic box support

# ToDo list

# Variables
#   title = 1st line of data file
#   names = dictionary with atom attributes as keys, col #s as values 
#   headers = dictionary with header name as key, value or tuple as values
#   sections = dictionary with section name as key, array of lines as values
#   nselect = 1 = # of snapshots

# Imports and external programs

from os import popen

try: tmp = PIZZA_GUNZIP
except: PIZZA_GUNZIP = "gunzip"

# Class definition

class data:

  # --------------------------------------------------------------------

  def __init__(self,*list):
    self.nselect = 1
    
    if len(list) == 0:
      self.title = "LAMMPS data file"
      self.names = {}
      self.headers = {}
      self.sections = {}
      return

    file = list[0]
    if file[-3:] == ".gz": f = popen("%s -c %s" % (PIZZA_GUNZIP,file),'r')
    else: f = open(file)

    self.title = f.readline()
    self.names = {}
    
    headers = {}
    while 1:
      line = f.readline()
      line = line.strip()
      if len(line) == 0:
        continue
      found = 0
      for keyword in hkeywords:
        if line.find(keyword) >= 0:
	  found = 1
	  words = line.split()
	  if keyword == "xlo xhi" or keyword == "ylo yhi" or \
	    keyword == "zlo zhi":
	    headers[keyword] = (float(words[0]),float(words[1]))
	  elif keyword == "xy xz yz":
	    headers[keyword] = \
              (float(words[0]),float(words[1]),float(words[2]))
          else:
	    headers[keyword] = int(words[0])
      if not found:
        break

    sections = {}
    while 1:
      found = 0
      for pair in skeywords:
        keyword,length = pair[0],pair[1]
        if keyword == line:
	  found = 1
          if not headers.has_key(length):
            raise StandardError, \
                  "data section %s has no matching header value" % line
	  f.readline()
          list = []
          for i in xrange(headers[length]): list.append(f.readline())
          sections[keyword] = list
      if not found:
        raise StandardError,"invalid section %s in data file" % line
      f.readline()
      line = f.readline()
      if not line:
        break
      line = line.strip()
      
    f.close()
    self.headers = headers
    self.sections = sections

  # --------------------------------------------------------------------
  # assign names to atom columns

  def map(self,*pairs):
    if len(pairs) % 2 != 0:
      raise StandardError, "data map() requires pairs of mappings"
    for i in range(0,len(pairs),2):
      j = i + 1
      self.names[pairs[j]] = pairs[i]-1

  # --------------------------------------------------------------------
  # extract info from data file fields

  def get(self,*list):
    if len(list) == 1:
      field = list[0]
      array = []
      lines = self.sections[field]
      for line in lines:
        words = line.split()
        values = map(float,words)
        array.append(values)
      return array
    elif len(list) == 2:
      field = list[0]
      n = list[1] - 1
      vec = []
      lines = self.sections[field]
      for line in lines:
        words = line.split()
        vec.append(float(words[n]))
      return vec
    else:
      raise StandardError, "invalid arguments for data.get()"

  # --------------------------------------------------------------------
  # reorder columns in a data file field

  def reorder(self,name,*order):
    n = len(order)
    natoms = len(self.sections[name])
    oldlines = self.sections[name]
    newlines = natoms*[""]
    for index in order:
      for i in xrange(len(newlines)):
        words = oldlines[i].split()
        newlines[i] += words[index-1] + " "
    for i in xrange(len(newlines)):
      newlines[i] += "\n"
    self.sections[name] = newlines

  # --------------------------------------------------------------------
  # replace a column of named section with vector of values

  def replace(self,name,icol,vector):
    lines = self.sections[name]
    newlines = []
    j = icol - 1
    for i in xrange(len(lines)):
      line = lines[i]
      words = line.split()
      words[j] = str(vector[i])
      newline = ' '.join(words) + '\n'
      newlines.append(newline)
    self.sections[name] = newlines

  # --------------------------------------------------------------------
  # replace x,y,z in Atoms with x,y,z values from snapshot ntime of dump object
  # assumes id,x,y,z are defined in both data and dump files
  # also replaces ix,iy,iz if they are defined

  def newxyz(self,dm,ntime):
    nsnap = dm.findtime(ntime)

    if dm.scaled(nsnap): scaleflag = 1
    else: scaleflag = 0
    dm.sort(ntime)
    
    if scaleflag: dm.unscale(ntime)
    x,y,z = dm.vecs(ntime,"x","y","z")
    if scaleflag: dm.scale(ntime)

    self.replace("Atoms",self.names['x']+1,x)
    self.replace("Atoms",self.names['y']+1,y)
    self.replace("Atoms",self.names['z']+1,z)
    
    if dm.names.has_key("ix") and self.names.has_key("ix"):
      ix,iy,iz = dm.vecs(ntime,"ix","iy","iz")
      self.replace("Atoms",self.names['ix']+1,ix)
      self.replace("Atoms",self.names['iy']+1,iy)
      self.replace("Atoms",self.names['iz']+1,iz)
      
  # --------------------------------------------------------------------
  # delete header value or section from data file

  def delete(self,keyword):

    if self.headers.has_key(keyword): del self.headers[keyword]
    elif self.sections.has_key(keyword): del self.sections[keyword]
    else: raise StandardError, "keyword not found in data object"

  # --------------------------------------------------------------------
  # write out a LAMMPS data file

  def write(self,file):
    f = open(file,"w")
    print >>f,self.title
    for keyword in hkeywords:
      if self.headers.has_key(keyword):
        if keyword == "xlo xhi" or keyword == "ylo yhi" or \
               keyword == "zlo zhi":
	  pair = self.headers[keyword]
	  print >>f,pair[0],pair[1],keyword
        elif keyword == "xy xz yz":
	  triple = self.headers[keyword]
	  print >>f,triple[0],triple[1],triple[2],keyword
        else:
	  print >>f,self.headers[keyword],keyword
    for pair in skeywords:
      keyword = pair[0]
      if self.sections.has_key(keyword):
        print >>f,"\n%s\n" % keyword
        for line in self.sections[keyword]:
	  print >>f,line,
    f.close()

  # --------------------------------------------------------------------
  # iterator called from other tools

  def iterator(self,flag):
    if flag == 0: return 0,0,1
    return 0,0,-1

  # --------------------------------------------------------------------
  # time query from other tools

  def findtime(self,n):
    if n == 0: return 0
    raise StandardError, "no step %d exists" % (n)
   
  # --------------------------------------------------------------------
  # return list of atoms and bonds to viz for data object

  def viz(self,isnap):
    if isnap: raise StandardError, "cannot call data.viz() with isnap != 0"
    
    id = self.names["id"]
    type = self.names["type"]
    x = self.names["x"]
    y = self.names["y"]
    z = self.names["z"]
    
    xlohi = self.headers["xlo xhi"]
    ylohi = self.headers["ylo yhi"]
    zlohi = self.headers["zlo zhi"]
    box = [xlohi[0],ylohi[0],zlohi[0],xlohi[1],ylohi[1],zlohi[1]]

    # create atom list needed by viz from id,type,x,y,z

    atoms = []
    atomlines = self.sections["Atoms"]
    for line in atomlines:
      words = line.split()
      atoms.append([int(words[id]),int(words[type]),
                    float(words[x]),float(words[y]),float(words[z])])

    # create list of current bond coords from list of bonds
    # assumes atoms are sorted so can lookup up the 2 atoms in each bond

    bonds = []
    if self.sections.has_key("Bonds"):
      bondlines = self.sections["Bonds"]
      for line in bondlines:
        words = line.split()
        bid,btype   = int(words[0]),int(words[1])
        atom1,atom2 = int(words[2]),int(words[3])
        atom1words  = atomlines[atom1-1].split()
        atom2words  = atomlines[atom2-1].split()
        bonds.append([bid,btype,
                      float(atom1words[x]),float(atom1words[y]),
                      float(atom1words[z]),
                      float(atom2words[x]),float(atom2words[y]),
                      float(atom2words[z]),
		      float(atom1words[type]),float(atom2words[type])])
      
    tris = []
    lines = []
    return 0,box,atoms,bonds,tris,lines

  # --------------------------------------------------------------------
  # return box size

  def maxbox(self):
    xlohi = self.headers["xlo xhi"]
    ylohi = self.headers["ylo yhi"]
    zlohi = self.headers["zlo zhi"]
    return [xlohi[0],ylohi[0],zlohi[0],xlohi[1],ylohi[1],zlohi[1]]

  # --------------------------------------------------------------------
  # return number of atom types

  def maxtype(self):
    return self.headers["atom types"]

# --------------------------------------------------------------------
# data file keywords, both header and main sections

hkeywords = ["atoms","lines","tris",
             "bonds","angles","dihedrals","impropers",
	     "atom types","bond types","angle types","dihedral types",
	     "improper types","xlo xhi","ylo yhi","zlo zhi","xy xz yz"]

skeywords = [["Masses","atom types"],
             ["Atoms","atoms"],["Lines","lines"],["Triangles","tris"],
             ["Bonds","bonds"],
	     ["Angles","angles"],["Dihedrals","dihedrals"],
	     ["Impropers","impropers"],["Velocities","atoms"],
             ["Pair Coeffs","atom types"],
	     ["Bond Coeffs","bond types"],["Angle Coeffs","angle types"],
	     ["Dihedral Coeffs","dihedral types"],
	     ["Improper Coeffs","improper types"],
             ["BondBond Coeffs","angle types"],
             ["BondAngle Coeffs","angle types"],
             ["MiddleBondTorsion Coeffs","dihedral types"],
             ["EndBondTorsion Coeffs","dihedral types"],
             ["AngleTorsion Coeffs","dihedral types"],
             ["AngleAngleTorsion Coeffs","dihedral types"],
             ["BondBond13 Coeffs","dihedral types"],
             ["AngleAngle Coeffs","improper types"]]
