# -*- coding: utf-8 -*-

"""
This module provides a simple interface into RProfile-01-????.bobaaa files
generated by the PPMstar code
"""
from __future__ import division
from __future__ import print_function

from builtins import str
from builtins import range
from past.utils import old_div
from builtins import object
import struct
import logging
import os
import re
import sys

try:
  import numpy as np
except ImportError:
  print("numpy is required for reading rprofiles")

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger('rp_info')
log.setLevel(logging.DEBUG)

center_types = [(k, 'f8') for k in ('phi', 'theta', 'x', 'y', 'z')]
normal_types = [(k, 'f8') for k in ('x', 'y', 'z')]

class rprofile_set(object):

  def __init__(self, source, **kwargs):
    """Load a set of RProfiles. Can pass a `path` to a set of files, or a list of `files`. Passing
    the `lazy=True` makes everything operate only from disk, without loading things into memory
    (for large datasets).

    `source` is the path to a directory containing profiles or a list of RProfile files to open

    `stride` the iteration increment. Every `stride` element is looked at during iteration. Default is 1.

    `first_dump` the first dump to iterate from or None

    `last_dump` the last dump to iterate through (inclusive)

    There is no `.get` method... you must iterate through the files like this:

    .. code-block:: python
      :linenos:

      rp_set = lcse.rprofile_set(path=targetpath, lazy=True, logging=True)
      for rp in rp_set:
        rp.get("fv_hi")
        rp.get_table("fv")
    """

    self.path = source
    self.files = source if isinstance(source, list) else []
    self.lazy = kwargs.get('lazy', True)
    self.stride = kwargs.get('stride', 1)
    self.first_dump = kwargs.get('first_dump')
    self.last_dump = kwargs.get('last_dump')

    self._logging = kwargs.get('logging')
    self.log = log if self._logging else None

    self.ray_profiles = {}

    self._current_ix = 0
    self._current = None

    if self.path:
      self.files = self.get_file_list_for_path(self.path) if os.path.isdir(self.path) else [self.path]

    dump_re = re.compile('(.*)-([\d]{4})\.bobaaa')

    self.dump_map = dict((int(dump_re.match(f).groups()[1]), f) for f in self.files if dump_re.match(f))
    self.file_map = dict((f, int(dump_re.match(f).groups()[1])) for f in self.files if dump_re.match(f))

    self.dumps = list(self.dump_map.keys())
    self.dumps.sort()

  def __iter__(self):
    self._current_ix = self.dumps.index(self.first_dump) if self.first_dump else 0
    return self

  # Python 3 bullshit
  def __next__(self):
    return next(self)

  def __next__(self):

    if self._current_ix < len(self.dumps):
      dump = self.dumps[self._current_ix]

      if self.last_dump and dump > self.last_dump:
        raise StopIteration()

      rp = self.ray_profiles.get(dump, rprofile(self.dump_map[dump], lazy=self.lazy, logging=self._logging))

      if not self.lazy and (dump not in self.ray_profiles):
          self.ray_profiles[dump] = rp

      self._current = rp
      self._current_ix += self.stride
      return rp

    else:
      raise StopIteration()

#  def reset(self):
#
#    if first_dump:
#      self.start = self.dumps.index(first_dump)
#
#    self._current_ix = self.start
#    self._current = None

  def get_dump(self, dump=None):
    """ Get a new `rprofile` instance for `dump`. These are NOT cached internally."""

    if self.dumps and dump is None:
      dump = self.dumps[-1]
    elif dump not in self.dump_map:
      return None

    return self.ray_profiles.get(dump, rprofile(self.dump_map[dump], lazy=self.lazy, logging=self._logging))

  def get_file_list_for_path(self, path):
    """ Return a list of RProfiles at the given path"""

    filenames = [os.path.join(path, f) for f in os.listdir(path) if f.startswith('RProfile') and f.endswith('.bobaaa')]
    filenames.sort()

    return filenames

