#
# Changes
# 
# 2022-09, Arjo Segers, Met-Norway
#   Updated documentation.
#
# 2022-02, Arjo Segers, Met-Norway
#   Create output directory if needed.
#

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

"""
.. _cso-regions:

**********************
``cso_regions`` module
**********************

The :py:mod:`cso_regions` module provides classes to create and apply region masks.


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

The classes and are defined according to the following hierchy:

* :py:class:`.UtopyaRc`

  * :py:class:`.CSO_CreateCountryMask`
  * :py:class:`.CSO_CatalogueBase`

      * :py:class:`.CSO_Catalogue_RegionsMaps`
      * :py:class:`.CSO_Catalogue_RegionsTimeSeries`

  * :py:class:`.CSO_Statistics_RegionsTables`

* :py:class:`.CreateCountryMap`
* :py:class:`.RegionMap`


Classes
=======


"""


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

# modules:
import logging

# tools:
import utopya
import cso_catalogue


########################################################################
###
### create country masks
###
########################################################################

class CSO_CreateCountryMask( utopya.UtopyaRc ) :

    """
    Create country mask at grid.
    
    Arguments:
    
    * ``rcfile``, ``rcbase``, ``env`` : settings file, prefix for keys, environment

    Specify how the grid is defined::

        ! how is grid defined?
        ! - "sample"   : file with 'longitude' and 'latitude' coordinates
        ! - "regular"  : define regular grid
        <rcbase>.grid.type            :   regular
        
    For grid type ``regular``, the definition should look like::

        ! define regular grid:
        <rcbase>.grid.regular.west    :   -180
        <rcbase>.grid.regular.east    :    180
        <rcbase>.grid.regular.south   :    -90
        <rcbase>.grid.regular.north   :     90
        <rcbase>.grid.regular.dlon    :   0.1
        <rcbase>.grid.regular.dlat    :   0.1

    For grid type ``sample``, define the sample file with::
        
        ! sample file that defines target grid:
        <rcbase>.grid.sample.file   :  /data/FLEXPART/grid_time_20170103000000.nc
        
    The country borders are taken from shape files shipped with the :py:mod:`cartopy` module.
    Specify the resolution to be used::
        
        ! required resolution of regions shape files:
        !   10m  110m
        <rcbase>.resolution          :  110m
        
    The region map will be written to a netCDF file together with the region names and codes.
    In addition, a csv with just the region names/codes will be written.
    Specify the basename of both files::
    
        ! basename of region map files:
        !  <basename>.nc    : netcdf file with map and region names etc
        !  <basename>.csv   : text file with region names etc (same as in nc file)
        <rcbase>.basename            :  /data/regions/regions-map_COSMO-7_countries
        
    Existing files are only replaced if the following flag is set::
    
        ! renew files if already present?
        <rcbase>.renew               :  True
        
    The csv table with region records has for example as content::
    
        code;name
        OCE;ocean
        ALB;Albania
        AUT;Austria
        BEL;Belgium
        :
        
    Based on this map it is posslible to create additional maps with aggregated countries.
    Provide a list with names of aggeragated maps to be created::

        ! create aggregated maps?
        <rcbase>.aggregates          :  aggr total

    For each aggregate, provide:
    
    * The name of csv table with codes and names for the aggregated regions.
      For example, if the aggreation should use "BNL" for combined Belgium/Netherlands/Luxemburg,
      then the table should contain::
      
        code;name
        BNL;BeNeLux
        :
        
    * The name of csv table with mapping from original to aggregated codes::
    
        code;newcode
        BEL;BNL
        :

    * The basename of the aggregated map to be created.
    
    These settings could look like::
    
        <rcbase>.aggregate.aggr.table      :  samples/regions_aggr.csv
        <rcbase>.aggregate.aggr.mapping    :  samples/regions_aggr-from-countries.csv
        <rcbase>.aggregate.aggr.basename   :  /data/regions/regions-map_COSMO-7_aggr

    """
    
    def __init__( self, rcfile, rcbase='', env={}, indent='' ) :
    
        """
        Create country mask.
        """
        
        # modules:
        import os
        import xarray
        import numpy
        import scipy.interpolate
        
        # init base object:
        utopya.UtopyaRc.__init__( self, rcfile, rcbase=rcbase, env=env )
        
        # info ...
        logging.info( indent+'' )
        logging.info( indent+'** Create country mask' )
        logging.info( indent+'' )
        
        # renew existing files?
        renew = self.GetSetting( 'renew', totype='bool' )
        
        # basename for target files:
        basename = self.GetSetting( 'basename' )
        # target files:
        filename_nc  = basename+'.nc'
        filename_csv = basename+'.csv'
        # no base map yet ..
        rm = None
        # create?
        if renew or (not os.path.isfile(filename_nc)) or (not os.path.isfile(filename_csv)) :
            # info ...
            logging.info( 'create %s.{nc,csv} ...' % basename )
            
            # how is grid defined?
            grid_type = self.GetSetting( 'grid.type' )
            # switch:
            # ...............................
            if grid_type == 'regular' :
            # ...............................
            
                # read domain:
                west  = self.GetSetting( 'grid.regular.west' , totype='float' )
                east  = self.GetSetting( 'grid.regular.east' , totype='float' )
                south = self.GetSetting( 'grid.regular.south', totype='float' )
                north = self.GetSetting( 'grid.regular.north', totype='float' )
                # spacing:
                dlon = self.GetSetting( 'grid.regular.dlon' , totype='float' )
                dlat = self.GetSetting( 'grid.regular.dlat' , totype='float' )
                # shape:
                nlon = numpy.round( (east-west)/dlon )
                nlat = numpy.round( (north-south)/dlat )
                
                # create regular coordinates:
                lons_da = xarray.DataArray( west+(numpy.arange(nlon)+0.5)*dlon, dims=('longitude',),
                                            attrs={ 'standard_name' : 'longitude', 'units' : 'degrees_east' } )
                lats_da = xarray.DataArray( south+(numpy.arange(nlat)+0.5)*dlat, dims=('latitude',),
                                            attrs={ 'standard_name' : 'latitude', 'units' : 'degrees_north' } )
            
            # ...............................
            elif grid_type == 'sample' :
            # ...............................
            
                # basename for target files:
                sample_file = self.GetSetting( 'grid.sample.file' )
                # info ...
                logging.info( '  read grid from: %s' % sample_file )
                # check ...
                if not os.path.isfile(sample_file) :
                    logging.error( 'file not found: %s' % sample_file )
                    raise Exception
                #endif
                # read:
                with xarray.open_dataset( sample_file ) as sample_ds :
                    # assume regular grid ...
                    lons_da = sample_ds['longitude']
                    lats_da = sample_ds['latitude']
                #endwith
                
            # ...............................
            else :
                logging.error( 'unsupported grid type "%s"' % grid_type )
                raise Exception
            #endif
            # ...............................

            ## testing ..
            #lons_da = lons_da[0:100]
            #lats_da = lats_da[0:100]

            # extract values:
            lons = lons_da.values
            lats = lats_da.values
            # corners, should be clockwise:
            #   3 o--o 0
            #     |  |
            #   2 o--o 1
            nlon = len(lons)
            nlat = len(lats)
            blons = scipy.interpolate.interp1d( numpy.arange(nlon), lons, fill_value='extrapolate' )( numpy.arange(nlon+1)-0.5 )
            blats = scipy.interpolate.interp1d( numpy.arange(nlat), lats, fill_value='extrapolate' )( numpy.arange(nlat+1)-0.5 )
            lon_bounds = numpy.zeros((nlat,nlon,4),float)
            lat_bounds = numpy.zeros((nlat,nlon,4),float)
            for j in range(nlat) :
                lon_bounds[j,:,0] = blons[ 1:]
                lon_bounds[j,:,1] = blons[ 1:]
                lon_bounds[j,:,2] = blons[:-1]
                lon_bounds[j,:,3] = blons[:-1]
            #endfor
            for i in range(nlon) :
                lat_bounds[:,i,0] = blats[ 1:]
                lat_bounds[:,i,1] = blats[:-1]
                lat_bounds[:,i,2] = blats[:-1]
                lat_bounds[:,i,3] = blats[ 1:]
            #endif
            
            # grid definition:
            grid_dims   = ('latitude','longitude')
            grid_coords = { 'latitude' : lats_da, 'longitude' : lons_da }
            grid_attrs  = {}
            
            # resolution of shape files:
            resolution = self.GetSetting( 'resolution' )
            # info ...
            logging.info( '  resolution of shape files: %s' % resolution )
        
            # info ...
            logging.info( '  create country map, might take a while ...' )
            # create country map object:
            cm = CreateCountryMap( lons,lats, lon_bounds,lat_bounds, resolution )
    
            # info ...
            logging.info( '  create country table ...' )
            # table with columns:
            #   icountry : original country index
            #   iregion  : new numbering
            #   code
            #   name
            ctab = cm.GetCountryTable()

            # info ..
            logging.info( '  create region map ...' )
            # new map following "iregion" numbering:
            cmap = cm.GetRegionMap( ctab )
    
            # info ..
            logging.info( '  create map object ...' )
            # store:
            rm = RegionMap()
            # global attributes:
            #history = 'country map on grid defined in %s' % os.path.basename(sample_file)
            #source = 'cartopy shapefiles at %s resolution' % resolution
            gattrs = {} #self.GetGlobalAttributes( {}, history=history, source=source )
            # fill:
            rm.Create( cmap, ctab, (grid_dims,grid_coords,grid_attrs), gattrs=gattrs )
            
            # save:
            rm.Write( basename, savetable=True )
            
        else :
            # info ...
            logging.info( 'keep existing %s.{nc,csv} ...' % basename )
        #endif # renew
        
        # * aggregated regions
        
        # aggregates?
        aggregates = self.GetSetting( 'aggregates' ).split()
        # loop:
        for aggregate in aggregates :
            # info ...
            logging.info( 'aggregate "%s" ...' % aggregate )

            # renew existing files?
            renew = self.GetSetting( 'aggregate.%s.renew' % aggregate, totype='bool' )

            # basename of target files:
            aggr_basename = self.GetSetting( 'aggregate.%s.basename' % aggregate )
            # target files:
            aggr_filename_nc  = aggr_basename+'.nc'
            aggr_filename_csv = aggr_basename+'.csv'
            # create?
            if renew or (not os.path.isfile(aggr_filename_nc)) or (not os.path.isfile(aggr_filename_csv)) :
                # info ...
                logging.info( '  create %s.{nc,csv} ...' % aggr_basename )
                
                # no basemap yet?
                if rm is None :
                    # info ..
                    logging.info( '    read basemap %s.{nc,csv} ...' % basename )
                    # read:
                    rm = RegionMap( basename=basename )
                #endif
                
                # table with new codes/names:
                ctab_new = self.GetSetting( 'aggregate.%s.table' % aggregate )
                # mapping from orignal codes to new:
                ctab_to_new = self.GetSetting( 'aggregate.%s.mapping' % aggregate )
                # info ...
                logging.info( '    aggregated regions: %s' % ctab_new )
                logging.info( '    mapping from base : %s' % ctab_to_new )
                # global attributes:
                #history = 'aggregated to %s using %s' % (os.path.basename(ctab_new),os.path.basename(ctab_to_new))
                gattrs = {}  #self.GetGlobalAttributes( rm.gattrs, history=history )
                # create new aggregated region map object:
                rma = rm.Remap( ctab_to_new, ctab_new,
                                     gattrs={ 'source' : 'aggregated regions' },
                                     indent='    ' )
                # save:
                rma.Write( aggr_basename, savetable=True )

            else :
                # info ...
                logging.info( '  keep existing %s.{nc,csv} ...' % basename )
            #endif # renew

        #endfor aggregates
        
        # * coarsenings to maps with fractions
        
        # coarsenings?
        coarsenings = self.GetSetting( 'coarsenings' ).split()
        # loop:
        for coarsening in coarsenings :
            # info ...
            logging.info( 'coarsening "%s" ...' % coarsening )

            # renew existing files?
            renew = self.GetSetting( 'coarsening.%s.renew' % coarsening, totype='bool' )

            # basename of target files:
            coarse_basename = self.GetSetting( 'coarsening.%s.basename' % coarsening )
            # target files:
            coarse_filename_nc  = coarse_basename+'.nc'
            coarse_filename_csv = coarse_basename+'.csv'
            # create?
            if renew or (not os.path.isfile(coarse_filename_nc)) or (not os.path.isfile(coarse_filename_csv)) :
                # info ...
                logging.info( '  create %s.{nc,csv} ...' % coarse_basename )
                
                # input base:
                basename = self.GetSetting( 'coarsening.%s.input' % coarsening )
                # info ..
                logging.info( '    read basemap %s.{nc,csv} ...' % basename )
                # read:
                rm = RegionMap( basename=basename )
                
                # read domain:
                west  = self.GetSetting( 'coarsening.%s.west'  % coarsening, totype='float' )
                east  = self.GetSetting( 'coarsening.%s.east'  % coarsening, totype='float' )
                south = self.GetSetting( 'coarsening.%s.south' % coarsening, totype='float' )
                north = self.GetSetting( 'coarsening.%s.north' % coarsening, totype='float' )
                # spacing:
                dlon = self.GetSetting( 'coarsening.%s.dlon' % coarsening, totype='float' )
                dlat = self.GetSetting( 'coarsening.%s.dlat' % coarsening, totype='float' )
                # assign some values to the rest? for example ocean to countries in same cell:
                distribute = self.GetSetting( 'coarsening.%s.distribute' % coarsening ).split()

                # global attributes:
                #history = 'coarsened from %s.{nc,csv}' % os.path.basename(basename)
                gattrs = {}  #self.GetGlobalAttributes( rm.gattrs, history=history )
                # create new coarsened region map object:
                rmc = rm.Coarsen( [west,east,south,north], [dlon,dlat], distribute=distribute,
                                     gattrs={ 'source' : 'coarsened regions' },
                                     indent='    ' )
                # save:
                rmc.Write( coarse_basename, savetable=True )

            else :
                # info ...
                logging.info( '  keep existing %s.{nc,csv} ...' % basename )
            #endif # renew

        #endfor coarsenings

        # *

        # info ...
        logging.info( indent+'' )
        logging.info( indent+'** end' )
        logging.info( indent+'' )
    
    #enddef __init__
    
