# Big endian vs Little endian pain in the neck
# <f4 means read as np.float32 type little endian byter order
# >f4 means read as np.float32 type big    endian byter order
# Had to do the latter for runs on my desktop (NBI)

# There is a VTK reader in standard python libraries but 
# it was not working with snoopy... so had to do it the old
# fashioned way: read binary file line by line (or as a byte chunk)

import numpy as np
import matplotlib.pyplot as plt

def readvtk_sno(fname='v0001.vtk'):
  '''
  Usage:
  tt,vx,vy,vz,bx,by,bz,N,D=readvtk_sno(fname='v1000.vtk')
  '''
  ff = open(fname,'rb')

  ff.seek(0,0) # start of file

  ff.readline() # Call it l1 Out[70]: b'# vtk DataFile Version 2.0\n'
  
  # Read time
  l2=ff.readline() # l2 Out[71]: b't= 5.025051671818386e-01 Snoopy Code v5.0\n'
  tt = np.float32(l2.split()[1])

  ff.readline() # l3 Out[72]: b'BINARY\n'
  ff.readline() # l4 Out[73]: b'DATASET STRUCTURED_POINTS\n'
  
  # Read dimensions
  l5=ff.readline() # l5 Out[74]: b'DIMENSIONS 32 64 32\n'
  dimensions=l5.split()
  nx,ny,nz = np.int32(dimensions[1]),np.int32(dimensions[2]),np.int32(dimensions[3])
  dims=np.array([nx,ny,nz],dtype=np.int32)

  ff.readline() # l6 Out[75]: b'ORIGIN -0.5 -1 -0.5\n'

  # Read spacing
  l7=ff.readline() # l7 Out[76]: b'SPACING 0.03125 0.03125 0.03125\n'
  spacing=l7.split()
  dx,dy,dz = np.float32(spacing[1]),np.float32(spacing[2]),np.float32(spacing[3])
  dgrid=np.array([dx,dy,dz],dtype=np.float32)

  ff.readline() # l8 Out[77]: b'POINT_DATA 65536\n'
  ff.readline() # l9 Out[78]: b'SCALARS vx float\n'
  ff.readline() # l10 Out[79]: b'LOOKUP_TABLE default\n'

  ###
  #VERY IMPORTANT to read it as FLOAT32 - FLOAT64 reads more than it should!!
  ###
  #vx = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  vx = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)
  ff.readline() # l11 Out[81]: b'FIELD FieldData 5\n'
  ff.readline() # l12 Out[82]: b'vy 1 65536 float\n'

  #vy = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  vy = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)
  ff.readline() # l13 Out[84]: b'vz 1 65536 float\n'

  #vz = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  vz = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)
  ff.readline() # Out[86]: b'bx 1 65536 float\n'

  #bx = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  bx = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)
  ff.readline() #Out[88]: b'by 1 65536 float\n'

  #by = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  by = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)
  ff.readline() #Out[90]: b'bz 1 65536 float\n'

  #bz = np.fromfile(ff,dtype=np.float32,count=nx*ny*nz).reshape(dims)
  bz = np.fromfile(ff,dtype='>f4',count=nx*ny*nz).reshape(dims)

  return (tt,vx,vy,vz,bx,by,bz,dims,dgrid)