# -*- coding: utf-8 -*-
"""
QGIS Processing Toolbox script for rescaling rasters using inflection points.
Based on r.fuzzy.set from GRASS GIS:
https://grass.osgeo.org/grass-stable/manuals/addons/r.fuzzy.set.html
Author: Paulo van Breugel
"""

from qgis.core import (
    QgsProcessingAlgorithm,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterString,
    QgsProcessingParameterNumber,
    QgsProcessingParameterEnum,
    QgsProcessingParameterRasterDestination,
    QgsProcessingException,
)
from qgis.PyQt.QtCore import QCoreApplication
import numpy as np
from osgeo import gdal


def tr(text):
    return QCoreApplication.translate("FuzzyRescale", text)


class FuzzyRescale(QgsProcessingAlgorithm):
    INPUT = "INPUT"
    LEFT = "LEFT"
    RIGHT = "RIGHT"
    SHAPE = "SHAPE"
    FUNCTION = "FUNCTION"
    OUTPUT = "OUTPUT"

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(self.INPUT, tr("Input Raster Layer"))
        )

        self.addParameter(
            QgsProcessingParameterString(
                self.LEFT,
                tr("Left side (A,B)"),
                defaultValue="",
                optional=True,
            )
        )
        self.parameterDefinition(self.LEFT).setHelp(
            tr(
                "Optional. Enter two comma-separated values A,B to define the left slope of the rescale function."
            )
        )

        self.addParameter(
            QgsProcessingParameterString(
                self.RIGHT, tr("Right side (C,D)"), optional=True
            )
        )
        self.parameterDefinition(self.RIGHT).setHelp(
            tr(
                "Optional. Enter two comma-separated values C,D to define the right slope of the rescale function."
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.SHAPE,
                tr("Shape Parameter (-1 to 1)"),
                QgsProcessingParameterNumber.Double,
                0.0,
                False,
                -1.0,
                1.0,
            )
        )
        self.parameterDefinition(self.SHAPE).setHelp(
            tr(
                "Controls the sharpness of the curve.\n"
                "-1 = soft transition\n"
                " 0 = default\n"
                "+1 = sharp transition"
            )
        )

        self.addParameter(
            QgsProcessingParameterEnum(
                self.FUNCTION,
                tr("Rescale function"),
                options=["Linear", "S-shaped", "G-shaped", "J-shaped"],
                defaultValue=0,
            )
        )
        self.parameterDefinition(self.FUNCTION).setHelp(
            tr(
                "Choose the type of rescale function to apply:\n"
                "- Linear: piecewise linear slope(s)\n"
                "- S-shaped: smooth curve with sinusoidal shape\n"
                "- G-shaped: gradual increase\n"
                "- J-shaped: steep increase"
            )
        )

        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT, tr("Output Rescaled Raster")
            )
        )

    def processAlgorithm(self, parameters, context, feedback):
        shape = self.parameterAsDouble(parameters, self.SHAPE, context)
        func_index = self.parameterAsEnum(parameters, self.FUNCTION, context)
        function_name = ["Linear", "S-shaped", "G-shaped", "J-shaped"][func_index]
        left_str = self.parameterAsString(parameters, self.LEFT, context)
        right_str = self.parameterAsString(parameters, self.RIGHT, context)

        def parse_pair(pair_str):
            try:
                values = [float(v.strip()) for v in pair_str.split(",")]
                if len(values) != 2:
                    raise ValueError
                return values
            except Exception:
                raise QgsProcessingException(
                    f"Invalid format for value pair: '{pair_str}'. Use format: value1,value2"
                )

        a, b = parse_pair(left_str) if left_str else (None, None)
        c, d = parse_pair(right_str) if right_str else (None, None)

        use_left = a is not None and b is not None
        use_right = c is not None and d is not None

        if not (use_left or use_right):
            raise QgsProcessingException(
                "Please provide at least one valid inflection point pair (left or right)."
            )

        mode = "both" if use_left and use_right else "left" if use_left else "right"

        input_layer = self.parameterAsRasterLayer(parameters, self.INPUT, context)
        output_path = self.parameterAsOutputLayer(parameters, self.OUTPUT, context)

        dataset = gdal.Open(input_layer.source(), gdal.GA_ReadOnly)
        band = dataset.GetRasterBand(1)
        arr = band.ReadAsArray().astype(np.float32)

        def normalize(val):
            return np.clip(val, 0, 1)

        def rescale_linear(x):
            result = np.zeros_like(x)
            if use_left:
                result[(x > a) & (x < b)] = (x[(x > a) & (x < b)] - a) / (b - a)
                if not use_right:
                    result[(x >= b)] = 1.0
            if use_right:
                result[(x > c) & (x < d)] = (d - x[(x > c) & (x < d)]) / (d - c)
                if use_left and b is not None and c is not None:
                    result[(x >= b) & (x <= c)] = 1.0
                elif not use_left:
                    result[(x <= c)] = 1.0
            return normalize(result)

        def rescale_s(x):
            m = 2 ** np.exp(2 * abs(shape)) if shape != 0 else 2
            linear_result = rescale_linear(x)
            output = np.ones_like(linear_result)
            mask = linear_result < 1.0

            if shape > 0:
                output[mask] = np.sin(linear_result[mask] * np.pi / 2) ** m
            elif shape < 0:
                x_clip = np.clip(linear_result[mask], 0, 1)
                cos_part = np.cos(x_clip * np.pi / 2)
                cos_part = np.where(cos_part < 1e-8, 0.0, cos_part)
                output[mask] = 1 - (cos_part**m)
            else:
                output[mask] = np.sin(linear_result[mask] * np.pi / 2) ** 2

            return normalize(output)

        def rescale_g(x):
            m = 2 ** np.exp(-2 * shape) if shape < 0 else 2 ** (1 - shape)
            linear_result = rescale_linear(x)
            output = np.ones_like(linear_result)
            mask = linear_result < 1.0
            safe_val = np.clip(linear_result[mask], 1e-6, 1)
            output[mask] = np.tan(safe_val * np.pi / 4) ** (1 / m)
            return normalize(output)

        def rescale_j(x):
            m = 2 ** np.exp(2 * shape) if shape > 0 else 2 ** (1 + shape)
            linear_result = rescale_linear(x)
            output = np.ones_like(linear_result)
            mask = linear_result < 1.0
            safe_val = np.clip(linear_result[mask], 1e-6, 1)
            output[mask] = np.tan(safe_val * np.pi / 4) ** m
            return normalize(output)

        if function_name == "Linear":
            out_arr = rescale_linear(arr)
        elif function_name == "S-shaped":
            out_arr = rescale_s(arr)
        elif function_name == "G-shaped":
            out_arr = rescale_g(arr)
        elif function_name == "J-shaped":
            out_arr = rescale_j(arr)
        else:
            raise QgsProcessingException(f"Function '{function_name}' not implemented.")

        driver = gdal.GetDriverByName("GTiff")
        out_ds = driver.Create(
            output_path, arr.shape[1], arr.shape[0], 1, gdal.GDT_Float32
        )
        out_ds.SetGeoTransform(dataset.GetGeoTransform())
        out_ds.SetProjection(dataset.GetProjection())
        out_ds.GetRasterBand(1).WriteArray(out_arr)
        out_ds.GetRasterBand(1).SetNoDataValue(np.nan)
        out_ds.FlushCache()

        return {self.OUTPUT: output_path}

    def name(self):
        return "rescale_rescale_inflection"

    def displayName(self):
        return tr("Rescale or fuzzy membership functions")

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

    def groupId(self):
        return "raster_analysis"

    def shortHelpString(self):
        return tr(
            "<h2>Description</h2>"
            "<p>In addition to the built-in Fuzzify function available in the QGIS toolbox, this custom rescale tool offers "
            "an extended set of rescale functions that you can use to transform your raster data and are arguably easier to use.</p>"
            "<h2>Shape</h2>"
            "<p>To determine how values of the input raster are transformed, four different functions are available "
            "(linear, S-shaped, G-shaped and J-shaped). Each function defines how values are transformed between 0 and 1. "
            "Choose the one that best reflects how suitability increases or decreases across your landscape. The functions can "
            "be applied to the left boundary (suitability increasing with input values) or the right boundary (suitability "
            "decreasing with input values).</p>"
            "<p>With the <b>shape parameter</b> (-1 to 1), you can fine-tune how gradual or abrupt the shift from unsuitable "
            "to suitable (or vice versa) should be. The default is 0. Positive values result in a steeper transition, while "
            "negative values result in a gentler transition.</p>"
            "<h2>Boundary</h2>"
            "<p>You can define either 2 or 4 inflection points, depending on whether you want to rescale just one side "
            "(lower or higher end of the distribution) or both sides of the distribution. Use the <b>Left side</b> parameter "
            "to enter two comma-separated values A,B and the <b>Right side</b> parameter to enter two comma-separated values C,D.</p>"
            "<p><b>Both A,B and C,D provided</b> (two-sided rescaling): values &lt; A or &gt; D become 0; values between B and C "
            "become 1; values between A and B or between C and D are rescaled between 0 and 1 using the selected function.</p>"
            "<p><b>Only A,B provided</b> (left-sided rescaling): values &lt; A become 0; values &gt; B become 1; values between "
            "A and B are rescaled using the selected function.</p>"
            "<p><b>Only C,D provided</b> (right-sided rescaling): values &lt; C become 1; values &gt; D become 0; values between "
            "C and D are rescaled using the selected function.</p>"
            "<p>The points do not have to be in map range, but this may lead to only 0 or 1 membership for the whole map.</p>"
            "<h2>Note</h2>"
            "<p>Based on the <em>r.fuzzy.set</em> addon for GRASS. See its "
            "<a href='https://grass.osgeo.org/grass-stable/manuals/addons/r.fuzzy.set.html'>manual page</a> for more information about the "
            "functions and a visual explanation of the parameters.</p>"
            "<p>The term 'rescale function' is used instead of 'fuzzy membership function' because it can be used in applications "
            "other than fuzzy logic, in contexts such as raster normalization and suitability mapping.</p>"
            "<h2>References</h2>"
            "<p>Jasiewicz, J. (2011). A new GRASS GIS fuzzy inference system for massive data analysis. Computers & Geosciences "
            "(37) 1525-1531. DOI https://doi.org/10.1016/j.cageo.2010.09.008</p>"
            "<p>r.fuzzy.set addon for GRASS <a href='https://grass.osgeo.org/grass-stable/manuals/addons/r.fuzzy.set.html'>manual page</a></p>"
            "<h2>Author</h2>"
            "<p><a href='https:ecodiv.earth'>Paulo van Breugel</a>, <a href='https://has.nl'>HAS green academy</a>, "
            "<a href='https://www.has.nl/en/research/professorships/innovative-bio-monitoring-professorship/'>Innovative Biomonitoring research group</a>, "
            "<a href='https://www.has.nl/en/research/professorships/climate-robust-landscapes-professorship/'>Climate-robust Landscapes research group</a></p>"
        )

    @staticmethod
    def createInstance():
        return FuzzyRescale()