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

def parseCNT(data, filename, format, 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())
    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 i in range(len(elements)):
      allElements += [elements[i]] * nElements[i]

    # 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, False)
    # 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, False);
    # Mandatory call at the end to update some internal dimensions
    # Related to the box.
    box.setMargin(data.getMaxElementSize() + data.getAllNodeExtens(box), True);
    # Set a commentry for this subset of data.
    data.setFileCommentary(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.FileFormat.new("VASP CNT files", ("*.cnt", "cnt.*"))
v_sim.Rendering.getByName("Atom visualisation").addFileFormat(0, fmt, 95, parseCNT)