#!/usr/bin/env python
from gi.repository import v_sim
import traceback

def parseCNT(format, filename, data, iset, cancel):
  print "test file: %s for VisuData: %s(%d)" % (filename, data, iset)
  valid = False
  try:
    # Read the given file.
    f = open(filename, "r")
    desc = f.readline()
    f.readline()
    box = map(float, f.readline().split()) + \
          map(float, f.readline().split()) + \
          map(float, f.readline().split())
    if len(box) != 9:
      raise ValueError
    elements = f.readline().split()
    nElements = map(int, f.readline().split())
    # From here, we consider that this file is indeed a VASP file.
    valid = True
    f.readline()
    f.readline()
    coords = []
    for line in f.xreadlines():
      coords.append(map(float, line.split()[:3]))
    f.close()
    allElements = []
    for (ele, n) in zip(elements, nElements):
      allElements += [ele, ] * n

    # Pass the variables to V_Sim.
    # Declare one set of data for this file.
    data.setNSubset(1)
    # Setup the box geomtry.
    box = v_sim.Box.new_full(box, v_sim.BoxBoundaries.PERIODIC)
    # Setting up the box geometry here make available to use the
    # converting routines to go from reduced coordinates to cartesian.
    data.setBox(box)
    # Allocate the internal arrays for these elemnts and number of nodes.
    data.allocateByNames(nElements, elements)
    for (coord, ele) in zip(coords, allElements):
      # For each node, add it with its coordinates
      data.addNodeFromElementName(ele, coord, False);
    # Set a commentry for this subset of data.
    data.setDescription(desc.strip(), 0)

    return True
  except Exception as error:
    traceback.print_exc()
    return valid
  return False

# Declare this new fileformat to V_Sim.
fmt = v_sim.DataLoader.new("VASP CNT files", ("*.cnt", "cnt.*"), False, parseCNT, 95)
v_sim.DataAtomicClass.addLoader(fmt)