s2toolkit_algorithm.py

# -*- coding: utf-8 -*-

"""
/***************************************************************************
 S2Toolkit
                                 A QGIS plugin
 Tools for the S2 Geometry.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2024-05-02
        copyright            : (C) 2024 by Mikhail Melnikov, Oleg Petrov
        email                : support@ibtgeo.ru
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

__author__ = 'Mikhail Melnikov, Oleg Petrov'
__date__ = '2024-05-02'
__copyright__ = '(C) 2024 by Mikhail Melnikov, Oleg Petrov'

# This will get replaced with a git SHA1 when you do a git archive

__revision__ = '$Format:%H$'


from qgis.PyQt.QtCore import QCoreApplication
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtCore import QVariant
from qgis.core import (QgsProcessing,
                       QgsFeatureSink,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterFeatureSink,
                       QgsProcessingException,
                       QgsFields,
                       QgsField,
                       QgsWkbTypes,
                       QgsFeatureRequest,
                       QgsRectangle,
                       QgsFeature,
                       QgsGeometry,
                       QgsCoordinateTransform,
                       QgsCoordinateReferenceSystem,
                       QgsProject,
                       QgsProcessingParameterNumber)
from .s2sphere import s2sphere
from . import resources
import json
import os
pluginPath = os.path.dirname(__file__)

class S2Toolkit_fromLayerExtent(QgsProcessingAlgorithm):
    """
    TODO: Add docstring
    """

    OUTPUT = 'OUTPUT'
    LVL = 'LVL'
    INPUT = 'INPUT'

    def initAlgorithm(self, config):
        
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT,
                self.tr('Input layer'),
                [QgsProcessing.TypeVectorAnyGeometry]
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.LVL,
                self.tr('S2 Geometry level'),
                type=QgsProcessingParameterNumber.Type.Integer,
                minValue = 1, defaultValue = 12
            )
        )
        
        self.addParameter(
            QgsProcessingParameterFeatureSink(
                self.OUTPUT,
                self.tr('Output layer')
            )
        )
        

    def processAlgorithm(self, parameters, context, feedback):
        
        source = self.parameterAsSource(parameters, self.INPUT, context)
        if source is None:
            raise QgsProcessingException(self.invalidSourceError(parameters,
                                                                 self.INPUT))
        fields = QgsFields()
        fields.append(QgsField('id', QVariant.Int, '', 20))
        fields.append(QgsField('level', QVariant.Int, '', 3))
        fields.append(QgsField('cellid_int', QVariant.String, 'string', 20))
        fields.append(QgsField('cellid_hex', QVariant.String, 'string', 19))
        
        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT,
                                    context, fields,
                                    QgsWkbTypes.Type.Polygon,
                                    QgsCoordinateReferenceSystem("EPSG:4326"))
        if sink is None:
            raise QgsProcessingException(self.invalidSinkError(parameters,
                                                               self.OUTPUT))

        # Bounding box
        features = source.getFeatures(QgsFeatureRequest().
                                      setSubsetOfAttributes([]))
        bounds = QgsRectangle()
        for current, f in enumerate(features):
            if feedback.isCanceled():
                break
            if not f.hasGeometry():
                continue
            # Bounding boxes, calculate on the fly for efficiency
            bounds.combineExtentWith(f.geometry().boundingBox())

        # Transform projected coordinates to WGS84
        transform = QgsCoordinateTransform(source.sourceCrs(),
                                    QgsCoordinateReferenceSystem("EPSG:4326"),
                                    QgsProject.instance())
        bounds_wgs84 = transform.transformBoundingBox(bounds)
        # Get coordinates of boundaries
        xmin, xmax, ymax, ymin = (bounds_wgs84.xMaximum(),
                                  bounds_wgs84.xMinimum(),
                                  bounds_wgs84.yMaximum(),
                                  bounds_wgs84.yMinimum())
        # Get S2 Rectangle
        region_rect = s2sphere.LatLngRect.from_point_pair(
            s2sphere.LatLng.from_degrees(ymin, xmin),
            s2sphere.LatLng.from_degrees(ymax, xmax))
        # Create coverer
        coverer = s2sphere.RegionCoverer()
        # Get min and max levels from parameters
        min_level = self.parameterAsInt(parameters, self.LVL, context)
        max_level = self.parameterAsInt(parameters, self.LVL, context)
        # Raise Error if min_level bigger than max_level
        coverer.min_level = min_level
        coverer.max_level = max_level
        covering = coverer.get_covering(region_rect)
        
        total = 100.0 / (len(covering)) if covering else 0
        
        # Cell id to geometry in WKT, then into sink
        for current, cell in enumerate(covering):
            cell_obj = s2sphere.Cell(s2sphere.CellId(cell.id()))
            vs = []
            for i in range(4):
                vertex = cell_obj.get_vertex(i)
                lat_lng = s2sphere.LatLng.from_point(vertex)
                vs.append((float(lat_lng.lng().degrees),
                           float(lat_lng.lat().degrees)))
            wkt = (f'POLYGON (({vs[0][0]} {vs[0][1]}, {vs[1][0]} {vs[1][1]}, '
                   f'{vs[2][0]} {vs[2][1]}, {vs[3][0]} {vs[3][1]}, '
                   f'{vs[0][0]} {vs[0][1]}))')
            if not feedback.isCanceled():
                s2feature = QgsFeature()
                s2feature.setGeometry(QgsGeometry.fromWkt(wkt))
                s2feature.setAttributes([current, min_level,
                                       cell.id(), hex(cell.id())])
                sink.addFeature(s2feature, QgsFeatureSink.Flag.FastInsert)
                # Update the progress bar
                feedback.setProgress(int(current * total))

        return {self.OUTPUT: dest_id}

    def name(self):
        return 'S2 from Layer Extent'

    def displayName(self):
        return self.tr(self.name())

    def group(self):
        return self.tr(self.groupId())

    def groupId(self):
        return ''

    def shortHelpString(self):
        help_str = """
        This tool generates an S2 geometry polygon vector layer from the \
        bounding box extent of a vector layer with any type of geometry \
        (points, lines, polygons). Unlike other methods that may use the \
        layer's predefined extent, this tool specifically calculates \
        the bounding box to capture the actual spatial extent of the \
        vector features.

        <b>Usage</b>
        <u>Input Layer</u>: Any vector layer.
        <u>S2 Geometry Level</u>: Desired S2 geometry level. \
        CAUTION: Levels higher than 17 may require extended processing time!
        <u>Output Layer</u>: Polygon layer projected in WGS84 (EPSG:4326)
        """

        return self.tr(help_str)

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def icon(self):
        return QIcon(os.path.join(pluginPath, "logo.png"))
    
    def createInstance(self):
        return S2Toolkit_fromLayerExtent()

class S2Toolkit_fromFeatureExtents(QgsProcessingAlgorithm):
    """
    TODO: Add docstring
    """

    OUTPUT = 'OUTPUT'
    LVL = 'LVL'
    INPUT = 'INPUT'

    def initAlgorithm(self, config):
        
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT,
                self.tr('Input layer'),
                [QgsProcessing.TypeVectorAnyGeometry]
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.LVL,
                self.tr('S2 Geometry level'),
                type=QgsProcessingParameterNumber.Type.Integer,
                minValue = 1, defaultValue = 12
            )
        )
        
        self.addParameter(
            QgsProcessingParameterFeatureSink(
                self.OUTPUT,
                self.tr('Output layer')
            )
        )
        

    def processAlgorithm(self, parameters, context, feedback):
        source = self.parameterAsSource(parameters, self.INPUT, context)
        if source is None:
            raise QgsProcessingException(self.invalidSourceError(parameters, self.INPUT))

        fields = QgsFields()
        fields.append(QgsField('id', QVariant.Int, '', 20))
        fields.append(QgsField('level', QVariant.Int, '', 3))
        fields.append(QgsField('cellid_int', QVariant.String, 'string', 20))
        fields.append(QgsField('cellid_hex', QVariant.String, 'string', 19))

        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT,
                                               context, fields,
                                               QgsWkbTypes.Polygon,
                                               QgsCoordinateReferenceSystem("EPSG:4326"))
        if sink is None:
            raise QgsProcessingException(self.invalidSinkError(parameters, self.OUTPUT))

        min_level = self.parameterAsInt(parameters, self.LVL, context)
        max_level = self.parameterAsInt(parameters, self.LVL, context)

        total = 100.0 / source.featureCount() if source.featureCount() else 0

        for current, feature in enumerate(source.getFeatures()):
            if feedback.isCanceled():
                break

            if not feature.hasGeometry():
                continue

            # Get feature bounding box and transform it to WGS84
            bbox = feature.geometry().boundingBox()
            transform = QgsCoordinateTransform(source.sourceCrs(),
                                               QgsCoordinateReferenceSystem("EPSG:4326"),
                                               QgsProject.instance())
            bbox_wgs84 = transform.transformBoundingBox(bbox)

            xmin, xmax, ymax, ymin = (bbox_wgs84.xMaximum(),
                                      bbox_wgs84.xMinimum(),
                                      bbox_wgs84.yMaximum(),
                                      bbox_wgs84.yMinimum())

            region_rect = s2sphere.LatLngRect.from_point_pair(
                s2sphere.LatLng.from_degrees(ymin, xmin),
                s2sphere.LatLng.from_degrees(ymax, xmax))

            coverer = s2sphere.RegionCoverer()
            coverer.min_level = min_level
            coverer.max_level = max_level
            covering = coverer.get_covering(region_rect)

            for cell in covering:
                cell_obj = s2sphere.Cell(s2sphere.CellId(cell.id()))
                vs = []
                for i in range(4):
                    vertex = cell_obj.get_vertex(i)
                    lat_lng = s2sphere.LatLng.from_point(vertex)
                    vs.append((float(lat_lng.lng().degrees),
                               float(lat_lng.lat().degrees)))
                wkt = (f'POLYGON (({vs[0][0]} {vs[0][1]}, {vs[1][0]} {vs[1][1]}, '
                       f'{vs[2][0]} {vs[2][1]}, {vs[3][0]} {vs[3][1]}, '
                       f'{vs[0][0]} {vs[0][1]}))')

                s2feature = QgsFeature()
                s2feature.setGeometry(QgsGeometry.fromWkt(wkt))
                s2feature.setAttributes([current, min_level,
                                         cell.id(), hex(cell.id())])
                sink.addFeature(s2feature, QgsFeatureSink.FastInsert)

            feedback.setProgress(int(current * total))

        return {self.OUTPUT: dest_id}

    def name(self):
        return 'S2 from Feature Extents'

    def displayName(self):
        return self.tr(self.name())

    def group(self):
        return self.tr(self.groupId())

    def groupId(self):
        return ''

    def shortHelpString(self):
        help_str = """
        This tool generates an S2 geometry polygon vector layer from the \
        bounding boxes of a vector layer features with any type of geometry \
        (points, lines, polygons). This tool specifically calculates bounding \
        box of every feature in input vector layer.

        <b>Usage</b>
        <u>Input Layer</u>: Any vector layer.
        <u>S2 Geometry Level</u>: Desired S2 geometry level. \
        CAUTION: Levels higher than 17 may require extended processing time!
        <u>Output Layer</u>: Polygon layer projected in WGS84 (EPSG:4326)
        """

        return self.tr(help_str)

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def icon(self):
        return QIcon(os.path.join(pluginPath, "logo.png"))
    
    def createInstance(self):
        return S2Toolkit_fromFeatureExtents()

class S2Toolkit_fromGeometry(QgsProcessingAlgorithm):
    """
    TODO: Add docstring
    """

    OUTPUT = 'OUTPUT'
    LVL = 'LVL'
    INPUT = 'INPUT'

    def initAlgorithm(self, config):
        
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT,
                self.tr('Input layer'),
                [QgsProcessing.TypeVectorAnyGeometry]
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.LVL,
                self.tr('S2 Geometry level'),
                type=QgsProcessingParameterNumber.Type.Integer,
                minValue = 1, defaultValue = 12
            )
        )
        
        self.addParameter(
            QgsProcessingParameterFeatureSink(
                self.OUTPUT,
                self.tr('Output layer')
            )
        )
        

    def processAlgorithm(self, parameters, context, feedback):
        source = self.parameterAsSource(parameters, self.INPUT, context)
        if source is None:
            raise QgsProcessingException(self.invalidSourceError(parameters, self.INPUT))

        fields = QgsFields()
        fields.append(QgsField('id', QVariant.Int, '', 20))
        fields.append(QgsField('level', QVariant.Int, '', 3))
        fields.append(QgsField('cellid_int', QVariant.String, 'string', 20))
        fields.append(QgsField('cellid_hex', QVariant.String, 'string', 19))

        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT,
                                               context, fields,
                                               QgsWkbTypes.Polygon,
                                               QgsCoordinateReferenceSystem("EPSG:4326"))
        if sink is None:
            raise QgsProcessingException(self.invalidSinkError(parameters, self.OUTPUT))

        features = source.getFeatures(QgsFeatureRequest().setSubsetOfAttributes([]))
        bounds = QgsRectangle()
        for current, f in enumerate(features):
            if feedback.isCanceled():
                break
            if not f.hasGeometry():
                continue
            bounds.combineExtentWith(f.geometry().boundingBox())

        transform = QgsCoordinateTransform(source.sourceCrs(),
                                           QgsCoordinateReferenceSystem("EPSG:4326"),
                                           QgsProject.instance())
        bounds_wgs84 = transform.transformBoundingBox(bounds)

        xmin, xmax, ymax, ymin = (bounds_wgs84.xMaximum(),
                                  bounds_wgs84.xMinimum(),
                                  bounds_wgs84.yMaximum(),
                                  bounds_wgs84.yMinimum())

        region_rect = s2sphere.LatLngRect.from_point_pair(
            s2sphere.LatLng.from_degrees(ymin, xmin),
            s2sphere.LatLng.from_degrees(ymax, xmax))

        coverer = s2sphere.RegionCoverer()
        min_level = self.parameterAsInt(parameters, self.LVL, context)
        max_level = self.parameterAsInt(parameters, self.LVL, context)

        coverer.min_level = min_level
        coverer.max_level = max_level
        covering = coverer.get_covering(region_rect)

        total = 100.0 / len(covering) if covering else 0

        input_features = list(source.getFeatures())

        for current, cell in enumerate(covering):
            if feedback.isCanceled():
                break

            cell_obj = s2sphere.Cell(s2sphere.CellId(cell.id()))
            vs = []
            for i in range(4):
                vertex = cell_obj.get_vertex(i)
                lat_lng = s2sphere.LatLng.from_point(vertex)
                vs.append((float(lat_lng.lng().degrees),
                           float(lat_lng.lat().degrees)))
            wkt = (f'POLYGON (({vs[0][0]} {vs[0][1]}, {vs[1][0]} {vs[1][1]}, '
                   f'{vs[2][0]} {vs[2][1]}, {vs[3][0]} {vs[3][1]}, '
                   f'{vs[0][0]} {vs[0][1]}))')
            
            cell_geometry = QgsGeometry.fromWkt(wkt)
            intersects = False
            
            for f in input_features:
                geom = f.geometry()
                geom.transform(transform)
                if cell_geometry.intersects(geom):
                    intersects = True
                    break

            if intersects:
                s2feature = QgsFeature()
                s2feature.setGeometry(cell_geometry)
                s2feature.setAttributes([current, min_level,
                                         cell.id(), hex(cell.id())])
                sink.addFeature(s2feature, QgsFeatureSink.FastInsert)
                feedback.setProgress(int(current * total))


        return {self.OUTPUT: dest_id}

    def name(self):
        return 'S2 from Geometry'

    def displayName(self):
        return self.tr(self.name())

    def group(self):
        return self.tr(self.groupId())

    def groupId(self):
        return ''

    def shortHelpString(self):
        help_str = """
        This tool generates an S2 geometry polygon vector layer from the \
        features of a vector layer with any type of geometry (points, lines, \
        polygons) that intersects S2 Cells.

        <b>Usage</b>
        <u>Input Layer</u>: Any vector layer.
        <u>S2 Geometry Level</u>: Desired S2 geometry level. \
        CAUTION: Levels higher than 17 may require extended processing time!
        <u>Output Layer</u>: Polygon layer projected in WGS84 (EPSG:4326)
        """

        return self.tr(help_str)

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def icon(self):
        return QIcon(os.path.join(pluginPath, "logo.png"))
    
    def createInstance(self):
        return S2Toolkit_fromGeometry()