content: # -*- coding: utf-8 -*-
  # Copyright (C) Duncan Macleod (2018-2020)
  #
  # This file is part of GWpy.
  #
  # GWpy is free software: you can redistribute it and/or modify
  # it under the terms of the GNU General Public License as published by
  # the Free Software Foundation, either version 3 of the License, or
  # (at your option) any later version.
  #
  # GWpy is distributed in the hope that it will be useful,
  # but WITHOUT ANY WARRANTY; without even the implied warranty of
  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  # GNU General Public License for more details.
  #
  # You should have received a copy of the GNU General Public License
  # along with GWpy.  If not, see <http://www.gnu.org/licenses/>.
  
  """Extension of `~matplotlib.axes.Axes` for gwpy
  """
  
  import warnings
  from functools import wraps
  from math import log
  from numbers import Number
  
  import numpy
  
  from astropy.time import Time
  
  from matplotlib import rcParams
  from matplotlib.artist import allow_rasterization
  from matplotlib.axes import Axes as _Axes
  from matplotlib.axes._base import _process_plot_var_args
  from matplotlib.collections import PolyCollection
  from matplotlib.lines import Line2D
  from matplotlib.projections import register_projection
  
  from . import (Plot, colorbar as gcbar)
  from .colors import format_norm
  from .gps import GPS_SCALES
  from .legend import HandlerLine2D
  from ..time import to_gps
  
  __author__ = 'Duncan Macleod <duncan.macleod@ligo.org>'
  
  
  def log_norm(func):
      """Wrap ``func`` to handle custom gwpy keywords for a LogNorm colouring
      """
      @wraps(func)
      def decorated_func(*args, **kwargs):
          norm, kwargs = format_norm(kwargs)
          kwargs['norm'] = norm
          return func(*args, **kwargs)
      return decorated_func
  
  
  def xlim_as_gps(func):
      """Wrap ``func`` to handle pass limit inputs through `gwpy.time.to_gps`
      """
      @wraps(func)
      def wrapped_func(self, left=None, right=None, **kw):
          if right is None and numpy.iterable(left):
              left, right = left
          kw['left'] = left
          kw['right'] = right
          gpsscale = self.get_xscale() in GPS_SCALES
          for key in ('left', 'right'):
              if gpsscale:
                  try:
                      kw[key] = numpy.longdouble(str(to_gps(kw[key])))
                  except TypeError:
                      pass
          return func(self, **kw)
      return wrapped_func
  
  
  def restore_grid(func):
      """Wrap ``func`` to preserve the Axes current grid settings.
      """
      @wraps(func)
      def wrapped_func(self, *args, **kwargs):
          try:
              grid = (
                  self.xaxis._minor_tick_kw["gridOn"],
                  self.xaxis._major_tick_kw["gridOn"],
                  self.yaxis._minor_tick_kw["gridOn"],
                  self.yaxis._major_tick_kw["gridOn"],
              )
          except KeyError:  # matplotlib < 3.3.3
              grid = (self.xaxis._gridOnMinor, self.xaxis._gridOnMajor,
                      self.yaxis._gridOnMinor, self.yaxis._gridOnMajor)
          try:
              return func(self, *args, **kwargs)
          finally:
              # reset grid
              self.xaxis.grid(grid[0], which="minor")
              self.xaxis.grid(grid[1], which="major")
              self.yaxis.grid(grid[2], which="minor")
              self.yaxis.grid(grid[3], which="major")
      return wrapped_func
  
  
  # -- new Axes -----------------------------------------------------------------
  
  class Axes(_Axes):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
  
          # handle Series in `ax.plot()`
          self._get_lines = PlotArgsProcessor(self)
  
          # reset data formatters (for interactive plots) to support
          # GPS time display
          self.fmt_xdata = self._fmt_xdata
          self.fmt_ydata = self._fmt_ydata
  
      @allow_rasterization
      def draw(self, *args, **kwargs):
          labels = {}
  
          for ax in (self.xaxis, self.yaxis):
              if ax.get_scale() in GPS_SCALES and ax.isDefault_label:
                  labels[ax] = ax.get_label_text()
                  trans = ax.get_transform()
                  epoch = float(trans.get_epoch())
                  unit = trans.get_unit_name()
                  iso = Time(epoch, format='gps', scale='utc').iso
                  utc = iso.rstrip('0').rstrip('.')
                  ax.set_label_text('Time [{0!s}] from {1!s} UTC ({2!r})'.format(
                      unit, utc, epoch))
  
          try:
              super().draw(*args, **kwargs)
          finally:
              for ax in labels:  # reset labels
                  ax.isDefault_label = True
  
      # -- auto-gps helpers -----------------------
  
      def _fmt_xdata(self, x):
          if self.get_xscale() in GPS_SCALES:
              return str(to_gps(x))
          return self.xaxis.get_major_formatter().format_data_short(x)
  
      def _fmt_ydata(self, y):
          if self.get_yscale() in GPS_SCALES:
              return str(to_gps(y))
          return self.yaxis.get_major_formatter().format_data_short(y)
  
      set_xlim = xlim_as_gps(_Axes.set_xlim)
  
      def set_epoch(self, epoch):
          """Set the epoch for the current GPS scale.
  
          This method will fail if the current X-axis scale isn't one of
          the GPS scales. See :ref:`gwpy-plot-gps` for more details.
  
          Parameters
          ----------
          epoch : `float`, `str`
              GPS-compatible time or date object, anything parseable by
              :func:`~gwpy.time.to_gps` is fine.
          """
          scale = self.get_xscale()
          return self.set_xscale(scale, epoch=epoch)
  
      def get_epoch(self):
          """Return the epoch for the current GPS scale/
  
          This method will fail if the current X-axis scale isn't one of
          the GPS scales. See :ref:`gwpy-plot-gps` for more details.
          """
          return self.get_xaxis().get_transform().get_epoch()
  
      # -- overloaded plotting methods ------------
  
      def scatter(self, x, y, c=None, **kwargs):
          # scatter with auto-sorting by colour
          try:
              if c is None:
                  raise ValueError
              c_array = numpy.asanyarray(c, dtype=float)
          except ValueError:  # no colour array
              pass
          else:
              c_sort = kwargs.pop('c_sort', True)
              if c_sort:
                  sortidx = c_array.argsort()
                  x = numpy.asarray(x)[sortidx]
                  y = numpy.asarray(y)[sortidx]
                  c = numpy.asarray(c)[sortidx]
  
          return super().scatter(x, y, c=c, **kwargs)
  
      scatter.__doc__ = _Axes.scatter.__doc__.replace(
          'marker :',
          'c_sort : `bool`, optional, default: True\n'
          '    Sort scatter points by `c` array value, if given.\n\n'
          'marker :',
      )
  
      @log_norm
      def imshow(self, array, *args, **kwargs):
          """Display an image, i.e. data on a 2D regular raster.
  
          If ``array`` is a :class:`~gwpy.types.Array2D` (e.g. a
          :class:`~gwpy.spectrogram.Spectrogram`), then the defaults are
          _different_ to those in the upstream
          :meth:`~matplotlib.axes.Axes.imshow` method. Namely, the defaults are
  
          - ``origin='lower'`` (coordinates start in lower-left corner)
          - ``aspect='auto'`` (pixels are not forced to be square)
          - ``interpolation='none'`` (no image interpolation is used)
  
          In all other usage, the defaults from the upstream matplotlib method
          are unchanged.
  
          Parameters
          ----------
          array : array-like or PIL image
              The image data.
  
          *args, **kwargs
              All arguments and keywords are passed to the inherited
              :meth:`~matplotlib.axes.Axes.imshow` method.
  
          See also
          --------
          matplotlib.axes.Axes.imshow
              for details of the image rendering
          """
          if hasattr(array, "yspan"):  # Array2D
              return self._imshow_array2d(array, *args, **kwargs)
  
          image = super().imshow(array, *args, **kwargs)
          self.autoscale(enable=None, axis='both', tight=None)
          return image
  
      def _imshow_array2d(self, array, origin='lower', interpolation='none',
                          aspect='auto', **kwargs):
          """Render an `~gwpy.types.Array2D` using `Axes.imshow`
          """
          # NOTE: If you change the defaults for this method, please update
          #       the docstring for `imshow` above.
  
          # calculate extent
          extent = tuple(array.xspan) + tuple(array.yspan)
          if self.get_xscale() == 'log' and extent[0] == 0.:
              extent = (1e-300,) + extent[1:]
          if self.get_yscale() == 'log' and extent[2] == 0.:
              extent = extent[:2] + (1e-300,) + extent[3:]
          kwargs.setdefault('extent', extent)
  
          return self.imshow(array.value.T, origin=origin, aspect=aspect,
                             interpolation=interpolation, **kwargs)
  
      @restore_grid
      @log_norm
      def pcolormesh(self, *args, **kwargs):
          """Create a pseudocolor plot with a non-regular rectangular grid.
  
          When using GWpy, this method can be called with a single argument
          that is an :class:`~gwpy.types.Array2D`, for which the ``X`` and ``Y``
          coordinate arrays will be determined from the indexing.
  
          In all other usage, all ``args`` and ``kwargs`` are passed directly
          to :meth:`~matplotlib.axes.Axes.pcolormesh`.
  
          Notes
          -----
          Unlike the upstream :meth:`matplotlib.axes.Axes.pcolormesh`,
          this method respects the current grid settings.
  
          See also
          --------
          matplotlib.axes.Axes.pcolormesh
          """
          if len(args) == 1 and hasattr(args[0], "yindex"):  # Array2D
              return self._pcolormesh_array2d(*args, **kwargs)
          return super().pcolormesh(*args, **kwargs)
  
      def _pcolormesh_array2d(self, array, *args, **kwargs):
          """Render an `~gwpy.types.Array2D` using `Axes.pcolormesh`
          """
          x = numpy.concatenate((array.xindex.value, array.xspan[-1:]))
          y = numpy.concatenate((array.yindex.value, array.yspan[-1:]))
          xcoord, ycoord = numpy.meshgrid(x, y, copy=False, sparse=True)
          return self.pcolormesh(xcoord, ycoord, array.value.T, *args, **kwargs)
  
      def hist(self, x, *args, **kwargs):
          x = numpy.asarray(x)
  
          # re-format weights as array if given as float
          weights = kwargs.get('weights', None)
          if isinstance(weights, Number):
              kwargs['weights'] = numpy.ones_like(x) * weights
  
          # calculate log-spaced bins on-the-fly
          if (kwargs.pop('logbins', False) and
                  not numpy.iterable(kwargs.get('bins', None))):
              nbins = kwargs.get('bins', None) or rcParams.get('hist.bins', 30)
              # get range
              hrange = kwargs.pop('range', None)
              if hrange is None:
                  try:
                      hrange = numpy.min(x), numpy.max(x)
                  except ValueError as exc:
                      if str(exc).startswith('zero-size array'):  # no data
                          exc.args = ('cannot generate log-spaced histogram '
                                      'bins for zero-size array, '
                                      'please pass `bins` or `range` manually',)
                      raise
              # log-scale the axis and extract the base
              if kwargs.get('orientation') == 'horizontal':
                  self.set_yscale('log', nonposy='clip')
                  logbase = self.yaxis._scale.base
              else:
                  self.set_xscale('log', nonposx='clip')
                  logbase = self.xaxis._scale.base
              # generate the bins
              kwargs['bins'] = numpy.logspace(
                  log(hrange[0], logbase), log(hrange[1], logbase),
                  nbins+1, endpoint=True)
  
          return super().hist(x, *args, **kwargs)
  
      hist.__doc__ = _Axes.hist.__doc__.replace(
          'color :',
          'logbins : boolean, optional\n'
          '    If ``True``, use logarithmically-spaced histogram bins.\n\n'
          '    Default is ``False``\n\n'
          'color :')
  
      # -- new plotting methods -------------------
  
      def plot_mmm(self, data, lower=None, upper=None, **kwargs):
          """Plot a `Series` as a line, with a shaded region around it.
  
          The ``data`` `Series` is drawn, while the ``lower`` and ``upper``
          `Series` are plotted lightly below and above, with a fill
          between them and the ``data``.
  
          All three `Series` should have the same `~Series.index` array.
  
          Parameters
          ----------
          data : `~gwpy.types.Series`
              Data to plot normally.
  
          lower : `~gwpy.types.Series`
              Lower boundary (on Y-axis) for shade.
  
          upper : `~gwpy.types.Series`
              Upper boundary (on Y-axis) for shade.
  
          **kwargs
              Any other keyword arguments acceptable for
              :meth:`~matplotlib.Axes.plot`.
  
          Returns
          -------
          artists : `tuple`
              All of the drawn artists:
  
              - `~matplotlib.lines.Line2d` for ``data``,
              - `~matplotlib.lines.Line2D` for ``lower``, if given
              - `~matplotlib.lines.Line2D` for ``upper``, if given
              - `~matplitlib.collections.PolyCollection` for shading
  
          See also
          --------
          matplotlib.axes.Axes.plot
              for a full description of acceptable ``*args`` and ``**kwargs``
          """
          alpha = kwargs.pop('alpha', .1)
  
          # plot mean
          line, = self.plot(data, **kwargs)
          out = [line]
  
          # modify keywords for shading
          kwargs.update({
              'label': '',
              'linewidth': line.get_linewidth() / 2,
              'color': line.get_color(),
              'alpha': alpha * 2,
          })
  
          # plot lower and upper Series
          fill = [data.xindex.value, data.value, data.value]
          for i, bound in enumerate((lower, upper)):
              if bound is not None:
                  out.extend(self.plot(bound, **kwargs))
                  fill[i+1] = bound.value
  
          # fill between
          out.append(self.fill_between(
              *fill, alpha=alpha, color=kwargs['color'],
              rasterized=kwargs.get('rasterized', True)))
  
          return out
  
      def tile(self, x, y, w, h, color=None,
               anchor='center', edgecolors='face', linewidth=0.8,
               **kwargs):
          """Plot rectanguler tiles based onto these `Axes`.
  
          ``x`` and ``y`` give the anchor point for each tile, with
          ``w`` and ``h`` giving the extent in the X and Y axis respectively.
  
          Parameters
          ----------
          x, y, w, h : `array_like`, shape (n, )
              Input data
  
          color : `array_like`, shape (n, )
              Array of amplitudes for tile color
  
          anchor : `str`, optional
              Anchor point for tiles relative to ``(x, y)`` coordinates, one of
  
              - ``'center'`` - center tile on ``(x, y)``
              - ``'ll'`` - ``(x, y)`` defines lower-left corner of tile
              - ``'lr'`` - ``(x, y)`` defines lower-right corner of tile
              - ``'ul'`` - ``(x, y)`` defines upper-left corner of tile
              - ``'ur'`` - ``(x, y)`` defines upper-right corner of tile
  
          **kwargs
              Other keywords are passed to
              :meth:`~matplotlib.collections.PolyCollection`
  
          Returns
          -------
          collection : `~matplotlib.collections.PolyCollection`
              the collection of tiles drawn
  
          Examples
          --------
          >>> import numpy
          >>> from matplotlib import pyplot
          >>> import gwpy.plot  # to get gwpy's Axes
  
          >>> x = numpy.arange(10)
          >>> y = numpy.arange(x.size)
          >>> w = numpy.ones_like(x) * .8
          >>> h = numpy.ones_like(x) * .8
  
          >>> fig = pyplot.figure()
          >>> ax = fig.gca()
          >>> ax.tile(x, y, w, h, anchor='ll')
          >>> pyplot.show()
          """
          # get color and sort
          if color is not None and kwargs.get('c_sort', True):
              sortidx = color.argsort()
              x = x[sortidx]
              y = y[sortidx]
              w = w[sortidx]
              h = h[sortidx]
              color = color[sortidx]
  
          # define how to make a polygon for each tile
          if anchor == 'll':
              def _poly(x, y, w, h):
                  return ((x, y), (x, y+h), (x+w, y+h), (x+w, y))
          elif anchor == 'lr':
              def _poly(x, y, w, h):
                  return ((x-w, y), (x-w, y+h), (x, y+h), (x, y))
          elif anchor == 'ul':
              def _poly(x, y, w, h):
                  return ((x, y-h), (x, y), (x+w, y), (x+w, y-h))
          elif anchor == 'ur':
              def _poly(x, y, w, h):
                  return ((x-w, y-h), (x-w, y), (x, y), (x, y-h))
          elif anchor == 'center':
              def _poly(x, y, w, h):
                  return ((x-w/2., y-h/2.), (x-w/2., y+h/2.),
                          (x+w/2., y+h/2.), (x+w/2., y-h/2.))
          else:
              raise ValueError("Unrecognised tile anchor {!r}".format(anchor))
  
          # build collection
          cmap = kwargs.pop('cmap', rcParams['image.cmap'])
          coll = PolyCollection((_poly(*tile) for tile in zip(x, y, w, h)),
                                edgecolors=edgecolors, linewidth=linewidth,
                                **kwargs)
          if color is not None:
              coll.set_array(color)
              coll.set_cmap(cmap)
  
          out = self.add_collection(coll)
          self.autoscale_view()
          return out
  
      # -- overloaded auxiliary methods -----------
  
      def legend(self, *args, **kwargs):
          # handle deprecated keywords
          linewidth = kwargs.pop("linewidth", None)
          if linewidth:
              warnings.warn(
                  "the linewidth keyword to gwpy.plot.Axes.legend has been "
                  "deprecated and will be removed in a future release; "
                  "please update your code to use a custom legend handler, "
                  "e.g. gwpy.plot.legend.HandlerLine2D.",
                  DeprecationWarning,
              )
          alpha = kwargs.pop("alpha", None)
          if alpha:
              kwargs.setdefault("framealpha", alpha)
              warnings.warn(
                  "the alpha keyword to gwpy.plot.Axes.legend has been "
                  "deprecated and will be removed in a future release; "
                  "use framealpha instead.",
                  DeprecationWarning,
              )
  
          # build custom handler
          handler_map = kwargs.setdefault("handler_map", dict())
          if isinstance(handler_map, dict):
              handler_map.setdefault(Line2D, HandlerLine2D(linewidth or 6))
  
          # create legend
          return super().legend(*args, **kwargs)
  
      legend.__doc__ = _Axes.legend.__doc__.replace(
          "Call signatures",
          """.. note::
  
     This method uses a custom default legend handler for
     `~matplotlib.lines.Line2D` objects, with increased linewidth relative
     to the upstream :meth:`~matplotlib.axes.Axes.legend` method.
     To disable this, pass ``handler_map=None``, or create and pass your
     own handler class.  See :ref:`gwpy-plot-legend` for more details.
  
  Call signatures""",
      )
  
      def colorbar(self, mappable=None, **kwargs):
          """Add a `~matplotlib.colorbar.Colorbar` to these `Axes`
  
          Parameters
          ----------
          mappable : matplotlib data collection, optional
              collection against which to map the colouring, default will
              be the last added mappable artist (collection or image)
  
          fraction : `float`, optional
              fraction of space to steal from these `Axes` to make space
              for the new axes, default is ``0.`` if ``use_axesgrid=True``
              is given (default), otherwise default is ``.15`` to match
              the upstream matplotlib default.
  
          **kwargs
              other keyword arguments to be passed to the
              :meth:`Plot.colorbar` generator
  
          Returns
          -------
          cbar : `~matplotlib.colorbar.Colorbar`
              the newly added `Colorbar`
  
          See also
          --------
          Plot.colorbar
          """
          fig = self.get_figure()
          if kwargs.get('use_axesgrid', True):
              kwargs.setdefault('fraction', 0.)
          if kwargs.get('fraction', 0.) == 0.:
              kwargs.setdefault('use_axesgrid', True)
          mappable, kwargs = gcbar.process_colorbar_kwargs(
              fig, mappable=mappable, ax=self, **kwargs)
          if isinstance(fig, Plot):
              # either we have created colorbar Axes using axesgrid1, or
              # the user already gave use_axesgrid=False, so we forcefully
              # disable axesgrid here in case fraction == 0., which causes
              # gridspec colorbars to fail.
              kwargs['use_axesgrid'] = False
          return fig.colorbar(mappable, **kwargs)
  
  
  # override default Axes with this one by registering a projection with the
  # same name
  
  register_projection(Axes)
  
  
  # -- overload Axes.plot() to handle Series ------------------------------------
  
  class PlotArgsProcessor(_process_plot_var_args):
      """This class controls how ax.plot() works
      """
      def __call__(self, *args, **kwargs):
          """Find `Series` data in `plot()` args and unwrap
          """
          newargs = []
          while args:
              # strip first argument
              this, args = args[:1], args[1:]
              # it its a 1-D Series, then parse it as (xindex, value)
              if hasattr(this[0], "xindex") and this[0].ndim == 1:
                  this = (this[0].xindex.value, this[0].value)
              # otherwise treat as normal (must be a second argument)
              else:
                  this += args[:1]
                  args = args[1:]
              # allow colour specs
              if args and isinstance(args[0], str):
                  this += args[0],
                  args = args[1:]
              newargs.extend(this)
  
          return super().__call__(*newargs, **kwargs)
  