#  def load_files(self, filenames):
#    """ Loads `filenames` """
#
#    # This should add to the existing
#
#    self.files = filenames
#    self.files.sort()
#    self.ray_profiles = [rprofile(f, lazy=self.lazy, logging=self._logging) for f in self.files]

  def check_for_new(self, path=None):
    """Check path for new files"""

    current_files = self.get_file_list_for_path(self.path or path)

    new_files = [f for f in current_files if f not in self.files]

    self.files.extend(new_files)
    self.files.sort()

    return len(new_files) > 0

class rprofile(object):
  """
   `rprofile.header_attrs` is a dictionary of header attributes

  """


  header_var_list = [
      dict(name='version', pos=0, type='i'),
      dict(name='cell_ct_low', pos=1, type='i'),
      dict(name='nbuckets', pos=2, type='i'),
      dict(name='dump', pos=3, type='i'),
      dict(name='sizeof_float', pos=4, type='i'),
      dict(name='has_centers', pos=5, type='i'),
      dict(name='has_corners', pos=6, type='i'),
      dict(name='has_normals', pos=7, type='i'),
      dict(name='isrestart', pos=8, type='i', min_ver=12),
      dict(name='var_ct_low', pos=9, type='i'),
      dict(name='var_ct_high', pos=10, type='i'),
      dict(name='cell_ct_high', pos=11, type='i'),
      dict(name='ncpucores', pos=12, type='i'),
      dict(name='ntxbricks', pos=13, type='i'),
      dict(name='ntybricks', pos=14, type='i'),
      dict(name='ntzbricks', pos=15, type='i'),
      dict(name='nxteams', pos=16, type='i'),
      dict(name='nyteams', pos=17, type='i'),
      dict(name='nzteams', pos=18, type='i'),
      dict(name='nx', pos=19, type='i'),
      dict(name='ny', pos=20, type='i'),
      dict(name='nz', pos=21, type='i'),
      dict(name='nsugar', pos=22, type='i'),
      dict(name='nbdy', pos=23, type='i'),
      dict(name='nfluids', pos=24, type='i'),
      dict(name='nvars', pos=25, type='i'),
      dict(name='nhalfwaves', pos=26, type='i'),
      dict(name='maxrad', pos=27, type='i'),
      dict(name='nteamsinbunch', pos=28, type='i'),
      dict(name='ndumps', pos=29, type='i'),
      dict(name='ndumpstodo', pos=30, type='i'),
      dict(name='nrminrad', pos=31, type='i'),
      dict(name='nrmaxrad', pos=32, type='i'),
      dict(name='iburn', pos=33, type='i', min_ver=12),
      dict(name='imuffledbdry', pos=34, type='i', min_ver=12),
      dict(name='ireflectbdry', pos=35, type='i', min_ver=12),

      # fheader (the offsets are in the fheader)
      dict(name='radin0', pos=0, type='f', help='Gravity completely off inside this radius'),
      dict(name='radinner', pos=1, type='f', help='Gravity starts turning off inside this radius'),
      dict(name='radbase', pos=2, type='f', help='Bot convect zone'),
      dict(name='radtop', pos=3, type='f', help='Top convect zone'),
      dict(name='radouter', pos=4, type='f', help='Grav starts turning off outside this radius'),
      dict(name='radout0', pos=5, type='f', help='Gravity completely off outside this radius'),
      dict(name='radmax', pos=6, type='f', help='distance from center of grid to nearest edge'),
      dict(name='dlayerbot', pos=7, type='f', help='thickness of flame zone'),
      dict(name='dlayertop', pos=8, type='f', help='thickness of transition @ top of convect zone'),
      dict(name='totallum', pos=9, type='f'),
      dict(name='grav00base', pos=10, type='f'),
      dict(name='rho00base', pos=11, type='f'),
      dict(name='prs00base', pos=12, type='f'),
      dict(name='gammaconv', pos=13, type='f'),
      dict(name='gammabelow', pos=14, type='f'),
      dict(name='gammaabove', pos=15, type='f'),
      dict(name='gravconst', pos=16, type='f'),
      dict(name='rhoconv', pos=17, type='f'),
      dict(name='rhoabove', pos=18, type='f'),
      dict(name='airmu', pos=19, type='f', min_ver=13),
      dict(name='cldmu', pos=20, type='f', min_ver=13),
      dict(name='fkair', pos=21, type='f', min_ver=13),
      dict(name='fkcld', pos=22, type='f', min_ver=13),
      dict(name='atomicnoair', pos=23, type='f', min_ver=13),
      dict(name='atomicnocld', pos=24, type='f', min_ver=13),

      # Global T-history
      dict(name='time', pos=31+0, type='f'),
      dict(name='timerescaled', pos=31+1, type='f'),
      dict(name='bubbleheight', pos=31+2, type='f'),
      dict(name='spikeheight', pos=31+3, type='f'),
      dict(name='cycl', pos=31+4, type='f', min_ver=12),
      dict(name='dt', pos=31+5, type='f'),
      dict(name='courmx', pos=31+6, type='f'),
      dict(name='urbubmx', pos=31+7, type='f'),
      dict(name='urspkmn', pos=31+8, type='f'),
      dict(name='ekmx', pos=31+9, type='f'),
      dict(name='ekrmx', pos=31+10, type='f'),
      dict(name='ektmx', pos=31+11, type='f'),
      dict(name='ekurmn', pos=31+12, type='f'),
      dict(name='ekurmx', pos=31+13, type='f'),
      dict(name='eiurmn', pos=31+14, type='f'),
      dict(name='eiurmx', pos=31+15, type='f'),
      dict(name='Hurmn', pos=31+16, type='f'),
      dict(name='Hurmx', pos=31+17, type='f'),
      dict(name='ekurspkmn', pos=31+18, type='f'),
      dict(name='ekurbubmx', pos=31+19, type='f'),
      dict(name='eiurspkmn', pos=31+20, type='f'),
      dict(name='eiurbubmx', pos=31+21, type='f'),
      dict(name='Hurspkmn', pos=31+22, type='f'),
      dict(name='Hurbubmx', pos=31+23, type='f'),
      dict(name='ekbubmx', pos=31+24, type='f'),
      dict(name='ekrbubmx', pos=31+25, type='f'),
      dict(name='ektbubmx', pos=31+26, type='f'),
      dict(name='ekspkmx', pos=31+27, type='f'),
      dict(name='ekrspkmx', pos=31+28, type='f'),
      dict(name='ektspkmx', pos=31+29, type='f'),

      # Args images
      dict(name='ai_vort', pos=64+0, type='f', len=2),
      dict(name='ai_divu', pos=64+2, type='f', len=2),
      dict(name='ai_s', pos=64+4, type='f', len=2),
      dict(name='ai_fv', pos=64+6, type='f', len=2),
      dict(name='ai_rho', pos=64+8, type='f', len=2),
      dict(name='ai_p', pos=64+10, type='f', len=2),
      dict(name='ai_ux', pos=64+12, type='f', len=2),
      dict(name='ai_uy', pos=64+14, type='f', len=2),
      dict(name='ai_uz', pos=64+16, type='f', len=2),
  ]

  def __init__(self, filename, lazy=True, **kwargs):
    """Create a ray profile reader object.

    `lazy` means only the header is loaded on open
    """

    logging = kwargs.get('logging')

    self._filename = filename
    self.lazy = lazy

    self.version = None
    self.bucket_count = 0

    self._centers = None
    self._corners = None
    self._normals = None

    self._cache = {}

    self._variable_map = {}

    self._names = []
    self._data = []

    self.header_attrs = {}

    if logging:
      import logging
      logging.basicConfig(level=logging.DEBUG)
      self.log = logging.getLogger('rp_info')
      self.log.setLevel(logging.DEBUG)
    else:
      self.log = None

    if str(filename).isdigit():
      filename = 'RProfile-01-%04i.bobaaa' % int(filename)

    if self.log: self.log.info("Opening %s" % filename)

    f = open(filename, 'rb')
    header = f.read(128)
    '''
    if header[:8] != 'LCSE:RPS':
      if self.log: self.log.warn('File %s is not a new Ray Profile, try an older rp_info.py' % filename)
      f.close()
      raise Exception('Unsupported file version')
    '''
    self.version = struct.unpack("i", header[8:12])[0]

    f.seek(0)

    if self.version < 8:
      raise Exception('Unsupported version %i' % self.version)
    elif self.version == 8:
      self._header_size = 128
      hlen = 8

      self.header_var_list = self.header_var_list[:8]

