# load_data module
# contains functions to load data from various formats of file

version = '29th April 2021'

# Nexus is a registered trademark of the Halliburton Company

import logging
log = logging.getLogger(__name__)
log.debug('load_data.py version %s', version)

# defs:
#    def file_exists(file_name):
#    def load_corp_array_from_file(file_name, extent_kji = None, comment_char = '!',
#                                  data_free_of_comments = False, use_binary = False,
#                                  eight_mode = False):
#    def load_array_from_file(file_name, extent, data_type = 'real', comment_char = '!',
#                             data_free_of_comments = False, use_binary = False):
#    def load_array_from_ascii_file(file_name, extent, data_type='real', comment_char = '!', data_free_of_comments = False):
#    def load_fault_mask(file_name, direction, fault_mask, use_binary = False):

import os
import numpy as np

import resqpy.olio.utility as pu
import resqpy.olio.box_utilities as box
import resqpy.olio.keyword_files as kf
import resqpy.olio.write_data as wd
import resqpy.olio.grid_functions as gf
import resqpy.olio.ab_toolbox as abt


def file_exists(file_name, must_be_more_recent_than_file = None):
   """Returns True if the file exists (and is more recent than other file, if given)."""

   if file_name is None or len(file_name) == 0: return False
   existe = os.path.exists(file_name)
   if not existe: return False
   if not must_be_more_recent_than_file or file_name == must_be_more_recent_than_file or  \
      not os.path.exists(must_be_more_recent_than_file): return True
   return os.path.getmtime(file_name) > os.path.getmtime(must_be_more_recent_than_file)


######################################################################################################
# load_corp_array_from_file():
# function to load a nexus corp array of data from a named file
# if the grid extent is not known, the file must be free of comments
# NB: extent_kji here is extent of grid, rather than that of the corp array
# returns a pagoda style 7D array, resequenced

def load_corp_array_from_file(file_name, extent_kji = None, corp_bin = False, swap_bytes = True,
                              max_lines_for_keyword = 100, comment_char = None,
                              data_free_of_comments = False, use_binary = False,
                              eight_mode = False, use_numbers_only = None):
   """Loads a nexus corner point (CORP) array from a file, returns a 7D numpy array in pagoda ordering.

      arguments:
         file_name: The name of an ascii file holding the CORP data (no other keywords with numeric data
                  should be in the file); write access to the directory is likely to be needed if
                  use_binary is True
         extent_kji: The extent of the grid as a list or a 3 element numpy array, in the order [NK, NJ, NI].
                  If extent_kji is None, the extent is figured out from the data.  It must be
                  given for 1D or 2D models
         corp_bin (boolean, default False): if True, input file is in bespoke corp binary format, otherwise ascii
         swap_bytes (boolean, default True): if True, byte ordering of corp bin data is reversed; only relevant
                  if corp_bin is True
         max_lines_for_keyword: the maximum number of lines to search for CORP keyword; set to zero if file is
                  known to be data only
         comment_char: A single character string which is interpreted as introducing a comment
         data_free_of_comments: If True, once the numeric data is encountered, it is assumed that there are no
                  further comments (allowing a faster load)
         use_binary: If True, a more recent file containing a pure binary copy of the data is looked for first,
                  in the same directory; if found, the data is loaded directly from that file; if not found, the
                  binary file is created after the ascii has been loaded (ready for next time)
         eight_mode: If True, the data is assumed to be in CORP EIGHT ordering; otherwise the normal ordering
                  (The code does not look for keywords.); this is not automatically determined from any keyword
                  in the file
         use_numbers_only: no longer in use, ignored

      returns:
         A numpy array containing the CORP data in 7D pagoda protocol ordering. The extent of the grid, and hence
         shape of the array is determined from the corner point data unless extent_kji has been specified.
   """

