"""simple_lines.py: functions for handling simple lines in relation to a resqml grid."""

version = '19th January 2021'

# line is represented as 2D numpy array of shape (NP, 3): [point index, xyz]

import logging
log = logging.getLogger(__name__)
log.debug('simple_lines.py version ' + version)

import numpy as np


def read_lines(filename):
   """Returns a list of line arrays, read from ascii file.

      argument:
         filename (string): the path of the ascii file holding a set of poly-lines

      returns:
         list of numpy arrays, each array representing one poly-line

      notes:
         each line in the file must contain 3 floating point numbers: x, y, z;
         each poly-line must be terminated with a null marker line: 999.0 999.0 999.0
         there is no handling of units; elsewhere they will implicitly be assumed to be
         those of a crs for a grid object
   """

   lines_list = []
   end_of_line = np.array([999.0, 999.0, 999.0])
   try:
      with open(filename, 'r') as fp:
         point_count = 0
         point_list = None
         while True:
            line = fp.readline()
            if len(line) == 0:
               assert point_count == 0, 'unterminated list of line points at end of file'
               break
            words = line.split()
            if len(words) == 0: continue # blank line
            assert len(words) == 3, 'badly formed line (expecting 3 reals)'
            point = np.empty(3)
            for p in range(3): point[p] = float(words[p])
            if np.all(np.isclose(point, end_of_line)):
               if point_count:
                  lines_list.append(point_list)
                  point_count = 0
            elif point_count:
               point_list = np.concatenate((point_list, point.reshape((1, 3))), axis = 0)
               point_count += 1
            else:
               point_list = np.empty((1, 3))
               point_list[0, :] = point
               point_count = 1
   except:
      log.exception('failed to read simple lines from ascii file: ' + filename)
      lines_list = []
   log.info(str(len(lines_list)) + ' line(s) read from file: ' + filename)
   return lines_list


def polygon_line(line, tolerance = 0.001):
   """Returns a copy of the line with the last vertex stripped off if it is close to the first vertex.

      arguments:
         line (numpy array of floats): representation of a poly-line which might be closed (last vertex
            matches first vertex)
         tolerance (float, default = 0.001): the maximum Manhatten distance between two points for them
            to be treated as coincident

      returns:
         numpy array of floats which is either a copy of line, or a copy of line with the last point
         removed
   """

   last = len(line) - 1
   if last < 1: return line
   dims = len(line[0])
   manhattan = 0.0
   for d in range(dims): manhattan += abs(line[0][d] - line[last][d])
   if manhattan > tolerance: return line
   return line[:-1]


def duplicate_vertices_removed(line, tolerance = 0.001):
   """Returns a copy of the line with neighbouring duplicate vertices removed.

      arguments:
         line (2D numpy array of floats): representation of a poly-line
         tolerance (float, default = 0.001): the maximum Manhatten distance between two points for them
            to be treated as coincident

      returns:
         numpy array of floats which is either a copy of line, or a copy of line with some points
         removed

      notes:
         does not treat the line as a closed polyline, use polygon_line() function as well to remove
         duplicated first/last point;
         always preserves first and last point, even if they are identical and there are no other vertices
   """

   assert line.ndim == 2
   if len(line) < 3: return line
   dims = line.shape[1]
   whittled = np.zeros(line.shape)
   whittled[0] = line[0]
   c_i = 0
   for i in range(1, len(line)):
      manhattan = 0.0
      for d in range(dims): manhattan += abs(line[i, d] - whittled[c_i, d])
      if manhattan > tolerance:
         c_i += 1
         whittled[c_i] = line[i]
   if c_i == len(line) - 1: return line
   whittled[c_i] = line[-1]
   return whittled[:c_i + 1]


def nearest_pillars(line_list, grid, ref_k = 0, ref_kp = 0):
   """Finds pillars nearest to each point on each line; returns list of lists of (j0, i0).

      arguments:
         line_list (list of numpy arrays of floats): set of poly-lines, the points of which are used
            to find the nearest pillars in grid
         grid (grid.Grid object): the grid whose pillars are compared with the poly-line vertices
         ref_k (integer, default 0): the reference layer in the grid to compare against the vertices;
            zero based
         ref_kp (integer, default 0): 0 to indicate the top corners of the reference layer, 1 for the base

      returns:
         a list of lists of pairs of integers, each being the (j, i) pillar indices of the nearest pillar
         to the corresponding vertex of the poly-line; zero based indexing

      notes:
         this is a 2D search in the x, y plane; z values are ignored;
         poly-line x, y values must implicitly be in the same crs as the grid's points data
   """

   pillar_list_list = []
   for line in line_list:
      pillar_list = []
      p_count = line.shape[0]
      for p in range(p_count):
         xy = line[p, 0:2]
         ji = grid.nearest_pillar(xy, ref_k0 = ref_k + ref_kp)
         pillar_list.append(ji)
      pillar_list_list.append(pillar_list)
   return pillar_list_list


