#
# Changes
#
# 2022-09, Arjo Segers
#   Added class to create super-observations.
#   Fixed output of coordinate variables.
#
# 2023-08, Arjo Segers
#   Reformatted using 'black'.
#


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

"""
.. _cso-superobs:

***********************
``cso_superobs`` module
***********************

The :py:mod:`cso_superobs` module provides classes to create *super observations*
as average of pixels over a regular grid.
The result is stored into a :py:class:`CSO_File` object as if it is
a collection of pixels.


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

The classes and are defined according to the following hierchy:


* :py:class:`.UtopyaRc`

  * :py:class:`.CSO_SuperObs`



Classes
=======


"""


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

# modules:
import logging

# tools:
import utopya


########################################################################
###
### super observations
###
########################################################################


class CSO_SuperObs(utopya.UtopyaRc):

    """
    Average pixels over regular grid using rcfile settings.

    The value assigned to a cell is a weighted average of the pixel values,
    where the weight is relative to the area of the overlapping part.
    Since a cell might be only partly covered by pixels, the normalization is done
    using the total overlapping area:
    
    .. math::
        x(i_k,j_k) ~=~ \left(\ \sum\limits_{p\in P_k} y_p\ w_{p,k}\ \\right)\ /\ \sum\limits_{p\in P_k}\ w_{p,k}
        
    where:
    
    * :math:`i_k,j_k` are the indices of grid cell :math:`k`;
    * :math:`P_k` is the set of pixels that overlap with cell :math:`k`;
    * :math:`y_p` is the data value of pixel :math:`p`;
    * :math:`w_p` is the footprint area [m\ :sup:`2`] of pixel :math:`p`
    * :math:`w_{p,k}` is the area [m\ :sup:`2`] of pixel :math:`p` that overlaps with the cell :math:`k`.
    
    The overlapping area is computed using the :py:meth:`LonLatPolygonCentroids <cso_mapping.LonLatPolygonCentroids>` method.
    This fractions a footprint area into a large number of triangles, and returns
    for each triangle the centroid and the area.
    The centroids are collected per grid cell, and sum over the associated triangle area's is
    then used as measure for the overlapping area.
    
    The target grid is regular longitude/latitude and defined using the following settings::
        
        ! domain definition:
        <rcbase>.west              :  40.0
        <rcbase>.east              :  60.0
        <rcbase>.south             :  20.0
        <rcbase>.north             :  35.0
        ! grid resolution:
        <rcbase>.dlon              :  0.1
        <rcbase>.dlat              :  0.1

    The level recursive splitting of footprint into triangles and assignment of centroids
    is defined with::

        ! for 4-corner footprints, number of centroids is:
        !  1 (levels=0), 4 (1), 8 (2), 16 (3), 64 (5), 256 (7)
        <rcbase>.mapping.levels         :  7

    A gridded average is created for each orbit within a time range.
    Specify the time range in the settings using::
    
      <rcbase>.timerange.start        :  2018-06-01 00:00
      <rcbase>.timerange.end          :  2018-06-03 23:59

    Input data could originate from different sources, for example a *data* file with footprints
    and retrieval, and a *state* file with simulations by a model.
    Specify a list with keywords describing the sources; the first source should be a file
    that contains the footprints::
    
        ! keywords for source files:
        <rcbase>.sources                :  data

    For each source, specify the input file(s) as a filename pattern::
    
        ! input file or filename pattern for each source type:
        <rcbase>.source.data         :  /scratch/CSO-data/S5p/OFFL/CH4/Middle-East/%Y/%m/S5p_OFFL_CH4_*.nc

    
    The input could be filtered by applying selections, for example on quality value;
    see :py:meth:`CSO_File.SelectPixels` method for supported filter types::

        ! keywords for filters:
        <rcbase>.filters                :  quality

        ! minimum quality value required:
        <rcbase>.filter.quality.var       :  qa_value
        <rcbase>.filter.quality.type      :  min
        <rcbase>.filter.quality.min       :  0.8
        <rcbase>.filter.quality.units     :  1

    Specify the name of the target file; could be useful to include a description of the target grid and the applied filters in the path::

        ! target file, might contain templates:
        !   %Y,%m,etc     : time values
        !   %{basename}   : basename (without extension) of first source file
        <rcbase>.output.file            :   /scratch/CSO-data/S5p/OFFL/CH4/Middle-East_r01x01_qa08/%Y/%m/%{basename}_gridded.nc

    By default all variables are written with ``zlib`` compression enabled to save storage size.
    The default compression level is 1 (out of 9); set the following flag to a higher level
    to have stronger compression (on expense of computation time), or set to 0 to disable compression::
    
        ! zlib compression level (default 1, 0 for no compression):
        <rcbase>.output.complevel           :  1

    Existing files are replaced if the following flag is set::
    
        ! renew existing files?
        <rcbase>.renew                  :  True

    """

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

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

        # modules:
        import os
        import numpy
        import xarray
        import datetime
        import glob
        import pandas

        # tools:
        import cso_file
        import cso_mapping

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

        # info ...
        logging.info(indent + "")
        logging.info(indent + "** Superobs from 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)))

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

        # variable data type:
        dtype = "f4"
        fill_value = numpy.finfo(dtype).max

        # regular grid:
        # ~ domain:
        west = self.GetSetting("grid.west", "float")
        east = self.GetSetting("grid.east", "float")
        south = self.GetSetting("grid.south", "float")
        north = self.GetSetting("grid.north", "float")
        # ~ resolution:
        dlon = self.GetSetting("grid.dlon", "float")
        dlat = self.GetSetting("grid.dlat", "float")
        # ~ shape:
        nlon = int(numpy.round((east - west) / dlon))
        nlat = int(numpy.round((north - south) / dlat))
        # info ...
        logging.info(indent + "define regular grid ...")
        logging.info(
            indent + "  domain     :  [%7.2f,%7.2f] x [%7.2f,%7.2f]" % (west, east, south, north)
        )
        logging.info(indent + "  resolution :  %5.2f x %5.2f" % (dlon, dlat))
        logging.info(indent + "  resolution :  %5i x %5i" % (nlon, nlat))

        # centers and corners:
        #    3 o---o 2
        #      | * |
        #    0 o---o 1
        lons = numpy.zeros((nlat, nlon), float)
        lon_bounds = numpy.zeros((nlat, nlon, 4), float)
        for j in range(nlat):
            lons[j, :] = west + (numpy.arange(nlon) + 0.5) * dlon
            lon_bounds[j, :, 0] = west + numpy.arange(0, nlon) * dlon
            lon_bounds[j, :, 1] = west + numpy.arange(1, nlon + 1) * dlon
            lon_bounds[j, :, 2] = west + numpy.arange(1, nlon + 1) * dlon
            lon_bounds[j, :, 3] = west + numpy.arange(0, nlon) * dlon
        # endfor
        lats = numpy.zeros((nlat, nlon), float)
        lat_bounds = numpy.zeros((nlat, nlon, 4), float)
        for i in range(nlon):
            lats[:, i] = south + (numpy.arange(nlat) + 0.5) * dlat
            lat_bounds[:, i, 0] = south + numpy.arange(0, nlat) * dlon
            lat_bounds[:, i, 1] = south + numpy.arange(0, nlat) * dlon
            lat_bounds[:, i, 2] = south + numpy.arange(1, nlat + 1) * dlon
            lat_bounds[:, i, 3] = south + numpy.arange(1, nlat + 1) * dlon
        # endfor

        # grid cell area:
        cell_area = cso_mapping.LonLatRectanglesArea(lon_bounds, lat_bounds)

        # number of regridding points:
        mapping_levels = self.GetSetting("mapping.levels", "int")
        # info ..
        logging.info("level of footprint mapping: %i" % mapping_levels)
        # how are pixels weighted?
        if mapping_levels == 0:
            # keyword for weight by number of pixels assigned to cell:
            wtype = "number"
            # threshold for number of pixels included in average:
            minimum_number = self.GetSetting("mapping.minimum_number", "int")
            # description:
            mapping_threshold = "at least %i original pixels" % minimum_number
            # info ..
            logging.info("  map by pixel center; %s" % mapping_threshold)
        else:
            # keyword for weight by footprint area covering cell:
            wtype = "area"
            # threshold for number of pixels included in average:
            minimum_coverage = self.GetSetting("mapping.minimum_coverage", "float")
            # description:
            mapping_threshold = (
                "at least %4.2f%% of area covered by original pixels" % minimum_coverage
            )
            # info ..
            logging.info("  map by pixel footprint coverage; %s" % mapping_threshold)
        # endif

        # input listing file:
        listing_file = self.GetSetting("source.listing")
        # info ..
        logging.info("input listing: %s" % listing_file)
        # read:
        listing = cso_file.CSO_Listing(listing_file, exist=True)

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

        # compression level:
        complevel = self.GetSetting("output.complevel", "int")

        # counter:
        nfile = 0

        # loop over records:
        for irec in range(len(listing)):
            # current:
            rec = listing.GetRecord(irec)

            # info ...
            logging.info("%s .." % rec["filename"])

            # filter on time range:
            if (rec["start_time"] > t2) or (rec["end_time"] < t1):
                # info ...
                logging.info("   skip; outside time range ...")
                # next record:
                continue
            # endif

            # target file:
            # ~ replace time template from orbit start time:
            outfile = rec["start_time"].strftime(outfile_template)
            # ~ loop over record elements, replace '%{orbit}' etc:
            for key in rec.keys():
                outfile = outfile.replace("%{" + key + "}", str(rec[key]))
            # endif
            # root of input file:
            root, ext = os.path.splitext(os.path.basename(rec["filename"]))
            outfile = outfile.replace("%{root}", root)

            # renew?
            if (not os.path.isfile(outfile)) or renew:

                # info ...
                logging.info("  create %s .." % outfile)
                # init history:
                history = []

                # full path:
                source_filename = os.path.join(listing.dirname, rec["filename"])
                # info ..
                logging.info("    read %s .." % source_filename)
                # read:
                source_file = cso_file.CSO_File(filename=source_filename)

                # filter:
                selected, history = source_file.SelectPixels(
                    self.rcf, rcbase, indent=indent + "  "
                )
                # info ..
                logging.info(
                    indent + "  selected %i / %i pixels .." % (selected.sum(), selected.size)
                )
                # selected indices:
                (iipix,) = numpy.where(selected)
                # number of selected pixels:
                npix = len(iipix)

                # footprint or points?
                if wtype == "number":
                    # info ...
                    logging.info(indent + "  use pixel centers as points ...")
                    # check ...
                    if ("longitude" not in source_file.ds) or ("latitude" not in source_file.ds):
                        logging.error(
                            "requested mapping using pixel centers (mapping level %i),"
                            % mapping_levels
                        )
                        logging.error('but no "longitude" and "latitude" arrays in source file:')
                        logging.error("  %s" % source_filename)
                        raise Exception
                    # endif
                    # copy:
                    xxp = source_file.ds["longitude"].values[iipix].reshape((1, npix))
                    yyp = source_file.ds["latitude"].values[iipix].reshape((1, npix))
                    wwp = numpy.ones((1, npix), dtype=float)
                    #
                elif wtype == "area":
                    # info ...
                    logging.info(indent + "  distribute points over footprints ...")
                    # check ...
                    if ("longitude_bounds" not in source_file.ds) or (
                        "latitude_bounds" not in source_file.ds
                    ):
                        logging.error(
                            "requested mapping using pixel footprints (mapping level %i),"
                            % mapping_levels
                        )
                        logging.error(
                            'but no "longitude_bounds" and "latitude_bounds" arrays in source file:'
                        )
                        logging.error("  %s" % source_filename)
                        raise Exception
                    # endif
                    # distribute points within footprint by deviding it recursively in triangles;
                    # get points and weights:
                    xxp, yyp, wwp = cso_mapping.LonLatPolygonCentroids(
                        source_file.ds["longitude_bounds"].values[iipix, :],
                        source_file.ds["latitude_bounds"].values[iipix, :],
                        maxlevel=mapping_levels,
                    )
                else:
                    logging.error('unsupported weighting type "%s"' % wtype)
                    raise Exception
                # endif

                # convert coordinates to cell indices:
                ii = numpy.array(numpy.floor((xxp - west) / dlon), int)
                jj = numpy.array(numpy.floor((yyp - south) / dlat), int)

                # store in dataframe:
                df = pandas.DataFrame()
                df["ii"] = ii.flatten()
                df["jj"] = jj.flatten()
                df["ww"] = wwp.flatten()

                # info ...
                logging.info("  collect input variables ...")
                # variables:
                vkeys = []
                # loop over source variables:
                for vkey in source_file.ds.keys():
                    # info ...
                    logging.info('    variable "%s" ...' % vkey)

                    # filter original locations ...
                    if vkey in ["longitude", "latitude", "longitude_bounds", "latitude_bounds"]:
                        logging.info("      skip original pixel locations ...")
                        continue
                    elif vkey in [
                        "track_longitude",
                        "track_latitude",
                        "track_longitude_bounds",
                        "track_latitude_bounds",
                    ]:
                        logging.info("      skip original track ...")
                        continue
                    # endif

                    # extract dims:
                    vdims = source_file.ds[vkey].dims
                    vshp = source_file.ds[vkey].shape
                    ## info ..
                    # logging.info( '      dimensions: %s' % str(vdims) )
                    # check ..
                    if vdims[0] != "pixel":
                        logging.error(
                            'expected first dim "pixel", found dims "%s" for variable "%s" in : %s'
                            % (str(vdims), vkey, source_filename)
                        )
                        raise Exception
                    # endif
                    # store:
                    vkeys.append(vkey)

                    # number of data values per pixels:
                    nv = int(numpy.product(vshp[1:]))
                    # selection, and reshape to have row of values per pixel;
                    # row should have length 1 for variables with dims: ('pixel')
                    if vkey == "time":
                        # convert time values:
                        times = pandas.to_datetime(source_file.ds[vkey].values)
                        # reference time:
                        tref = pandas.Timestamp("%i-01-01 00:00:00" % times[0].year)
                        tunits = tref.strftime("seconds since %Y-%m-%d %H:%M:%S")
                        # convert to timestep:
                        values = numpy.array((times - tref).total_seconds()).reshape((npix, nv))
                    elif len(vshp) == 1:
                        vshp = (vshp[0], 1)
                        values = source_file.ds[vkey].values[iipix].reshape((npix, nv))
                    else:
                        values = source_file.ds[vkey].values[iipix, :].reshape((npix, nv))
                    # endif
                    # loop over data values:
                    for iv in range(nv):
                        # column name:
                        ckey = "%s_%i" % (vkey, iv)
                        # store weights*values in data frame for sumation;
                        # note that '*' gives an element-wise product of wwp (np,npix) and the values (npix):
                        df[ckey] = (wwp * values[:, iv]).flatten()
                    # endfor

                # endfor # source variables

                # info ...
                logging.info(indent + "  average %i pixels over grid cells ..." % len(iipix))

                # group per target grid cell, sum over weights:
                df2 = df.groupby(["jj", "ii"]).sum()
                # extract target indices within grid:
                ii2 = df2.index.get_level_values("ii").values
                jj2 = df2.index.get_level_values("jj").values

                # mask for super pixels with enough contributions(by number or area)
                # of original pixels:
                if wtype == "number":
                    # minimum number:
                    enough_input = df2["ww"].values >= minimum_number
                #
                elif wtype == "area":
                    # minimum area fraction:
                    enough_input = df2["ww"].values / cell_area[jj2, ii2] >= minimum_coverage
                #
                else:
                    logging.error('unsupported weighting type "%s"' % wtype)
                    raise Exception
                # endif

                # within grid, and enough pixels (number or area) that contribute to average:
                (kk,) = numpy.where(
                    (ii2 >= 0) & (ii2 <= nlon - 1) & (jj2 >= 0) & (jj2 <= nlat - 1) & enough_input
                )

                # number of super observations:
                nsup = len(kk)
                # check ...
                if nsup == 0:
                    logging.info(
                        "  no super observations could be formed matching coverage threshold; skip orbit ..."
                    )
                    continue
                # endif
                # info ...
                logging.info(indent + "    number of super obs: %i" % nsup)

                # init output file:
                outf = cso_file.CSO_File()

                # pixel number is compressed index in grid:
                pixel_indx = jj2[kk] * nlon + ii2[kk]
                # super pixel coordinate:
                pixel_coords = {}
                pixel_coords["pixel"] = xarray.DataArray(
                    pixel_indx,
                    dims=("pixel",),
                    attrs={
                        "long_name": "pixel index",
                        "units": "1",
                        "compress": "track_scan track_pixel",
                    },
                )

                # pixel centers:
                outf.ds["longitude"] = xarray.DataArray(
                    lons[jj2[kk], ii2[kk]],
                    dims=("pixel",),
                    coords=pixel_coords,
                    attrs={
                        "standard_name": "longitude",
                        "units": "degrees_east",
                        "bounds": "longitude_bounds",
                    },
                )
                outf.ds["latitude"] = xarray.DataArray(
                    lats[jj2[kk], ii2[kk]],
                    dims=("pixel",),
                    coords=pixel_coords,
                    attrs={
                        "standard_name": "latitude",
                        "units": "degrees_north",
                        "bounds": "latitude_bounds",
                    },
                )
                # pixel corners;
                # strange: when using units 'degrees_east' and 'degrees_north',
                #   the "units" attributes for these variables disapear from the output file ...
                outf.ds["longitude_bounds"] = xarray.DataArray(
                    lon_bounds[jj2[kk], ii2[kk], :],
                    dims=("pixel", "corner"),
                    coords=pixel_coords,
                    attrs={"long_name": "longitude_bounds", "units": "degrees"},
                )
                outf.ds["latitude_bounds"] = xarray.DataArray(
                    lat_bounds[jj2[kk], ii2[kk], :],
                    dims=("pixel", "corner"),
                    coords=pixel_coords,
                    attrs={"long_name": "latitude_bounds", "units": "degrees"},
                )
                # "track" is whole grid:
                track_coords = {}
                track_coords["track_scan"] = xarray.DataArray(
                    numpy.arange(nlat),
                    dims=("track_scan",),
                    attrs={"long_name": "scan index in grid (0-based)", "units": "1"},
                )
                track_coords["track_pixel"] = xarray.DataArray(
                    numpy.arange(nlon),
                    dims=("track_pixel",),
                    attrs={"long_name": "pixel index in grid (0-based)", "units": "1"},
                )
                # centers:
                outf.ds["track_longitude"] = xarray.DataArray(
                    lons,
                    dims=("track_scan", "track_pixel"),
                    coords=track_coords,
                    attrs={"long_name": "track center longitudes", "units": "degrees_east"},
                )
                outf.ds["track_latitude"] = xarray.DataArray(
                    lats,
                    dims=("track_scan", "track_pixel"),
                    coords=track_coords,
                    attrs={"long_name": "track center latitudes", "units": "degrees_north"},
                )
                # corners:
                outf.ds["track_longitude_bounds"] = xarray.DataArray(
                    lon_bounds,
                    dims=("track_scan", "track_pixel", "corner"),
                    coords=track_coords,
                    attrs={"long_name": "track longitude bounds", "units": "degrees_east"},
                )
                outf.ds["track_latitude_bounds"] = xarray.DataArray(
                    lat_bounds,
                    dims=("track_scan", "track_pixel", "corner"),
                    coords=track_coords,
                    attrs={"long_name": "track latitude bounds", "units": "degrees_north"},
                )

                # pixel ara:
                outf.ds["area"] = xarray.DataArray(
                    cell_area[jj2[kk], ii2[kk]],
                    dims=("pixel",),
                    coords=pixel_coords,
                    attrs={"long_name": "super-pixel area", "units": "m2"},
                )

                # total pixel weight per superobs:
                sup_ww = df2["ww"].values[kk]

                # switch:
                if wtype == "number":
                    # number of averaged pixels:
                    outf.ds["pixel_number"] = xarray.DataArray(
                        sup_ww,
                        dims=("pixel"),
                        coords=pixel_coords,
                        attrs={"long_name": "number of averaged pixels", "units": "1"},
                    )
                #
                elif wtype == "area":
                    # area covered by pixels:
                    outf.ds["pixel_area"] = xarray.DataArray(
                        sup_ww,
                        dims=("pixel"),
                        coords=pixel_coords,
                        attrs={"long_name": "area covered by averaged pixels", "units": "m2"},
                    )
                #
                else:
                    logging.error('unsupported weighting type "%s"' % wtype)
                    raise Exception
                # endif

                # loop over output variables:
                for vkey in vkeys:

                    # extract shape:
                    vdims = source_file.ds[vkey].dims
                    vshp = source_file.ds[vkey].shape
                    # number of data values per pixels:
                    nv = int(numpy.product(vshp[1:]))

                    # set attributes:
                    attrs = source_file.ds[vkey].attrs
                    if "long_name" in source_file.ds[vkey].attrs:
                        attrs["long_name"] = (
                            "averaged %s" % source_file.ds[vkey].attrs["long_name"]
                        )
                    # endif

                    # adhoc: reconstruct time values ..
                    if vkey == "time":
                        # init list:
                        values = []
                        # loop over super pixels:
                        for k in kk:
                            # average time:
                            values.append(tref + pandas.Timedelta(df2[ckey].values[k], "seconds"))
                        # endfor

                        # update attributes, conversion has not always set correct long_name ...
                        attrs["long_name"] = "averaged time"
                        # attrs['units'    ] = tunits
                        # store:
                        outf.ds[vkey] = xarray.DataArray(values, dims=vdims, attrs=attrs)
                        # ensure correct output:
                        outf.ds[vkey].encoding["units"] = tunits
                        outf.ds[vkey].encoding["calendar"] = "standard"
                        outf.ds[vkey].encoding["_FillValue"] = None

                    else:

                        # storage for super-obs:
                        values = numpy.zeros((nsup, nv), dtype=dtype) + fill_value
                        # loop over data values:
                        for iv in range(nv):
                            # column name:
                            ckey = "%s_%i" % (vkey, iv)
                            # copy sum per grid cells from dataframe into field:
                            values[:, iv] = df2[ckey].values[kk]
                            # normalize with total weight (non-zero weights only):
                            values[:, iv] /= sup_ww
                        # endfor

                        # store:
                        outf.ds[vkey] = xarray.DataArray(
                            values.reshape((nsup,) + vshp[1:]),
                            dims=vdims,
                            coords=pixel_coords,
                            attrs=attrs,
                        )

                    # endif # time or values

                # endfor # vkeys

                # update history list:
                history.append("averaged %i pixels over grid cells" % len(iipix))
                history.append("selected %i pixels with %s" % (nsup, mapping_threshold))

                # write:
                outf.Write(
                    outfile, attrs=source_file.ds.attrs, history=history, complevel=complevel
                )

                # increase counter:
                nfile += 1

            else:
                # info ...
                logging.info("  keep  %s .." % outfile)
            # endif # renew

            ## testing ..
            # if nfile > 0 :
            #    logging.warning( 'break after first day' )
            #    break
            ##endif

        # endfor # listing records

        # info ...
        logging.info(indent + "")
        logging.info(indent + "processed %i files" % nfile)
        logging.info(indent + "")
        logging.info(indent + "** end")
        logging.info(indent + "")

    # enddef __init__


# endclass CSO_SuperObs


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