import numpy as np
from qgis.core import (
    Qgis,
    QgsCategorizedSymbolRenderer,
    QgsColorRampShader,
    QgsGeometryGeneratorSymbolLayer,
    QgsHashedLineSymbolLayer,
    QgsInterpolatedLineColor,
    QgsInterpolatedLineSymbolLayer,
    QgsInterpolatedLineWidth,
    QgsLineSymbol,
    QgsMarkerSymbol,
    QgsProperty,
    QgsRendererCategory,
    QgsRuleBasedLabeling,
    QgsSimpleLineSymbolLayer,
    QgsSimpleMarkerSymbolLayer,
    QgsSymbol,
    QgsSymbolLayer,
    QgsVectorLayer,
)
from qgis.PyQt.QtCore import QVariant
from qgis.PyQt.QtGui import QColor

from openlog.core.pint_utilities import convert_to_base_unit
from openlog.datamodel.assay.generic_assay import AssayDatabaseDefinition
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)


class TraceSymbology:
    """
    Convert symbology from AssayColumnVisualizationConfig and apply to a QgsVectorLayer
    """

    def __init__(
        self,
        config: AssayColumnVisualizationConfig = None,
        assay_def: AssayDatabaseDefinition = None,
    ) -> None:
        self.config = config
        self.assay_def = assay_def

    def _get_extended_num_layer(
        self, col_list: [str], layer: QgsVectorLayer = None
    ) -> QgsSimpleLineSymbolLayer:
        """
        Generic method creating a QgsSimpleLineSymbolLayer for :
        - extended assay
        Color is mapped through expression.
        """
        # construct expression
        col_name = col_list[0]
        params = self.config.color_ramp_handler.get_parameters()

        if params.get("discretize_method") == "Linear":
            discrete = "false"
            from_, to_, colors = params.get("lut")
        else:
            discrete = "true"
            from_, to_, colors = self.config.color_ramp_handler.get_lut(fill_gaps=True)

        qcolors = [QColor(color) for color in colors]
        min_ = params.get("lut_min_value")
        max_ = params.get("lut_max_value")
        color_ramp_str = [
            f"{(value - min_)/(max_ - min_)},'{qcolor.red()},{qcolor.green()},{qcolor.blue()},{qcolor.alpha()}'"
            for value, qcolor in zip(from_, qcolors)
        ]
        color_ramp_str = ",".join(color_ramp_str)
        exp = f"""ramp_color(create_ramp(map({color_ramp_str}), {discrete}), ("{col_name}" - {min_})/({max_} - {min_}) )"""

        symbol_layer = QgsSimpleLineSymbolLayer()
        symbol_layer.setDataDefinedProperty(
            QgsSymbolLayer.Property.StrokeColor, QgsProperty.fromExpression(exp)
        )
        symbol_layer.setWidth(1.26)
        symbol_layer.setPenCapStyle(Qgis.EndCapStyle.Flat)

        return symbol_layer

    def _get_gradient_symbol_layer(
        self, col_list: [str], layer: QgsVectorLayer = None
    ) -> QgsInterpolatedLineSymbolLayer:
        """
        Generic method creating a QgsInterpolatedLineSymbolLayer for :
        - discrete assay
        """
        log = False
        if len(col_list) == 2:
            from_col, to_col = col_list
            log = self.config.log_param.value()
        else:
            from_col = to_col = col_list[0]

        # numpy array discribing pyqtgraph palette : rows (breakpoints), columns (rgba)
        is_discretized = (
            self.config.color_ramp_handler.symbology_widget._is_dicretized()
        )
        from_, to_, colors_ = self.config.color_ramp_handler.get_lut(fill_gaps=True)
        ramp_pg = np.array(
            [
                [
                    QColor(hex).red(),
                    QColor(hex).green(),
                    QColor(hex).blue(),
                    QColor(hex).alpha(),
                ]
                for hex in colors_
            ]
        )
        n_colors = ramp_pg.shape[0]
        # get min and max values
        fields = layer.fields()
        from_index = fields.indexFromName(from_col)
        to_index = fields.indexFromName(to_col)
        # get min/max values from column config
        min_val = self.config.color_ramp_handler.symbology_widget.min_val_sb.value()
        max_val = self.config.color_ramp_handler.symbology_widget.max_val_sb.value()
        # convert to base unit
        if min_val is not None:
            min_val = convert_to_base_unit(self.config.column.unit, min_val)
        if max_val is not None:
            max_val = convert_to_base_unit(self.config.column.unit, max_val)

        # get all layer min/max if invalid values
        if min_val is None or max_val is None:

            min_val = np.nanmin(
                [layer.minimumValue(index) for index in [from_index, to_index]]
            )

            max_val = np.nanmax(
                [layer.maximumValue(index) for index in [from_index, to_index]]
            )

        # transform into log if necessary
        if log:
            min_val = np.log10(min_val)
            max_val = np.log10(max_val)
            from_col = f"log10({from_col})"
            to_col = f"log10({to_col})"
            # if infinity, check for smallest/biggest value among non-zero values
            all_values = list(layer.uniqueValues(from_index)) + list(
                layer.uniqueValues(to_index)
            )
            # remove NULL values
            all_values = np.array(
                [elt for elt in all_values if not isinstance(elt, QVariant)]
            )
            if np.isinf(min_val):
                min_val = np.log10(np.nanmin(all_values[all_values != 0]))
            if np.isinf(max_val):
                max_val = np.log10(np.nanmax(all_values[all_values != 0]))

        # define qgis palette breakpoints
        bp = to_ if is_discretized else from_
        lst = [
            QgsColorRampShader.ColorRampItem(
                value, QColor(*ramp_pg[color_index, :].tolist())
            )
            for value, color_index in zip(bp, range(n_colors))
        ]

        # symbol layer
        symbol_layer = QgsInterpolatedLineSymbolLayer()

        # ramp shader
        ramp = QgsColorRampShader()
        ramp.setColorRampItemList(lst)
        # add source color ramp : if not set, junctions between lines are black on QGIS canvas
        ramp.setSourceColorRamp(ramp.createColorRamp())
        # set discrete or linear
        ramp.setColorRampType(
            Qgis.ShaderInterpolationMethod.Discrete
            if is_discretized
            else Qgis.ShaderInterpolationMethod.Linear
        )
        # interpolated line color instantiated with color ramp shader
        int_line = QgsInterpolatedLineColor(ramp)
        int_line.setColoringMethod(1)
        # fixed line width
        int_width = QgsInterpolatedLineWidth()
        int_width.setFixedStrokeWidth(1.26)

        symbol_layer.setInterpolatedColor(int_line)
        symbol_layer.setExpressionsStringForColor(from_col, to_col)
        symbol_layer.setInterpolatedWidth(int_width)

        return symbol_layer

    def _set_discrete_symbology(self, layer: QgsVectorLayer = None) -> QgsVectorLayer:
        """
        Return a QgsVectorLayer with associated symbology.

        Args:
            layer: QgsVectorLayer for which we need to compute symbology

        """
        col_name = self.config.column.name
        # if switched
        if self.config.get_default_config() != self.config:
            col_list = [col_name]
        else:
            col_list = [col_name + "_from_val", col_name + "_to_val"]
        # gradient color
        symbol_layer = self._get_gradient_symbol_layer(col_list=col_list, layer=layer)
        # hachure
        hachure_layer = QgsGeometryGeneratorSymbolLayer.create({})
        symbol = QgsLineSymbol.createSimple({})
        symbol.deleteSymbolLayer(0)
        symbol.appendSymbolLayer(QgsHashedLineSymbolLayer.create({}))
        hachure_layer.setSymbolType(QgsSymbol.Line)
        hachure_layer.setSubSymbol(symbol)
        exp = f"""
        case when "from_p" >= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' and "{col_list[-1]}" is not null then $geometry
        when "from_p" <= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' then line_substring($geometry, get_projected_length_of_3D_line($geometry, ("seuil" - "from_p")/("to_p" - "from_p") * length3D($geometry)), length3D($geometry))
        else NULL
        end
        """
        hachure_layer.setGeometryExpression(exp)

        # delete default symbol layer and append the new one
        layer.renderer().symbol().deleteSymbolLayer(0)
        layer.renderer().symbol().appendSymbolLayer(symbol_layer)
        layer.renderer().symbol().appendSymbolLayer(hachure_layer)

        return layer

    def _set_extended_num_symbology(self, layer: QgsVectorLayer = None):

        col_name = self.config.column.name
        # if switched
        if self.config.get_default_config() != self.config:
            col_list = [col_name + "_from_val", col_name + "_to_val"]
        else:
            col_list = [col_name]
        symbol_layer = self._get_extended_num_layer(col_list=col_list, layer=layer)

        # hachure
        hachure_layer = QgsGeometryGeneratorSymbolLayer.create({})
        symbol = QgsLineSymbol.createSimple({})
        symbol.deleteSymbolLayer(0)
        symbol.appendSymbolLayer(QgsHashedLineSymbolLayer.create({}))
        hachure_layer.setSymbolType(QgsSymbol.Line)
        hachure_layer.setSubSymbol(symbol)
        exp = f"""
        case when "from_p" >= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' and "{col_name}" is not null then $geometry
        when "from_p" <= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' then line_substring($geometry, get_projected_length_of_3D_line($geometry, ("seuil" - "from_p")/("to_p" - "from_p") * length3D($geometry)), length3D($geometry))
        else NULL
        end
        """
        hachure_layer.setGeometryExpression(exp)

        layer.renderer().symbol().deleteSymbolLayer(0)
        layer.renderer().symbol().appendSymbolLayer(symbol_layer)
        layer.renderer().symbol().appendSymbolLayer(hachure_layer)

        return layer

    def _set_extended_cat_symbology(self, layer: QgsVectorLayer = None):
        """
        Return a QgsVectorLayer with associated symbology.
        Palette is extracted from config.bar_symbology.

        Args:
            layer: QgsVectorLayer for which we need to compute symbology

        """
        # logviewer symbology
        col_map = self.config.bar_symbology.color_map
        col_names = self.config.bar_symbology.color_col
        col_name_db = self.assay_def.y_col[col_names]

        categories = []

        exp = f"""
        case when "from_p" >= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' and "{col_name_db}" is not null then $geometry
        when "from_p" <= "seuil" and "to_p" >= "seuil" and "mode" = 'planned_geom' then line_substring($geometry, get_projected_length_of_3D_line($geometry, ("seuil" - "from_p")/("to_p" - "from_p") * length3D($geometry)), length3D($geometry))
        else NULL
        end
        """

        for cat, color in col_map.items():
            symbol = QgsLineSymbol()
            symbol.setColor(QColor(color))
            symbol.setWidth(1.26)
            # hached if above gemetry limit
            hachure_layer = QgsGeometryGeneratorSymbolLayer.create({})
            hashed_symbol = QgsLineSymbol.createSimple({})
            hashed_symbol.deleteSymbolLayer(0)
            hashed_symbol.appendSymbolLayer(QgsHashedLineSymbolLayer.create({}))
            hachure_layer.setSymbolType(QgsSymbol.Line)
            hachure_layer.setSubSymbol(hashed_symbol)
            hachure_layer.setGeometryExpression(exp)

            symbol.appendSymbolLayer(hachure_layer)

            category = QgsRendererCategory(cat, symbol, cat, True)
            categories.append(category)

        symbol_layer = QgsCategorizedSymbolRenderer(col_name_db, categories)
        layer.setRenderer(symbol_layer)

        return layer

    def _set_discrete_cat_symbology(
        self, layer: QgsVectorLayer = None
    ) -> QgsVectorLayer:
        """
        Return a QgsVectorLayer with associated symbology.
        Palette is extracted from config.bar_symbology.
        Size, symbols and contour color is extracted from config.

        Args:
            layer: QgsVectorLayer for which we need to compute symbology

        """
        symbol_map = {
            "o": QgsSimpleMarkerSymbolLayer.Circle,
            "x": QgsSimpleMarkerSymbolLayer.CrossFill,
            "s": QgsSimpleMarkerSymbolLayer.Square,
            "d": QgsSimpleMarkerSymbolLayer.Diamond,
        }

        # logviewer symbology
        col_map = self.config.point_symbology.color_map
        col_names = self.config.point_symbology.color_col
        col_name_db = self.assay_def.y_col[col_names]
        size = self.config.symbol_size_parameter.value()
        pg_symbol = self.config.symbol_parameter.value()
        border_color = self.config.symbol_color_parameter.value()

        categories = []

        for cat, color in col_map.items():
            symbol_marker_layer = QgsSimpleMarkerSymbolLayer()
            symbol_marker_layer.setShape(symbol_map.get(pg_symbol))
            symbol_marker_layer.setStrokeColor(border_color)
            symbol = QgsMarkerSymbol()
            symbol.changeSymbolLayer(0, symbol_marker_layer)
            symbol.setColor(QColor(color))
            symbol.setSize(size / 3)

            category = QgsRendererCategory(cat, symbol, cat, True)
            categories.append(category)

        symbol_layer = QgsCategorizedSymbolRenderer(col_name_db, categories)
        layer.setRenderer(symbol_layer)

        return layer

    def set_symbology(self, layer: QgsVectorLayer = None):
        """
        Use appropriate method according to configuration class
        """
        # remove custom tolerance (since 3.42)
        try:
            props = layer.elevationProperties()
            props.setCustomToleranceEnabled(False)
        except Exception:
            pass

        if self.config.is_discrete and self.config.is_minmax_registered:
            return self._set_discrete_symbology(layer)
        elif self.config.is_discrete and self.config.is_categorical:
            return self._set_discrete_cat_symbology(layer)
        elif self.config.is_categorical and not self.config.is_discrete:
            return self._set_extended_cat_symbology(layer)
        elif self.config.is_minmax_registered and not self.config.is_discrete:
            return self._set_extended_num_symbology(layer)
        else:
            raise NotImplementedError