#   assert(extent_kji is not None or data_free_of_comments)   # no longer a requirement

   if extent_kji is None:
      extent = None
   else:
      extent = [extent_kji[0] * extent_kji[1] * extent_kji[2], 8, 3]

   if corp_bin:
      bin_file_size = os.path.getsize(file_name)
      bin_cell_count, remainder = divmod(bin_file_size, 26 * 4)   # corp bin format has records of 24 + 2 (head, tail) 32 bit words
      if remainder:
         log.error('corp binary file ' + str(file_name) + ' is not a whole number of records')
         return None
      dt = np.dtype('float32')
      if swap_bytes: dt = dt.newbyteorder()
      cp_array = np.fromfile(file_name, dtype = dt, count = 26 * bin_cell_count).reshape((-1, 26))[:, 1:25]
      if extent_kji is not None:
         if bin_cell_count != (extent_kji[0] * extent_kji[1] * extent_kji[2]):
            log.error('corp binary file ' + str(file_name) + ' contains data for ' + str(bin_cell_count) +
                      ' cells; extent requires ' + str(extent_kji[0] * extent_kji[1] * extent_kji[2]))

   else:
      cp_array = load_array_from_file(file_name, extent = extent, data_type = 'real',
                                      keyword = 'CORP', max_lines_for_keyword = max_lines_for_keyword,
                                      comment_char = comment_char,
                                      data_free_of_comments = data_free_of_comments,
                                      use_binary = use_binary)

   cell_count, remainder = divmod(cp_array.size, 24)
   if remainder:
      log.error('file ' + file_name + ' contains ' + str(cp_array.size) + ' data, which is not a multiple of 24')
      return None

   cp_array = cp_array.reshape(1, 1, cell_count, 2, 2, 2, 3)  # pagoda 7D, temporarily with all cells on I axis
   log.debug('resequencing corner point data')
   gf.resequence_nexus_corp(cp_array, eight_mode = eight_mode)  # switches IP points where JP = 1 (unless eight_mode)

   if extent_kji is None:
      log.info('determining grid extent from corner points')
      extent_kji = gf.determine_corp_extent(cp_array)
      if extent_kji is None:
         log.error('failed to determine extent of grid from corner points')
         return cp_array

   return cp_array.reshape(extent_kji[0], extent_kji[1], extent_kji[2], 2, 2, 2, 3)


######################################################################################################
# load_array_from_file():
# function to load an array of data from a named file
# if the extent is not known (None), the file must be free of comments
# a new numpy array is returned
# binary data file creation & reuse is supported

def load_array_from_file(file_name, extent = None, data_type = 'real', keyword = None,
                         max_lines_for_keyword = None, comment_char = None,
                         data_free_of_comments = False, use_binary = False, use_numbers_only = None):
   """Load an array from an ascii (or pure binary) file.

      Arguments are similar to those for load_corp_array_from_file().
   """

   if not use_binary:
      return load_array_from_ascii_file(file_name, extent = extent, data_type = data_type,
                                        keyword = keyword, max_lines_for_keyword = max_lines_for_keyword,
                                        comment_char = comment_char,
                                        data_free_of_comments = data_free_of_comments)

   (extension, ab_type) = abt.binary_file_extension_and_np_type_for_data_type(data_type)

   if extent is None:
       cell_count = -1   # np.fromfile interprets this as 'read everything'
       log.debug('Loading unknown number of array data elements from file ' + file_name)
   else:
       cell_count = pu.cell_count_from_extent(extent)
       log.debug('Loading %1d array data elements from file %s', cell_count, file_name)

   ascii_file_name = file_name
   binary_file_name = file_name
   if len(binary_file_name) < 4 or binary_file_name[-3:] != extension: binary_file_name += extension
   else: ascii_file_name = ascii_file_name[:-3]   # strip off '.db' or similar

   try:   # tentatively try to read data from an existing binary file, if present
      if file_exists(binary_file_name, must_be_more_recent_than_file = ascii_file_name):
         with open(binary_file_name, 'rb') as binary_file_in:
            result = np.fromfile(binary_file_in, dtype = ab_type, count = cell_count)
            if extent is not None: result = result.reshape(extent)
            # check that end of file has been reached, ie. not too much data in file
            try: # expected to return null
               c = binary_file_in.read(1)
               if len(c): log.warning('binary file contains more data than expected: ' + binary_file_name)
            except:
               pass
            log.info('Data loaded from binary file %s', binary_file_name)
            return result
   except:
      pass

   # read from ascii file
   result = load_array_from_ascii_file(file_name, extent = extent, data_type = data_type,
                                       keyword = keyword, max_lines_for_keyword = max_lines_for_keyword,
                                       comment_char = comment_char,
                                       data_free_of_comments = data_free_of_comments)

   # create a binary file (to be used next time)
   try:
      wd.write_pure_binary_data(binary_file_name, result)
   except:
      log.warn('Failed to write data to binary file %s', binary_file_name)
      # todo: could delete the binary file in case a corrupt file is left for use next time

   return result

