import numpy as np
import pandas as pd
import pyqtgraph as pq
from pyqtgraph import mkPen
from qgis.PyQt.QtCore import QRectF, QTimer
from qgis.PyQt.QtGui import QBrush, QColor, QImage, QPainter, QTransform
from qgis.PyQt.QtSvg import QSvgRenderer

from openlog.datamodel.assay.generic_assay import GenericAssay
from openlog.gui.assay_visualization.categorical_symbology import CategoricalSymbology
from openlog.gui.assay_visualization.extended.bar_symbology import BarSymbology
from openlog.gui.assay_visualization.extended.text_bar_graph_items import (
    TextBarGraphItems,
)
from openlog.gui.assay_visualization.item_tags import CategoricalLegendTag, LegendTag
from openlog.gui.assay_visualization.plot_item_utils import (
    _convert_x_values,
    disable_tick_and_value_display,
)
from openlog.gui.pyqtgraph.ItemSwitcher import (
    DiscreteItemSwitcher,
    ExtendedItemSwitcher,
)


# EXTENDED
class SvgBarGraphItem(pq.BarGraphItem):
    def __init__(self, **opts):
        """
        Override pq.BarGraphItem to display svg pattern in each graph bars

        To maintain a correct display of wanted svg size, brush transform ratio is computed at each view transform changes

        Additional keyword options are:
        svg_files, svg_sizes, brushes

        Example uses:

            SvgBarGraphItem(x=range(3), height=[1,5,2], width=0.5,
                            svg_files= ['pattern1.svg','pattern1.svg','pattern3.svg'],
                            svg_sizes= [100,150,200],
                            brushes= ['red', 'blue', 'yellow'])

        Args:
            **opts:
        """
        pq.BarGraphItem.__init__(self, **opts)

        # Map of images of svg files render with different option (size,color)
        self.images = {}

        # Define options
        self.svg_sizes = None
        self.svg_files = None

        if "svg_files" in opts:
            self.svg_files = opts["svg_files"]
        if "svg_sizes" in opts:
            self.svg_sizes = opts["svg_sizes"]

        # Add timer for brush update on viewTransformChanged
        self._update_brush_timer = QTimer()
        self._update_brush_timer.setSingleShot(True)
        self._update_brush_timer.timeout.connect(self._update_brush_transform)
        # Update brush texture to use svg files
        if self.svg_files is not None:
            self.update_brush_texture()

    def update_brush_texture(self) -> None:
        """
        Update brush texture with current svg files

        """
        default_brush = self.opts["brushes"]
        brushes = []

        for i in range(len(self.svg_files)):
            color = default_brush[i]
            svg_file = self.svg_files[i]
            svg_size = self._get_item_svg_size(i)

            # Render svg file if not available at this size and color
            if (svg_file, svg_size, color) not in self.images:
                self._render_svg_image(svg_file, svg_size, color)

            # Update bar brush
            brush = QBrush()
            brush.setTextureImage(self.images[(svg_file, svg_size, color)])
            brushes.append(brush)
        self.setOpts(brushes=brushes)
        self._update_brush_transform()

    def viewTransformChanged(self) -> None:
        """
        Update transform for each brush depending on svg size and current device

        """
        # Use timer to avoid multiple brush update
        self._update_brush_timer.stop()
        self._update_brush_timer.start(75)

    def _update_brush_transform(self):
        """
        Update brush transform for correct svg display

        """
        reset_picture = False
        for i in range(len(self.opts["brushes"])):
            wanted_pix_w = self._get_item_svg_size(i)
            wanted_pix_h = wanted_pix_w

            svg_rect = QRectF(0, 0, wanted_pix_w, wanted_pix_h)
            svg_rect_device = self.mapRectFromDevice(svg_rect)
            if svg_rect_device is not None:
                ratio_w = wanted_pix_w / svg_rect_device.width()
                ratio_h = wanted_pix_h / svg_rect_device.height()

                brush = self.opts["brushes"][i]
                current_transform = brush.transform()
                transform = QTransform()
                transform = transform.scale(1 / ratio_w, 1 / ratio_h)
                if current_transform != transform:
                    brush.setTransform(transform)
                    self.opts["brushes"][i] = brush
                    reset_picture = True

        if reset_picture:
            self.prepareGeometryChange()
            self.picture = None

    def _render_svg_image(self, svg_file: str, svg_size: int, color: str) -> None:
        """
        Render svg image and store it in images map

        Args:
            svg_file: (str) svg file path
            svg_size: (int) wanted svg size
            color: (str) color as html string
        """
        renderer = QSvgRenderer(svg_file)
        default_size = renderer.defaultSize()
        ratio = default_size.height() / default_size.width()
        hSize = int(svg_size * ratio)
        image = QImage(svg_size, hSize, QImage.Format_ARGB32_Premultiplied)
        image.fill(QColor(color))
        impainter = QPainter(image)
        renderer.render(impainter)
        impainter.end()
        self.images[(svg_file, svg_size, color)] = image

    def _get_item_svg_size(self, i: int) -> str:
        """
        Get item svg size : if not defined use 150

        Args:
            i: bar graph item index

        Returns: (int) svg size

        """
        svg_sizes = 150
        if self.svg_sizes is not None:
            svg_sizes = self.svg_sizes[i]
        return svg_sizes


