QGIS/python/plugins/processing/algs/qgis/Relief.py

"""
***************************************************************************
    Relief.py
    ---------------------
    Date                 : December 2016
    Copyright            : (C) 2016 by Alexander Bruy
    Email                : alexander dot bruy at gmail dot com
***************************************************************************
*                                                                         *
*   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__ = "Alexander Bruy"
__date__ = "December 2016"
__copyright__ = "(C) 2016, Alexander Bruy"

import os

from qgis.PyQt.QtGui import QIcon, QColor

from qgis.analysis import QgsRelief
from qgis.core import (
    QgsProcessingParameterDefinition,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterNumber,
    QgsProcessingParameterBoolean,
    QgsProcessingParameterRasterDestination,
    QgsProcessingParameterFileDestination,
    QgsRasterFileWriter,
    QgsProcessingException,
)
from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm

pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]


class ParameterReliefColors(QgsProcessingParameterDefinition):

    def __init__(self, name="", description="", parent=None, optional=True):
        super().__init__(name, description, None, optional)
        self.parent = parent
        self.setMetadata(
            {
                "widget_wrapper": "processing.algs.qgis.ui.ReliefColorsWidget.ReliefColorsWidgetWrapper"
            }
        )

    def type(self):
        return "relief_colors"

    def clone(self):
        return ParameterReliefColors(
            self.name(),
            self.description(),
            self.parent,
            self.flags() & QgsProcessingParameterDefinition.Flag.FlagOptional,
        )

    @staticmethod
    def valueToColors(value):
        if value is None:
            return None

        if value == "":
            return None

        if isinstance(value, str):
            return value.split(";")
        else:
            return ParameterReliefColors.colorsToString(value)

    @staticmethod
    def colorsToString(colors):
        return ";".join(
            f"{c[0]:f}, {c[1]:f}, {c[2]:d}, {c[3]:d}, {c[4]:d}" for c in colors
        )


class Relief(QgisAlgorithm):
    INPUT = "INPUT"
    Z_FACTOR = "Z_FACTOR"
    AUTO_COLORS = "AUTO_COLORS"
    COLORS = "COLORS"
    OUTPUT = "OUTPUT"
    FREQUENCY_DISTRIBUTION = "FREQUENCY_DISTRIBUTION"

    def icon(self):
        return QIcon(os.path.join(pluginPath, "images", "dem.png"))

    def group(self):
        return self.tr("Raster terrain analysis")

    def groupId(self):
        return "rasterterrainanalysis"

    def __init__(self):
        super().__init__()

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(self.INPUT, self.tr("Elevation layer"))
        )
        self.addParameter(
            QgsProcessingParameterNumber(
                self.Z_FACTOR,
                self.tr("Z factor"),
                type=QgsProcessingParameterNumber.Type.Double,
                minValue=0.00,
                defaultValue=1.0,
            )
        )
        self.addParameter(
            QgsProcessingParameterBoolean(
                self.AUTO_COLORS,
                self.tr("Generate relief classes automatically"),
                defaultValue=False,
            )
        )
        self.addParameter(
            ParameterReliefColors(
                self.COLORS, self.tr("Relief colors"), self.INPUT, True
            )
        )
        self.addParameter(
            QgsProcessingParameterRasterDestination(self.OUTPUT, self.tr("Relief"))
        )
        self.addParameter(
            QgsProcessingParameterFileDestination(
                self.FREQUENCY_DISTRIBUTION,
                self.tr("Frequency distribution"),
                "CSV files (*.csv)",
                optional=True,
                createByDefault=False,
            )
        )

    def name(self):
        return "relief"

    def displayName(self):
        return self.tr("Relief")

    def processAlgorithm(self, parameters, context, feedback):
        inputFile = self.parameterAsRasterLayer(
            parameters, self.INPUT, context
        ).source()
        zFactor = self.parameterAsDouble(parameters, self.Z_FACTOR, context)
        automaticColors = self.parameterAsBoolean(parameters, self.AUTO_COLORS, context)
        outputFile = self.parameterAsOutputLayer(parameters, self.OUTPUT, context)
        frequencyDistribution = self.parameterAsFileOutput(
            parameters, self.FREQUENCY_DISTRIBUTION, context
        )

        outputFormat = QgsRasterFileWriter.driverForExtension(
            os.path.splitext(outputFile)[1]
        )

        relief = QgsRelief(inputFile, outputFile, outputFormat)

        if automaticColors:
            reliefColors = relief.calculateOptimizedReliefClasses()
        else:
            colors = ParameterReliefColors.valueToColors(parameters[self.COLORS])
            if colors is None or len(colors) == 0:
                raise QgsProcessingException(
                    self.tr(
                        'Specify relief colors or activate "Generate relief classes automatically" option.'
                    )
                )

            reliefColors = []
            for c in colors:
                v = c.split(",")
                color = QgsRelief.ReliefColor(
                    QColor(int(v[2]), int(v[3]), int(v[4])), float(v[0]), float(v[1])
                )
                reliefColors.append(color)

        relief.setReliefColors(reliefColors)
        relief.setZFactor(zFactor)
        if frequencyDistribution:
            relief.exportFrequencyDistributionToCsv(frequencyDistribution)
        res = relief.processRaster(feedback)
        if res == 1:
            raise QgsProcessingException(self.tr("Can not open input file."))
        elif res == 2:
            raise QgsProcessingException(
                self.tr("Can not get GDAL driver for output file.")
            )
        elif res == 3:
            raise QgsProcessingException(self.tr("Can not create output file."))
        elif res == 4:
            raise QgsProcessingException(self.tr("Can not get input band."))
        elif res == 5:
            raise QgsProcessingException(self.tr("Can not create output bands."))
        elif res == 6:
            raise QgsProcessingException(
                self.tr("Output raster size is too small (at least 3 rows needed).")
            )
        elif res == 7:
            feedback.pushInfo(self.tr("Canceled."))

        return {
            self.OUTPUT: outputFile,
            self.FREQUENCY_DISTRIBUTION: frequencyDistribution,
        }