text: # -*- coding: utf-8 -*-
  """
  
      mslib.mswms.mpl_hsec
      ~~~~~~~~~~~~~~~~~~~~
  
      Horizontal section style super classes for use with the
      HorizontalSectionDriver class.
  
      This file is part of MSS.
  
      :copyright: Copyright 2008-2014 Deutsches Zentrum fuer Luft- und Raumfahrt e.V.
      :copyright: Copyright 2011-2014 Marc Rautenhaus (mr)
      :copyright: Copyright 2016-2022 by the MSS team, see AUTHORS.
      :license: APACHE-2.0, see LICENSE for details.
  
      Licensed under the Apache License, Version 2.0 (the "License");
      you may not use this file except in compliance with the License.
      You may obtain a copy of the License at
  
         http://www.apache.org/licenses/LICENSE-2.0
  
      Unless required by applicable law or agreed to in writing, software
      distributed under the License is distributed on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      See the License for the specific language governing permissions and
      limitations under the License.
  """
  # style definitions should be put in mpl_hsec_styles.py
  
  
  import io
  import logging
  from abc import abstractmethod
  import mswms_settings
  
  import matplotlib as mpl
  from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
  import mpl_toolkits.basemap as basemap
  import mpl_toolkits.axes_grid1
  import numpy as np
  import PIL.Image
  
  from mslib.mswms import mss_2D_sections
  from mslib.utils.coordinate import get_projection_params
  from mslib.utils.units import convert_to
  from mslib.mswms.utils import make_cbar_labels_readable
  
  
  BASEMAP_CACHE = {}
  BASEMAP_REQUESTS = []
  
  
  class AbstractHorizontalSectionStyle(mss_2D_sections.Abstract2DSectionStyle):
      """
      Abstract horizontal section super class. Use this class as a parent
      to classes implementing different plotting backends. For example,
      to derive a Matplotlib-based style class, or a Magics++-based
      style class.
      """
  
      @abstractmethod
      def plot_hsection(self):
          """
          Re-implement this function to perform the actual plotting.
          """
          pass
  
      def add_colorbar(self, contour, label=None, tick_levels=None, width="3%", height="30%", cb_format=None,
                       fraction=0.05, pad=0.08, shrink=0.7, loc=4, extend="both", tick_position="left"):
          if not self.noframe:
              cbar = self.fig.colorbar(contour, fraction=fraction, pad=pad, shrink=shrink)
              cbar.set_label(label)
          else:
              axins1 = mpl_toolkits.axes_grid1.inset_locator.inset_axes(
                  self.bm.ax, width=width, height=height, loc=loc)
              self.fig.colorbar(
                  contour, cax=axins1, orientation="vertical",
                  ticks=tick_levels, extend=extend, format=cb_format)
              axins1.yaxis.set_ticks_position(tick_position)
              make_cbar_labels_readable(self.fig, axins1)
  
  
  class MPLBasemapHorizontalSectionStyle(AbstractHorizontalSectionStyle):
      """
      Matplotlib-based super class for all horizontal section styles.
      Sets up the map projection and draws a basemap.
      """
      name = "BASEMAP"
      title = "Matplotlib basemap"
      _plot_countries = True  # set to False in derived class to disable country plotting
  
      def _plot_style(self):
          """
          Overwrite this method to plot style-specific data on the map.
          """
          pass
  
      def supported_epsg_codes(self):
          return list(mswms_settings.epsg_to_mpl_basemap_table.keys())
  
      def support_epsg_code(self, crs):
          """
          Returns a list of supported EPSG codes.
          """
          try:
              get_projection_params(crs)
          except ValueError:
              return False
          return True
  
      def supported_crs(self):
          """
          Returns a list of the coordinate reference systems supported by
          this style.
          """
          crs_list = set([
              "EPSG:3031",  # WGS 84 / Antarctic Polar Stereographic
              "EPSG:3995",  # WGS 84 / Arctic Polar Stereographic
              "EPSG:3857",  # WGS 84 / Spherical Mercator
              "EPSG:4326",  # WGS 84 / cylindric
              "MSS:stere"])
          for code in self.supported_epsg_codes():
              crs_list.add(f"EPSG:{code:d}")
          return sorted(crs_list)
  
      def _draw_auto_graticule(self, bm):
          """
          """
          # Compute some map coordinates that are required below for the automatic
          # determination of which meridians and parallels to draw.
          axis = bm.ax.axis()
          upperLeftCornerLon, upperLeftCornerLat = bm(axis[0], axis[3], inverse=True)
          lowerRightCornerLon, lowerRightCornerLat = bm(axis[1], axis[2], inverse=True)
          middleUpperBoundaryLon, middleUpperBoundaryLat = bm(np.mean([axis[0], axis[1]]), axis[3], inverse=True)
          middleLowerBoundaryLon, middleLowerBoundaryLat = bm(np.mean([axis[0], axis[1]]), axis[2], inverse=True)
  
          # Determine which parallels and meridians should be drawn.
          #   a) determine which are the minimum and maximum visible
          #      longitudes and latitudes, respectively. These
          #      values depend on the map projection.
          if bm.projection in ['npstere', 'spstere', 'stere', 'lcc']:
              # For stereographic projections: Draw meridians from the minimum
              # longitude contained in the map at one of the four corners to the
              # maximum longitude at one of these corner points. If
              # the map centre is contained in the map, draw all meridians
              # around the globe.
              # FIXME? This is only correct for polar stereographic projections.
              # Draw parallels from the min latitude contained in the map at
              # one of the four corners OR the middle top or bottom to the
              # maximum latitude at one of these six points.
              # If the map centre in contained in the map, replace either
              # start or end latitude by the centre latitude (whichever is
              # smaller/larger).
  
              # check if centre point of projection is contained in the map,
              # use projection coordinates for this test
              centre_x = bm.projparams["x_0"]
              centre_y = bm.projparams["y_0"]
              contains_centre = (centre_x < bm.urcrnrx) and (centre_y < bm.urcrnry)
              # merdidians
              if contains_centre:
                  mapLonStart = -180.
                  mapLonStop = 180.
              else:
                  mapLonStart = min(upperLeftCornerLon, bm.llcrnrlon,
                                    bm.urcrnrlon, lowerRightCornerLon)
                  mapLonStop = max(upperLeftCornerLon, bm.llcrnrlon,
                                   bm.urcrnrlon, lowerRightCornerLon)
              # parallels
              mapLatStart = min(middleLowerBoundaryLat, lowerRightCornerLat,
                                bm.llcrnrlat,
                                middleUpperBoundaryLat, upperLeftCornerLat,
                                bm.urcrnrlat)
              mapLatStop = max(middleLowerBoundaryLat, lowerRightCornerLat,
                               bm.llcrnrlat,
                               middleUpperBoundaryLat, upperLeftCornerLat,
                               bm.urcrnrlat)
              if contains_centre:
                  centre_lat = bm.projparams["lat_0"]
                  mapLatStart = min(mapLatStart, centre_lat)
                  mapLatStop = max(mapLatStop, centre_lat)
          else:
              # for other projections (preliminary): difference between the
              # lower left and the upper right corner.
              mapLonStart = bm.llcrnrlon
              mapLonStop = bm.urcrnrlon
              mapLatStart = bm.llcrnrlat
              mapLatStop = bm.urcrnrlat
  
          # b) parallels and meridians can be drawn with a spacing of
          #      >spacingValues< degrees. Determine the appropriate
          #      spacing for the lon/lat differences: about 10 lines
          #      should be drawn in each direction. (The following lines
          #      filter the spacingValues list for all spacing values
          #      that are larger than lat/lon difference / 10, then
          #      take the first value (first values that's larger)).
          spacingValues = [0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 20, 30, 40]
          deltaLon = mapLonStop - mapLonStart
          deltaLat = mapLatStop - mapLatStart
          spacingLon = [i for i in spacingValues if i > (deltaLon / 10.)][0]
          spacingLat = [i for i in spacingValues if i > (deltaLat / 10.)][0]
  
          #   c) parallels and meridians start at the first value in the
          #      spacingLon/Lat grid that's smaller than the lon/lat of the
          #      lower left corner; they stop at the first values in the
          #      grid that's larger than the lon/lat of the upper right corner.
          lonStart = np.floor((mapLonStart / spacingLon)) * spacingLon
          lonStop = np.ceil((mapLonStop / spacingLon)) * spacingLon
          latStart = np.floor((mapLatStart / spacingLat)) * spacingLat
          latStop = np.ceil((mapLatStop / spacingLat)) * spacingLat
  
          #   d) call the basemap methods to draw the lines in the determined
          #      range.
          bm.map_parallels = bm.drawparallels(np.arange(latStart, latStop,
                                                        spacingLat),
                                              labels=[1, 1, 0, 0],
                                              color='0.5', dashes=[5, 5])
          bm.map_meridians = bm.drawmeridians(np.arange(lonStart, lonStop,
                                                        spacingLon),
                                              labels=[0, 0, 0, 1],
                                              color='0.5', dashes=[5, 5])
  
      def plot_hsection(self, data, lats, lons, bbox=(-180, -90, 180, 90),
                        level=None, figsize=(960, 640), crs=None,
                        proj_params=None,
                        valid_time=None, init_time=None, style=None,
                        resolution=-1, noframe=False, show=False,
                        transparent=False):
          """
          EPSG overrides proj_params!
          """
          if proj_params is None:
              proj_params = {"projection": "cyl"}
              bbox_units = "latlon"
          # Projection parameters from EPSG code.
          if crs is not None:
              proj_params, bbox_units = [get_projection_params(crs)[_x] for _x in ("basemap", "bbox")]
  
          logging.debug("plotting data..")
  
          # Check if required data is available.
          self.data_units = self.driver.data_units.copy()
          for datatype, dataitem, dataunit in self.required_datafields:
              if dataitem not in data:
                  raise KeyError(f"required data field '{dataitem}' not found")
              origunit = self.driver.data_units[dataitem]
              if dataunit is not None:
                  data[dataitem] = convert_to(data[dataitem], origunit, dataunit)
                  self.data_units[dataitem] = dataunit
              else:
                  logging.debug("Please add units to plot variables")
  
          # Copy parameters to properties.
          self.data = data
          self.lats = lats
          self.lons = lons
          self.level = level
          self.valid_time = valid_time
          self.init_time = init_time
          self.style = style
          self.resolution = resolution
          self.noframe = noframe
          self.crs = crs
  
          # Derive additional data fields and make the plot.
          logging.debug("preparing additional data fields..")
          self._prepare_datafields()
  
          logging.debug("creating figure..")
          dpi = 80
          figsize = (figsize[0] / dpi), (figsize[1] / dpi)
          facecolor = "white"
          fig = mpl.figure.Figure(figsize=figsize, dpi=dpi, facecolor=facecolor)
          logging.debug("\twith frame and legends" if not noframe else
                        "\twithout frame")
          if noframe:
              ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])
          else:
              ax = fig.add_axes([0.05, 0.05, 0.9, 0.88])
  
          # The basemap instance is created with a fixed aspect ratio for framed
          # plots; the aspect ratio is not fixed for frameless plots (standard
          # WMS). This means that for WMS plots, the map will always fill the
          # entire image area, no matter of whether this stretches or shears the
          # map. This is the behaviour specified by the WMS standard (WMS Spec
          # 1.1.1, Section 7.2.3.8):
          # "The returned picture, regardless of its return format, shall have
          # exactly the specified width and height in pixels. In the case where
          # the aspect ratio of the BBOX and the ratio width/height are different,
          # the WMS shall stretch the returned map so that the resulting pixels
          # could themselves be rendered in the aspect ratio of the BBOX.  In
          # other words, it should be possible using this definition to request a
          # map for a device whose output pixels are themselves non-square, or to
          # stretch a map into an image area of a different aspect ratio."
          # NOTE: While the MSUI always requests image sizes that match the aspect
          # ratio, for instance the Metview 4 client does not (mr, 2011Dec16).
  
          # Some additional code to store the last 20 coastlines in memory for quicker
          # access.
          key = repr((proj_params, bbox, bbox_units))
          basemap_use_cache = getattr(mswms_settings, "basemap_use_cache", False)
          basemap_request_size = getattr(mswms_settings, "basemap_request_size ", 200)
          basemap_cache_size = getattr(mswms_settings, "basemap_cache_size", 20)
          bm_params = {"area_thresh": 1000., "ax": ax, "fix_aspect": (not noframe)}
          bm_params.update(proj_params)
          if bbox_units == "degree":
              bm_params.update({"llcrnrlon": bbox[0], "llcrnrlat": bbox[1],
                                "urcrnrlon": bbox[2], "urcrnrlat": bbox[3]})
          elif bbox_units.startswith("meter"):
              # convert meters to degrees
              try:
                  bm_p = basemap.Basemap(resolution=None, **bm_params)
              except ValueError:  # projection requires some extent
                  bm_p = basemap.Basemap(resolution=None, width=1e7, height=1e7, **bm_params)
              bm_center = [float(_x) for _x in bbox_units[6:-1].split(",")]
              center_x, center_y = bm_p(*bm_center)
              bbox_0, bbox_1 = bm_p(bbox[0] + center_x, bbox[1] + center_y, inverse=True)
              bbox_2, bbox_3 = bm_p(bbox[2] + center_x, bbox[3] + center_y, inverse=True)
              bm_params.update({"llcrnrlon": bbox_0, "llcrnrlat": bbox_1,
                                "urcrnrlon": bbox_2, "urcrnrlat": bbox_3})
          elif bbox_units == "no":
              pass
          else:
              raise ValueError(f"bbox_units '{bbox_units}' not known.")
          if basemap_use_cache and key in BASEMAP_CACHE:
              bm = basemap.Basemap(resolution=None, **bm_params)
              (bm.resolution, bm.coastsegs, bm.coastpolygontypes, bm.coastpolygons,
               bm.coastsegs, bm.landpolygons, bm.lakepolygons, bm.cntrysegs) = BASEMAP_CACHE[key]
              logging.debug("Loaded '%s' from basemap cache", key)
          else:
              bm = basemap.Basemap(resolution='l', **bm_params)
              # read in countries manually, as those are laoded only on demand
              bm.cntrysegs, _ = bm._readboundarydata("countries")
              if basemap_use_cache:
                  BASEMAP_CACHE[key] = (bm.resolution, bm.coastsegs, bm.coastpolygontypes, bm.coastpolygons,
                                        bm.coastsegs, bm.landpolygons, bm.lakepolygons, bm.cntrysegs)
          if basemap_use_cache:
              BASEMAP_REQUESTS.append(key)
              BASEMAP_REQUESTS[:] = BASEMAP_REQUESTS[-basemap_request_size:]
  
              if len(BASEMAP_CACHE) > basemap_cache_size:
                  useful = {}
                  for idx, key in enumerate(BASEMAP_REQUESTS):
                      useful[key] = useful.get(key, 0) + idx
                  least_useful = sorted([(value, key) for key, value in useful.items()])[:-basemap_cache_size]
                  for _, key in least_useful:
                      del BASEMAP_CACHE[key]
                      BASEMAP_REQUESTS[:] = [_x for _x in BASEMAP_REQUESTS if key != _x]
  
          if self._plot_countries:
              # Set up the map appearance.
              try:
                  bm.drawcoastlines(color='0.25')
              except ValueError as ex:
                  logging.error("Error in basemap/matplotlib call of drawcoastlines: %s", ex)
              bm.drawcountries(color='0.5')
              bm.drawmapboundary(fill_color='white')
  
              # zorder = 0 is necessary to paint over the filled continents with
              # scatter() for drawing the flight tracks and trajectories.
              # Curiously, plot() works fine without this setting, but scatter()
              # doesn't.
              bm.fillcontinents(color='0.98', lake_color='white', zorder=0)
              self._draw_auto_graticule(bm)
  
          if noframe:
              ax.axis('off')
  
          self.bm = bm  # !! BETTER PASS EVERYTHING AS PARAMETERS?
          self.fig = fig
          self.shift_data()
          self.mask_data()
          self.lonmesh, self.latmesh = bm(*np.meshgrid(self.lons, self.lats))
          self._plot_style()
  
          # Set transparency for the output image.
          if transparent:
              fig.patch.set_alpha(0.)
  
          # Return the image as png embedded in a StringIO stream.
          canvas = FigureCanvas(fig)
          output = io.BytesIO()
          canvas.print_png(output)
  
          if show:
              logging.debug("saving figure to mpl_hsec.png ..")
              canvas.print_png("mpl_hsec.png")
  
          # Convert the image to an 8bit palette image with a significantly
          # smaller file size (~factor 4, from RGBA to one 8bit value, plus the
          # space to store the palette colours).
          # NOTE: PIL at the current time can only create an adaptive palette for
          # RGB images, hence alpha values are lost here. If transparency is
          # requested, the figure face colour is stored as the "transparent"
          # colour in the image. This works in most cases, but might lead to
          # visible artefacts in some cases.
          logging.debug("converting image to indexed palette.")
          # Read the above stored png into a PIL image and create an adaptive
          # colour palette.
          output.seek(0)  # necessary for PIL.Image.open()
          palette_img = PIL.Image.open(output).convert(mode="RGB").convert("P", palette=PIL.Image.ADAPTIVE)
          output = io.BytesIO()
          if not transparent:
              logging.debug("saving figure as non-transparent PNG.")
              palette_img.save(output, format="PNG")  # using optimize=True doesn't change much
          else:
              # If the image has a transparent background, we need to find the
              # index of the background colour in the palette. See the
              # documentation for PIL's ImagePalette module
              # (http://www.pythonware.com/library/pil/handbook/imagepalette.htm). The
              # idea is to create a 256 pixel image with the same colour palette
              # as the original image and use it as a lookup-table. Converting the
              # lut image back to RGB gives us a list of all colours in the
              # palette. (Why doesn't PIL provide a method to directly access the
              # colours in a palette??)
              lut = palette_img.resize((256, 1))
              lut.putdata(list(range(256)))
              lut = [c[1] for c in lut.convert("RGB").getcolors()]
              facecolor_rgb = list(mpl.colors.hex2color(mpl.colors.cnames[facecolor]))
              for i in [0, 1, 2]:
                  facecolor_rgb[i] = int(facecolor_rgb[i] * 255)
              try:
                  facecolor_index = lut.index(tuple(facecolor_rgb))
                  logging.debug("saving figure as transparent PNG with transparency index %s.", facecolor_index)
                  palette_img.save(output, format="PNG", transparency=facecolor_index)
              except ValueError:
                  logging.debug("transparency requested but not possible, saving non-transparent instead")
                  palette_img.save(output, format="PNG")
  
          logging.debug("returning figure..")
          return output.getvalue()
  
      def shift_data(self):
          """Shift the data fields such that the longitudes are in the range
          left_longitude .. left_longitude+360, where left_longitude is the
          leftmost longitude appearing in the plot.
  
          Necessary to prevent data cut-offs in situations where the requested
          map covers a domain that crosses the data longitude boundaries
          (e.g. data is stored on a -180..180 grid, but a map in the range
          -200..-100 is requested).
          """
          # Shifting makes the grid irregular for stereographic and other projections
          # in case the data is not global (i.e. covers 360 degrees).
          if self.bm.projection in ['npstere', 'spstere', 'stere', 'lcc']:
              return
          # Determine the leftmost longitude in the plot.
          axis = self.bm.ax.axis()
          ulcrnrlon, ulcrnrlat = self.bm(axis[0], axis[3], inverse=True)
          left_longitude = min(self.bm.llcrnrlon, ulcrnrlon)
          logging.debug("shifting data grid to leftmost longitude in map %2.f..", left_longitude)
  
          # Add a column of nans for non-global data to keep it from being
          # interpolated over long gaps in cylindrical
          dlon = self.lons[1] - self.lons[0]
          # not global!
          if abs(((180 + self.lons[-1] + dlon - self.lons[0]) % 360) - 180) > 1e-4:
              self.lons = np.concatenate([self.lons, [self.lons[-1] + dlon]])
              for key in self.data:
                  self.data[key] = np.hstack(
                      [self.data[key], np.full((len(self.lats), 1), np.nan)])
  
          # Shift the longitude field such that the data is in the range
          # left_longitude .. left_longitude+360.
          self.lons = ((self.lons - left_longitude) % 360) + left_longitude
          self.lon_indices = self.lons.argsort()
          self.lons = self.lons[self.lon_indices]
  
          # Shift data fields correspondingly.
          for key in self.data:
              self.data[key] = self.data[key][:, self.lon_indices]
  
      def mask_data(self):
          """Mask data arrays so that all values outside the map domain
             are masked. This is required for clabel to work correctly.
  
          See:
          http://www.mail-archive.com/matplotlib-users@lists.sourceforge.net/msg02892.html
          (Re: [Matplotlib-users] clabel and basemap).
  
          (mr, 2011-01-18)
          """
          # compute native map projection coordinates of lat/lon grid.
          lonmesh_, latmesh_ = np.meshgrid(self.lons, self.lats)
          x, y = self.bm(lonmesh_, latmesh_)
          # test which coordinates are outside the map domain.
  
          add_x = (self.bm.xmax - self.bm.xmin) / 10
          add_y = (self.bm.ymax - self.bm.ymin) / 10
  
          mask1 = x < self.bm.xmin - add_x
          mask2 = x > self.bm.xmax + add_x
          mask3 = y > self.bm.ymax + add_y
          mask4 = y < self.bm.ymin - add_y
          mask = mask1 & mask2 & mask3 & mask4
          # mask data arrays.
          for key in self.data:
              self.data[key] = np.ma.masked_array(self.data[key], mask=mask)
  