"""fine_coarse.py: Module providing support for grid refinement and coarsening."""

version = '29th April 2021'

# Nexus is a registered trademark of the Halliburton Company

# todo: cater for parallel to top & parallel to base style refinement proportions in k

import logging

log = logging.getLogger(__name__)

import numpy as np

import resqpy.olio.box_utilities as box


class FineCoarse:
   """Class for holding a mapping between fine and coarse grids."""

   def __init__(self, fine_extent_kji, coarse_extent_kji, within_fine_box = None, within_coarse_box = None):
      """Partial initialisation function, call other methods to assign ratios and proportions.

      arguments:
         fine_extent_kji (triple int): the local extent of the fine grid in k, j, i axes, ie (nk, nj, ni).
         coarse_extent_kji (triple int): the local extent of the coarse grid in k, j, i axes.
         within_fine_box (numpy int array of shape (2, 3), optional): if present, the subset of a larger
            fine grid that the mapping refers to; axes are min,max and k,j,i; values are zero based indices;
            max values are included in box (un-pythonesque); use this in the case of a local grid coarsening
         within_coarse_box (numpy int array of shape (2, 3), optional): if present, the subset of a larger
            coarse grid that the mapping refers to; axes are min,max and k,j,i; values are zero based indices;
            max values are included in box (un-pythonesque); use this in the case of a local grid refinement;
            required for write_cartref() method to work

      returns:
         newly formed FineCoarse object awaiting determination of ratios and proportions by axes.

      notes:
         at most one of within_fine_box and within_coarse_box may be passed; this information is not really
         used internally by the FineCoarse class but is noted in order to support local grid refinement and
         local grid coarsening applications;
         after intialisation, set_* methods should be called to establish the mapping
      """

      assert len(fine_extent_kji) == 3 and len(coarse_extent_kji) == 3
      assert within_fine_box is None or within_coarse_box is None

      self.fine_extent_kji = tuple(fine_extent_kji)  #: fine extent
      self.coarse_extent_kji = tuple(coarse_extent_kji)  #: coarse extent
      self._assert_extents_valid()
      if within_fine_box is not None:
         _assert_valid_box(fine_extent_kji, within_fine_box)
      if within_coarse_box is not None:
         _assert_valid_box(coarse_extent_kji, within_coarse_box)
      self.within_fine_box = within_fine_box  #: if not None, a box within an unidentified larger fine grid
      self.within_coarse_box = within_coarse_box  #: if not None, a box within an unidentified larger coarse grid

      self.constant_ratios = [None, None, None]  #: list for 3 axes kji, each None or int
      self.vector_ratios = [None, None, None]  #: list for 3 axes kji, each numpy vector of int or None
      self.equal_proportions = [True, True, True]  #: list for 3 axes kji, each boolean defaulting to equal proportions
      self.vector_proportions = [None, None, None
                                ]  #: list for 3 axes kji, each None or list of numpy vectors of float summing to 1.0

      self.fine_to_coarse_mapping = None  # derived triplet of int vectors holding coarse cell index for each fine cell index

   def assert_valid(self):
      """Checks consistency of everything within the fine coarse mapping; raises assertion error if not valid."""

      self._assert_extents_valid()

      assert self.constant_ratios is not None and len(self.constant_ratios) == 3
      assert self.vector_ratios is not None and len(self.vector_ratios) == 3
      assert self.equal_proportions is not None and len(self.equal_proportions) == 3
      assert self.vector_proportions is not None and len(self.vector_proportions) == 3

      for axis in range(3):
         assert (self.constant_ratios[axis] is None or
                 _is_int(self.constant_ratios[axis]) and self.constant_ratios[axis] > 0)
         if self.constant_ratios[axis] is None:
            assert (self.vector_ratios[axis] is not None and isinstance(self.vector_ratios[axis], np.ndarray) and
                    self.vector_ratios[axis].ndim == 1 and str(self.vector_ratios[axis].dtype).startswith('int'))
            assert len(self.vector_ratios[axis]) == self.coarse_extent_kji[axis]
            assert np.all(self.vector_ratios[axis] > 0)
            assert np.sum(self.vector_ratios[axis]) == self.fine_extent_kji[axis]
         else:
            assert self.vector_ratios[axis] is None
            assert self.coarse_extent_kji[axis] * self.constant_ratios[axis] == self.fine_extent_kji[axis]
         assert isinstance(self.equal_proportions[axis], bool)
         if self.equal_proportions[axis]:
            assert self.vector_proportions[axis] is None
         else:
            assert (self.vector_proportions[axis] is not None and
                    isinstance(self.vector_proportions[axis], list) and  # could allow a tuple
                    len(self.vector_proportions[axis]) == self.coarse_extent_kji[axis])
            for c0 in range(self.coarse_extent_kji[axis]):
               assert (isinstance(self.vector_proportions[axis][c0], np.ndarray) and
                       self.vector_proportions[axis][c0].ndim == 1 and
                       str(self.vector_proportions[axis][c0].dtype).startswith('float') and
                       len(self.vector_proportions[axis][c0]) == self.ratio(axis, c0))
               assert np.min(self.vector_proportions[axis][c0]) > 0.0
               assert abs(np.sum(self.vector_proportions[axis][c0]) - 1.0) < 1.0e-6

   def ratio(self, axis, c0):
      """Return fine:coarse ratio in given axis and coarse slice."""

      if self.constant_ratios[axis] is not None:
         return self.constant_ratios[axis]
      return self.vector_ratios[axis][c0]

   def ratios(self, c_kji0):
      """Return find:coarse ratios triplet for coarse cell."""

      return tuple(self.ratio(axis, c_kji0[axis]) for axis in range(3))

   def coarse_for_fine(self):
      """Returns triplet of numpy int vectors being the axial coarse cell indices for the axial fine cell indices."""

      if self.fine_to_coarse_mapping is None:
         self._set_fine_to_coarse_mapping()
      return self.fine_to_coarse_mapping

   def coarse_for_fine_kji0(self, fine_kji0):
      """Returns the index of the coarse cell which the given fine cell falls within."""

      if self.fine_to_coarse_mapping is None:
         self._set_fine_to_coarse_mapping()
      return (self.fine_to_coarse_mapping[0][fine_kji0[0]], self.fine_to_coarse_mapping[1][fine_kji0[1]],
              self.fine_to_coarse_mapping[2][fine_kji0[2]])

   def coarse_for_fine_axial(self, axis, f0):
      """Returns the index, for a single axis, of the coarse cell which the given fine cell falls within."""

      if self.fine_to_coarse_mapping is None:
         self._set_fine_to_coarse_mapping()
      return self.fine_to_coarse_mapping[axis][f0]

   def coarse_for_fine_axial_vector(self, axis):
      """Returns a numpy int vector, for a single axis, of the coarse cell index which each fine cell falls within."""

      if self.fine_to_coarse_mapping is None:
         self._set_fine_to_coarse_mapping()
      return self.fine_to_coarse_mapping[axis]

   def fine_base_for_coarse_axial(self, axis, c0):
      """Returns the index, for a single axis, of the 'first' fine cell within the coarse cell (lowest fine index)."""

      if self.constant_ratios[axis] is not None:
         return c0 * self.constant_ratios[axis]
      return np.sum(self.vector_ratios[axis][:c0])  # todo: check this returns zero for c0 == 0

   def fine_base_for_coarse(self, c_kji0):
      """Returns a 3-tuple being the 'first' (min) k, j, i0 in fine grid for given coarse cell."""

      base = []
      for axis in range(3):
         base.append(self.fine_base_for_coarse_axial(axis, c_kji0[axis]))
      return tuple(base)

   def fine_box_for_coarse(self, c_kji0):
      """Returns a numpy int array of shape (2, 3) being the min, max for k, j, i0 in fine grid for given coarse cell."""

      box = np.empty((2, 3), dtype = int)
      box[0] = self.fine_base_for_coarse(c_kji0)
      box[1] = box[0] + self.ratios(c_kji0) - 1
      return box

   def proportion(self, axis, c0):
      """Return a numpy vector of floats (summing to one) being the axial relative proportions of fine within coarse."""

      if self.equal_proportions[axis]:
         count = self.ratio(axis, c0)
         fraction = 1.0 / float(count)
         return np.full((count,), fraction)
      return self.vector_proportions[axis][c0]

   def interpolation(self, axis, c0):
      """Return a numpy vector of floats starting at zero and increasing monotonically to less than one, ready for interpolation."""

      count = self.ratio(axis, c0)
      fractions = np.zeros((count,))
      proport = self.proportion(axis, c0)
      for f0 in range(1, count):
         fractions[f0] = fractions[f0 - 1] + proport[f0 - 1]
      return fractions

   def proportions(self, c_kji0):
      """Return triplet of axial proportions for refinement of coarse cell."""

      return (self.proportion(axis, c_kji0[axis]) for axis in range(3))

   def set_constant_ratio(self, axis):
      """Set the refinement ratio for axis based on the ratio of the fine to coarse extents."""

      assert 0 <= axis < 3
      extent_ratio, remainder = divmod(self.fine_extent_kji[axis], self.coarse_extent_kji[axis])
      assert remainder == 0, 'coarse extent in ' + 'KJI'[axis] + ' is not a multiple of fine extent'
      self.constant_ratios[axis] = extent_ratio
      self.vector_ratios[axis] = None

   def set_ij_ratios_constant(self):
      """Set the refinement ratio for I & J axes based on the ratio of the fine to coarse extents."""

      for axis in (1, 2):
         self.set_constant_ratio(axis)

   def set_all_ratios_constant(self):
      """Set all refinement ratios constant based on the ratio of the fine to coarse extents."""

      for axis in range(3):
         self.set_constant_ratio(axis)

   def set_ratio_vector(self, axis, vector):
      """Set fine:coarse ratios for axis from numpy int vector of length matching coarse extent."""

      assert 0 <= axis < 3
      if isinstance(vector, list) or isinstance(vector, tuple):
         vector = np.array(vector)
      assert isinstance(vector, np.ndarray) and vector.ndim == 1
      assert str(vector.dtype).startswith('int')
      assert len(vector) == self.coarse_extent_kji[axis],  \
             'length of vector of refinement ratios does not match coarse extent'
      assert np.sum(vector) == self.fine_extent_kji[axis],  \
             'sum of refinement ratios in vector does not match fine extent'
      minimum = np.min(vector)
      assert minimum > 0
      if minimum == np.max(vector):
         self.set_constant_ratio(axis)
         assert self.constant_ratios[axis] == minimum
      else:
         self.vector_ratios[axis] = vector.copy()
         self.constant_ratios[axis] = None

   def set_equal_proportions(self, axis):
      """Set proportions equal for axis."""

      assert 0 <= axis < 3
      self.equal_proportions[axis] = True
      self.vector_proportions[axis] = None

   def set_all_proprtions_equal(self):
      """Sets proportions equal in all 3 axes."""

      for axis in range(3):
         self.set_equal_proportions(axis)

   def set_proportions_list_of_vectors(self, axis, list_of_vectors):
      """Sets the proportions for given axis, with one vector for each coarse slice in the axis."""

      assert 0 <= axis < 3
      assert len(list_of_vectors) == self.coarse_extent_kji[axis]
      if self.vector_proportions is None:
         self.vector_proportions = [None, None, None]
      fractions_list_of_vectors = []
      all_equal = True
      for c0 in range(self.coarse_extent_kji[axis]):
         vector = list_of_vectors[c0]
         count = self.ratio(axis, c0)
         if isinstance(vector, list) or isinstance(vector, tuple):
            vector = np.array(vector)
         assert isinstance(vector, np.ndarray) and vector.ndim == 1
         assert len(vector) == count, 'wrong number of proportions for axis: ' + str(axis) + ' slice(0): ' + c0
         assert not np.any(np.isnan(vector))
         assert np.all(vector > 0)
         total = float(np.sum(vector))
         fractions = np.empty((count,), dtype = float)
         fractions[:] = vector
         fractions /= total
         fractions_list_of_vectors.append(fractions)
         if np.max(fractions) - np.min(fractions) > 1.0e-6:
            all_equal = False
      if all_equal:
         self.equal_proportions[axis] = True
         self.vector_proportions[axis] = None
      else:
         self.equal_proportions[axis] = False
         self.vector_proportions[axis] = fractions_list_of_vectors

   def fine_for_coarse_natural_column_index(self, coarse_col):
      """Returns the fine equivalent natural (first) column index for coarse natural column index."""

      j, i = divmod(coarse_col, self.coarse_extent_kji[2])
      ratio_j, ratio_i = self.constant_ratios[1], self.constant_ratios[2]
      if ratio_j is None:
         if j:
            j = np.sum(self.vector_ratios[1][0:j])
      else:
         j *= ratio_j
      if ratio_i is None:
         if i:
            i = np.sum(self.vector_ratios[2][0:i])
      else:
         i *= ratio_i
      return j * self.fine_extent_kji[2] + i

   def fine_for_coarse_natural_pillar_index(self, coarse_p):
      """Returns the fine equivalent natural (first) pillar index for coarse natural pillar index."""

      p_j, p_i = divmod(coarse_p, self.coarse_extent_kji[2] + 1)
      ratio_j, ratio_i = self.constant_ratios[1], self.constant_ratios[2]
      if ratio_j is None:
         if p_j:
            p_j = np.sum(self.vector_ratios[1][0:p_j])
      else:
         p_j *= ratio_j
      if ratio_i is None:
         if p_i:
            p_i = np.sum(self.vector_ratios[2][0:p_i])
      else:
         p_i *= ratio_i
      return p_j * (self.fine_extent_kji[2] + 1) + p_i

   def write_cartref(self,
                     filename,
                     lgr_name,
                     mode = 'a',
                     root_name = None,
                     preceeding_blank_lines = 0,
                     trailing_blank_lines = 0):
      """Write Nexus ascii input format CARTREF; within_coarse_box must have been set."""

      assert self.within_coarse_box is not None, 'box within larger coarse grid required for cartref'
      assert filename and lgr_name, 'filename or lgr name missing for cartref'
      assert ' ' not in lgr_name, 'lgr name may not contain spaces'
      if len(lgr_name) > 8 or (root_name and len(root_name) > 8):
         log.warning('Nexus might only differentiate first 8 characters of grid names')
      with open(filename, mode) as fp:
         for _ in range(preceeding_blank_lines):
            fp.write('\n')
         if root_name:
            fp.write('LGR ' + str(root_name) + '\n')
         fp.write('LGR\n')
         fp.write('CARTREF ' + str(lgr_name) + '\n')
         fp.write('   ' + box.spaced_string_iijjkk1_for_box_kji0(self.within_coarse_box) + '\n')
         for axis in range(2, -1, -1):
            fp.write('  ')
            for c0 in range(self.coarse_extent_kji[axis]):
               fp.write(' ' + str(self.ratio(axis, c0)))
            fp.write('\n')
            if not self.equal_proportions[axis]:
               log.warning('unequal propertions in axis ' + 'KJI'[axis] + ': define Lgr corp separately')
         fp.write('ENDREF\n')
         fp.write('ENDLGR\n')
         for _ in range(trailing_blank_lines):
            fp.write('\n')

   def _assert_extents_valid(self):

      for axis in range(3):
         assert _is_int(self.fine_extent_kji[axis]) and _is_int(self.coarse_extent_kji[axis])
         assert self.fine_extent_kji[axis] > 0
         assert self.coarse_extent_kji[axis] <= self.fine_extent_kji[axis]

   def _set_fine_to_coarse_mapping(self):
      self.fine_to_coarse_mapping = [None, None, None]
      for axis in range(3):
         self.fine_to_coarse_mapping[axis] = np.empty(self.fine_extent_kji[axis], dtype = int)
         f0 = 0
         for c0 in range(self.coarse_extent_kji[axis]):
            ratio = self.ratio(axis, c0)
            self.fine_to_coarse_mapping[axis][f0:f0 + ratio] = c0
            f0 += ratio
         assert f0 == self.fine_extent_kji[axis]