class CategoricalBarGraphItem(SvgBarGraphItem):

    WIDTH = 10

    def __init__(
        self,
        assay: GenericAssay,
        column: str,
        symbology: BarSymbology,
        coordinates: np.ndarray = None,
        planned_coordinates: np.ndarray = None,
        planned_eoh=None,
    ):
        """
        Override SvgBarGraphItem to display assay data as categorical series

        Args:
            assay: GenericAssay
            column: assay column name (should have an extended data extent and categorical series type)
            symbology: BarSymbology for category display
        """
        self._assay = assay
        self._column = column
        self.symbology_config = symbology

        # For now display all values available
        (x, y) = assay.get_all_values(column, remove_none=True)
        x = x.astype(float)

        if assay.get_dimension(x.shape) == 2:
            y0_val = x[:, 0]
            height_val = [interval[1] - interval[0] for interval in x]

        else:
            y0_val = []
            height_val = []

        x0_val = np.full(len(x), 0)
        width_val = np.full(len(x), self.WIDTH)

        super().__init__(height=height_val, width=width_val, x0=x0_val, y0=y0_val)
        self.set_symbology(self.symbology_config)
        self.switcher = ExtendedItemSwitcher(
            self,
            coordinates=coordinates,
            planned_coordinates=planned_coordinates,
            planned_eoh=planned_eoh,
        )

    def copy_item(self):
        """
        For minimap displaying.
        """
        s = self.symbology_config

        # no pattern
        s.pattern_map = {
            key: str(s.USGS_SYMB_ROOT_ALIAS + "/usgs691.svg")
            for key, _ in s.pattern_map.items()
        }
        new_item = CategoricalBarGraphItem(self._assay, self._column, s)
        o = self.opts
        new_item.setOpts(
            pen=None, x0=o["x0"], y0=o["y0"], height=o["height"], width=o["width"]
        )
        new_item.setOpacity(0.7)

        return new_item

    def set_symbology(self, symbology: BarSymbology):
        """
        Define used bar symbology map

        Args:
            symbology_map: (BarSymbologyMap)
        """
        self.symbology_config = symbology

        # For now display all values available
        val_dict = {}
        if self.symbology_config.pattern_col not in val_dict:
            (_, y) = self._assay.get_all_values(
                self.symbology_config.pattern_col, remove_none=True
            )
            val_dict[self.symbology_config.pattern_col] = y
        self.svg_files = [
            self.symbology_config.get_pattern_file(str(cat))
            for cat in val_dict[self.symbology_config.pattern_col]
        ]

        if self.symbology_config.scale_col not in val_dict:
            (_, y) = self._assay.get_all_values(
                self.symbology_config.scale_col, remove_none=True
            )
            val_dict[self.symbology_config.scale_col] = y
        self.svg_sizes = [
            self.symbology_config.get_scale(str(cat))
            for cat in val_dict[self.symbology_config.scale_col]
        ]

        if self.symbology_config.color_col not in val_dict:
            (_, y) = self._assay.get_all_values(
                self.symbology_config.color_col, remove_none=True
            )
            val_dict[self.symbology_config.color_col] = y
        brushes = [
            self.symbology_config.get_color(str(cat))
            for cat in val_dict[self.symbology_config.color_col]
        ]
        self.setOpts(brushes=brushes)
        self.update_brush_texture()

    def setPen(self, *args, **kargs):
        """
        Sets the pen used to draw graph line.
        The argument can be a :class:`QtGui.QPen` or any combination of arguments accepted by
        :func:`pyqtgraph.mkPen() <pyqtgraph.mkPen>`.
        """
        self.setOpts(pen=None)