#endclass CSO_CreateCountryMask



    
#=======================================================================
#===
#=== CreateCountryMap
#===
#=======================================================================
    
# modules:
import numba

@numba.jit(nopython=True)
def minmax(array):
    """Find both min and max value of arr in one pass"""
    minval = maxval = array[0]

    for e in array[1:]:
        if e < minval:
            minval = e
        #endif
        if e > maxval:
            maxval = e
        #endif
    #endfor

    return minval, maxval

#enddef minmax

# *

@numba.jit(nopython=True)
def intersection_possible( points, c_minx, c_miny, c_maxx, c_maxy ) :

    """
    Determine if country could contain any of the points.

    The maximum and minimum x and y values of the country and the points
    are compared, to determine if any of the points lie in the bounding box
    of country.

    For rectangles oriented along the coordinate axes, Polygon.intersects()
    is faster than this function. However, we're working with arbitrarily
    oriented (and possibly distorted) gridcells, for which
    Polygon.intersects() is considerably slower.

    Arguments:
    
    * ``points`` : array of shape (4,3) containing the 4 corners of a cell, ordered clockwise
    * ``c_minx,c_miny,c_maxx,c_maxy`` : country bounds

    Returns:
    
    * ``True`` if an intersection is possible, ``False`` otherwise.
    """

    # check x-coordinates
    p_minx,p_maxx = minmax( points[:,0] )

    if (c_minx >= p_maxx) or (c_maxx <= p_minx) :
        return False

    # check y-coordinates (if this still needs to be done)
    p_miny, p_maxy = minmax( points[:,1] )

    return (c_miny < p_maxy) and (c_maxy > p_miny)

