#
# Changes
#
# 2022-09, Arjo Segers
#   Trap non-existing or empty listing files when creating catalogue of converted data.
#
# 2022-09, Arjo Segers
#   Specify output file template for catalogue figures rather than output directory.
#   Support steps over year when plotting series of gridded observations.
#   Define full output filename in settings instead of path only.
#
# 2022-02, Arjo Segers
#   Added option to specify which layer of gridded fields should be plotted.
#
# 2022-10, Arjo Segers
#   Adhoc fix to trap nans in lon/lat arrays in S5p/glyoxal data.
#
# 2023, Lewis Blake, Arjo Segers
#   Formatted using "black".
#
# 2024-01, Arjo Segers
#   Use new plotting routines based on cartopy.
#

########################################################################
###
### help
###
########################################################################

"""
.. _cso-catalogue:

************************
``cso_catalogue`` module
************************

The :py:mod:`cso_catalogue` module provides classes to create
plots of CSO data.


Class hierchy
=============

The classes and are defined according to the following hierchy:


* :py:class:`.UtopyaRc`

  * :py:class:`.CSO_CatalogueBase`

    * :py:class:`.CSO_Catalogue`
    * :py:class:`.CSO_SimCatalogue`
    * :py:class:`.CSO_GriddedCatalogue`



Classes
=======


"""


########################################################################
###
### modules
###
########################################################################

# modules:
import logging

# tools:
import utopya


########################################################################
###
### catalogue base
###
########################################################################


class CSO_CatalogueBase(utopya.UtopyaRc):

    """
    Base class for catalogues.

    """

    def _ConvertUnits(self, values, attrs, vunits):

        """
        Convert units of values.

        Arguments:

        * ``values` : input data
        * ``attrs`` : attributes dictionairy, should contain at least a value for ``units``,
            but eventually also conversion factors as provided with original data
        * ``vunits`` : target units, ``'None'`` if units do not need to be changed

        Return:

        * ``values``
        * ``vunits`` : current units of ``values``
        """

        # modules:
        import numpy

        # tools:
        import cso_constants

        # no target units defined?
        if (vunits is None) or (vunits == "None"):

            # unchanged:
            vunits = attrs["units"]

        # different units?
        elif attrs["units"] != vunits:

            # which conversion?
            conversion = "%s -> %s" % (attrs["units"], vunits)
            # switch:
            if conversion in ["1e-9 -> ppb"]:

                # just other notation ...
                pass

            # error covariance, convert to std.dev.:
            elif conversion in ["(1e-9)**2 -> ppb"]:

                # other notation, only apply sqrt:
                values = numpy.sqrt(values)

            # pressures:
            elif conversion == "Pa -> hPa":

                # apply factor:
                values = values / 1e2

            # vertical column density:
            elif conversion == "mol m-2 -> umol/m2":

                # apply factor:
                values = values * 1e6

            # vertical column density:
            elif conversion == "mol m-2 -> mmol/m2":

                # apply factor:
                values = values * 1e3

            # vertical column density:
            elif conversion == "mol m-2 -> 1e15 mlc/cm2":

                # apply factor:
                #        mol/m2        mlc/mol         m2/cm2
                values = values * cso_constants.Avog * 1e-4 / 1e15  # 1e15 mlc/cm2

            # vertical column density error covariance:
            elif conversion == "(mol m-2)**2 -> 1e15 mlc/cm2":

                # take square root and apply factor:
                #              mol/m2              mlc/mol         m2/cm2
                values = numpy.sqrt(values) * cso_constants.Avog * 1e-4 / 1e15  # 1e15 mlc/cm2

            # vertical column density error covariance:
            elif conversion == "(mol m-2)**2 -> umol/m2":

                # take square root and apply factor:
                values = numpy.sqrt(values) * 1e6

            else:
                logging.error('unsupported conversion "%s"' % conversion)
                raise Exception
            # endif

        # endif

        # ok
        return values, vunits

    # enddef _ConvertUnits


# endclass CSO_CatalogueBase


########################################################################
###
### catalogue
###
########################################################################


