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

"""
This module provides a simple interface into AProfile-01-????.bobaaa files
generated by the ANP tool
"""

import struct
import logging
import os
import re
import sys

logging.basicConfig(level=logging.DEBUG)

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

class aprofile_set(object):

  def __init__(self, source, lazy=False, verbose=False, **kwargs):
    """Load a set of AProfiles. 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 AProfile files to open

    `lazy` if True data variables and grid state are not loaded from disk until they are needed

    `verbose` enable debug output

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

    .. code-block:: python
      :linenos:

      rp_set = lcse.aprofile_set(targetpath, lazy=True, verbose=True)
      for rp in rp_set:
        rp.get("fv_hi")
        rp.get_table("fv")
    """

    self.path = source if isinstance(source, str) else None
    self.files = source if isinstance(source, list) else []
    self.lazy = lazy

    self.stride = kwargs.get('stride', 1)
    self.first_dump = kwargs.get('first_dump')
    self.last_dump = kwargs.get('last_dump')

    self.verbose = verbose
    self.log = logging.getLogger('aprofile_set')

    if verbose:
      self.log.setLevel(logging.DEBUG)
    else:
      self.log.setLevel(logging.ERROR)

    self.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('(.*)AProfile-01-([\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 = 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 self.next()

  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.profiles.get(dump, aprofile(self.dump_map[dump], lazy=self.lazy, verbose=self.verbose))

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

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

    else:
      raise StopIteration()

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

    if self.dumps and dump is None:
      dump = self.dumps[-1]

    if dump in self.profiles:
      return self.profiles[dump]
    elif dump not in self.dump_map:
      log.error("dump %i not found" % dump)
      return None

    rp = aprofile(self.dump_map[dump], lazy=self.lazy, verbose=self.verbose)

    if not self.lazy:
      self.profiles[dump] = rp

    return rp

  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('AProfile') and f.endswith('.bobaaa')]
    filenames.sort()

    return filenames


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

  header_var_list = [
      dict(name='version', pos=0, type='i'),
      dict(name='species', pos=1, type='i'),
      dict(name='cells', pos=2, type='i'),
      dict(name='buckets', pos=3, type='i'),
      dict(name='dump', pos=4, type='i'),
      dict(name='time', pos=5, type='f'),
      dict(name='timestamp', pos=6, type='i'),
      dict(name='pid', pos=7, type='i'),
  ]

  def __init__(self, filename, verbose=False, lazy=False, names=None, **kwargs):
    """ Load an AProfile

    `verbose` if True turn on debug logging (sets `.log.setLevel(logging.DEBUG)`)

    `lazy`

    `names` list of length `.species` corresponding to the contained
    """

    self._filename = filename
    self.lazy = lazy

    # Map of variable names to numbers (user must load this)
    self.names = names if names else []
    self.header = {}

    self.log = logging.getLogger('aprofile')

    if verbose:
      self.log.setLevel(logging.DEBUG)
    else:
      self.log.setLevel(logging.ERROR)

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

    self._header_size = 128

    f = open(filename, 'rb')
    header = f.read(self._header_size)

    if header[:8] != 'LCSE:ANP':
      raise Exception('Unsupported file')

    var_types = ''.join(v['type'] for v in self.header_var_list)
    hmap = struct.unpack(var_types, header[8 : 8 + 4 * len(self.header_var_list)])

    for var in self.header_var_list:
      name = var['name']
      pos = var['pos']

      self.header[name] = hmap[pos]

    # Add instance variables
    self.buckets = self.header.get('buckets')
    self.cells = self.header.get('cells')
    self.dump = self.header.get('dump')
    self.species = self.header.get('species')
    self.version = self.header.get('version')

    # Load grid information
    self.normals = np.fromfile(f, dtype=np.float32, count=9 * self.buckets).reshape((3, 3, -1), order='F')
    self.centers = np.fromfile(f, dtype=np.float32, count=3 * self.buckets).reshape((3, -1), order='F')
    self.corners = np.fromfile(f, dtype=np.float32, count=9 * self.buckets).reshape((3, 3, -1), order='F')
    self.neighbors = np.fromfile(f, dtype=np.int32, count=6 * self.buckets).reshape((2, 3, -1), order='F')
    self.r = np.fromfile(f, dtype=np.float32, count=self.cells)

    if not self.lazy:
      self._data = f.read()

    f.close()

  def get_attributes(self):
    attrs = self.header_attrs.keys()
    attrs.sort()
    return attrs

  def get_state(self):
    '''Return a table with dimension (6, cells, buckets). The first dimension contains:
       rho, p, U_r, U_t (3 components, one for each bucket side).
    '''

    shape = (6, self.cells, self.buckets)
    count = 6 * self.cells * self.buckets

    if self.lazy:
      with open(self._filename, 'r') as f:
        f.seek(self._header_size + 108 * self.buckets)
        data = np.fromfile(f, dtype=np.float32, count=count).reshape(shape, order='F')
    else:
      data = np.frombuffer(self._data, dtype=np.float32,
                           count=count).reshape(shape, order='F')
    return data

  def get_table(self, index):
    '''Read and return the abundances for variable given by `index`. Where `index` is either a
       an integer representing the 0-based index of the given variable or if `.names` contains
       a list of names then `index` can be staring from `names`.
    '''

    if isinstance(index, str):
      index = self.names.get(index, -1)

    if index > self.species:
      self.log.error("Index %i out of range" % index)
      return None

    count = self.cells * self.buckets
    data_offset = 4 * (6*count + index * count)
    shape = (self.cells, self.buckets)

    if self.lazy:
      with open(self._filename, 'r') as f:
        f.seek(self._header_size + 108 * self.buckets + data_offset)
        data = np.fromfile(f, dtype=np.float32, count=count).reshape(shape, order='F')
    else:
      data = np.frombuffer(self._data[data_offset:], dtype=np.float32,
                           count=count).reshape(shape, order='F')

    return data

  def get_centers(self): return self.centers
  def get_normals(self): return self.normals

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

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

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

  path = sys.argv[1]

  if os.path.isdir(path):
    ap_set = aprofile_set(sys.argv[1])

    for ap in ap_set:
      print ap.dump

    print(ap_set.dumps)
  else:
    ap = aprofile(sys.argv[1], logging=True)
    log.info('File version %i, with %i buckets, %i radial cells, %i species, dump %i' %
             (ap.version, ap.buckets, ap.cells, ap.species, ap.dump))

    header_keys = ap.header.keys()
    header_keys.sort()

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

    print "Normals", rp.normals
    print "Centers", rp.centers
    print "Edges", rp.corners
    print "Neighbors", rp.neighbors
    print "State", rp.get_state()
    print "r", rp.r


if __name__ == "__main__":
    main()
