#!/usr/bin/env python2
import os
import xml.etree.ElementTree as ET
 
#A helpful function to load compressed or uncompressed XML files
def loadXMLFile(filename):
    #Check if the file is compressed or not, and 
    if (os.path.splitext(filename)[1][1:].strip() == "bz2"):
        import bz2
        f = bz2.BZ2File(filename)
        doc = ET.parse(f)
        f.close()
        return doc
    else:
        return ET.parse(filename)

import sys
if len(sys.argv) == 1:
    print "dynamo2xyz CONFIG-FILE-NAME1 [CONFIG-FILE-NAME2]"
    print " This program converts a dynamo configuration file to xyz format"
    print " and prints it on the screen. To save this conversion, just "
    print "redirect it to a file. For example,"
    print "  dynamo2xyz config.out.xml.bz2 > config.xyz"
    print "If multiple file names are given, they are stitched together,"
    print "in order, to make an animation file."
    exit(1);


particlecount=-1
for filename in sys.argv[1:]:
    XMLDoc=loadXMLFile(filename)
    RootElement=XMLDoc.getroot()

    #All files need the same number of particles inside them, first
    #store the count if we're on the first file
    currentparticlecount = len(RootElement.findall('.//Pt'))
    if particlecount == -1:
        particlecount = currentparticlecount

    #Now check that all files have the same particle count
    if particlecount != currentparticlecount:
        print "input file", filename, "has",currentparticlecount,"particles, but all files must have the same number of particles. The first file has",particlecount
        exit(1)
        
    #Start of the XYZ file format. We skip the optional comment line
    #line 1: number of particles
    print particlecount
    #line 2: molecule name
    print "DynamOdata"
    #line 3-onwards: atom_number x y z vx vy vz
    for particleNode in RootElement.findall('.//Pt'):
        posNode = particleNode.find('P')
        velNode = particleNode.find('V')
        #First column is the atom name/type. Note the trailing comma to prevent a newline!
        print "H", 
        #Next three columns are the position
        print posNode.attrib['x'], posNode.attrib['y'], posNode.attrib['z'],
        #Next three columns are the velocity
        print velNode.attrib['x'], velNode.attrib['y'], velNode.attrib['z']