#enddef intersection_possible


#
# *
#
    
class CreateCountryMap( object ) :

    """
    Create map with indices in country list.
    
    Arguments:
    
    * ``lons``, ``lats`` : 1D arrays with grid cell centers
    * ``lon_bounds``, ``lat_bounds`` : 3D arrays ``(nlat,nlon,ncorner)`` with grid cell corners
    * ``resolution`` : shape file resolution to be used: '10m', '110m'
    
    Attributes created:
    
    * ``map`` : 2D map with indices
    * ``attrs`` : dict with lists for indices
    * ``used_country_index`` : list with country indices
    """
    
    def __init__( self, lons,lats, lon_bounds,lat_bounds, resolution ) :
    
        """
        Create map with index values.
        """
        
        # modules:
        import cartopy
        import numpy
        import shapely

        # source file:
        shpfilename = cartopy.io.shapereader.natural_earth( resolution=resolution, 
                                     category="cultural", name="admin_0_countries" )
        # read list with items for all countries in the world:
        countries = list( cartopy.io.shapereader.Reader( shpfilename ).records() )

        # Keep track of countries with no code
        no_country_indx = []
        self.used_country_index = []

        ## extract list of country bounds:
        #country_bounds = [ country.bounds for country in countries ]

        # Name of the country attribute holding the ISO3 country abbreviation
        iso3 = "ADM0_A3"

        #
        # create dictionairy for selected attributes,
        # with in each a list of values:
        #
        #   attrs['code'] = ['OCE', 'FJI', 'TZA', 'SAH', 'CAN', ...]
        #        ['name'] = ['ocean', 'Fiji', 'Tanzania', 'W. Sahara', 'Canada', ...]
        #
        self.attrs = { 'code' : ['OCE'], 
                       'name' : ['ocean'] }
        # original record names:
        crec = {}
        crec['code'] = iso3
        crec['name'] = 'NAME'
        # loop over initialized attributes:
        for aname in self.attrs.keys() :
            # loop over countries:
            for k in range(len(countries)) :
                # extract:
                avalue = countries[k].attributes[crec[aname]]
                # extend list:
                self.attrs[aname].append( avalue )
            #endfor # country code range
        #endfor # attributes

        # projection of the shapefile: WGS84, which the PlateCarree defaults to    
        projection_in  = cartopy.crs.PlateCarree()
        # projection of the target grid:
        #projection_out = grid.get_projection()
        projection_out = cartopy.crs.PlateCarree()
        
        # size:
        nx = len(lons)
        ny = len(lats)

        # init result:
        self.map = numpy.empty( (ny,nx), dtype=int )
        
        # for info: total number of grid points:
        np = nx * ny
        # when to print:
        npx = int( np / 100 )        

        # loop over grid cells:
        for j in range(ny) :
            for i in range(nx) :
            
                # info ...
                ip = j*nx + i+1 
                if ip % npx == 0 : logging.info( '  %3i %% ..' % (ip/np*100) )

                # corners:
                cell_x = lon_bounds[j,i,:]
                cell_y = lat_bounds[j,i,:]
                # project:
                projected_corners = projection_in.transform_points( projection_out, cell_x, cell_y )
                # create polygon:
                projected_cell = shapely.geometry.Polygon( projected_corners )

                #intersected_countries = [
                #    country
                #    for country, bounds in zip(countries,country_bounds)
                #    if intersection_possible(projected_corners, *bounds)
                #    and projected_cell.intersects(country.geometry)
                #]

                # intersections:
                intersected_countries = []
                for country in countries :
                    if intersection_possible( projected_corners, *country.bounds ) and \
                                  projected_cell.intersects(country.geometry) :
                        intersected_countries.append( country )
                    #endif
                #endfor

                try:
                    # Multiple countries might be intersected, 
                    # assign the one with the largest area in the cell:
                    cell_country = max( intersected_countries,
                                        key=lambda c: projected_cell.intersection(c.geometry).area )
                    # extract code:
                    cell_country_code = cell_country.attributes[iso3]
                    # extract index:
                    cell_country_indx = self.attrs['code'].index( cell_country_code )
                except ValueError:
                    # max on empty sequence raises ValueError
                    # If no intersection was found, country code 0 is assigned. It
                    # corresponds to ocean.
                    # Even if a cell is predominantly ocean, it will still get assigned
                    # the country code of the largest country.
                    cell_country_indx = 0
                except KeyError:
                    # Didn't find country name in country_codes
                    no_country_indx.append(cell_country_code)
                    cell_country_indx = -1
                #endtry

                # update list of used codes:
                if cell_country_indx not in self.used_country_index :
                    self.used_country_index.append( cell_country_indx )
                #endif

                # store:
                self.map[j,i] = cell_country_indx

            #endfor i

            ## testing ...
            #break

        #endfor j

        if len(no_country_indx) > 0:
            print( "The following countries were found, but didn't have a corresponding code:" )
            for cname in no_country_indx :
                print( '  %s' % cname )
            #endfor
        #endif
    
    #enddef __init__
    
    # *
    
    def GetCountryTable( self ) :
    
        """
        Return DataFrame with used countries;
        ``icountry`` is the original country number,
        ``iregion`` a new numbering::
        
            icountry;iregion;code;name
             0;0;OCE;ocean
            12;1;ALB;Albania
            :
        
        """
        
        # modules:
        import pandas

        # collect used country codes (except "OCE"), store original number:
        #   ccode_info['ALB'] = { 'icountry' : 1, 'name' : 'Albania' }
        #                :
        ccode_info = {}
        for icountry in self.used_country_index :
            # current:
            ccode = self.attrs['code'][icountry]
            name  = self.attrs['name'][icountry]
            # skip oceans, prepended later on ...
            if ccode == 'OCE' : continue
            # add:
            ccode_info[ccode] = { 'icountry' : icountry, 'name' : name }
        #endfor

        # exract used codes:
        ccodes = list(ccode_info.keys())
        # sort inplace:
        ccodes.sort()
        # prepend oceans:
        ccodes = ['OCE']+ccodes
        ccode_info['OCE'] = { 'icountry' : 0, 'name' : 'ocean' }

        # init table for info on used country codes:
        ctab = pandas.DataFrame( columns=['icountry','iregion','code','name'] )
        # loop over used codes:
        for icode in range(len(ccodes)) :
            # current:
            ccode = ccodes[icode]
            # store values:
            ctab.loc[icode,'icountry' ] = ccode_info[ccode]['icountry']
            ctab.loc[icode,'iregion'  ] = icode
            ctab.loc[icode,'code'     ] = ccode
            ctab.loc[icode,'name'     ] = ccode_info[ccode]['name']
        #endfor

        # ok
        return ctab
        
    #enddef GetCountryTable
    
    # * 
    
    def GetRegionMap( self, ctab ) :
    
        """
        Return 2D array with region numbers as in the ``iregion`` column 
        of ``ctab``, the DataFrame returned by :py:meth`GetCountryTable`.
        """
        
        # modules:
        import numpy

        # new map:
        cmap = numpy.zeros( self.map.shape, int ) - 999
        # loop over countries:
        for ccode,rec in ctab.iterrows() :
            # mask:
            jj,ii = numpy.where( self.map == rec.icountry )
            # new:
            cmap[jj,ii] = rec.iregion
        #endfor
        
        # ok
        return cmap
        
    #enddef GetRegionMap