class CSO_Catalogue(CSO_CatalogueBase):

    """
    CSO task to create catalogue of satellite data figures.

    .. figure:: figs/NO2/CSO_NO2_catalogue.png
       :scale: 50 %
       :align: center
       :alt: image catalogue

    The catalogue is created in an output directory specified with::

        <rcbase>.output.dir     :  /Scratch/CSO/catalogue

    A time range for which images should be created is specified with::

      ! specifiy timerange:
      <rcbase>.timerange.start  :  2012-01-01 00:00
      <rcbase>.timerange.end    :  2012-12-31 23:59

    The time range is used to select oribts from a listing file created
    during conversion from raw data::

      <rcbase>.input.listing        :  /Scratch/CSO/S5p/RPRO/NO2/Europe/listing.csv

    Specify a list of variables to be plotted::

      <rcbase>.vars                   :  vcd qa_flag

    Each ``var`` should be a variable in the orbit file,
    unless a different variable name is specified with::

      <rcbase>.var.vcd.source             :  vcd_trop

    Optionally specify target units that are  different from the input::

        ! convert units:
        <rcbase>.var.vcd.units             :  1e15 mlc/cm2

    The value range of the colorbar could be tunes using (default limits are based on data values)::

        <rcbase>.var.vcd.vmin              :   0.0
        <rcbase>.var.vcd.vmax              :  10.0

    The colors in the colorbar could be changed using::

        <rcbase>.var.qa_flag.colors        :   ['red','yellow','green']

    The label below the colorbar will by default show the variable name,
    unless a ``long_name`` is defined::

        <rcbase>.var.vcd.long_name         :   retrieved vertical column density

    The created images will have filenames including subdirectories for year/month/day::

        2018/06/01/S5p_RPRO_NO2_03278__vcd.png
                   S5p_RPRO_NO2_03278__qa_value.png

    Enable the following flag to re-create existing files,
    by default only non-existing files are created::

      <rcbase>.renew                  :  False

    The figures size could be specified to fine tune the shape::

        ! figure size (inches), default is (8,6):
        <rcbase>.figsize             :  (6,6)

    Specify the domain of the map, projection is regular longitude/latitude:

        ! map domain (west east south north):
        <rcbase>.domain       :  -30 45 35 75

    """

    def __init__(self, rcfile, rcbase="", env={}, indent=""):

        """
        Create catalogue of CSO figures.
        """

        # modules:
        import os
        import pandas
        import numpy
        import matplotlib
        import matplotlib.pyplot as plt

        # tools:
        import cso_file
        import cso_plot

        # setup graphics back-end:
        matplotlib.use("Agg")

        # init base object:
        CSO_CatalogueBase.__init__(self, rcfile, rcbase=rcbase, env=env)

        # info ...
        logging.info(indent + "")
        logging.info(indent + "** plot CSO data")
        logging.info(indent + "")

        # time range:
        t1 = self.GetSetting("timerange.start", totype="datetime")
        t2 = self.GetSetting("timerange.end", totype="datetime")
        # info ...
        tfmt = "%Y-%m-%d %H:%M"
        logging.info(indent + "timerange: [%s,%s]" % (t1.strftime(tfmt), t2.strftime(tfmt)))

        # target files:
        figfile_template = self.GetSetting("output.file")

        # listing file:
        lstfile = self.GetSetting("input.listing")
        # info ...
        logging.info("listing file:")
        logging.info("  %s" % lstfile)
        # should exist ...
        if not os.path.isfile(lstfile):
            logging.error("file not found: %s" % lstfile)
            raise Exception
        # endif
        # read listing:
        lst = cso_file.CSO_Listing(lstfile, indent=indent)
        # subset of files in timerange:
        df = lst.Select(tr=(t1, t2))
        # number of records:
        nrec = len(df)

        # info ...
        logging.info("  records within timerange: %i" % nrec)

        # any data?
        if nrec > 0:

            # variables to be plotted:
            varnames = self.GetSetting("vars").split()

            # renew existing figures?
            renew = self.GetSetting("renew", totype="bool")

            # plot domain:
            domain = list(map(float, self.GetSetting("domain").split()))

            # mid of domain, used for local time:
            domain_lon_mid = 0.5 * (domain[0] + domain[1])
            # time zone shift based on bands of 360/24=15 deg
            tzone_dhour = int(numpy.round(domain_lon_mid / 15.0))
            # increment:
            tzone_dt = pandas.Timedelta(value=tzone_dhour, unit="hour")
            # annote:
            if tzone_dhour < 0:
                tzone_label = "UTC%i" % tzone_dhour
            elif tzone_dhour > 0:
                tzone_label = "UTC+%i" % tzone_dhour
            else:
                tzone_label = "UTC"
            # endif

            # figure size:
            figsize = eval(self.GetSetting("figsize"))

            # no-data color:
            color_nan = self.GetSetting("color_nan", default="0.80")
            # extra map properties:
            bmp_kwargs = self.GetSetting("map", totype="dict", default=dict())

            # storage for csv files with orbits per subdir:
            orbits = {}

            # base path, filenames are relative to listing file:
            inputdir = os.path.dirname(lstfile)

            # info ..
            logging.info(indent + "loop over orbit files ...")
            # loop:
            for fname in df.index:
                # info ..
                logging.info(indent + "  %s ..." % fname)
                # full path:
                filename = os.path.join(inputdir, fname)

                # orbit key:
                orbit = df.at[fname, "orbit"]

                ## testing ...
                # if '1200' not in filename : continue

                # time range:
                tr1 = df.at[fname, "start_time"]
                tr2 = df.at[fname, "end_time"]
                # mid time:
                taver = tr1 + 0.5 * (tr2 - tr1)

                # subdir:
                figdir = os.path.dirname(taver.strftime(figfile_template))
                ## create if necessary:
                # if not os.path.isdir( figdir ) : os.makedirs( figdir )

                # create storage for orbit lists:
                if figdir not in orbits.keys():
                    orbits[figdir] = []
                # store current:
                orbits[figdir].append(orbit)

                ## basename of file without extension:
                # bname,ext = os.path.splitext( os.path.basename(filename) )

                # no orbit file read yet ..
                orb = None

                # loop over variables ...
                for varname in varnames:

                    # target file:
                    # figfile = '%s/%s__%s.png' % (figdir,bname,varname)
                    figfile = taver.strftime(figfile_template)
                    figfile = figfile.replace("%{orbit}", orbit)
                    figfile = figfile.replace("%{varname}", varname)

                    # renew?
                    if (not os.path.isfile(figfile)) or renew:
                        # info ..
                        logging.info(indent + "    create %s ..." % figfile)

                        # read?
                        if orb is None:
                            # info ...
                            logging.info(indent + "      read ...")
                            # init storage and read:
                            orb = cso_file.CSO_File(filename=filename)
                        # endif

                        # display as local standard time:
                        taver_lst = taver + tzone_dt
                        # annote:
                        title = "%s (%s) %s" % (
                            taver_lst.strftime("%Y-%m-%d %H:%M"),
                            tzone_label,
                            orbit,
                        )

                        # settings for this variable:
                        vkey = "var.%s" % varname
                        # plot type:
                        ptype = self.GetSetting(vkey + ".type", default="map")
                        # input variable:
                        vname = self.GetSetting(vkey + ".source", default=varname)
                        # target units:
                        vunits = self.GetSetting(vkey + ".units", default="None")
                        # long name used in labels:
                        long_name = self.GetSetting(vkey + ".long_name", default=varname)

                        # switch:
                        if ptype == "map":

                            # style:
                            vmin = eval(self.GetSetting(vkey + ".vmin", default="None"))
                            vmax = eval(self.GetSetting(vkey + ".vmax", default="None"))
                            colors = eval(self.GetSetting(vkey + ".colors", default="None"))

                            # extract corner grids and values:
                            xx, yy, values, attrs = orb.GetTrack(vname)

                            # adhoc, some chocho files have nans in lon/lat bound arrays,
                            # these are still present because the scan line has some valid values ...
                            if numpy.any(numpy.isnan(xx)):
                                logging.warning("    found nans in track corners, skip ...")
                                break  # leave loop over variables
                            # endif

                            # squeeze
                            if len(values.shape) > 2:
                                # remove lenth-1 dim "retr" from dims: (scan,pixel,retr)
                                if len(values.shape) == 3:
                                    if values.shape[2] == 1:
                                        values = values[:, :, 0]
                                    else:
                                        logging.error(
                                            "could not squeeze array of shape %s to 2D field"
                                            % str(values.shape)
                                        )
                                        raise Exception
                                    # endif
                                elif len(values.shape) == 4:
                                    if (values.shape[2] == 1) and (values.shape[3] == 1):
                                        values = values[:, :, 0, 0]
                                    else:
                                        logging.error(
                                            "could not squeeze array of shape %s to 2D field"
                                            % str(values.shape)
                                        )
                                        raise Exception
                                    # endif
                                else:
                                    logging.error(
                                        "could not squeeze array of shape %s to 2D field"
                                        % str(values.shape)
                                    )
                                    raise Exception
                                # endif
                            # endif

                            # convert if necessary:
                            values, units = self._ConvertUnits(values, attrs, vunits)

                            # annote:
                            label = "%s [%s]" % (long_name, units)

                            # map properties:
                            bmp = dict(countries=True)
                            bmp.update(bmp_kwargs)

                            # create map figure:
                            fig = cso_plot.QuickPat(
                                values,
                                xx=xx,
                                yy=yy,
                                vmin=vmin,
                                vmax=vmax,
                                cmap=dict(colors=colors),  # color_bad=color_nan),
                                bmp=bmp,
                                domain=domain,
                                cbar=dict(label=label),
                                title=title,
                                figsize=figsize,
                            )
                            # save:
                            fig.Export(figfile)

                            # close:
                            fig.Close()

                        else:
                            logging.error('unsupported dimensions "%s"' % (str(dims)))
                            raise Exception
                        # endif  # dim types

                        ## testing ...
                        # logging.warning( 'break after first variable' )
                        # break

                    # endif # renew

                # endfor # variables

                ## testing ...
                # logging.warning( 'break after first orbit' )
                # break

            # endfor # orbit files

            # write orbit lists:
            for figdir in orbits.keys():
                # new:
                df = pandas.DataFrame({"orbit": orbits[figdir]})
                # write:
                df.to_csv(os.path.join(figdir, "orbits.csv"), index=False, sep=";")
            # endfor

        # endif # nrec > 0

    # enddef __init__