def tartan_refinement(coarse_extent_kji,
                      coarse_fovea_box,
                      fovea_ratios_kji,
                      decay_rates_kji = None,
                      decay_mode = 'exponential',
                      within_coarse_box = None):
   """Returns a new FineCoarse object set to a tartan grid refinement; fine extent is determined from arguments.

   arguments:
      coarse_extent_kji (triple int): the extent of the coarse grid being refined
      coarse_fovea_box (numpy int array of shape (2, 3)): the central box within the coarse grid to receive maximum refinement
      fovea_ratios_kji (triple int): the maximum refinement ratios, to be applied in the coarse_fovea_box
      decay_rates_kji (triple float or triple int, optional): controls how quickly refinement ratio reduces in slices away from fovea;
         if None then default values will be generated; see notes for more details
      decay_mode (str, default 'exponential'): 'exponential' or 'linear'; see notes
      within_coarse_box (numpy int array of shape (2, 3), optional): if present, is preserved in FineCoarse for possible use
         in setting resqml ParentWindow or generating Nexus CARTREF

   returns:
      FineCoarse object holding the tartan refinement mapping

   notes:
      each axis is treated independently;
      the fovea (box of maximum refinement) may be a column of cells (for a vertical well) or any other logical cuboid;
      the refinement factor is reduced monotonically in slices moving away from the fovea;
      two refinement factor decay functions are available: 'exponential' and 'linear', with different meaning to decay_rates_kji;
      for exponential decay, each decay rate should be a float in the range 0.0 to 1.0, with 0.0 causing immediate change
      to no refinement (factor 1), and 1.0 causing no decay (constant refinement at fovea factor);
      for linear decay, each decay rate should typically be a non-negative integer (though float is also okay), with 0 causing
      no decay, 1 causing a reduction in refinement factor of 1 per coarse slice, 2 meaning refinement factor reduces by 2 with
      each coarse slice etc.;
      in all cases, the refinement factor is given a lower limit of 1;
      the factor is rounded to an int for each slice, when working with floats;
      if decay rates are not passed as arguments, suitable values are generated to give a gradual reduction in refinement to a
      ratio of one at the boundary of the grid
   """

   assert box.valid_box(coarse_fovea_box, coarse_extent_kji)
   assert decay_mode in ['exponential', 'linear']
   assert len(fovea_ratios_kji) == 3

   fovea_ratios_kji = np.array(fovea_ratios_kji)
   assert np.all(fovea_ratios_kji >= 1)

   # generate decay rates if not passed in
   rates = np.zeros((2, 3), dtype = float)  # -/+ side of fovea, in each axis
   if decay_rates_kji is None:
      border = np.zeros((2, 3), dtype = float)
      border[0] = coarse_fovea_box[0].astype(float) + 0.5
      border[1] = (coarse_extent_kji - coarse_fovea_box[1]).astype(float) - 0.5
      for side in (0, 1):
         if decay_mode == 'linear':
            rates[side] = (fovea_ratios_kji - 1) / border[side]
         else:  # exponential decay
            rates[side] = np.power(1.5 / fovea_ratios_kji, 1.0 / border[side])
   else:
      assert len(decay_rates_kji) == 3
      rates[:] = decay_rates_kji

   log.debug(f'decay rates: {rates}')

   ratios_list = []
   fine_extent_kji = np.zeros(3, dtype = int)

   for axis in range(3):
      ratios = np.ones(coarse_extent_kji[axis], dtype = int)
      slice_a, slice_b = coarse_fovea_box[:, axis]
      ratios[slice_a:slice_b + 1] = fovea_ratios_kji[axis]
      floating_ratio = np.ones(2, dtype = float)
      if decay_mode == 'linear':
         floating_ratio[:] = fovea_ratios_kji[axis] - rates[:, axis]
      else:
         floating_ratio[:] = fovea_ratios_kji[axis] * rates[:, axis]  # exponential decay
      rounded_ratio = np.around(floating_ratio).astype(int)
      while np.any(rounded_ratio > 1) and (slice_a > 0 or slice_b < coarse_extent_kji[axis] - 1):
         slice_a -= 1
         slice_b += 1
         if slice_a >= 0:
            ratios[slice_a] = rounded_ratio[0]
         if slice_b < coarse_extent_kji[axis]:
            ratios[slice_b] = rounded_ratio[1]
         if decay_mode == 'linear':
            floating_ratio[:] -= rates[:, axis]
         else:
            floating_ratio[:] *= rates[:, axis]  # exponential decay
         rounded_ratio = np.around(floating_ratio).astype(int)
      ratios_list.append(ratios)
      fine_extent_kji[axis] = np.sum(ratios)