#      header = struct.unpack(hlen * "i", header[8:8+4*hlen])

#      self.header_attrs['version'] = header[0]
#      self.header_attrs['cell_ct_low'] = header[1]
#      self.header_attrs['nbuckets'] = header[2]
#      self.header_attrs['dump'] = header[3]
#      self.header_attrs['sizeof_float'] = header[4]
#      self.header_attrs['has_centers'] = header[5]
#      self.header_attrs['has_corners'] = header[6]
#      self.header_attrs['has_normals'] = header[7]

    elif self.version > 8:
      self._header_size = 1024
      hlen = 127

    header = f.read(self._header_size)
      # Bug fixes
      # Using the header info from v9
     # if self.version < 11:
     #   self._init_v9()
     #   self._init_legacy()

     #   raw_header = struct.unpack(hlen * "i", header[8:8+4*hlen])
     #   raw_fheader = struct.unpack(hlen * "f", header[8+4*hlen:8+8*hlen])
     #   self.header_attrs.update([(k, raw_header[i]) for i, k in enumerate(self._header_names)])
     #   self.header_attrs.update([(k, raw_fheader[i]) for i, k in enumerate(self._fheader_names)])
     #   self.header_attrs.update([(k, raw_fheader[32 + i]) for i, k in enumerate(self._fheader_names2)])
     #   self.header_attrs.update([(k, (raw_fheader[64 + 2*i], raw_fheader[64 + 2*i + 1] )) for i, k in enumerate(self._argsimg_names)])
     #elif self.version <= 12:

    hmap = dict(i=struct.unpack(hlen * "i", header[8 : 8 + 4 * hlen]),
                f=struct.unpack(hlen * "f", header[8 + 4 * hlen : 8 * (1 + hlen)]))

    for var in self.header_var_list:

      name = var['name']
      pos = var['pos']
      var_type = var['type']
      var_len = var.get('len', 1)
      min_ver = var.get('min_ver', 0)

      if self.version < min_ver:
        continue

      # A slight offset problem
      if self.version == 11 and var_type == 'f' and pos > 30: # and pos < 64:
        pos = pos + 1

      attr = hmap[var_type][pos] if var_len == 1 else hmap[var_type][pos : pos + var_len]
      self.header_attrs[name] = attr

    # Fix header problems
    if self.version == 8:
      self.header_attrs['cell_ct_high'] = 2 * self.header_attrs['cell_ct_low']
      self.header_attrs['var_ct_high'] = 1
      self.header_attrs['var_ct_low'] = 14

    if self.version == 9:
      self.header_attrs['cell_ct_low'] -= 2
      self.header_attrs['cell_ct_high'] -= 4

    if self.version < 12:
      self.header_attrs['isreflectbdry'] = 1
      self.header_attrs['ismuffledbdry'] = 0

    if self.version > 13:
      self.header_attrs['has_corners'] = False

    self.bucket_count = self.header_attrs['nbuckets']
    self.dump = self.header_attrs['dump']
    self.buckets = self.header_attrs['nbuckets']

    if self.version > 10:
      self._init_v11()

      if not self.lazy:
        f = open(self._filename, 'r')
        self._data = f.read()
        f.close()

    else:
      self._init_legacy()

      if self.version == 8:
        self._init_v8()
      else:
        self._init_v9()

      for k in ['has_centers', 'has_corners', 'has_normals']:
        self.header_attrs[k] = self.header_attrs.get(k, 0) == 1

      float_type = 'f8' if self.header_attrs.get('sizeof_float') == 8 else 'f4'

      self._dtypes_hi = [('j_hi', 'i4')]
      self._dtypes_hi.extend([(n, float_type) for n in self._names_hi])
      self._col_names_hi = ['j_hi'] + self._names_hi

      self._dtypes = [('j', 'i4')]
      self._dtypes.extend([(n, float_type) for n in self._names])

      self._col_names = ['j'] + self._names

      if self.lazy:
        log.warn("Lazy Loading not supported for v %i" % self.version)

      self._load(f)

    f.close()

  def _load(self, f):

    nbuckets = self.header_attrs.get('nbuckets')
    cell_ct_low = self.header_attrs.get('cell_ct_low')
    cell_ct_high = self.header_attrs.get('cell_ct_high')

    # Read the high resolution table
    self._data_hi = np.fromfile(f, dtype=self._dtypes_hi, count=cell_ct_high*(nbuckets+1))

    # Read the low resolution table
    self._data_low = np.fromfile(f, dtype=self._dtypes, count=cell_ct_low*(nbuckets+1))

    if self.header_attrs.get('has_centers'):
      vals = 3 if self.version > 12 else 5
      self._centers = np.fromfile(f, dtype=np.float64, count=5 * nbuckets).reshape((vals, -1), order='F')

    if self.header_attrs.get('has_normals'):
      self._normals = np.fromfile(f, dtype=np.float64, count=9*nbuckets).reshape((3, 3, -1), order='F')

    if self.header_attrs.get('has_corners'):
      self._corners = np.fromfile(f, dtype=np.float64, count=9*nbuckets).reshape((3, 3, -1), order='F')

  def get_centers(self):
    """ Get centers of the buckets as an array of x, y, z """

    if self._centers is None and self.version >= 11:
      self._centers = self._get_array('centers')

      if self.version < 13:
        centers = self._centers[2:]
        self._centers = old_div(-centers, np.sqrt((centers * centers).sum(0)))

    return self._centers

  def get_corners(self):
    """ Get corners of the buckets as an array of (xyz,side #, bucket #) """

    if self._corners is None and self.version >= 11:
      normals = self.get_normals()

      self._corners = np.zeros((3, 3, normals.shape[2]))
      self._corners[:,0,:] = np.cross(normals[:,0,:], normals[:,1,:], axis=0)
      self._corners[:,1,:] = np.cross(normals[:,1,:], normals[:,2,:], axis=0)
      self._corners[:,2,:] = np.cross(normals[:,2,:], normals[:,0,:], axis=0)
      self._corners[:,:,:] /= np.sqrt((self._corners * self._corners).sum(0))

    return self._corners

  def get_normals(self):
    """ Get normals of the buckets as an array of (x/y/z coordinate, side #, bucket #).
    """

    if self._normals is None and self.version >= 11:
      self._normals = self._get_array('normals')
      normals_len = np.sqrt((self._normals * self._normals).sum(0))
      normals_ix = normals_len > 0.0
      self._normals[:, normals_ix] = old_div(self._normals[:, normals_ix], normals_len[normals_ix])

    return self._normals

  def get_cell_volumes(self):
    """
    Get an array of dimension (`bucket_ct`, `cell_ct_low`) containing the volume of each
    totopo shaped cell.
    """

    volumes = []

    ys = self.get('y')

    normals = self.get_normals()

    a, b, c = normals[:,0,:], normals[:,1,:], normals[:,2,:]

    ang_a = np.pi - np.arccos((a * b).sum(0))
    ang_b = np.pi - np.arccos((b * c).sum(0))
    ang_c = np.pi - np.arccos((c * a).sum(0))

    angles = np.vstack([ang_a, ang_b, ang_c])

    dr = (ys[1] - ys[0])

    r_sq = dr * (ys**2)

    bucket_angles = angles.sum(0) - np.pi

    for i in range(0, self.bucket_count):
        vols = r_sq * bucket_angles[i]
        volumes.append(vols.reshape((-1,1)))

    return np.hstack(volumes)


  def get_table(self, var):
    """Get a table of dimension (4, ncells, nbuckets+1) containg all  buckets
    (including the global average bucket zero). The first dimension contains the
    statastical information: average, min, max, sd.
    """

    if var not in self._variable_names:
      print('Variable %s not found in table. Available variables are %s' % (var, self._variable_names))
      return

    if self.version >= 11:
      return self._get_array(var)
    else:
      return self._get_legacy(var)

  def get(self, var):
    """Get the global bucket for variable `var` or  get header attribute `var`.

    Use `get_table(self, var)` to get the same variable but for all buckets.

    If the global bucket is returned an array of dimension (4, ncells) is returned.
    The first dimension contains avg, min, max, sd.
    """

    if var in self.header_attrs:
      return self.header_attrs.get(var)

    if self.version >= 11:
      return self._get_array(var, global_only=True)
    else:
      return self._get_legacy(var, global_only=True)

  def _get_array(self, var, global_only=False):

    if var not in self._variable_map:
      return None

    if var in self._cache:
      return self._cache[var]

    offset, dtype, count, shape = self._variable_map[var]