# end of load_array_from_file() def
######################################################################################################


######################################################################################################
# load_array_from_ascii_file():
# function to load an array of data from a named ascii file
# if the extent is not known (None), the file must be free of comments
# a new numpy array is returned

def load_array_from_ascii_file(file_name, extent = None, data_type = 'real', keyword = None,
                               max_lines_for_keyword = None, comment_char = None,
                               data_free_of_comments = False, skip_c_space = True, use_numbers_only = None):
   """Returns a numpy array with data loaded from an ascii file.

   arguments:
      file_name: string holding name of the existing ascii data file, eg. 'Cell_depth.dat'
      extent: a python list of integers specifying the extent (shape) of the array, eg. [148, 270, 103];
         if None, all data is read and returned as a 'flat' 1D array (data must be free from comments)
      data_type: the type of individual data elements, one of 'real','float','int', 'integer', 'bool' or 'boolean'
      keyword: if present, an attempt is made to find the keyword before reading data, if keyword is None or is
         not found, data is read from the start of the file
      max_lines_for_keyword: can be used to limit the search for keyword (for speed efficiency)
      comment_char: a single character string being the character used to introduce a comment in the file
      data_free_of_comments: if set to True, a faster load is used once any header line comments have been
         skipped
      skip_c_space: if True then a line starting 'C' followed by white space is skipped as a comment
      use_numbers_only: this argument is no longer in use and is ignored

   returns:
      a numpy array of shape specified in extent argument with dtype matching data_type

   example call:
      depth_array = load_data.load_array_from_ascii_file('Cell_depth.dat', [148, 270, 103])

   notes:
      In all use cases, this function is designed to load a single array of data from an ascii file
      that DOES NOT CONTAIN OTHER ARRAYS as well, ie. data for a single simulation keyword in the file.

      If skip_c_space is True, lines starting 'C ' are
      also treated as comments.  If data_free_of_comments is True, there must be
      at least one blank line before the data begins, and no further comments are permitted.
      (This format is designed to handle data files generated by a commonly used geomodelling package.)

      Repeat counts must not be present in the ascii data.

      The extent, if present, can contain any number of dimensions, typically 3 for reservoir modelling work.
      The total number of numbers in the file must match the number of elements in the given extent
      (ie. the product of the list of numbers in the extent argument).
      The order of indices in extent should be 'slowest changing' first, eg.: k,j,i

      The data_type defaults to 'real'
      'real' and 'float' are synonymous; 'int' and 'integer' are synonymous;
      'bool' and 'boolean' are synonymous; default is 'real'
      The numpy data type will be the default 64 bit float or 64 bit int
   """
   # todo: Code enhancement could cater for 32 bit options (and 8 bit for bool) if needed to reduce memory usage

   if extent is None:
      cell_count = -1   # np.fromfile interprets this as 'read everything'
      log.debug('Loading unknown number of array data elements from ascii file ' + file_name)
   else:
      cell_count = pu.cell_count_from_extent(extent)
      log.debug('Loading %1d array data elements from ascii file %s', cell_count, file_name)

   if data_type in ['real', 'float', float]: d_type = 'float'
   elif data_type in ['int', 'integer', int]: d_type = 'int'
   elif data_type in ['bool', 'boolean', bool]: d_type = 'int'  # read booleans as 0 or 1 and convert after
   else:
      assert False, 'Unknown data_type passed to load_array_from_ascii_file' + str(data_type)

   read_file_name = file_name

   with open(read_file_name, 'r') as data_file:

      if not comment_char and not data_free_of_comments:
         comment_char = kf.guess_comment_char(data_file)
         if not comment_char: comment_char = '!'

      if keyword:
         keyword_found = kf.find_keyword(data_file, keyword, max_lines = max_lines_for_keyword)
         if keyword_found: data_file.readline()  # skip keyword line

      kf.skip_blank_lines_and_comments(data_file, comment_char = comment_char, skip_c_space = skip_c_space)

      result = None

      if data_free_of_comments:   # use numpy fromfile function after passing header comments

         result = np.fromfile(data_file, dtype = d_type, count = cell_count, sep = ' ')

         if result is None:
            data_file.seek(0)

      if result is None:

         # builds one very big string, stripping trailing comments

