#!/usr/bin/env python

import sys
import common

# Check command line

try:
   x, filename, grid_type = sys.argv
except ValueError:
   print("usage: "+sys.argv[0]+" filename grid_type")
   sys.exit(2)


from ezfio import ezfio

def main():
  ezfio.set_file(filename)
  file = open("%s_%s.cube"%(filename,grid_type),"w")
  print(" Cube File", file=file)
  print(filename, grid_type, file=file)
  data = [ ezfio.nuclei_nucl_num ]+ezfio.grid_origin
  print("%5d%12.6f%12.6f%12.6f" % tuple(data), file=file)
  data = [ ezfio.grid_num_x, ezfio.grid_step_size[0],0.,0. ]
  print("%5d%12.6f%12.6f%12.6f" % tuple(data), file=file)
  data = [ ezfio.grid_num_y, 0., ezfio.grid_step_size[1],0. ]
  print("%5d%12.6f%12.6f%12.6f" % tuple(data), file=file)
  data = [ ezfio.grid_num_z, 0.,0., ezfio.grid_step_size[2] ]
  print("%5d%12.6f%12.6f%12.6f" % tuple(data), file=file)

  charge = ezfio.nuclei_nucl_charge
  coord = ezfio.nuclei_nucl_coord
  for i in range(ezfio.nuclei_nucl_num):
    data = [ charge[i], charge[i], coord[0][i], coord[1][i], coord[2][i] ]
    print("%5d%12.6f%12.6f%12.6f%12.6f" % tuple(data), file=file)

  data = getattr(ezfio,"grid_data_%s"%grid_type)
  for i in range(0,ezfio.grid_num_x):
   for j in range(0,ezfio.grid_num_y):
    for k in range(0,ezfio.grid_num_z,6):
     for l in range(k,min(ezfio.grid_num_z,k+6)):
      file.write (" %12.5E"%data[l][j][i])
     file.write ("\n")
  file.close()

main()