import numpy as np
from qgis.core import (
    QgsCategorizedSymbolRenderer,
    QgsColorRampShader,
    QgsInterpolatedLineColor,
    QgsInterpolatedLineSymbolLayer,
    QgsInterpolatedLineWidth,
    QgsLineSymbol,
    QgsMarkerSymbol,
    QgsRendererCategory,
    QgsSimpleMarkerSymbolLayer,
    QgsVectorLayer,
)
from qgis.PyQt.QtCore import QVariant
from qgis.PyQt.QtGui import QColor

from openlog.datamodel.assay.generic_assay import AssayDatabaseDefinition
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.discrete.config import (
    DiscreteAssayColumnVisualizationConfig,
    DiscreteCategoricalAssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.extended.config import (
    ExtendedCategoricalAssayColumnVisualizationConfig,
    ExtendedNumericalAssayColumnVisualizationConfig,
)


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_gradient_symbol_layer(
        self, col_list: [str], layer: QgsVectorLayer = None
    ) -> QgsInterpolatedLineSymbolLayer:
        """
        Generic method creating a QgsInterpolatedLineSymbolLayer for :
        - discrete assay
        - extended numerical assay

        Palette and min/max range are extracted from config.color_ramp_parameter.
        For extended numerical, start and end fields are the same column.

        """
        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)
        ramp_pg = self.config.color_ramp_parameter.value().color
        ramp_pg = (255 * ramp_pg).astype(int)
        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.min_color_ramp_param.value()
        max_val = self.config.max_color_ramp_param.value()

        # 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]))

        # create breakpoints
        bp = np.linspace(min_val, max_val, n_colors)

        # define qgis palette breakpoints
        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())
        # 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

        """
        symbol_layer = self._get_gradient_symbol_layer(
            col_list=["from_val", "to_val"], layer=layer
        )

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

        return layer

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

        col_name = self.config.column.name
        symbol_layer = self._get_gradient_symbol_layer(col_list=[col_name], layer=layer)
        layer.renderer().symbol().deleteSymbolLayer(0)
        layer.renderer().symbol().appendSymbolLayer(symbol_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 = []

        for cat, color in col_map.items():
            symbol = QgsLineSymbol()
            symbol.setColor(QColor(color))
            symbol.setWidth(1.26)
            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
        """
        if isinstance(self.config, DiscreteAssayColumnVisualizationConfig):
            return self._set_discrete_symbology(layer)
        elif isinstance(self.config, DiscreteCategoricalAssayColumnVisualizationConfig):
            return self._set_discrete_cat_symbology(layer)
        elif isinstance(self.config, ExtendedCategoricalAssayColumnVisualizationConfig):
            return self._set_extended_cat_symbology(layer)
        elif isinstance(self.config, ExtendedNumericalAssayColumnVisualizationConfig):
            return self._set_extended_num_symbology(layer)
        else:
            raise NotImplementedError