def nearest_rods(line_list, projection, grid, axis, ref_slice0 = 0, plus_face = False):
   """Finds rods nearest to each point on each line; returns list of lists of (k0, j0) or (k0, i0).

      arguments:
         line_list (list of numpy arrays of floats): set of poly-lines, the points of which are used
            to find the nearest rods in grid
         projection (string): 'xz' or 'yz'
         grid (grid.Grid object): the grid whose cross section points are compared with the poly-line vertices
         axis (string): 'I' or 'J' being the axis removed during slicing
         ref_slice0 (integer, default 0): the reference slice in the grid to compare against the vertices;
            zero based
         plus_face (boolean, default False): which face of the reference slice to use

      returns:
         a list of numpy arrays of pairs of integers, each being the (k, j) or (k, i) indices of the
         nearest rod to the corresponding vertex of the poly-line under projection; zero based indexing

      notes:
         this is a 2D search in the x, z or y, z plane;
         currently limited to unsplit grids without k gaps;
         poly-line x, y, z values must implicitly be in the same crs as the grid's points data
   """

   assert projection in ['xz', 'yz']
   assert axis.upper() in ['J', 'I']

   rod_list_list = []
   for line in line_list:
      line_a = np.array(line)
      rod_list_list.append(grid.nearest_rod(line_a, projection, axis, ref_slice0 = ref_slice0, plus_face = plus_face))
   return rod_list_list


def drape_lines(line_list, pillar_list_list, grid, ref_k = 0, ref_kp = 0, offset = -1.0, snap = False):
   """Roughly drapes lines over grid horizon; draped lines are suitable for 3D visualisation.

      arguments:
         line_list (list of numpy arrays of floats): the undraped poly-lines
         pillar_list_list: (list of lists of pairs of integers): as returned by nearest_pillars()
         grid: (grid.Grid object): the grid to which the poly-lines are to be draped
         ref_k (integer, default 0): the reference layer in the grid to which the lines will be
            draped; zero based
         ref_kp (integer, default 0): 0 to indicate the top corners of the reference layer, 1 for the base
         offset (float, default -1.0): the vertical offset to add to the z value of pillar points; positive
            drapes lines deeper, negative shallower (assumes grid's crs has z increasing downwards)
         snap (boolean, default False): if True, the x & y values of each vertex in the lines are moved to
            match the pillar point; if False, only the z values are adjusted

      returns:
         list of numpy arrays of floats, being the draped equivalents of the line_list

      notes:
         the units of the line list must implicitly be those of the crs for the grid;
         for grids with split coordinate lines (faults), only the primary pillars are currently used;
         the results of this function are intended for 3D visualisation of an indicative nature;
         resulting draped lines may penetrate the grid layer faces depending on the 3D geometry and
         the spacing of the vertices compared to cell sizes
   """

   assert len(line_list) == len(pillar_list_list)

   # note: currently uses primary pillars only: expect draping disasters by faults
   # pe_j = grid.extent_kji[1] + 1
   # pe_i = grid.extent_kji[2] + 1
   pillar_points = grid.horizon_points(ref_k0 = ref_k + ref_kp)

   draped_list = []
   for line, pillar_list in zip(line_list, pillar_list_list):
      drape = line.copy()
      p_count = line.shape[0]
      for p in range(p_count):
         ji = pillar_list[p]
         if snap:
            drape[p, :] = pillar_points[tuple(ji)]
         else:
            drape[p, 2] = pillar_points[tuple(ji)][2]
      drape[:, 2] += offset
      draped_list.append(drape)

   return draped_list


def drape_lines_to_rods(line_list, rod_list_list, projection,
                        grid, axis, ref_slice0 = 0, plus_face = False,
                        offset = -1.0, snap = False):
   """Roughly drapes lines near grid cross section; draped lines are suitable for 3D visualisation.

      arguments:
         line_list (list of numpy arrays of floats): the undraped poly-lines
         rod_list_list: (list of arrays of pairs of integers): as returned by nearest_rods()
         grid: (grid.Grid object): the grid to which the poly-lines are to be draped
         axis (string): 'I' or 'J' being the axis removed during slicing
         ref_slice0 (integer, default 0): the reference slice in the grid to drape to; zero based
         plus_face (boolean, default False): which face of the reference slice to use
         offset (float, default -1.0): the horzontal offset to add to the y or x value of rod points
         snap (boolean, default False): if True, the x & z or y & z values of each vertex in the lines
            are moved to match the rod point; if False, only the y or x values are adjusted

      returns:
         list of numpy arrays of floats, being the draped equivalents of the line_list

      notes:
         the units of the line list must implicitly be those of the crs for the grid;
         currently limited to unsplit grids with no k gaps;
         the results of this function are intended for 3D visualisation of an indicative nature;
         resulting draped lines may penetrate the grid layer faces depending on the 3D geometry and
         the spacing of the vertices compared to cell sizes
   """

   assert projection in ['xz', 'yz']
   assert axis.upper() in ['I', 'J']
   assert len(line_list) == len(rod_list_list)

   x_sect = grid.x_section_points(axis, ref_slice0 = ref_slice0, plus_face = plus_face)

   draped_list = []
   for line, rod_list in zip(line_list, rod_list_list):
      drape = line.copy()
      p_count = line.shape[0]
      for p in range(p_count):
         kj_or_ki = rod_list[p]
         if snap:
            drape[p, :] = x_sect[tuple(kj_or_ki)]
         elif projection == 'xz':
            drape[p, 1] = x_sect[tuple(kj_or_ki)][1]
         else:
            drape[p, 0] = x_sect[tuple(kj_or_ki)][0]
      if projection == 'xz':
         drape[:, 1] += offset
      else:
         drape[:, 0] += offset
      draped_list.append(drape)

   return draped_list
