import numpy as np

""" Example usage:

import numpy as np
from writeDx import writeDx

x = y = z = np.arange(-10,10,1)
X,Y,Z = np.meshgrid(x,y,z,indexing='ij')

origin = [a[0] for a in (x,y,z)]
delta  = [a[1]-a[0] for a in (x,y,z)] 

R = np.sqrt(X**2 + Y**2 + Z**2)

U = (R-5)**2
U[R<5] = 0

writeDx("harmonic-sphere.dx", U, origin, delta, fmt="%.2f")

"""

def writeDx(outfile, data, origin, delta, fmt="%.3f"):
  shape = np.shape(data)
  num = np.prod(shape)
  assert( len(shape) == 3 )
  assert( len(origin) == 3 )
  assert( len(delta) == 3 )
  headerInfo = dict( nx=shape[0], ny=shape[1], nz=shape[2],
                     ox=origin[0], oy=origin[1], oz=origin[2],
                     dx=delta[0], dy=delta[1], dz=delta[2],
                     num=num
                   )
  data = data.flatten(order='C')
  header = """# OpenDX density file
# File format: http://opendx.sdsc.edu/docs/html/pages/usrgu068.htm#HDREDF
object 1 class gridpositions counts  {nx} {ny} {nz}
origin {ox} {oy} {oz}
delta  {dx} 0.000000 0.000000
delta  0.000000 {dy} 0.000000
delta  0.000000 0.000000 {dz}
object 2 class gridconnections counts  {nx} {ny} {nz}
object 3 class array type double rank 0 items {num} data follows""".format(**headerInfo)
  len(data)

  if num == 3*(num//3):
    footer = ""
  else:
    footer = " ".join([fmt % x for x in data[3*(num//3):]]) # last line of data
    footer += "\n"

  footer += """attribute "dep" string "positions"
object "density" class field 
component "positions" value 1
component "connections" value 2
component "data" value 3
"""
  np.savetxt( outfile, np.reshape(data[:3*(num//3)], (num//3,3), order='C'), 
              fmt=fmt,
              header=header, comments='', footer=footer )