# endclass CSO_Catalogue


########################################################################
###
### catalogue of simulations
###
########################################################################


class CSO_SimCatalogue(CSO_CatalogueBase):

    """
    CSO task to create catalogue of satellite observations and simulations.

    .. figure:: figs/NO2/CSO_NO2_sim-catalogue.png
       :scale: 50 %
       :align: center
       :alt: image catalogue

    The catalogue is created in an output directory specified with::

        <rcbase>.output.dir     :  /Scratch/CSO/catalogue

    A time range for which images should be created is specified with::

      ! specifiy timerange:
      <rcbase>.timerange.start  :  2012-01-01 00:00
      <rcbase>.timerange.end    :  2012-12-31 23:59
      ! step is one of: hour | day | month
      <rcbase>.timerange.step   :  hour

    The time range is used to scan for output files from the satellite observation operator.
    Both a *data* file with for example footprints and observations,
    as well as a *state* file with simulations is needed::

        <rcbase>.input.data.file        :  /Scratch/model/output/CSO_output_%Y%m%d_%H%M_data.nc
        <rcbase>.input.state.file       :  /Scratch/model/output/CSO_output_%Y%m%d_%H%M_state.nc

    Specify a list of variables to be plotted, for example the retrieved and simulated column::

      <rcbase>.vars                   :  yr ys

    By default the ``var`` is expected to be a variable in the *data* file;
    use the following source specifications to explicitly define the origin::

      <rcbase>.var.yr.source                 ::  data:vcd
      <rcbase>.var.ys.source                 ::  state:vcd

    Optionally specify target units that are  different from the input::

        ! convert units:
        <rcbase>.var.vcd.units             :  1e15 mlc/cm2

    The value range of the colorbar could be tunes using (default limits are based on data values)::

        <rcbase>.var.vcd.vmin              :   0.0
        <rcbase>.var.vcd.vmax              :  10.0

    The colors in the colorbar could be changed using::

        <rcbase>.var.qa_flag.colors        :   ['red','yellow','green']

    The label below the colorbar will by default show the variable name,
    unless a ``long_name`` is defined::

        <rcbase>.var.vcd.long_name         :   retrieved vertical column density

    The name of the created image files is read from::

        ! target files, time tempates are replaced:
        <rcbase>.output.file            :  %Y/%m/%d/S5p_RPRO_NO2_%Y%m%d_%H%M_%{var}.png

        2018/06/01/S5p_RPRO_NO2_03278__vcd.png
                   S5p_RPRO_NO2_03278__qa_value.png

    Enable the following flag to re-create existing files,
    by default only non-existing files are created::

      <rcbase>.renew                  :  False

    The figures size could be specified to fine tune the shape::

        ! figure size (inches), default is (8,6):
        <rcbase>.figsize             :  (6,6)

    Specify the domain of the map, projection is regular longitude/latitude:

        ! map domain (west east south north):
        <rcbase>.domain       :  -30 45 35 75

    """

    def __init__(self, rcfile, rcbase="", env={}, indent=""):

        """
        Create catalogue of CSO figures.
        """

        # modules:
        import os
        import pandas
        import numpy
        import datetime
        import calendar
        import matplotlib
        import matplotlib.pyplot as plt

        # tools:
        import cso_file
        import cso_plot

        # setup graphics back-end:
        matplotlib.use("Agg")

        # init base object:
        CSO_CatalogueBase.__init__(self, rcfile, rcbase=rcbase, env=env)

        # info ...
        logging.info(indent + "")
        logging.info(indent + "** plot CSO simulations")
        logging.info(indent + "")

        # time range:
        t1 = self.GetSetting("timerange.start", totype="datetime")
        t2 = self.GetSetting("timerange.end", totype="datetime")
        tstep = self.GetSetting("timerange.step")
        # info ...
        tfmt = "%Y-%m-%d %H:%M"
        logging.info(
            indent + "timerange: [%s,%s] by %s" % (t1.strftime(tfmt), t2.strftime(tfmt), tstep)
        )

        # target file template:
        figfile_template = self.GetSetting("output.file")

        # template for data files:
        datafile_template = self.GetSetting("input.data.file")
        statefile_template = self.GetSetting("input.state.file")

        # variables to be plotted:
        varnames = self.GetSetting("vars").split()

        # renew existing figures?
        renew = self.GetSetting("renew", totype="bool")

        # plot domain:
        domain = list(map(float, self.GetSetting("domain").split()))

        # figure size:
        figsize = eval(self.GetSetting("figsize"))

        # no-data color:
        color_nan = self.GetSetting("color_nan", default="0.80")
        # extra map properties:
        bmp_kwargs = self.GetSetting("map", totype="dict", default=dict())

        # info ..
        logging.info(indent + "time loop ...")
        # time loop:
        t = t1
        while t <= t2:
            # info ..
            logging.info(indent + "  %s ..." % t.strftime("%Y-%m-%d %H:%M"))

            # timestep:
            if tstep in ["hour", "hourly"]:
                long_timestamp = t.strftime("%Y-%m-%d %H:%M")
                dt = datetime.timedelta(0, 3600)
            elif tstep in ["day", "daily"]:
                long_timestamp = t.strftime("%Y-%m-%d")
                dt = datetime.timedelta(1)
            elif tstep in ["month", "monthly"]:
                t = datetime.datetime(t.year, t.month, 1, 0, 0)
                long_timestamp = t.strftime("%Y-%m")
                weekday, nday = calendar.monthrange(t.year, t.month)
                dt = datetime.timedelta(nday)
            else:
                logging.error('unsupported time step "%s"' % tstep)
                raise Exception
            # endif

            # input files:
            datafile = t.strftime(datafile_template)
            statefile = t.strftime(statefile_template)

            # state file present?
            if os.path.isfile(statefile):

                # info ...
                logging.info(indent + "    found state file: %s" % statefile)

                # check ...
                if not os.path.isfile(datafile):
                    logging.error("data file not found: %s" % datafile)
                    raise Exception
                # endif

                # no data file read yet ..
                sdata = None

                # loop over variables ...
                for varname in varnames:

                    # target file:
                    figfile = t.strftime(figfile_template)
                    figfile = figfile.replace("%{var}", varname)

                    # renew?
                    if (not os.path.isfile(figfile)) or renew:
                        # info ..
                        logging.info(indent + "    create %s ..." % figfile)

                        # read?
                        if sdata is None:
                            # info ...
                            logging.info(indent + "      read ...")
                            # init storage and read:
                            sdata = cso_file.CSO_File(filename=datafile)
                            # init storage and read:
                            sstate = cso_file.CSO_File(filename=statefile)
                        # endif

                        # annote:
                        title = long_timestamp

                        # settings for this variable:
                        vkey = "var.%s" % varname
                        # originating variable:
                        vsource = self.GetSetting(vkey + ".source", default="data:" + varname)
                        # target units:
                        vunits = self.GetSetting(vkey + ".units", default="None")
                        # plot type:
                        ptype = self.GetSetting(vkey + ".type", default="map")
                        # long name used in labels:
                        long_name = self.GetSetting(vkey + ".long_name", default=varname)

                        # switch:
                        if ptype == "map":

                            # style:
                            vmin = eval(self.GetSetting(vkey + ".vmin", default="None"))
                            vmax = eval(self.GetSetting(vkey + ".vmax", default="None"))
                            colors = eval(self.GetSetting(vkey + ".colors", default="None"))

                            # variable source:
                            #    [data:]vname
                            #    state:vname
                            if ":" in vsource:
                                vfile, vname = vsource.split(":")
                            else:
                                vfile, vname = "data", vsource
                            # endif
                            # extract corner grids and values:
                            if vfile == "state":
                                # check ...
                                if vname not in sstate.ds.keys():
                                    logging.error(
                                        "variable `%s` not found in state from file: %s"
                                        % (vname, statefile)
                                    )
                                    raise Exception
                                # endif
                                # extract corner grids and values:
                                xx, yy, values, attrs = sdata.GetTrack(sstate.ds[vname])
                            elif vfile == "data":
                                # extract corner grids and values from data file:
                                xx, yy, values, attrs = sdata.GetTrack(vname)
                            else:
                                logging.error(
                                    'unsupported type keyword "%s" in source description "%s" for variable "%s"'
                                    % (vfile, vsource, varname)
                                )
                                raise Exception
                            # endif

                            # adhoc, some chocho files have nans in lon/lat bound arrays,
                            # these are still present because the scan line has some valid values ...
                            if numpy.any(numpy.isnan(xx)):
                                logging.warning("    found nans in track corners, skip ...")
                                break  # leave loop over variables
                            # endif

                            # squeeze
                            shp = values.shape
                            if len(shp) != 2:
                                # remove 1-dims:
                                if (len(shp) == 3) and (shp[2] == 1):
                                    values = values[:, :, 0]
                                else:
                                    logging.error(
                                        'could not squeeze "%s" %s to 2D track' % (vname, str(shp))
                                    )
                                    raise Exception
                                # endif
                            # endif

                            # convert if necessary:
                            values, units = self._ConvertUnits(values, attrs, vunits)

                            # annote:
                            label = "%s [%s]" % (long_name, units)

                            # map properties:
                            bmp = dict(countries=True)
                            bmp.update(bmp_kwargs)

                            # create map figure:
                            fig = cso_plot.QuickPat(
                                values,
                                xx=xx,
                                yy=yy,
                                vmin=vmin,
                                vmax=vmax,
                                cmap=dict(colors=colors, color_bad=color_nan),
                                bmp=bmp,
                                domain=domain,
                                cbar=dict(label=label),
                                title=title,
                                figsize=figsize,
                            )
                            # save:
                            fig.Export(figfile)

                            # close:
                            fig.Close()

                        else:
                            logging.error('unsupported plot type "%s"' % (ptype))
                            raise Exception
                        # endif  # ptype

                    else:
                        # info ..
                        logging.info(indent + "    keep  %s ..." % figfile)
                    # endif # renew

                # endfor # variables

                ## testing ...
                # break

            else:
                logging.info(indent + "    no    state file: %s" % statefile)
            # endif # state file present

            # next:
            t = t + dt

        # endfor # time loop

    # enddef __init__