#         s = ''
#         while True:
#            r = data_file.readline()
#            if len(r) == 0: break
#            s += r.partition(comment_char)[0] + ' '
#         result = np.fromstring(s, dtype = d_type, count = cell_count, sep = ' ')
#         # note: extra data will go unnoticed if cell count known!
#         del s

         b = bytearray(data_file.read().encode())
         bc = comment_char.encode()
         nl = b'\n'
         sp = b' '
         i = 0
         while True:
            i = b.find(bc, i)
            if i < 0: break
            eol = b.find(nl, i)
            if eol < 0: eol = len(b)
            b[i : eol] = sp * (eol - i)

         result = np.fromstring(b.decode(), dtype = d_type, count = cell_count, sep = ' ')
         # note: extra data will go unnoticed if cell count known!
         del b

      if result is None:

         assert(extent is not None)   # TODO: remove this restriction by dynamically extending a flat array

         view_1D = np.zeros(extent, dtype = d_type).flatten()
         elements = view_1D.size
         start_of_line = True

         for index in range(elements):

            while True:
               ch = data_file.read(1)
               assert(ch != '')        # premature end of file
               if ch == comment_char:
                  data_file.readline()
                  start_of_line = True
                  continue
               if skip_c_space and start_of_line and ch in ['C', 'c']:
                  next_ch = data_file.read(1)
                  if next_ch in ' \t\n':
                     if next_ch != '\n': data_file.readline()
                     continue
                  ch += next_ch
                  break
               if ch not in ' \t\n':
                  break

            start_of_line = False
            word = ch
            while True:
               ch = data_file.read(1)
               if ch == '':            # end of file (this assumes a whitespace character after last datum)
                  log.error('Not enough data in file %s: %1d of %1d numbers read', file_name, index, elements)
                  assert False, 'not enough data in file'
               if ch in ' \t\n': break
               word += ch
            if ch == '\n': start_of_line = True
            if d_type == 'float': view_1D[index] = float(word)
            else: view_1D[index] = int(word)

         result = view_1D

      kf.skip_blank_lines_and_comments(data_file, comment_char = comment_char, skip_c_space = skip_c_space)
      if not kf.end_of_file(data_file):
         assert False, 'too much data in ascii file: ' + file_name

   if result is not None and extent is not None: result = result.reshape(extent)

   if data_type in ['bool', 'boolean', bool]:
      return result != 0   # convert to boolean
   else:
      return result

# end of load_array_from_ascii_file() def
######################################################################################################


######################################################################################################
# load_fault_mask():
# function to load fault indicator data for cell faces from nexus input deck format ascii file
# indicator array is a numpy array with data type 'bool' and must have been initialised to false
# only faces specified in the input file are modified
# the indicator array is set to true for every face mentioned in the MULT data, regardless of value
# binary data file creation & reuse is supported

