text: """Visualize an absorption lookup table.
  
  Author: oliver.lemke@uni-hamburg.de
  """
  import re
  from itertools import zip_longest
  
  import matplotlib.pyplot as plt
  import numpy as np
  from cycler import cycler
  from matplotlib.lines import Line2D
  from scipy.interpolate import interp1d
  
  import typhon.constants
  from typhon.plots import (ScalingFormatter, set_xaxis_formatter)
  
  __all__ = [
      'plot_arts_lookup',
  ]
  
  
  def _calc_lookup_species_count(lookup):
      """Calculate number of cross sections per species.
  
      Usually one, except for the nonlinear species.
      """
      nlsspecies = lookup.nonlinearspecies
      speciescount = np.ones_like(lookup.speciestags, dtype=int)
      if nlsspecies is not None:
          speciescount[nlsspecies] = lookup.nonlinearspeciesvmrperturbations.size
      return speciescount
  
  
  def _get_lookup_species_index(lookup, species, vmrpert):
      """Get index of given species in lookup table."""
      ret = 0
      spindex = lookup.speciestags.index(species)
      nlsspecies = lookup.nonlinearspecies
      speciescount = _calc_lookup_species_count(lookup)
      if nlsspecies is not None and spindex in nlsspecies:
          if vmrpert >= speciescount[spindex]:
              raise RuntimeError(
                  'Nonlinear species VMR perturbation index too large')
          ret = vmrpert
  
      return ret + (np.sum(speciescount[0:spindex]) if spindex > 0 else 0)
  
  
  def plot_lookup_xsec(lookup, ipressures, species=None, ax=None, tpert=0,
                       vmrpert=0):
      """Plot the cross section for one or more species of an ARTS lookup table.
  
      Parameters:
          lookup (typhon.arts.catalogues.GasAbsLookup): ARTS lookup table.
          ipressures (ndarray(int)): Indices of pressure levels to plot.
          species (list(list(str)), optional):
              ARTS species tags e.g. [['H2O-*-*0*']]. If none is given, plots
              all species in the lookup table for the given vmr perturbation.
          ax (AxesSubplot, optional): Axes to plot in.
          vmrpert (int): Index of vmr perturbation for nonlinear species to plot.
          tpert (int): Index of temperature perturbation to plot.
      """
      if ax is None:
          ax = plt.gca()
  
      ax.set_yscale('log')
      if species is None:
          species = lookup.speciestags
  
      for tag in species:
          ax.set_prop_cycle(
              cycler('color', [plt.cm.viridis(i) for i in
                               np.linspace(0, 1, len(ipressures))]))
          for pi in ipressures:
              xsec = lookup.absorptioncrosssection[
                     tpert,
                     _get_lookup_species_index(lookup, tag, vmrpert), :, pi]
              ax.plot(lookup.frequencygrid, xsec)
  
      if len(species) > 1:
          ax.legend(fontsize='xx-small', frameon=False)
      else:
          ax.set_title(
              ',\n'.join(re.sub('(-\*)+$', '', s) for s in species[0]),
              y=1. - len(species[0]) * 0.05,
              fontsize='xx-small')
  
      set_xaxis_formatter(ScalingFormatter('giga', r'${x:g}$'), ax=ax)
      ax.tick_params(axis='both', which='major', labelsize='xx-small')
      ax.spines['right'].set_visible(False)
      ax.spines['top'].set_visible(False)
  
  
  def plot_lookup_opacity(lookup, opacity, species=None, vmrpert=0, ax=None,
                          oneline=False, total=False):
      """Plot the opacity for one or more species of an ARTS lookup table.
  
      Parameters:
          lookup (typhon.arts.catalogues.GasAbsLookup): ARTS lookup table.
          opacity (ndarray): Opacity per species in lookup table as generated by
              `calc_opacity_from_lookup`.
          species (list(list(str)), optional):
              ARTS species tags e.g. [['H2O-*-*0*']]. If none is given, plots
              all species in the lookup table for the given vmr perturbation.
          vmrpert (int): Index of vmr perturbation for nonlinear species to plot.
          ax (AxesSubplot, optional): Axes to plot in.
          oneline (bool, optional): Draw a line where opacity == 1.
          total (bool, optional): Additionally plot the sum of opacities of all
              species.
      """
      if ax is None:
          ax = plt.gca()
  
      ax.set_yscale('log')
      if species is None:
          species = lookup.speciestags
  
      for tag in species:
          ax.plot(lookup.frequencygrid,
                  opacity[_get_lookup_species_index(lookup, tag, vmrpert), :],
                  label=',\n'.join(tag))
      if oneline:
          ax.plot(lookup.frequencygrid, np.ones_like(lookup.frequencygrid),
                  linewidth=1, linestyle='--', color='k')
      if total:
          if lookup.nonlinearspecies is not None:
              speciescount = _calc_lookup_species_count(lookup)
              spindex = np.cumsum(speciescount)
              spindex[1:] = spindex[0:-1]
              spindex[0] = 0
              spindex[lookup.nonlinearspecies] += vmrpert
              o = opacity[spindex]
          else:
              o = opacity
          ax.plot(lookup.frequencygrid, np.sum(o, axis=0),
                  linewidth=1, color='k')
  
      if len(species) > 1:
          ax.legend(fontsize='xx-small', frameon=False)
      else:
          ax.set_title(',\n'.join(re.sub('(-\*)+$', '', s) for s in species[0]),
                       y=1. - len(species[0]) * 0.05,
                       fontsize='xx-small')
  
      set_xaxis_formatter(ScalingFormatter('giga', r'${x:g}$'), ax=ax)
      ax.tick_params(axis='both', which='major', labelsize='xx-small')
      ax.tick_params(axis='both', which='minor', labelsize='xx-small')
      ax.spines['right'].set_visible(False)
      ax.spines['top'].set_visible(False)
  
  
  def calc_opacity_from_lookup(lookup, z=None, g=typhon.constants.g,
                               r=typhon.constants.gas_constant_dry_air, tpert=0):
      """Calculate the opacity from an ARTS lookup table.
  
      Parameters:
          lookup (typhon.arts.catalogues.GasAbsLookup): ARTS lookup table.
          z (ndarray, optional): Altitude profile. If not given, the layer
              thicknesses are calculated based on the hypsometric formula.
          g (float, optional): Gravity constant. Default uses Earth's gravity.
          r (float, optional): Gas constant for dry air. Default uses constant
              for Earth.
          tpert (int, optional): Index of temperature perturbation to plot.
  
      Returns:
          ndarray: Opacity per species in lookup table.
      """
      speciescount = _calc_lookup_species_count(lookup)
      vmrs = (np.repeat(lookup.referencevmrprofiles, speciescount, axis=0)
              if lookup.nonlinearspecies is not None
              else lookup.referencevmrprofiles)
  
      ni = (lookup.pressuregrid * vmrs
            / lookup.referencetemperatureprofile / typhon.constants.boltzmann
            ).reshape(np.sum(speciescount), 1, len(lookup.pressuregrid))
  
      alpha = ni * lookup.absorptioncrosssection[tpert, :, :, :]
  
      if z is not None:
          z = interp1d(z.grids[0], z.data[:, 0, 0])(lookup.pressuregrid)
      else:
          # Calculate z from hypsometric formula
          pgrid = lookup.pressuregrid
          z = [r * t / g * np.log(p1 / p2)
               for p1, p2, t in zip(pgrid[:-1], pgrid[1:], (
                      lookup.referencetemperatureprofile[
                      1:] + lookup.referencetemperatureprofile[:-1]) / 2.)]
          z = np.cumsum(z)
          p = (pgrid[1:] + pgrid[:-1]) / 2.
          z = interp1d(p, z, fill_value='extrapolate')(lookup.pressuregrid)
  
      return np.vstack(np.trapz(ialpha, z, axis=1) for ialpha in alpha)
  
  
  def _add_opacity_legend(ax=None):
      """Add legend to an opacity lookup table plot."""
      if ax is None:
          ax = plt.gca()
  
      blue_line = Line2D([], [], label='species opacity')
      black_line = Line2D([], [], color='k', linewidth=1.,
                          label='total opacity')
      dashed_line = Line2D([], [], color='k', linestyle='--',
                           linewidth=1., label='opacity=1')
  
      handles = [blue_line, black_line, dashed_line]
      labels = [h.get_label() for h in handles]
  
      ax.legend(handles=handles, labels=labels, fontsize='xx-small',
                loc='upper left', ncol=6)
  
  
  def _add_xsec_legend(lookup, ipressures, ax=None):
      """Add legend to a cross section lookup table plot."""
      if ax is None:
          ax = plt.gca()
  
      pgrid = lookup.pressuregrid
      colors = [plt.cm.viridis(i) for i in np.linspace(0, 1, len(ipressures))]
      handles = [Line2D([], [],
                        color=colors[i],
                        label=f'{pgrid[ip]/100.:8.3f} hPa')
                 for i, ip in enumerate(ipressures)]
  
      labels = [h.get_label() for h in handles]
  
      ax.legend(handles=handles, labels=labels, fontsize='xx-small',
                loc='upper left', ncol=6)
  
  
  def _setup_lookup_figure(lookup, cols=3):
      """Create the figure and axes objects for the lookup table plot."""
      rows = int(np.ceil(len(lookup.speciestags) / cols))
      fig, ax = plt.subplots(rows + 1, cols, figsize=(4*cols, (rows + 1) * 2))
      fig.tight_layout()
  
      return rows, cols, fig, ax
  
  
  def plot_arts_lookup(lookup, opacity=True, z=None, g=typhon.constants.g,
                       r=typhon.constants.gas_constant_dry_air, tpert=0,
                       vmrpert=0, pressures=None, cols=3):
      """Visualize an ARTS lookup table.
  
      Plots the opacity or the absorption cross sections based on an
      ARTS lookup table.
  
      Parameters:
          lookup (typhon.arts.catalogues.GasAbsLookup): ARTS lookup table object.
          opacity (bool, optional): Set to False to plot the absorption cross
              sections.
          z (ndarray, optional): Altitude profile. Optional input for opacity
              calculation. If not given, the layer thicknesses are calculated
              based on the hypsometric formula.
          g (float, optional): Gravity constant. Uses Earth's gravity by default.
          r (float, optional): Gas constant for dry air.
              Uses constant for Earth by default.
          tpert (int, optional): Index of temperature perturbation to plot.
          vmrpert (int, optional): Index of vmr perturbation for nonlinear
              species to plot.
          pressures (ndarray(int), optional): Pressure levels to plot. If not
              given, up to 6 pressure levels are selected.
          cols (int, optional): Species to plot per row.
  
      Returns:
          matplotlib.figure.Figure, ndarray(AxesSubplot):
              Matplotlib Figure and Axes objects.
  
      Examples:
  
      .. plot::
          :include-source:
  
          from os.path import join, dirname
          import matplotlib.pyplot as plt
          import typhon as ty
  
          lookup_file = join(dirname(ty.__file__), 'tests', 'plots',
                             'reference', 'abs_lookup_small.xml')
          fig, ax = ty.plots.plot_arts_lookup(ty.arts.xml.load(lookup_file))
  
          fig.suptitle('Lookup table opacities')
          fig.subplots_adjust(top=0.88)
          plt.show()
  
      .. plot::
          :include-source:
  
          from os.path import join, dirname
          import matplotlib.pyplot as plt
          import typhon as ty
  
          lookup_file = join(dirname(ty.__file__), 'tests', 'plots',
                             'reference', 'abs_lookup_small.xml')
          fig, ax = ty.plots.plot_arts_lookup(ty.arts.xml.load(lookup_file),
                                              opacity=False)
  
          fig.suptitle('Lookup table absorption cross sections [m$^2$]')
          fig.subplots_adjust(top=0.88)
          plt.show()
  
      """
      rows, cols, fig, ax = _setup_lookup_figure(lookup, cols)
  
      if opacity:
          lookup_opacity = calc_opacity_from_lookup(lookup, z, g, r, tpert)
  
      for cax, species in zip_longest(ax.flatten() if len(ax.shape) == 2
                                      else ax.reshape(ax.size, 1),
                                      lookup.speciestags):
          if species is None:
              cax.axis('off')
              continue
  
          if opacity:
              plot_lookup_opacity(lookup, lookup_opacity, vmrpert=vmrpert,
                                  oneline=True, total=True, species=[species],
                                  ax=cax)
          else:
              psize = lookup.pressuregrid.size
              if pressures is not None:
                  ipressures = [np.abs(lookup.pressuregrid - p).argmin() for p
                                in eval(pressures)]
              else:
                  ipressures = (lookup.pressuregrid.size - 1 - (
                      range(psize) if psize <= 5
                      else np.linspace(0, lookup.pressuregrid.size,
                                       num=6,
                                       endpoint=False,
                                       dtype=int)))
              plot_lookup_xsec(lookup, ipressures, species=[species], ax=cax,
                               tpert=tpert, vmrpert=vmrpert)
  
      if opacity:
          _add_opacity_legend(ax[-1, 0])
      else:
          _add_xsec_legend(lookup, ipressures, ax[-1, 0])
  
      for cax in ax[-2, :]:
          cax.set_xlabel('Frequency [GHz]', fontsize='xx-small')
  
      return fig, ax
  