# endclass CSO_SimCatalogue


########################################################################
###
### catalogue of gridded averages
###
########################################################################


class CSO_GriddedCatalogue(CSO_CatalogueBase):

    """
    CSO task to create catalogue of gridded averages of retrievals and simulations.

    .. figure:: figs/NO2/CSO_NO2_gridded-catalogue.png
       :scale: 50 %
       :align: center
       :alt: image catalogue

    The catalogue is created in an output directory specified with::

        <rcbase>.output.dir     :  /Scratch/CSO/catalogue

    A time range for which images should be created is specified with::

      ! specifiy timerange:
      <rcbase>.timerange.start  :  2012-01-01 00:00
      <rcbase>.timerange.end    :  2012-12-31 23:59
      ! step is one of: hour | day | month
      <rcbase>.timerange.step   :  hour

    The time range is used to scan for output files from the :py:class:`CSO_GriddedAverage` class::

        <rcbase>.input.file        :  /Scratch/model/output/CSO_output_%Y%m%d_%H%M_gridded.nc

    Specify a list of variables to be plotted, for example the retrieved and simulated column::

      <rcbase>.vars                   :  yr ys

    For each variable, define the name of the source variable::

      <rcbase>.var.yr.source               :  yr
      <rcbase>.var.yr.source               :  ys

    Optionally specify a source layer for profiles, default is first layer::

      <rcbase>.var.yr.source_layer         :  0

    Optionally specify target units that are  different from the input::

        ! convert units:
        <rcbase>.var.yr.units              :  1e15 mlc/cm2

    The value range of the colorbar could be tunes using (default limits are based on data values)::

        <rcbase>.var.vcd.vmin              :   0.0
        <rcbase>.var.vcd.vmax              :  10.0

    The colors in the colorbar could be changed using::

        <rcbase>.var.qa_flag.colors        :   ['red','yellow','green']

    The label below the colorbar will by default show the variable name,
    unless a ``long_name`` is defined::

        <rcbase>.var.vcd.long_name         :   retrieved vertical column density

    The name of the created image files is read from::

        ! target files, time tempates are replaced:
        <rcbase>.output.file            :  %Y/%m/%d/S5p_RPRO_NO2_%Y%m%d_%H%M_gridded_%{var}.png

    which will give::

        2020/01/01/CSO_output_20200101_0100_gridded_yr.png
                   CSO_output_20200101_0100_gridded_ys.png
                                                    :

    Enable the following flag to re-create existing files,
    by default only non-existing files are created::

        <rcbase>.renew                  :  False

    The figures size could be specified to fine tune the shape::

        ! figure size (inches), default is (8,6):
        <rcbase>.figsize             :  (6,6)

    """

    def __init__(self, rcfile, rcbase="", env={}, indent=""):

        """
        Create catalogue of CSO figures.
        """

        # modules:
        import os
        import pandas
        import numpy
        import datetime
        import calendar
        import matplotlib
        import matplotlib.pyplot as plt

        # tools:
        import cso_file
        import cso_plot

        # setup graphics back-end:
        matplotlib.use("Agg")

        # init base object:
        CSO_CatalogueBase.__init__(self, rcfile, rcbase=rcbase, env=env)

        # info ...
        logging.info(indent + "")
        logging.info(indent + "** plot CSO simulations")
        logging.info(indent + "")

        # time range:
        t1 = self.GetSetting("timerange.start", totype="datetime")
        t2 = self.GetSetting("timerange.end", totype="datetime")
        tstep = self.GetSetting("timerange.step")
        # info ...
        tfmt = "%Y-%m-%d %H:%M"
        logging.info(
            indent + "timerange: [%s,%s] by %s" % (t1.strftime(tfmt), t2.strftime(tfmt), tstep)
        )

        # target file template:
        figfile_template = self.GetSetting("output.file")

        # variables to be plotted:
        #    y ys
        varnames = self.GetSetting("vars").split()

        # renew existing figures?
        renew = self.GetSetting("renew", totype="bool")

        # figure size:
        figsize = eval(self.GetSetting("figsize"))

        # extra map properties:
        bmp_kwargs = self.GetSetting("map", totype="dict", default=dict())

        # info ..
        logging.info(indent + "time loop ...")
        # time loop:
        t = t1
        while t <= t2:
            # info ..
            logging.info(indent + "  %s ..." % t.strftime("%Y-%m-%d %H:%M"))

            # timestep:
            if tstep in ["hour", "hourly"]:
                long_timestamp = t.strftime("%Y-%m-%d %H:%M")
                dt = datetime.timedelta(0, 3600)
            elif tstep in ["day", "daily"]:
                long_timestamp = t.strftime("%Y-%m-%d")
                dt = datetime.timedelta(1)
            elif tstep in ["month", "monthly"]:
                t = datetime.datetime(t.year, t.month, 1, 0, 0)
                long_timestamp = t.strftime("%Y-%m")
                weekday, nday = calendar.monthrange(t.year, t.month)
                dt = datetime.timedelta(nday)
            elif tstep in ["year", "yearly"]:
                t = datetime.datetime(t.year, 1, 1, 0, 0)
                long_timestamp = t.strftime("%Y")
                dt = datetime.timedelta(366)
            else:
                logging.error('unsupported time step "%s"' % tstep)
                raise Exception
            # endif

            # input file:
            infile = self.GetSetting("input.file")
            infile = t.strftime(infile)

            # input file present?
            if os.path.isfile(infile):

                # info ...
                logging.info(indent + "    found input file: %s" % infile)

                # no data read yet:
                indata = None

                # loop over variables ...
                for varname in varnames:

                    # target file:
                    figfile = t.strftime(figfile_template)
                    figfile = figfile.replace("%{var}", varname)

                    # renew?
                    if (not os.path.isfile(figfile)) or renew:
                        # info ..
                        logging.info(indent + "    create %s ..." % figfile)

                        # read?
                        if indata is None:
                            # info ...
                            logging.info(indent + "      read ...")
                            # init storage and read:
                            indata = cso_file.CSO_File(filename=infile)
                        # endif

                        # annote:
                        title = long_timestamp

                        # settings for this variable:
                        vkey = "var.%s" % varname
                        # originating variable:
                        vsource = self.GetSetting(vkey + ".source")
                        # optional layer selection:
                        vsource_layer = self.GetSetting(
                            vkey + ".source_layer", totype="int", default=0
                        )
                        # target units:
                        vunits = self.GetSetting(vkey + ".units", default="None")
                        # plot type:
                        ptype = self.GetSetting(vkey + ".type", default="map")
                        # long name used in labels:
                        long_name = self.GetSetting(vkey + ".long_name", default=varname)

                        # switch:
                        if ptype == "map":

                            # style:
                            vmin = eval(self.GetSetting(vkey + ".vmin", default="None"))
                            vmax = eval(self.GetSetting(vkey + ".vmax", default="None"))
                            colors = eval(self.GetSetting(vkey + ".colors", default="None"))

                            # extract values:
                            vda = indata.ds[vsource]

                            # grid:
                            if "longitude" in vda.coords.keys():
                                xm = vda.coords["longitude"].values
                                ym = vda.coords["latitude"].values
                            else:
                                logging.error("could not guess the grid coordinates yet from:")
                                for cname in vda.coords.keys():
                                    logging.error("  %s" % cname)
                                raise Exception
                            # endif

                            # data:
                            values = vda.values.squeeze()
                            # select layer if needed ..
                            if len(values.shape) == 3:
                                # check ..
                                if (vsource_layer < 0) or (values.shape[2] < vsource_layer + 1):
                                    logging.error(
                                        "could not extract source layer %i from values with shape (%s)"
                                        % (vsource_layer, str(vvalues.shape))
                                    )
                                    raise Exception
                                # endif
                                # extract:
                                values = values[:, :, vsource_layer]
                            #
                            elif len(values.shape) != 2:
                                logging.error(
                                    "could not plot map of array with shape (%s)"
                                    % (str(values.shape))
                                )
                                raise Exception
                            # endif

                            # adhoc fix, convert to masked array ...
                            values = numpy.ma.array(data=values, mask=numpy.isnan(values))

                            # convert if necessary:
                            values, units = self._ConvertUnits(values, vda.attrs, vunits)

                            # annote:
                            label = "%s [%s]" % (long_name, units)

                            # map properties:
                            bmp = dict(countries=True)
                            bmp.update(bmp_kwargs)

                            # create map figure:
                            fig = cso_plot.QuickPat(
                                values,
                                xm=xm,
                                ym=ym,
                                vmin=vmin,
                                vmax=vmax,
                                cmap=dict(colors=colors),
                                bmp=bmp,
                                cbar=dict(label=label),
                                title=title,
                                figsize=figsize,
                            )
                            # save:
                            fig.Export(figfile)

                            # close:
                            fig.Close()

                        else:
                            logging.error('unsupported plot type "%s"' % (ptype))
                            raise Exception
                        # endif  # ptype

                    # endif # renew

                    ## testing ...
                    # break

                # endfor # variables

            else:
                logging.info(indent + "    no    input file: %s" % infile)
            # endif # state file present

            # next:
            t = t + dt

        # endfor # time loop

    # enddef __init__


# endclass CSO_GriddedCatalogue


########################################################################
###
### end
###
########################################################################
