#  Copyright (C) 2012,2013
#      Max Planck Institute for Polymer Research
#  Copyright (C) 2008,2009,2010,2011
#      Max-Planck-Institute for Polymer Research & Fraunhofer SCAI
#  
#  This file is part of ESPResSo++.
#  
#  ESPResSo++ is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#  
#  ESPResSo++ is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>. 


"""
***************************************************
**lammpsfilewrite** - method to write a LAMMPS file
***************************************************

This method creates a LAMMPS file with the data from a specific system.

The method needs the following parameters:

* filename
    name of the file where the table schould be saved in
* system
    ESPResSo system which creates the data e.g.:
    >>>system, integrator = espresso.standard_system.LennardJones(100,(10,10,10))
* append
    =False
      the data in the file will be overwritten
    =True
      the data will be appended
"""
import espresso
from math import sqrt
from espresso import Real3D

def lammpsfilewrite(filename, system, append=False):
  if append:
    file = open(filename, 'a')
  else:    
    file = open(filename,'w')
  maxParticleID = int(espresso.analysis.MaxPID(system).compute())
  pid = 0
  atomtype = []
  typeA = []
  q = []
  xpos = []
  ypos = []
  zpos = []
  vx = []
  vy = []
  vz = []
  particleID = []
  mass = 0
  masses = {}
  atomtypes = 0
  angletypes = 0
  bondtypes = 0   
  dihedrals = 0   
  impropers = 0   
  bonds = 0       
  angles = 0      
  dihedrals = 0   
  impropers = 0   
  number = 0
  boxsizex = system.bc.boxL[0]  
  boxsizey = system.bc.boxL[1]
  boxsizez = system.bc.boxL[2]
  boxsize = 0    

  while pid <= maxParticleID:
    if system.storage.particleExists(pid):
      particle = system.storage.getParticle(pid)

      particleID.append(pid)
      typeA.append(particle.type)
      xpos.append(particle.pos[0])
      ypos.append(particle.pos[1])
      zpos.append(particle.pos[2])
      q.append(particle.q)
      vx.append(particle.v[0])
      vy.append(particle.v[1])
      vz.append(particle.v[2])

      #Masses
      if typeA[number] not in atomtype:
        atomtype.append(typeA[number])
        masses[atomtype[atomtypes]] = particle.mass
        atomtypes += 1

      number += 1
      pid += 1

    else:
      pid += 1
      
  #Head
  st = "LAMMPS data file\nREMARK generated by ESPResSo++\n%d atoms\n%d bonds\n%d angles\n%d dihedrals\n%d imropers\n"%(maxParticleID, bonds, angles, dihedrals, impropers)
  file.write(st)

  st = "\n%d atom types\n%d bond types\n%d angle types\n%d dihedral types\n%d improper types\n"%(atomtypes, bondtypes, angles, dihedrals, impropers)
  file.write(st)

  st = "\n%d %d xlo xhi\n%d %d ylo yhi\n%d %d zlo zhi\n"%(boxsize, boxsizex, boxsize, boxsizey, boxsize, boxsizez)
  file.write(st)

  #Masses
  st = "\nMasses\n\n"
  file.write(st)
  atomtype.sort()
  pid = 0
  while pid < atomtypes:
    mass = masses[atomtype[pid]]
    st = "%d %15.10f\n"%(atomtype[pid], mass)
    file.write(st)
    pid += 1

  #Atoms
  st = "\nAtoms\n\n"
  file.write(st)
  pid = 0           
  while pid < number:
    st = "%d %d %d %15.10f %15.10f %15.10f\n"%(particleID[pid], typeA[pid], q[pid], xpos[pid], ypos[pid], zpos[pid])#moleculetag should be at the second poition, but does not exist in ESPResSo yet.
    file.write(st)
    pid += 1

  #Velocities
  st = "\nVelocities\n\n"
  file.write(st)
  pid = 0  
  while pid < number:
    st =  "%d %15.10f %15.10f %15.10f\n"%(particleID[pid], vx[pid], vy[pid], vz[pid])
    file.write(st)
    pid += 1

  file.close()
