content: """
  This is a set of classes for reading budget information out of MODFLOW-style
  listing files.  Cumulative and incremental budgets are returned as numpy
  recarrays, which can then be easily plotted.
  
  """
  
  import collections
  import os
  import re
  import sys
  from datetime import timedelta
  import numpy as np
  
  from ..utils.utils_def import totim_to_datetime
  
  
  class ListBudget(object):
      """
      MODFLOW family list file handling
  
      Parameters
      ----------
      file_name : str
          the list file name
      budgetkey : str
          the text string identifying the budget table. (default is None)
      timeunit : str
          the time unit to return in the recarray. (default is 'days')
  
      Notes
      -----
      The ListBudget class should not be instantiated directly.  Access is
      through derived classes: MfListBudget (MODFLOW), SwtListBudget (SEAWAT)
      and SwrListBudget (MODFLOW with the SWR process)
  
      Examples
      --------
      >>> mf_list = MfListBudget("my_model.list")
      >>> incremental, cumulative = mf_list.get_budget()
      >>> df_in, df_out = mf_list.get_dataframes(start_datetime="10-21-2015")
  
      """
  
      def __init__(self, file_name, budgetkey=None, timeunit='days'):
  
          # Set up file reading
          assert os.path.exists(file_name)
          self.file_name = file_name
          if sys.version_info[0] == 2:
              self.f = open(file_name, 'r')
          elif sys.version_info[0] == 3:
              self.f = open(file_name, 'r', encoding='ascii', errors='replace')
  
          self.tssp_lines = 0
  
          # Assign the budgetkey, which should have been overriden
          if budgetkey is None:
              self.set_budget_key()
          else:
              self.budgetkey = budgetkey
  
          self.totim = []
          self.timeunit = timeunit
          self.idx_map = []
          self.entries = []
          self.null_entries = []
  
          self.time_line_idx = 20
          if timeunit.upper() == 'SECONDS':
              self.timeunit = 'S'
              self.time_idx = 0
          elif timeunit.upper() == 'MINUTES':
              self.timeunit = 'M'
              self.time_idx = 1
          elif timeunit.upper() == 'HOURS':
              self.timeunit = 'H'
              self.time_idx = 2
          elif timeunit.upper() == 'DAYS':
              self.timeunit = 'D'
              self.time_idx = 3
          elif timeunit.upper() == 'YEARS':
              self.timeunit = 'Y'
              self.time_idx = 4
          else:
              raise Exception('need to reset time_idxs attribute to '
                              'use units other than days and check usage of '
                              'timedelta')
  
          # Fill budget recarrays
          self._load()
          self._isvalid = False
          if len(self.idx_map) > 0:
              self._isvalid = True
  
          # Close the open file
          self.f.close()
  
          # return
          return
  
      def set_budget_key(self):
          raise Exception('Must be overridden...')
  
      def isvalid(self):
          """
          Get a boolean indicating if budget data are available in the file.
  
          Returns
          -------
          out : boolean
              Boolean indicating if budget data are available in the file.
  
          Examples
          --------
          >>> mf_list = MfListBudget('my_model.list')
          >>> valid = mf_list.isvalid()
  
          """
          return self._isvalid
  
      def get_record_names(self):
          """
          Get a list of water budget record names in the file.
  
          Returns
          -------
          out : list of strings
              List of unique text names in the binary file.
  
          Examples
          --------
          >>> mf_list = MfListBudget('my_model.list')
          >>> names = mf_list.get_record_names()
  
          """
          if not self._isvalid:
              return None
          return self.inc.dtype.names
  
      def get_times(self):
          """
          Get a list of unique water budget times in the list file.
  
          Returns
          -------
          out : list of floats
              List contains unique water budget simulation times (totim) in list file.
  
          Examples
          --------
          >>> mf_list = MfListBudget('my_model.list')
          >>> times = mf_list.get_times()
  
          """
          if not self._isvalid:
              return None
          return self.inc['totim'].tolist()
  
      def get_kstpkper(self):
          """
          Get a list of unique stress periods and time steps in the list file
          water budgets.
  
          Returns
          ----------
          out : list of (kstp, kper) tuples
              List of unique kstp, kper combinations in list file.  kstp and
              kper values are zero-based.
  
          Examples
          --------
          >>> mf_list = MfListBudget("my_model.list")
          >>> kstpkper = mf_list.get_kstpkper()
  
          """
          if not self._isvalid:
              return None
          kstpkper = []
          for kstp, kper in zip(self.inc['time_step'],
                                self.inc['stress_period']):
              kstpkper.append((kstp, kper))
          return kstpkper
  
      def get_incremental(self, names=None):
          """
          Get a recarray with the incremental water budget items in the list file.
  
          Parameters
          ----------
          names : str or list of strings
              Selection of column names to return.  If names is not None then
              totim, time_step, stress_period, and selection(s) will be returned.
              (default is None).
  
          Returns
          -------
          out : recarray
              Numpy recarray with the water budget items in list file. The
              recarray also includes totim, time_step, and stress_period.
  
          Examples
          --------
          >>> mf_list = MfListBudget("my_model.list")
          >>> incremental = mf_list.get_incremental()
  
          """
          if not self._isvalid:
              return None
          if names is None:
              return self.inc
          else:
              if not isinstance(names, list):
                  names = [names]
              names.insert(0, 'stress_period')
              names.insert(0, 'time_step')
              names.insert(0, 'totim')
              return self.inc[names].view(np.recarray)
  
      def get_cumulative(self, names=None):
          """
          Get a recarray with the cumulative water budget items in the list file.
  
          Parameters
          ----------
          names : str or list of strings
              Selection of column names to return.  If names is not None then
              totim, time_step, stress_period, and selection(s) will be returned.
              (default is None).
  
          Returns
          -------
          out : recarray
              Numpy recarray with the water budget items in list file. The
              recarray also includes totim, time_step, and stress_period.
  
          Examples
          --------
          >>> mf_list = MfListBudget("my_model.list")
          >>> cumulative = mf_list.get_cumulative()
  
         """
          if not self._isvalid:
              return None
          if names is None:
              return self.cum
          else:
              if not isinstance(names, list):
                  names = [names]
              names.insert(0, 'stress_period')
              names.insert(0, 'time_step')
              names.insert(0, 'totim')
              return self.cum[names].view(np.recarray)
  
      def get_budget(self, names=None):
          """
          Get the recarrays with the incremental and cumulative water budget items
          in the list file.
  
          Parameters
          ----------
          names : str or list of strings
              Selection of column names to return.  If names is not None then
              totim, time_step, stress_period, and selection(s) will be returned.
              (default is None).
  
          Returns
          -------
          out : recarrays
              Numpy recarrays with the water budget items in list file. The
              recarray also includes totim, time_step, and stress_period. A
              separate recarray is returned for the incremental and cumulative
              water budget entries.
  
          Examples
          --------
          >>> mf_list = MfListBudget("my_model.list")
          >>> budget = mf_list.get_budget()
  
          """
          if not self._isvalid:
              return None
          if names is None:
              return self.inc, self.cum
          else:
              if not isinstance(names, list):
                  names = [names]
              names.insert(0, 'stress_period')
              names.insert(0, 'time_step')
              names.insert(0, 'totim')
              return self.inc[names].view(np.recarray), self.cum[names].view(
                      np.recarray)
  
      def get_data(self, kstpkper=None, idx=None, totim=None, incremental=False):
          """
          Get water budget data from the list file for the specified conditions.
  
          Parameters
          ----------
          idx : int
              The zero-based record number.  The first record is record 0.
              (default is None).
          kstpkper : tuple of ints
              A tuple containing the time step and stress period (kstp, kper).
              These are zero-based kstp and kper values. (default is None).
          totim : float
              The simulation time. (default is None).
          incremental : bool
              Boolean flag used to determine if incremental or cumulative water
              budget data for the specified conditions will be returned. If
              incremental=True, incremental water budget data will be returned.
              If incremental=False, cumulative water budget data will be
              returned. (default is False).
  
          Returns
          -------
          data : numpy recarray
              Array has size (number of budget items, 3). Recarray names are 'index',
              'value', 'name'.
  
          See Also
          --------
  
          Notes
          -----
          if both kstpkper and totim are None, will return the last entry
  
          Examples
          --------
          >>> import matplotlib.pyplot as plt
          >>> import flopy
          >>> mf_list = flopy.utils.MfListBudget("my_model.list")
          >>> data = mf_list.get_data(kstpkper=(0,0))
          >>> plt.bar(data['index'], data['value'])
          >>> plt.xticks(data['index'], data['name'], rotation=45, size=6)
          >>> plt.show()
  
          """
          if not self._isvalid:
              return None
          ipos = None
          if kstpkper is not None:
              try:
                  ipos = self.get_kstpkper().index(kstpkper)
              except:
                  pass
          elif totim is not None:
              try:
                  ipos = self.get_times().index(totim)
              except:
                  pass
          elif idx is not None:
              ipos = idx
          else:
              ipos = -1
  
          if ipos is None:
              print('Could not find specified condition.')
              print('  kstpkper = {}'.format(kstpkper))
              print('  totim = {}'.format(totim))
              return None
  
          if incremental:
              t = self.inc[ipos]
          else:
              t = self.cum[ipos]
  
          dtype = np.dtype(
                  [('index', np.int32), ('value', np.float32), ('name', '|S25')])
          v = np.recarray(shape=(len(self.inc.dtype.names[3:])), dtype=dtype)
          for i, name in enumerate(self.inc.dtype.names[3:]):
              mult = 1.
              if '_OUT' in name:
                  mult = -1.
              v[i]['index'] = i
              v[i]['value'] = mult * t[name]
              v[i]['name'] = name
          return v
  
      def get_dataframes(self, start_datetime='1-1-1970',diff=False):
          """
          Get pandas dataframes with the incremental and cumulative water budget
          items in the list file.
  
          Parameters
          ----------
          start_datetime : str
              If start_datetime is passed as None, the rows are indexed on totim.
              Otherwise, a DatetimeIndex is set. (default is 1-1-1970).
  
          Returns
          -------
          out : panda dataframes
              Pandas dataframes with the incremental and cumulative water budget
              items in list file. A separate pandas dataframe is returned for the
              incremental and cumulative water budget entries.
  
          Examples
          --------
          >>> mf_list = MfListBudget("my_model.list")
          >>> incrementaldf, cumulativedf = mf_list.get_dataframes()
  
          """
  
          try:
              import pandas as pd
          except Exception as e:
              raise Exception(
                      "ListBudget.get_dataframe() error import pandas: " + \
                      str(e))
  
          if not self._isvalid:
              return None
          totim = self.get_times()
          if start_datetime is not None:
              totim = totim_to_datetime(totim,
                                        start=pd.to_datetime(start_datetime),
                                        timeunit=self.timeunit)
  
          df_flux = pd.DataFrame(self.inc, index=totim).loc[:, self.entries]
          df_vol = pd.DataFrame(self.cum, index=totim).loc[:, self.entries]
  
          if not diff:
              return df_flux, df_vol
  
          else:
              in_names = [col for col in df_flux.columns if col.endswith("_IN")]
              out_names = [col for col in df_flux.columns if col.endswith("_OUT")]
              #print(in_names,out_names)
              #print(df_flux.columns)
              base_names = [name.replace("_IN",'') for name in in_names]
              for name in base_names:
                  in_name = name + "_IN"
                  out_name = name + "_OUT"
                  df_flux.loc[:,name.lower()] = df_flux.loc[:,in_name] - df_flux.loc[:,out_name]
                  df_flux.pop(in_name)
                  df_flux.pop(out_name)
                  df_vol.loc[:,name.lower()] = df_vol.loc[:,in_name] - df_vol.loc[:,out_name]
                  df_vol.pop(in_name)
                  df_vol.pop(out_name)
              cols = list(df_flux.columns)
              cols.sort()
              cols = [col.lower() for col in cols]
              df_flux.columns = cols
              df_vol.columns = cols
              return df_flux, df_vol
      def _build_index(self, maxentries):
          self.idx_map = self._get_index(maxentries)
          return
  
      def _get_index(self, maxentries):
          # --parse through the file looking for matches and parsing ts and sp
          idxs = []
          l_count = 1
          while True:
              seekpoint = self.f.tell()
              line = self.f.readline()
              if line == '':
                  break
              if self.budgetkey in line:
                  for l in range(self.tssp_lines):
                      line = self.f.readline()
                  try:
                      ts, sp = self._get_ts_sp(line)
                  except:
                      print('unable to cast ts,sp on line number', l_count,
                            ' line: ', line)
                      break
                  # print('info found for timestep stress period',ts,sp)
  
                  idxs.append([ts, sp, seekpoint])
  
                  if maxentries and len(idxs) >= maxentries:
                      break
  
          return idxs
  
      def _seek_to_string(self, s):
          """
          Parameters
          ----------
          s : str
              Seek through the file to the next occurrence of s.  Return the
              seek location when found.
  
          Returns
          -------
          seekpoint : int
              Next location of the string
  
          """
          while True:
              seekpoint = self.f.tell()
              line = self.f.readline()
              if line == '':
                  break
              if s in line:
                  break
          return seekpoint
  
      def _get_ts_sp(self, line):
          """
          From the line string, extract the time step and stress period numbers.
  
          """
  
          # Old method.  Was not generic enough.
          # ts = int(line[self.ts_idxs[0]:self.ts_idxs[1]])
          # sp = int(line[self.sp_idxs[0]:self.sp_idxs[1]])
  
          # Get rid of nasty things
          line = line.replace(',', '')
  
          searchstring = 'TIME STEP'
          idx = line.index(searchstring) + len(searchstring)
          ll = line[idx:].strip().split()
          ts = int(ll[0])
  
          searchstring = 'STRESS PERIOD'
          idx = line.index(searchstring) + len(searchstring)
          ll = line[idx:].strip().split()
          sp = int(ll[0])
  
          return ts, sp
  
      def _set_entries(self):
          if len(self.idx_map) < 1:
              return None, None
          if len(self.entries) > 0:
              raise Exception('entries already set:' + str(self.entries))
          if not self.idx_map:
              raise Exception('must call build_index before call set_entries')
          try:
              incdict, cumdict = self._get_sp(self.idx_map[0][0],
                                              self.idx_map[0][1],
                                              self.idx_map[0][2])
          except:
              raise Exception('unable to read budget information from first '
                              'entry in list file')
          self.entries = incdict.keys()
          null_entries = collections.OrderedDict()
          incdict = collections.OrderedDict()
          cumdict = collections.OrderedDict()
          for entry in self.entries:
              incdict[entry] = []
              cumdict[entry] = []
              null_entries[entry] = np.NaN
          self.null_entries = [null_entries, null_entries]
          return incdict, cumdict
  
      def _load(self, maxentries=None):
          self._build_index(maxentries)
          incdict, cumdict = self._set_entries()
          if incdict is None and cumdict is None:
              return
          totim = []
          for ts, sp, seekpoint in self.idx_map:
              tinc, tcum = self._get_sp(ts, sp, seekpoint)
              for entry in self.entries:
                  incdict[entry].append(tinc[entry])
                  cumdict[entry].append(tcum[entry])
  
              # Get the time for this record
              seekpoint = self._seek_to_string('TIME SUMMARY AT END')
              tslen, sptim, tt = self._get_totim(ts, sp, seekpoint)
              totim.append(tt)
  
          # get kstp and kper
          idx_array = np.array(self.idx_map)
  
          # build dtype for recarray
          dtype_tups = [('totim', np.float32), ("time_step", np.int32),
                        ("stress_period", np.int32)]
          for entry in self.entries:
              dtype_tups.append((entry, np.float32))
          dtype = np.dtype(dtype_tups)
  
          # create recarray
          nentries = len(incdict[entry])
          self.inc = np.recarray(shape=(nentries,), dtype=dtype)
          self.cum = np.recarray(shape=(nentries,), dtype=dtype)
  
          # fill each column of the recarray
          for entry in self.entries:
              self.inc[entry] = incdict[entry]
              self.cum[entry] = cumdict[entry]
  
          # file the totim, time_step, and stress_period columns for the
          # incremental and cumulative recarrays (zero-based kstp,kper)
          self.inc['totim'] = np.array(totim)[:]
          self.inc["time_step"] = idx_array[:, 0] - 1
          self.inc["stress_period"] = idx_array[:, 1] - 1
  
          self.cum['totim'] = np.array(totim)[:]
          self.cum["time_step"] = idx_array[:, 0] - 1
          self.cum["stress_period"] = idx_array[:, 1] - 1
  
          return
  
      def _get_sp(self, ts, sp, seekpoint):
          self.f.seek(seekpoint)
          # --read to the start of the "in" budget information
          while True:
              line = self.f.readline()
              if line == '':
                  print(
                          'end of file found while seeking budget information for ts,sp',
                          ts, sp)
                  return self.null_entries
  
              # --if there are two '=' in this line, then it is a budget line
              if len(re.findall('=', line)) == 2:
                  break
  
          tag = 'IN'
          incdict = collections.OrderedDict()
          cumdict = collections.OrderedDict()
          while True:
  
              if line == '':
                  # raise Exception('end of file found while seeking budget information')
                  print(
                          'end of file found while seeking budget information for ts,sp',
                          ts, sp)
                  return self.null_entries
              if len(re.findall('=', line)) == 2:
                  try:
                      entry, flux, cumu = self._parse_budget_line(line)
                  except e:
                      print('error parsing budget line in ts,sp', ts, sp)
                      return self.null_entries
                  if flux is None:
                      print(
                              'error casting in flux for', entry,
                              ' to float in ts,sp',
                              ts, sp)
                      return self.null_entries
                  if cumu is None:
                      print(
                              'error casting in cumu for', entry,
                              ' to float in ts,sp',
                              ts, sp)
                      return self.null_entries
                  if entry.endswith(tag.upper()):
                      if ' - ' in entry.upper():
                          key = entry.replace(' ', '')
                      else:
                          key = entry.replace(' ', '_')
                  elif 'PERCENT DISCREPANCY' in entry.upper():
                      key = entry.replace(' ', '_')
                  else:
                      key = '{}_{}'.format(entry.replace(' ', '_'), tag)
                  incdict[key] = flux
                  cumdict[key] = cumu
              else:
                  if 'OUT:' in line.upper():
                      tag = 'OUT'
              line = self.f.readline()
              if entry.upper() == 'PERCENT DISCREPANCY':
                  break
  
          return incdict, cumdict
  
      def _parse_budget_line(self, line):
  
          # get the budget item name
          entry = line.strip().split('=')[0].strip()
  
          # get the cumulative string
          idx = line.index('=') + 1
          line2 = line[idx:]
          ll = line2.strip().split()
          cu_str = ll[0]
  
          idx = line2.index('=') + 1
          fx_str = line2[idx:].strip()
  
          #
          # cu_str = line[self.cumu_idxs[0]:self.cumu_idxs[1]]
          # fx_str = line[self.flux_idxs[0]:self.flux_idxs[1]]
  
  
          flux, cumu = None, None
          try:
              cumu = float(cu_str)
          except:
              if 'NAN' in cu_str.strip().upper():
                  cumu = np.NaN
          try:
              flux = float(fx_str)
          except:
              if 'NAN' in fx_str.strip().upper():
                  flux = np.NaN
          return entry, flux, cumu
  
      def _get_totim(self, ts, sp, seekpoint):
          self.f.seek(seekpoint)
          # --read header lines
          ihead = 0
          while True:
              line = self.f.readline()
              ihead += 1
              if line == '':
                  print(
                          'end of file found while seeking time information for ts,sp',
                          ts, sp)
                  return np.NaN, np.NaN, np.Nan
              elif ihead == 2 and 'SECONDS     MINUTES      HOURS       DAYS        YEARS' not in line:
                  break
              elif '-----------------------------------------------------------' in line:
                  line = self.f.readline()
                  break
          tslen = self._parse_time_line(line)
          if tslen == None:
              print('error parsing tslen for ts,sp', ts, sp)
              return np.NaN, np.NaN, np.Nan
  
          sptim = self._parse_time_line(self.f.readline())
          if sptim == None:
              print('error parsing sptim for ts,sp', ts, sp)
              return np.NaN, np.NaN, np.Nan
  
          totim = self._parse_time_line(self.f.readline())
          if totim == None:
              print('error parsing totim for ts,sp', ts, sp)
              return np.NaN, np.NaN, np.Nan
          return tslen, sptim, totim
  
      def _parse_time_line(self, line):
          if line == '':
              print('end of file found while parsing time information')
              return None
          try:
              time_str = line[self.time_line_idx:]
              raw = time_str.split()
              idx = self.time_idx
              # catch case where itmuni is undefined
              # in this case, the table format is different
              try:
                  v = float(raw[0])
              except:
                  time_str = line[45:]
                  raw = time_str.split()
                  idx = 0
              tval = float(raw[idx])
          except:
              print('error parsing tslen information', time_str)
              return None
          return tval
  
  
  class SwtListBudget(ListBudget):
      """
  
      """
  
      def set_budget_key(self):
          self.budgetkey = 'MASS BUDGET FOR ENTIRE MODEL'
          return
  
  
  class MfListBudget(ListBudget):
      """
  
      """
  
      def set_budget_key(self):
          self.budgetkey = 'VOLUMETRIC BUDGET FOR ENTIRE MODEL'
          return
  
  
  class MfusgListBudget(ListBudget):
      """
  
      """
  
      def set_budget_key(self):
          self.budgetkey = 'VOLUMETRIC BUDGET FOR ENTIRE MODEL'
          return
  
  
  class SwrListBudget(ListBudget):
      """
  
      """
  
      def set_budget_key(self):
          self.budgetkey = 'VOLUMETRIC SURFACE WATER BUDGET FOR ENTIRE MODEL'
          self.tssp_lines = 1
          return
  