#      log.debug(f'retio vector [{axis}]: {ratios}')

   fc = FineCoarse(fine_extent_kji, coarse_extent_kji, within_coarse_box = within_coarse_box)

   for axis in range(3):
      fc.set_ratio_vector(axis, ratios_list[axis])
      # todo: set proportions ?

   return fc


def axis_for_letter(letter):
   """Returns 0, 1 or 2 for 'K', 'J', or 'I'; as required for axis arguments in FineCoarse methods."""

   assert isinstance(letter, str) and len(letter) == 1
   u = letter.upper()
   return 'KJI'.index(u)


def letter_for_axis(axis):
   """Returns K, J, or I for axis; axis as required for axis arguments in FineCoarse methods."""

   assert isinstance(axis, int) and 0 <= axis < 3
   return 'KJI'[axis]


def _assert_valid_box(extent_kji, box):
   assert isinstance(box, np.ndarray) and box.shape == (2, 3) and str(box.dtype).startswith('int')
   for axis in range(3):
      assert 0 <= box[0, axis] <= box[1, axis]
      assert box[1, axis] - box[0, axis] + 1 == extent_kji[axis]


def _is_int(v):
   t = type(v)
   if t is int:
      return True
   if t in [np.int64, np.int32, np.int16, np.int8]:
      return True
   s = str(t)
   if s.startswith('int'):
      return True
   return False