#endclass CreateCountryMap

    
#=======================================================================
#===
#=== RegionMap
#===
#=======================================================================

class RegionMap( object ) :

    """
    Storage for region map.
    
    The map and auxilary data are stored in the attributes:
    
    * ``ds`` : :py:class:`xarray.Dataset` with coordinate:
    
      * ``latitude``, ``latitude`` : coordinates of regular grid
      * ``region`` : coordinate for region values
      
      and variables:
      
      * ``map(latitude,longitude)`` : integer region indices
      * ``region_index`` : region indices used in map
      * ``region_code`` : short character code
      * ``region_name`` : character strings with region names
      
    * ``gattrs`` : dict with global attributes
    * ``regions`` : :py:class:`pandas.DataFrame` with columns:

      * ``iregion`` : int index, same as ``region_index`` variable
      * ``code`` : str with ``region_code``
      * ``name`` : str with ``region_name``
    """
    
    def __init__( self, basename=None ) :
    
        """
        Initialize region map.
        """
        
        # no data set yet:
        self.ds = None
        self.gattrs = {}
        
        # read?
        if basename is not None :
            # read:
            self.Read( basename )
        #endif
        
    #enddef
   
    # *
    
    def __len__( self ) :
    
        """
        Returnn number of regions.
        """
        
        # defined?
        if self.ds is None :
            nregion = 0
        else :
            nregion = self.ds.dims['region']
        #endif
        
        # ok
        return nregion
        
    #enddef __len__
   
    # *
    
    def GetRegionInfo( self, iregion ) :
    
        """
        Return info on region:
        
        * ``code`` : region code
        * ``name`` : region name
        """
        
        # check ..
        if (iregion < 0) or (iregion > len(self)-1) :
            logging.error( 'region index %i should be in 0,..,%i-1' % (iregion,len(self)) )
            raise Exception
        #endif
        
        # ok
        return self.ds['region_code'].values[iregion],self.ds['region_name'].values[iregion]
        
    #enddef GetRegionInfo
   
    # *
    
    def GetRegionMask( self, iregion ) :
    
        """
        Return bool array of shape ``(latitude,longiutde)`` that is True for 
        the specified region.
        """
        
        # modules:
        import numpy
        
        # check ..
        if (iregion < 0) or (iregion > len(self)-1) :
            logging.error( 'region index %i should be in 0,..,%i-1' % (iregion,len(self)) )
            raise Exception
        #endif
        
        # init result:
        mask = numpy.zeros((self.ds.dims['latitude'],self.ds.dims['longitude']),bool)
        # cell indices for region:
        jj,ii = numpy.where( self.ds['map'] == self.ds['region_index'][iregion] )
        # reset in mask:
        if len(jj) > 0 :
            mask[jj,ii] = True
        #endif
        
        # ok
        return mask
        
    #enddef GetRegionMask
   
    # *
    
    def GetRegionDomain( self, iregion, edgefactor=0.0 ) :
    
        """
        Return bounding box ``[west,easth,south,north]`` around region,
        with eventually ``edgefactor`` extra edge.
        """
        
        # modules:
        import numpy
        
        # get mask for region:
        mask = self.GetRegionMask( iregion )
        # shape:
        nlat,nlon = mask.shape
        # cell indices for region:
        jj,ii = numpy.where( mask )
        # range:
        i0,i1 = min(ii),max(ii)
        j0,j1 = min(jj),max(jj)
        # extra length at edges:
        di = int( (i1-i0+1) * edgefactor/2 )
        dj = int( (j1-j0+1) * edgefactor/2 )
        # extend:
        i0 = max(      0, i0-di )
        i1 = min( nlon-1, i1+di )
        j0 = max(      0, j0-dj )
        j1 = min( nlat-1, j1+dj )
        # box:
        domain = [ self.ds.coords['longitude'].values[i0],
                   self.ds.coords['longitude'].values[i1],
                   self.ds.coords['latitude' ].values[j0],
                   self.ds.coords['latitude' ].values[j1] ]
        
        # ok
        return domain
        
    #enddef GetRegionDomain
        
    
    # *
    
    def Create( self, rmap, regions, grid, gattrs={} ) :
    
        """
        Create new region map.
        
        Arguments:

        * ``rmap[:,:]`` : 2D int array with region indices
        * ``regions`` : DataFrame with columns:

          * ``iregion`` : int index as used in ``rmap``
          * ``code`` : str region code
          * ``name`` : str region name

        * ``grid`` : 3-element tupple with info for grid variables:

          * ``('rlat','rlon')`` : tupple with dimension names
          * ``{ 'rlon' : .., 'rlat' : .. }`` : dict with :py:class:`xarray.DataArray` objects
          * ``{'grid_mapping' : 'rotated_pole', ... }`` : dict with attributes to be added to variables
          
        * ``gattrs`` : dict with global attributes
        """
        
        # modules:
        import xarray
        
        # expand tupple:
        grid_dims,grid_coords,grid_attrs = grid
        # store:
        self.gattrs = gattrs
        
        # init data arrays with grid coordinates:
        das = grid_coords.copy()
        
        # map attributes:
        attrs = {}
        attrs['long_name'] = 'region map'
        attrs['units'] = '1'
        # add grid description:
        attrs.update( grid_attrs )
        # define data array:
        da = xarray.DataArray( data=rmap, dims=grid_dims, attrs=attrs )
        # limit storage:
        #da.encoding['dtype'] = 'int8'
        da.encoding['dtype'] = 'int16'
        # no mask:
        da.encoding['_FillValue'] = None
        # store:
        das['map'] = da
        
        # add data arrays with region codes etc:
        das.update( self._GetRegionArrays(regions) )
        
        # create data set:
        self.ds = xarray.Dataset( das, attrs=gattrs )
        
        # create "self.regions" table from dataset,
        # should be similar to "region" argument:
        self._CreateTable()
        
    #enddef __init__
    
    # *

    def _GetRegionArrays( self, regions ) :
    
        """
        Return list of data arrays with region variables.
        
        Arguments:
        
        * ``regions`` : DataFrame with columns:

          * ``iregion`` : in index as used in ``rmap``
          * ``code`` : str region code
          * ``name`` : str region name
        """
        
        # modules:
        import xarray
        import numpy
        
        # init result:
        das = {}
        
        # number of regions:
        nregion = len(regions)
        # coordinate variable:
        da = xarray.DataArray( data=numpy.arange(1,nregion+1), dims=('region'), 
                                 attrs={ 'long_name' : 'region', 'units' : '1' } )
        da.encoding['dtype'] = 'int8'
        da.encoding['_FillValue'] = None
        # store:
        das['region'] = da
        
        # region index:
        da = xarray.DataArray( data=numpy.array(regions['iregion'].values,int), dims=('region'), 
                                 attrs={ 'long_name' : 'region index', 'units' : '1' } )
        da.encoding['dtype'] = 'int8'
        da.encoding['_FillValue'] = None
        # store:
        das['region_index'] = da
        
        # region code:
        da = xarray.DataArray( data=regions['code'].values, dims=('region'), 
                                 attrs={ 'long_name' : 'region code', 'units' : '1' } )
        # store:
        das['region_code'] = da
        
        # region name:
        da = xarray.DataArray( data=regions['name'].values, dims=('region'), 
                                 attrs={ 'long_name' : 'region name', 'units' : '1' } )
        # store:
        das['region_name'] = da
        
        # ok
        return das
        
    #enddef _GetRegionArrays
    
    # *
    
    def Write( self, basename, savetable=False ) :
    
        """
        Write region map to netCDF file (extension ".nc"),
        and optionally the regions table to a ".csv" file.
        """
        
        # modules:
        import os
        import xarray
        
        # target directory:
        dname = os.path.dirname(basename)
        # create?
        if len(dname) > 0 :
            if not os.path.isdir(dname) : os.makedirs( dname )
        #endif
        
        # save:
        self.ds.to_netcdf( basename+'.nc' )
        
        # save table?
        if savetable :
            # write csv table:
            self.regions.to_csv( basename+'.csv', sep=';', index=False,
                                  columns=['code','name'] )
        #endif
        
    #enddef Write
    
    # * 
    
    def Read( self, basename ) :
    
        """
        Read region map from specified netcdf filename.
        """
        
        # modules:
        import os
        import xarray
        import pandas
        
        # full name:
        filename = basename+'.nc'
        # check ...
        if not os.path.isfile(filename) :
            logging.error( 'file not found: %s' % filename )
            raise Exception
        #endif
        
        # load entire set:
        self.ds = xarray.load_dataset(filename)
        
        # no masks in coordinates ...
        for vname in self.ds.coords.keys() :
            self.ds.coords[vname].encoding['_FillValue'] = None
        #endfor
        # no masks in variables ...
        for vname in self.ds.keys() :
            self.ds[vname].encoding['_FillValue'] = None
        #endfor
        
        # create table:
        self._CreateTable()

        ## copy map:
        #self.rmap = ds['map'].values
        #
        ## map dimensions:
        #self.grid_dims = ds['map'].dims
        #
        ## copy known coordinate attribues if present:
        #self.grid_attrs = {}
        #for aname in ['coordinates','grid_mapping','cell_measures'] :
        #    if aname in ds['map'].attrs.keys() :
        #        self.grid_attrs[aname] = ds['map'].attrs[aname]
        #    #endif
        ##endfor
        #
        ## storage for coordinate arrays:
        #self.grid_coords = {}
        ## loop over dimensions:
        #for vname in self.grid_dims :
        #    # copy coordinate array:
        #    self.grid_coords[vname] = ds[vname]
        ##endfor
        ## loop over grid attributes:
        #for aname in self.grid_attrs.keys() :
        #    # content contains variables names or "area:" etc,
        #    # loop over all parts:
        #    for vname in self.grid_attrs[aname].split() :
        #        # part defines a variable?
        #        if vname in ds.keys() :
        #            # copy array:
        #            self.grid_coords[vname] = ds[vname]
        #        #endif
        #    #endfor
        ##endfor
        ## no masks ...
        #for vname in self.grid_coords.keys() :
        #    self.grid_coords[vname].encoding['_FillValue'] = None
        ##endfor
        #
        ## global attributes:
        #self.gattrs = ds.attrs.copy()
    
    #enddef Read
    
    # *
    
    def _CreateTable( self ) :
    
        """
        Create attribute 'regions' with DataFrame with columns:

          * ``iregion`` : in index as used in ``rmap``
          * ``code`` : str region code
          * ``name`` : str region name
        """
        
        # modules:
        import pandas
        
        # columns for dataframe:
        iregions   = self.ds['region_index'    ].values
        codes      = self.ds['region_code'     ].values
        names      = self.ds['region_name'     ].values
        
        # store:
        self.regions = pandas.DataFrame( { 'iregion' : iregions, 
                                           'code'    : codes,
                                           'name'    : names } )
        
    #enddef _CreateTable
    
    # *
    
    def Remap( self, tonew, new, gattrs={}, indent='' ) :
    
        """
        Assign regions to new names.
        
        Arguments:
        
        * ``tonew`` : csv file with old-to-new mapping; shoud have at least the columns:
          
          * ``code`` : original region codes
          * ``newcode`` : new codes
          
        * ``new`` : csv file with new regions; shoud have at least the columns:
          
          * ``code`` : region codes
          * ``name`` : region names
          
        * ``gattrs`` : dict with global attributes
        """
        
        # modules:
        import os
        import pandas
        import numpy
        import xarray
        
        # check ..
        if not os.path.isfile(tonew) :
            logging.error( 'file not found: %s' % tonew )
            raise Exception
        #endif
        # read:
        df_to = pandas.read_csv( tonew, sep=';', comment='#' )
        # check ..
        for key in ['code','newcode'] :
            if key not in df_to.keys() :
                logging.error( 'column "%s" not found in: %s' % (key,tonew) )
                raise Exception
            #endif
        #endfor
        
        # init list of new codes:
        newcodes = []
        # check if all codes are present ..
        for iregion in self.ds['region_index'].values :
            # code:
            code = self.ds['region_code'].values[iregion]
            # number of values:
            nval = (self.ds['map'].values == iregion).sum()
            # info ...
            logging.info( indent+'found %6i values with index %3i (%s)' % (nval,iregion,code) )
            # any values?
            if nval > 0 :
                # not in table?
                if not (code in df_to['code'].values) :
                    logging.error( 'code "%s" present in map but not found in: %s' % (code,tonew) )
                    raise Exception
                #endif
                # new code:
                newcodes.append( df_to.set_index('code').at[code,'newcode'] )
            #endif
        #endfor

        # check ..
        if not os.path.isfile(new) :
            logging.error( 'file not found: %s' % new )
            raise Exception
        #endif
        # read:
        df_new = pandas.read_csv( new, sep=';', comment='#' )
        # check ..
        for key in ['code','name'] :
            if key not in df_new.keys() :
                logging.error( 'column "%s" not found in: %s' % (key,tonew) )
                raise Exception
            #endif
        #endfor
        # check if all new codes are present ...
        for newcode in newcodes :
            if newcode not in df_new['code'].values :
                logging.error( 'new code "%s" not found in: %s' % (newcode,new) )
                raise Exception
            #endif
        #endfor
        # size:
        nregion_new = len(df_new)
        # new region codes
        df_new['iregion'] = numpy.arange(nregion_new)
        
        # init result:
        rm = RegionMap()
        # get arrays:
        das = self._GetRegionArrays( df_new )
        # copy attributes, replace if necessary:
        attrs = self.ds.attrs.copy()
        attrs.update( gattrs )
        # new dataset, copy current attributes:
        rm.ds = xarray.Dataset( das, attrs=attrs )
        # add map and grid variables:
        for vname in self.ds.keys() :
            # region arrays have been defined already ...
            if vname.startswith('region_') : continue
            # copy:
            rm.ds[vname] = self.ds[vname].copy()
        #endfor
        # create table:
        rm._CreateTable()
        
        # reset index variables:
        df_to  = df_to.set_index('code')
        df_new = df_new.set_index('code')
        
        # info ..
        logging.info( indent+'assign new indices ...' )
        # reset target map:
        rm.ds['map'].values *=  0
        # loop over source indices:
        for iregion in self.ds['region_index'].values :
            # find cells:
            jj,ii = numpy.where( self.ds['map'].values == iregion )
            # any?
            if len(jj) > 0 :
                # source code:
                code = self.ds['region_code'].values[iregion]
                # target code:
                code_new = df_to.at[code,'newcode']
                # store new region index:
                rm.ds['map'].values[jj,ii] = df_new.at[code_new,'iregion']
            #endif
        #endfor
        # info ..
        logging.info( indent+'  done ...' )
        
        # ok
        return rm
        
    #enddef Remap
    
    # *
    
    def Coarsen( self, domain, dres, distribute=None, gattrs={}, indent='' ) :
    
        """
        Coarsen a 2D ``(lat,lon)`` region map to a 3D ``(region,lat,lon)`` array.
        
        Arguments:
        
        * ``domain=[west,east,south,north]`` : target domain
        * ``dres=[dlon,dlat]`` : target resolution
        
        Optional arguments:
        
        * ``distribute`` : list with region codes to be distributed over other regions, 
          for example ``['OCE']`` to assign ocean area to surrounding regions
        * ``gattrs`` : dict with global attributes
        """
        
        # modules:
        import os
        import pandas
        import numpy
        import xarray
        
        # extract:
        west,east,south,north = domain
        dlon,dlat = dres
        # shape:
        nlon = int(numpy.round( (east -west )/dlon ))
        nlat = int(numpy.round( (north-south)/dlat ))
        # info ...
        logging.info( indent+'coarsen region map to fractions:' )
        logging.info( indent+'  domain     : [%f,%f] x [ %f,%f]' % (west,east,south,north) )
        logging.info( indent+'  resolution : %f x %f' % (dlon,dlat) )
        logging.info( indent+'  shape      : %i x %i' % (nlat,nlon) )
        
        # create regular coordinates:
        coords = {}
        coords['longitude'] = xarray.DataArray( west+(numpy.arange(nlon)+0.5)*dlon, 
                                    dims=('longitude',),
                                    attrs={ 'standard_name' : 'longitude', 'units' : 'degrees_east' } )
        coords['latitude'] = xarray.DataArray( south+(numpy.arange(nlat)+0.5)*dlat, 
                                    dims=('latitude',),
                                    attrs={ 'standard_name' : 'latitude', 'units' : 'degrees_north' } )
        coords['region'] = self.ds.coords['region']
        # number of regions:
        nregion = self.ds.dims['region']
        
        # storage for fraction maps:
        values = numpy.zeros( (nregion,nlat,nlon), float )
        # loop over regions:
        for iregion in range(nregion) :
            # region index in maps:
            indx = self.ds['region_index'].values[iregion]
            # points with this index:
            mjj,mii = numpy.where( self.ds['map'] == indx )
            # any?
            if len(mii) > 0 :
                # locations:
                mlons = self.ds['longitude'].values[mii]
                mlats = self.ds['latitude' ].values[mjj]
                # target indices:
                ii = numpy.array( numpy.floor( (mlons-west )/dlon ), dtype=int )
                jj = numpy.array( numpy.floor( (mlats-south)/dlat ), dtype=int )
                # add increments
                for k in range(len(ii)) :
                    # target cell:
                    j,i = jj[k],ii[k]
                    # in range?
                    if (i >= 0) and (i <= nlon-1) and (j >= 0) and (j <= nlat-1) :
                        values[iregion,j,i] += 1
                    #endif
                #endfor
            #endif
        #endfor # regions
        
        # distribute some?
        if distribute is not None :
            # loop:
            for dst_code in distribute :
                # info ...
                logging.info( indent+'  distribute "%s" over other regions ...' % dst_code )
                # region coordinate:
                iregs, = numpy.where( self.ds['region_code'].values == dst_code )
                # check ..
                if len(iregs) != 1 :
                    logging.error( 'distribute region "%s" found %i times in region names' % (dst_code,iregs) )
                    raise Exception
                #endif
                # extract:
                dst_ireg = iregs[0]
                # sum of all other weights:
                oth_sum = values.sum(axis=0) - values[dst_ireg,:,:]
                # find cells where other regions are present too:
                jj,ii = numpy.where( oth_sum > 0 )
                # any?
                if len(ii) > 0 :
                    # loop over cells:
                    for k in range(len(ii)) :
                        # cell:
                        j,i = jj[k],ii[k]
                        # source present?
                        if values[dst_ireg,j,i] > 0.0 :
                            # loop over other regions:
                            for ireg in range(nregion) :
                                # not the region to be distributed ..
                                if ireg == dst_ireg : continue
                                # any value here?
                                if values[ireg,j,i] > 0 :
                                    # add contribution:
                                    values[ireg,j,i] += values[dst_ireg,j,i] * values[ireg,j,i]/oth_sum[j,i]
                                #endif
                            #endfor # other regions
                            # reset source:
                            values[dst_ireg,j,i] = 0.0
                        #endif
                    #endfor # cells with other regions too
                #endif # any cells with other regions?
            #endfor # regions to be distributed
        #endif # distribute regions?
        
        # normalize:
        all_sum = values.sum(axis=0)             
        # where defined?
        jj,ii = numpy.where( all_sum > 0.0 )
        # any defined?
        if len(ii) > 0 :
            # loop over regions:
            for ireg in range(nregion) :
                # normalize:
                values[ireg,:,:] /= all_sum
            #endfor
        #endif
        
        # init result:
        rm = RegionMap()
        # init dataset:
        rm.ds = xarray.Dataset()
        # copy some:
        for key in self.ds.keys() :
            if key == 'map' : continue
            rm.ds[key] = self.ds[key]
        #endfor
        # add new variable:
        rm.ds['fraction_map'] = xarray.DataArray( values, dims=('region','latitude','longitude'),
                                                   coords=coords,
                                                   attrs={ 'units' : '1', 'long_name' : 'region fraction maps' } )
        # create table:
        rm._CreateTable()
            
        # ok
        return rm
        
    #enddef Coarsen
                
    