#    print self._variable_map[var], global_only

    if global_only and len(shape) == 3 and shape[2] == self.bucket_count + 1:
      count = shape[0] * shape[1]
      shape = shape[:2]

    if self.lazy:
      f = open(self._filename, 'r')
      f.seek(offset)
      data = np.fromfile(f, dtype=dtype, count=count).reshape(shape, order='F')
      f.close()
    else:
      data = np.frombuffer(self._data[offset:], dtype=dtype,
                           count=count).reshape(shape, order='F')

    if not global_only:
      self._cache[var] = data

    return data

  def _get_legacy(self, var, global_only=False):

    if var in self._col_names_hi:
      data_array = self._data_hi
      radial = self.header_attrs.get('cell_ct_high')
    elif var in self._col_names:
      radial = self.header_attrs.get('cell_ct_low')
      data_array = self._data_low
    else:
      #print var, self._col_names_hi, self._col_names, self.header_attrs
      raise Exception("Attribute '%s' not found, look in .get_variables() and .get_attributes()" % var)

    if var in self._legacy_remap:
      remap_vars = self._legacy_remap[var]

      data_out = np.zeros((len(remap_vars), radial, self.bucket_count + 1), order='F')

      var_array = []

      for i, v in enumerate(remap_vars):
        data_out[i,:,:] = data_array[:][v].reshape((radial, -1), order='F')

      data = data_out

    else:
      data = data_array[:][var].reshape((radial, -1), order='F')

    if var in ['y', 'j', 'j_hi', 'y_hi']:
       data = data[:,0]
    elif global_only:
      data = data[:,:,0]

    # The old format was backwards
    if self.version < 9:
      if len(data.shape) == 3:
        data = data[:,::-1,:]
      elif len(data.shape) == 2:
        data = data[:,::-1]
      else:
        data = data[::-1]

    return data

  def get_attributes(self):

    attrs = list(self.header_attrs.keys())
    attrs.sort()
    return attrs

  def get_variables(self):
    return self._variable_names

  def _init_v8(self):
    self._header_names = ['version', 'nradial_low', 'nbuckets', 'dump',
                          'sizeof_float', 'has_centers', 'has_corners', 'has_normals']

    self._names_hi = ['y_hi', 'fv_hi', 'fvmn_hi', 'fvmx_hi', 'fvsd_hi']

    self._names = ['counts', 'y',
                   'fv', 'fvmn','fvmx','fvsd',
                   'rho', 'rhomn', 'rhomx', 'rhosd',
                   'rhourbubble', 'rhourbubblemn', 'rhourbubblemx', 'rhourbubblesd',
                   'rhourspike', 'rhourspikemn', 'rhourspikemx', 'rhourspikesd',
                   'p', 'pmn', 'pmx', 'psd',
                   'ux','uxmn','uxmx','uxsd',
                   'uy', 'uymn', 'uymx', 'uysd',
                   'uz', 'uzmn', 'uzmx', 'uzsd',
                   'ekr', 'ekrmn','ekrmx','ekrsd',
                   'ekt','ektmn','ektmx','ektsd',
                   'ek', 'ekmn','ekmx','eksd',
                   'ekur','ekurmn','ekurmx','ekursd',
                   'eiur', 'eiurmn', 'eiurmx', 'eiursd',
                   'hur', 'hurmn', 'hurmx', 'hursd']

    self._header_arrays = ['normals', 'centers', 'corners']

  def _init_v9(self):
    self._init_v8()


    self._names.extend(['ceul', 'ceulmn', 'ceulmx', 'ceulsd',
                        'mach', 'machmn', 'machmx', 'machsd',
                        'enuc', 'enucmn', 'enucmx', 'enucsd',
                        'fnuc', 'fnucmn', 'fnucmx', 'fnucsd',
                        'dy', 'dymn', 'dymx', 'dysd'])

  def _init_legacy(self):
    ''' Initialize internals for old versions (< 11)'''

    buckets_total = 1 + self.bucket_count

    var_ct_high = self.header_attrs.get('var_ct_high')
    var_ct_low = self.header_attrs.get('var_ct_low')
    cell_ct_high = self.header_attrs.get('cell_ct_high')
    cell_ct_low = self.header_attrs.get('cell_ct_low')

    sizeof_float = self.header_attrs.get('sizeof_float')
    float_type = np.float64 if sizeof_float == 8 else np.float32

    self._variable_list = [('centers', float_type, sizeof_float, (5, self.bucket_count)),
                           ('normals', float_type, sizeof_float, (3, 3, self.bucket_count)),
                           ('corners', float_type, sizeof_float, (3, 3, self.bucket_count))]

    offset = self._header_size

    # Integer Arrays (j, h_hi)
    offset += 4 * buckets_total * (cell_ct_high + cell_ct_low)
    # Float array, yav + counts + 4 * nvars. No high counts
    offset += 8 * buckets_total * ((1 + 4 * var_ct_high) * cell_ct_high + (2 + 4 * var_ct_low) * cell_ct_low)

    for name, dtype, sizeof, shape in self._variable_list:
      count = np.prod(shape)
      size = sizeof * count
      self._variable_map[name] = (offset, dtype, count, shape)
      offset += size

    self._variable_names = list(self._variable_map.keys())
    self._variable_names.sort()

    # Variable meta
    self._legacy_remap = dict(fv_hi=('fv_hi', 'fvmn_hi','fvmx_hi','fvsd_hi'),
                              fv=('fv', 'fvmn','fvmx','fvsd'),
                              rho=('rho', 'rhomn', 'rhomx', 'rhosd'),
                              rhourbubble=('rhourbubble', 'rhourbubblemn', 'rhourbubblemx', 'rhourbubblesd'),
                              rhourspike=('rhourspike', 'rhourspikemn', 'rhourspikemx', 'rhourspikesd'),
                              p=('p', 'pmn', 'pmx', 'psd'),
                              ux=('ux','uxmn','uxmx','uxsd'),
                              uy=('uy', 'uymn', 'uymx', 'uysd'),
                              uz=('uz', 'uzmn', 'uzmx', 'uzsd'),
                              ekr=('ekr', 'ekrmn','ekrmx','ekrsd'),
                              ekt=('ekt','ektmn','ektmx','ektsd'),
                              ek=('ek', 'ekmn','ekmx','eksd'),
                              ekur=('ekur','ekurmn','ekurmx','ekursd'),
                              eiur=('eiur', 'eiurmn', 'eiurmx', 'eiursd'),
                              hur=('hur', 'hurmn', 'hurmx', 'hursd'),
                              ceul=('ceul', 'ceulmn', 'ceulmx', 'ceulsd'),
                              mach=('mach', 'machmn', 'machmx', 'machsd'),
                              enuc=('enuc', 'enucmn', 'enucmx', 'enucsd'),
                              fnuc=('fnuc', 'fnucmn', 'fnucmx', 'fnucsd'),
                              dy=('dy', 'dymn', 'dymx', 'dysd'))

    self._legacy_order = ['counts', 'y', 'fv', 'rho', 'rhourbubble', 'rhourspike',
                          'p', 'ux', 'uy', 'uz', 'ekr', 'ekt', 'ek', 'ekur', 'eiur', 'hur',
                          'ceul', 'mach', 'enuc', 'fnuc', 'dy',]

    self._variable_names = list(self._variable_map.keys()) + self._legacy_order + ['fv_hi', 'y_hi']
    self._variable_names.sort()

  def _init_v11(self):

    cell_ct_high = self.header_attrs.get('cell_ct_high')
    cell_ct_low = self.header_attrs.get('cell_ct_low')

    buckets_total = 1 + self.bucket_count

    sizeof_float = self.header_attrs.get('sizeof_float')
    float_type = np.float64 if sizeof_float == 8 else np.float32
    int_type = np.int32

    vals = 3 if self.version > 12 else 5

    # name, size_in_bytes, <array dimensions>
    self._variable_list = [('centers', float_type, sizeof_float, (vals, self.bucket_count)),
                           ('normals', float_type, sizeof_float, (3, 3, self.bucket_count)),
                           ('corners', float_type, sizeof_float, (3, 3, self.bucket_count)),
                           ('j_hi', int_type, 4, (cell_ct_high,)),
                           ('y_hi', float_type, sizeof_float, (cell_ct_high,)),
                           ('fv_hi', float_type, sizeof_float, (4, cell_ct_high, buckets_total)),
                           ('j', int_type, 4, (cell_ct_low,)),
                           ('y', float_type, sizeof_float, (cell_ct_low,)),
                           ('counts', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('fv', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('rho', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('rhobubble', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('rhospike', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('rhourbubble', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('rhourspike', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('p', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ux', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('uy', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('uz', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ceul', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('mach', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('enuc', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('fnuc', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('dy', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ekr', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ekt', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ek', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('ekur', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('eiur', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                           ('hur', float_type, sizeof_float, (4, cell_ct_low, buckets_total)),
                          ]

#    present_vars = [v['name'] for v in self.header_var_list if self.version >= v.get('min_ver', self.version)]

    # TODO: HACK UGH
    skip_vars = []
    if self.version > 12:
      skip_vars.append('corners')

    offset = self._header_size

    for name, dtype, sizeof, shape in self._variable_list:

      if name in skip_vars:
        continue

#      print (name, offset, dtype, sizeof, shape)

      count = np.prod(shape)
      size = sizeof * count
      self._variable_map[name] = (offset, dtype, count, shape)
      offset += size
#      print (name, offset, dtype, count, shape, sizeof)

    self._variable_names = list(self._variable_map.keys())
    self._variable_names.sort()

def main():
  '''Simple demo main function'''

  if len(sys.argv) < 2:
    print("Specify filename")
    return

  path = sys.argv[1]

  if os.path.isdir(path):
    rp_set = rprofile_set(sys.argv[1])
    print(rp_set.ray_profiles)
  else:
    rp = rprofile(sys.argv[1], logging=True)
    log.info('File version %i (real%i), with %i buckets and %i radial bins for dump %i' %
             (rp.version, rp.get('sizeof_float'), rp.get('nbuckets'), rp.get('cell_ct_low'), rp.get('dump')))

    header_keys = list(rp.header_attrs.keys())
    header_keys.sort()

    for k in header_keys:
      print("%s: %s" % (k, rp.header_attrs[k]))

    d = rp.get('fv')
    print(d)

#  print rp.get('j_hi', bucket=0)
  print(rp.get_table('y'))
  return

  print("ceul")
  print(rp.get('ceul'))
  print("mach")
  print(rp.get('mach'))
  print("enuc")
  print(rp.get('enuc'))
  print("fnuc")
  print(rp.get('fnucmx'))
  print("dy")
  print(rp.get('dy'))


if __name__ == "__main__":
    main()