def load_fault_mask(file_name, direction, fault_mask, use_binary = False):
   """Modifies a 3D numpy boolean array representing fault locations on cell faces in i,j or k.

   arguments:
      file_name: string giving name of existing ascii file holding nexus MULT keywords in input deck format
      direction: a single character giving direction of interest, must be 'i', 'j' or 'k'
      fault_mask: an existing 3D boolean numpy array which must have been initialised to False everywhere
      use_binary: if True, will attempt to load the boolean mask from a pre-written pure binary file,
                 which will be created if it does not exist

   returns:
      None

   notes :
      see nexus keyword reference for format of MULT keyword

      fault_mask element is set to True if corresponding face appears in the MULT data set,
      whatever function is being applied, and whatever value is specified the function returns None (nothing)
   """

   trans_keyword_dict = {'i':'TX', 'j':'TY', 'k':'TZ'}

   assert(direction in 'ijk')
   assert(fault_mask.ndim == 3)
   mask_extent_kji = fault_mask.shape
   ascii_file_name = file_name
   binary_file_name = file_name

   face_count = pu.cell_count_from_extent(mask_extent_kji)
   log.debug('Loading %1d fault mask elements from file %s', face_count, file_name)

   if use_binary:
      # try reading mask from binary file
      if len(binary_file_name) > 5 and binary_file_name[-5:] == '.' + direction.lower() + '.bb':
         ascii_file_name = ascii_file_name[:-5]
      else:
         binary_file_name = binary_file_name + '.' + direction.lower() + '.bb'
      try:
         with open(binary_file_name, 'rb') as binary_data_file:
            result = np.fromfile(binary_data_file, dtype = 'bool', count = face_count).reshape(mask_extent_kji)
            fault_mask.data = result.data
            # todo: check that end of file has been reached, ie. not too much data in file
            log.info('Fault mask data loaded from binary file %s', binary_file_name)
            return
      except:
         pass

   with open(ascii_file_name, 'r') as data_file:

      while kf.find_keyword_pair(data_file, 'MULT', trans_keyword_dict[direction]):

         line = data_file.readline().upper()
         line = line.partition('!')[0]         # removes trailing comment
         plus_minus_offset = [0,0,0]           # only possibly set to 1 in direction of interest
         if ('MINUS' not in line.split()):     # indicates data apply to 'negative' faces of specified cells
            plus_minus_offset['kji'.index(direction)] = 1

         if not kf.find_keyword(data_file, 'FNAME'): continue  # will mess up if not present for this MULT
         line = data_file.readline().upper()
         words = line.split()

         if not kf.find_number(data_file): continue            # will mess up if no numbers before next MULT

         count = 0
         while kf.number_next(data_file):
            # fetch a new box from the input file (usually the box contains just one cell)
            line = data_file.readline().upper()
            line = line.partition('!')[0]         # removes trailing comment
            words = line.split()
            box_kji0 = box.box_kji0_from_words_iijjkk1(words)
            box_kji0[0] = box_kji0[0] + plus_minus_offset   # adjust min indices
            box_kji0[1] = box_kji0[1] + plus_minus_offset   # adjust max indices
            # check that box is within extent of fault_mask
            for dir in range(2):
               assert(box_kji0[0, dir] >= 0)
               assert(box_kji0[1, dir] < mask_extent_kji[dir])
            # set fault_mask to True for face of every cell in box (usually just one cell, or column)
#            for k in range(box_kji0[0, 0], box_kji0[1, 0] + 1):
#               for j in range(box_kji0[0, 1], box_kji0[1, 1] + 1):
#                  for i in range(box_kji0[0, 2], box_kji0[1, 2] + 1):
#                     if not fault_mask[k,j,i]:
#                        fault_mask[k, j, i] = True
#                        count += 1
            count += (box.volume_of_box(box_kji0)  -
                      np.count_nonzero(fault_mask[box_kji0[0, 0] : box_kji0[1, 0] + 1,
                                                  box_kji0[0, 1] : box_kji0[1, 1] + 1,
                                                  box_kji0[0, 2] : box_kji0[1, 2] + 1]))
            fault_mask[box_kji0[0, 0] : box_kji0[1, 0] + 1,
                       box_kji0[0, 1] : box_kji0[1, 1] + 1,
                       box_kji0[0, 2] : box_kji0[1, 2] + 1] = True


   if use_binary:
      # create binary file for use next time
      try:
         with open(binary_file_name, 'wb') as binary_file_out:
            fault_mask.tofile(binary_file_out)
         log.info('Binary fault mask data file %s created', binary_file_name)
      except:
         log.warn('Failed to write data to binary fault mask file %s', binary_file_name)
         # todo: could delete the binary file in case a corrupt file is left for use next time

# end of load_fault_mask() def
######################################################################################################