#endclass RegionMap

    
########################################################################
###
### catalogue of gridded averages over regions
###
########################################################################


class CSO_Catalogue_RegionsMaps( cso_catalogue.CSO_CatalogueBase ) :

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

    .. NEEDS UPDATE
       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
      
    Define the basename of the region definition files (netcdf and csv),
    as created with the :py:class:`CSO_CreateCountryMask` class::

        ! basename for region definition files:
        <rcbase>.regions.basename       :  ${WORK}/projects/EMEP-S5p/regions/EMEP-002_country-mask_50m

    If some problematic countries should be skipped, specify their codes::

        ! skip some ...
        <rcbase>.regions.skip            :  AND
      
    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 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}_%{region}.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)

        ! extra keyword arguments for map:
        <rcbase>.map                 :  resolution='l'
        
    A list of region codes that where actually used to create plots (regions within the domain)
    is written to a file ``regions.csv`` to facilitate index generation.

    """
    
    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
        import cso_regions
        
        # setup graphics back-end:
        matplotlib.use('Agg') 

        # init base object:
        cso_catalogue.CSO_CatalogueBase.__init__( self, rcfile, rcbase=rcbase, env=env )
        
        # info ...
        logging.info( indent+'' )
        logging.info( indent+'** plot CSO simulations over regions' )
        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 directory:
        workdir = self.GetSetting( 'output.dir' )

        # info ..
        logging.info( indent+'work directory:' )
        logging.info( indent+'  %s' % workdir )
        # create?
        if not os.path.isdir(workdir) : os.makedirs( workdir )
        # store current:
        owd = os.getcwd()
        # change to:
        os.chdir( workdir )
        
        # basename for region definition files:
        regmap_basename = self.GetSetting( 'regions.basename' )
        # read:
        regmap = cso_regions.RegionMap( regmap_basename )
        # count:
        nreg = len(regmap)
        # skip lists:
        reg_skip = self.GetSetting( 'regions.skip' ).split()
        # storage for collected region codes, used for index:
        reg_used = pandas.DataFrame()
        
        # 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)
            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 regions:
                for ireg in range(nreg) :
                
                    # info:
                    reg_code,reg_name = regmap.GetRegionInfo( ireg )
                    # filter ..
                    if reg_code in reg_skip : continue
                    # store:
                    if (len(reg_used) == 0) or (reg_code not in reg_used['code'].values) :
                        reg_used = reg_used.append( { 'code' : reg_code, 'name' : reg_name }, ignore_index=True )
                    #endif

                    # no mask yet:
                    reg_mask = None

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

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

                        # 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

                            # setup mask?
                            if reg_mask is None :
                                # info ...
                                logging.info( indent+'      setup region mask ...' )
                                # mask:
                                reg_mask = regmap.GetRegionMask( ireg )
                                # domain, use 10% extra at edge:
                                reg_domain = regmap.GetRegionDomain( ireg, edgefactor=0.10 )
                            #endif
                            
                            # annote:
                            title = '%s  %s' % (reg_name,long_timestamp)

                            # settings for this variable:
                            vkey = 'var.%s' % varname
                            # originating variable:
                            vsource  = self.GetSetting( vkey+'.source' )
                            # 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()
                                # adhoc fix, convert to masked array ...
                                values = numpy.ma.array( data=values, \
                                                         mask=(numpy.isnan(values) | (~ reg_mask)) )

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

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

                                # map properties:
                                bmp = dict( resolution='i', countries=True, title=title )
                                bmp.update( bmp_kwargs )
                                bmp['domain'] = reg_domain

                                # create map figure:
                                fig = cso_plot.QuickMap( values, xm=xm, ym=ym, vmin=vmin, vmax=vmax,
                                                          cmap=dict(colors=colors),
                                                          bmp=bmp, cbar=dict(label=label),
                                                          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
                    
                    ## testing ...
                    #if len(reg_used) >= 5 : break
                    
                #endfor # regions

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

            # next:
            t = t + dt
            
        #endfor # time loop
        
        # write list of codes:
        reg_used.to_csv( 'regions.csv', sep=';', index=False )

        # back:
        os.chdir( owd )
    
    #enddef __init__


#endclass CSO_Catalogue_RegionsMaps


#
# ***
#



class CSO_Catalogue_RegionsTimeSeries( cso_catalogue.CSO_CatalogueBase ) :

    """
    CSO task to create catalogue of timeseries of region aggregates.

    """
    
    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
        import matplotlib.dates  as mdates
        import xarray
        
        # tools:
        import cso_file
        import cso_plot
        import cso_regions
        
        # setup graphics back-end:
        matplotlib.use('Agg') 

        # init base object:
        cso_catalogue.CSO_CatalogueBase.__init__( self, rcfile, rcbase=rcbase, env=env )
        
        # info ...
        logging.info( indent+'' )
        logging.info( indent+'** plot CSO timeseries over regions' )
        logging.info( indent+'' )

        # time range:
        t1    = self.GetSetting( 'timerange.start', totype='datetime' )
        t2    = self.GetSetting( 'timerange.end'  , totype='datetime' )
        
        # steps:
        tsteps = self.GetSetting( 'timerange.steps' ).split()
        # loop over steps:
        for tstep in tsteps :

            # info ...
            tfmt = '%Y-%m-%d %H:%M'
            logging.info( indent+'timerange: [%s,%s] by %s' % (t1.strftime(tfmt),t2.strftime(tfmt),tstep) )

            # target directory:
            workdir = self.GetSetting( 'output.dir' )

            # info ..
            logging.info( indent+'work directory:' )
            logging.info( indent+'  %s' % workdir )
            # create?
            if not os.path.isdir(workdir) : os.makedirs( workdir )
            # store current:
            owd = os.getcwd()
            # change to:
            os.chdir( workdir )

            # extract file for timeseries:
            regts_renew = self.GetSetting( 'regts.renew', totype='bool' )
            # filename:
            regts_file = self.GetSetting( 'regts.file.%s' % tstep )

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

            # target file template:
            figfile_template = self.GetSetting( 'output.file.%s' % tstep )
            # renew existing figures?
            renew = self.GetSetting( 'renew', totype='bool' )
            # figure size:
            figsize = eval( self.GetSetting('figsize') )

            # create extract?
            if (not os.path.isfile(regts_file)) or regts_renew :

                # info ...
                logging.info( indent+'  create %s ...' % regts_file )

                # basename for region definition files:
                regmap_basename = self.GetSetting( 'regions.basename' )
                # read:
                regmap = cso_regions.RegionMap( regmap_basename )
                # count:
                nreg = len(regmap)
                # skip list:
                reg_skip = self.GetSetting( 'regions.skip' ).split()
                # storage for collected region codes, used for index:
                reg_used = pandas.DataFrame()

                # storage:
                tsvalues = { 'tt' : [], 'region' : {} }

                # 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'] :
                        long_timestamp = t.strftime('%Y')
                        dt = datetime.datetime(t.year+1,1,1,0,0) - t
                    else :
                        logging.error( 'unsupported time step "%s"' % tstep )
                        raise Exception
                    #endif

                    # input file:
                    infile = self.GetSetting( 'input.file.%s' % tstep )
                    infile = t.strftime( infile )

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

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

                        # store time:
                        tsvalues['tt'].append( t )

                        # no data read yet:
                        indata = None

                        # loop over regions:
                        for ireg in range(nreg) :

                            # info:
                            reg_code,reg_name = regmap.GetRegionInfo( ireg )
                            # filter ..
                            if reg_code in reg_skip : continue
                            # store:
                            if (len(reg_used) == 0) or (reg_code not in reg_used['code'].values) :
                                reg_used = reg_used.append( { 'code' : reg_code, 'name' : reg_name }, ignore_index=True )
                            #endif

                            # mask:
                            reg_mask = regmap.GetRegionMask( ireg )

                            # storage:
                            if reg_code not in tsvalues['region'].keys() :
                                tsvalues['region'][reg_code] = { 'name' : reg_name }
                            #endif

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

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

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

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

                                # data:
                                values = vda.values.squeeze()
                                # adhoc fix, convert to masked array ...
                                values = numpy.ma.array( data=values, \
                                                         mask=(numpy.isnan(values) | (~ reg_mask)) )

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

                                # storage:
                                if varname not in tsvalues['region'][reg_code].keys() :
                                     tsvalues['region'][reg_code][varname] = { 'data' : [], 'units' : units, 'long_name' : long_name }
                                #endif
                                # store:
                                tsvalues['region'][reg_code][varname]['data'].append( values.mean() )

                                ## testing ...
                                #break

                            #endfor # variables

                            ## testing ...
                            #if len(reg_used) >= 5 : break

                        #endfor # regions

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

                    ## testing ...
                    #if len(tsvalues['tt']) >= 2 : break

                    # next:
                    t = t + dt

                #endfor # time loop

                # * save

                # new:
                dst = xarray.Dataset()
                # add countries:
                dst['region_code'] = xarray.DataArray( reg_used['code'], dims=('region') )
                dst['region_name'] = xarray.DataArray( reg_used['name'], dims=('region') )
                # timevalues:
                dst['time'] = xarray.DataArray( tsvalues['tt'], dims=('time') )
                # dims:
                nregion = dst.dims['region']
                ntime   = dst.dims['time']
                # timeseries:
                for varname in varnames :
                    # collect:
                    values = numpy.zeros((ntime,nregion),float)
                    for iregion in range(nregion) :
                        reg_code = reg_used.at[iregion,'code']
                        values[:,iregion] = tsvalues['region'][reg_code][varname]['data']
                    #endfor
                    # extract using latest region code:
                    long_name = tsvalues['region'][reg_code][varname]['long_name']
                    units     = tsvalues['region'][reg_code][varname]['units']
                    # store:
                    dst[varname] = xarray.DataArray( values, dims=('time','region'),
                                                      attrs={ 'units' : units, 'long_name' : long_name } )
                #endfor
                # save:
                dst.to_netcdf( regts_file )

            else :

                # read:
                dst = xarray.load_dataset( regts_file )

            #endif # renew

            # * plot

            # loop over region:
            for ireg in range(dst.dims['region']) :
                # current:
                reg_code = dst['region_code'].values[ireg]
                reg_name = dst['region_name'].values[ireg]

                # target file:
                figfile = figfile_template
                figfile = figfile.replace( '%{region}', reg_code )

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

                    # annote:
                    title = reg_name

                    # new:
                    fig = plt.figure( figsize=figsize )
                    ax = fig.add_axes([0.07,0.1,0.85,0.85])

                    # loop over variables:
                    for varname in varnames :

                        # annote:
                        label = dst[varname].attrs['long_name']
                        units = dst[varname].attrs['units']

                        # data:
                        tt = dst['time'].values
                        values = dst[varname].values[:,ireg]

                        # add:
                        ax.plot( tt, values, label=label )

                    #endfor

                    # annote:
                    ax.legend()
                    # time axis:
                    ax.set_xlim((t1,t2))
                    ax.xaxis.set_minor_locator( mdates.MonthLocator() )
                    ax.xaxis.set_major_locator( mdates.YearLocator() )
                    ax.xaxis.set_major_formatter( mdates.DateFormatter('%Y'))
                    ax.grid(axis='x')
                    # y-axis:
                    ax.set_ylabel( units )
                    ax.set_ylim((0.0,None))

                    # dir:
                    dname = os.path.dirname(figfile)
                    if len(dname) > 0 :
                        if not os.path.isdir(dname) : os.makedirs( dname )
                    #endif
                    # save:
                    fig.savefig( figfile )

                    # close:
                    plt.close( fig )

                #endif # renew

            #endfor # regions

            # create csv file with station codes, used for index generation:
            reg_used = pandas.DataFrame( { 'code' : dst['region_code'].values,
                                           'name' : dst['region_name'].values } )
            # write:
            reg_used.to_csv( 'regions.csv', sep=';', index=False )
            
        #endfor # tstep

        # * done

        # back:
        os.chdir( owd )
    
    #enddef __init__


#endclass CSO_Catalogue_RegionsTimeSeries





########################################################################
###
### table with statistics
###
########################################################################


class CSO_Statistics_RegionsTables( utopya.UtopyaRc ) :

    """
    CSO task to create tables with statistics over aggregations over regions.

    """
    
    def __init__( self, rcfile, rcbase='', env={}, indent='' ) :
    
        """
        Create statistics.
        """
                        
        # modules:
        import os
        import pandas
        import datetime
        import calendar
        import xarray
        
        # init base object:
        utopya.UtopyaRc.__init__( self, rcfile, rcbase=rcbase, env=env )
        
        # info ...
        logging.info( indent+'' )
        logging.info( indent+'** collect statistics over regions' )
        logging.info( indent+'' )

        # time range:
        t1    = self.GetSetting( 'timerange.start', totype='datetime' )
        t2    = self.GetSetting( 'timerange.end'  , totype='datetime' )

        # time steps:
        tsteps = self.GetSetting( 'timerange.steps' ).split()
        # loop over frequencies:
        for tstep in tsteps :

            # info ...
            tfmt = '%Y-%m-%d %H:%M'
            logging.info( indent+'timerange: [%s,%s] by %s' % (t1.strftime(tfmt),t2.strftime(tfmt),tstep) )

            # target directory:
            workdir = self.GetSetting( 'output.dir' )

            # info ..
            logging.info( indent+'work directory:' )
            logging.info( indent+'  %s' % workdir )
            # create?
            if not os.path.isdir(workdir) : os.makedirs( workdir )
            # store current:
            owd = os.getcwd()
            # change to:
            os.chdir( workdir )

            # filename:
            regts_file = self.GetSetting( 'regts.file.%s' % tstep )
            # create extract?
            if not os.path.isfile(regts_file) :
                logging.error( 'region timeseries files not found: %s' % regts_file )
                raise Exception
            #endif
            # read:
            dst = xarray.load_dataset( regts_file )

            # target file for table:
            outfile = self.GetSetting( 'output.file.%s' % tstep )

            # variable names:
            vname_obs = 'yr_m'
            vname_sim = 'ys_m'

            # table for iso2 column:
            ctab_file = self.GetSetting( 'ctab.file' )
            # read:
            ctab = pandas.read_csv( ctab_file, sep=';' )

            # info ..
            logging.info( indent+'collect statistics ...' )
            # target table:
            df = pandas.DataFrame()
            # loop over regions:
            for ireg in range(dst.dims['region']) :
                # current:
                reg_code = dst['region_code'].values[ireg]
                reg_name = dst['region_name'].values[ireg]

                # iso2 code:
                if reg_code not in ctab['Alpha-3 code'].values :
                    logging.warning( indent+'  iso3 code "%s" not found in table: %s' % (reg_code,ctab_file) )
                    continue
                    #raise Exception
                #endif
                reg_code2 = ctab[ ctab['Alpha-3 code'] == reg_code ].iloc[0]['Alpha-2 code']

                # loop over time values:
                for itime in range(dst.dims['time']) :

                    # time value:
                    t = pandas.Timestamp( dst['time'].values[itime] )
                    # formatted:
                    if tstep in ['month'] :
                        tlab = t.strftime('%Y-%m')
                    elif tstep in ['year'] :
                        tlab = t.strftime('%Y')
                    else :
                        logging.error( 'could nt set time label for tstep "%s"' % tstep )
                        raise Exception
                    #endif

                    # extract:
                    obs = dst[vname_obs].values[itime,ireg] ; obs_label = 'obs'
                    sim = dst[vname_sim].values[itime,ireg] ; sim_label = 'sim'
                    # bias relative to observation:
                    rbias = (sim-obs)/obs ; rbias_label = '(sim-obs)/obs'
                    # add record:
                    df = df.append( pandas.DataFrame( { 'iso2'      : [reg_code2],
                                                        'iso3'      : [reg_code],
                                                        'name'      : [reg_name],
                                                        'time'      : [tlab],
                                                        obs_label   : [obs],
                                                        sim_label   : [sim],
                                                        rbias_label : [rbias] } ),
                                    ignore_index=True )

                #endfor # itime
            #endfor # ireg

            # target file:
            outfile = 'EMEP_S5p-NO2_stats_%s.csv' % tstep
            # info ..
            logging.info( 'save to: %s' % outfile )
            # save ..
            df.to_csv( outfile, sep=';', index=False )

        #endfor # tstep

        # back:
        os.chdir( owd )
    
    #enddef __init__


#endclass CSO_Statistics_RegionsTables

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