def create_extended_categorical_plot_items(
    plot: pq.PlotWidget, assay: GenericAssay, column: str, config
) -> None:

    # Create item for categorical bar
    item = CategoricalBarGraphItem(
        assay,
        column,
        symbology=config.bar_symbology,
        coordinates=assay.get_coordinates(),
        planned_coordinates=assay.get_coordinates(planned=True),
        planned_eoh=assay.get_planned_eoh(),
    )
    plot.addItem(item)
    config.set_plot_item(item)

    # Create items for categorical text bar
    text_item = TextBarGraphItems(
        assay,
        column,
        CategoricalBarGraphItem.WIDTH,
        CategoricalBarGraphItem.WIDTH * 0.5,
        assay.get_altitude(),
        assay.get_altitude(planned=True),
    )
    text_item.add_to_plot(plot)
    config.set_category_text_bar_graph_item(text_item)

    # Disable auto range because of pyqtgraph issue :
    # If there is not enough space do display text, pyqtgraph has an infinite loop for auto range calculation
    plot.getViewBox().enableAutoRange(x=False)

    plot.setXRange(0, CategoricalBarGraphItem.WIDTH * 1.5)
    disable_tick_and_value_display(plot, "top")
    disable_tick_and_value_display(plot, "bottom")


# DISCRETE


class CategoricalScatterPlotItem(pq.ScatterPlotItem, CategoricalLegendTag):
    def __init__(
        self,
        assay: GenericAssay,
        column: str,
        symbology: CategoricalSymbology,
        coordinates: np.ndarray = None,
        planned_coordinates: np.ndarray = None,
        *args,
        **kargs,
    ):
        super().__init__(*args, **kargs)

        self.assay = assay
        self.column = column
        self.symbology = symbology
        (x, cat) = self.assay.get_all_values(self.column)
        self.x = x.astype(float)
        self.cat = cat.astype(str)
        self.legend_items = self.get_legend_items()

        self.set_symbology(self.symbology)
        self.switcher = DiscreteItemSwitcher(
            self, coordinates=coordinates, planned_coordinates=planned_coordinates
        )
        self.switcher.set_planned_eoh(assay.get_planned_eoh())

    def set_symbology(self, symbology: CategoricalSymbology):
        """
        Set a color by category.
        """

        self.symbology_config = symbology

        # For now display all values available
        val_dict = {}

        if self.symbology_config.color_col not in val_dict:
            (_, y) = self.assay.get_all_values(
                self.symbology_config.color_col, remove_none=True
            )
            val_dict[self.symbology_config.color_col] = y
        brushes = [
            self.symbology_config.get_color(str(cat))
            for cat in val_dict[self.symbology_config.color_col]
        ]

        self.setBrush(brushes)
        self.set_legend_brush(symbology)

    @classmethod
    def generate_single_brush_item(cls, brush: str = "") -> pq.ScatterPlotItem:
        """
        Given a color, return a pq.ScatterPlotItem with single brush
        """
        res = pq.ScatterPlotItem()
        res.setBrush(brush)
        res.setSize(15)
        return res

    def get_legend_items(self) -> list:
        """
        Return a list of item-category_name pairs for legend displaying.
        Only categories present in self are considered
        """
        result = []
        cats = set(self.cat)
        for cat in cats:
            brush = self.symbology.get_color(cat)
            obj = self.generate_single_brush_item(brush=brush)
            result.append((obj, cat))
        return result

    def set_legend_brush(self, symbology: CategoricalSymbology) -> None:
        """
        Update legend brushes according to self.symbology
        """
        for item, cat in self.legend_items:
            item.setBrush(symbology.get_color(cat))


def create_discrete_categorical_plot_items(plot, assay, column: str, config):

    """
    Create an AssayPlotDataItem from discrete assay

    Args:
        plot: pq.PlotWidget
        assay : GenericAssay
        column: assay column
        config: AssayColumnVisualizationConfig

    Returns: (AssayPlotDataItem)

    """
    # For now display all values available

    (x, _) = assay.get_all_values(column)
    x_val = _convert_x_values(assay, x)

    item = CategoricalScatterPlotItem(
        assay,
        column,
        symbology=config.point_symbology,
        x=np.full(x_val.shape, 0.0),
        y=x_val,
        coordinates=assay.get_coordinates(),
        planned_coordinates=assay.get_coordinates(planned=True),
    )

    plot.addItem(item)
    config.set_plot_item(item)
    config.update_symbol()
