import time

import numpy as np
import pyqtgraph as pg
from pyqtgraph import mkBrush, mkPen

from openlog.core.category_merger import CategoryMerger
from openlog.datamodel.assay.generic_assay import GenericAssay
from openlog.gui.pyqtgraph.CustomBarGraphItem import CustomBarGraphItem
from openlog.gui.pyqtgraph.ItemSwitcher import TextItemSwitcher


class CustomTextItem(pg.TextItem):
    def __init__(
        self,
        text="",
        color=(200, 200, 200),
        html=None,
        anchor=(0, 0),
        border=None,
        fill=None,
        angle=0,
        rotateAxis=None,
        altitude: tuple[np.ndarray, np.ndarray] = None,
        planned_altitude: tuple[np.ndarray, np.ndarray] = None,
    ):
        super().__init__(text, color, html, anchor, border, fill, angle, rotateAxis)
        self.switcher = TextItemSwitcher(self, altitude, planned_altitude)

    def set_switcher_XY(self):
        """
        Store initial position (drillhole length) in switcher.
        """
        x, y = self.pos()
        self.switcher.x = x
        self.switcher.y = y


class TextBarGraphItems:
    def __init__(
        self,
        assay: GenericAssay,
        column: str,
        offset: int = 0,
        width: int = 5,
        altitude: tuple[np.ndarray, np.ndarray] = None,
        planned_altitude: tuple[np.ndarray, np.ndarray] = None,
    ):
        """
        Store multiple plot items to display assay data as nominal series for extended data :

        - one CustomBarGraphItem to display rectangle for min/max x data
        - multiple pg.TextItem to display nominal series as string

        Args:
            assay: GenericAssay
            column: assay column name (should have an extended data extent and nominal or categorical series type)
            offset : offset for rectangle paint
            width: rectangle width
        """
        self._assay = assay
        self._column = column
        self._bar_item = None
        self._text_items = []
        self.merger = CategoryMerger(assay=assay, columns=[column])

        # For now display all values available
        (x, y) = self.merger.get_values()
        y = y[:, 0]

        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), offset)
        width_val = np.full(len(x), width)

        new_altitude = self.merger.get_altitude(columns=[column], planned=False)
        new_planned_altitude = self.merger.get_altitude(columns=[column], planned=True)

        self._bar_item = CustomBarGraphItem(
            height=height_val,
            width=width_val,
            x0=x0_val,
            y0=y0_val,
            altitude=new_altitude,
            planned_altitude=new_planned_altitude,
        )

        i = 0
        for text in y:
            x = x0_val[i]
            y = y0_val[i] + height_val[i] / 2.0

            if new_altitude is None:
                y0_alt = y1_alt = None
            else:
                y0_alt = new_altitude[0][i]
                y1_alt = new_altitude[1][i]

            if new_planned_altitude is None:
                y0_planned_alt = y1_planned_alt = None
            else:
                y0_planned_alt = new_planned_altitude[0][i]
                y1_planned_alt = new_planned_altitude[1][i]

            text_item = CustomTextItem(
                str(text),
                rotateAxis=(1, 0),
                angle=0,
                anchor=(0, 0.5),
                altitude=(y0_alt, y1_alt),
                planned_altitude=(y0_planned_alt, y1_planned_alt),
            )
            text_item.setPos(x, y)
            # store initial text position
            text_item.set_switcher_XY()
            self._text_items.append(text_item)
            i = i + 1

    def add_to_plot(self, plot: pg.PlotWidget) -> None:
        """
        Add stored items to plot

        Args:
            plot: pg.PlotWidget
        """
        if self._bar_item:
            plot.addItem(self._bar_item)
        for text_item in self._text_items:
            # Add ignoreBounds option so auto range is not computed with TextItem, very expensive in pyqtgraph
            plot.addItem(text_item, ignoreBounds=True)

    def setVisible(self, visible: bool) -> None:
        """
        Define store object visibility

        Args:
            visible: True to display items, False otherwise
        """
        if self._bar_item:
            self._bar_item.setVisible(visible)
        for text_item in self._text_items:
            text_item.setVisible(visible)

    def setPen(self, *args, **kargs) -> None:
        """
        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>`.
        """
        pen = mkPen(*args, **kargs)
        if self._bar_item:
            self._bar_item.setPen(pen)

    def setTextColor(self, color) -> None:
        """
        Set the color for pg.TextItem text.

        """
        for text_item in self._text_items:
            text_item.setColor(color)

    def setBrush(self, *args, **kargs):
        """
        Sets the brush used to draw bar rect.
        The argument can be a :class:`QBrush` or argument to :func:`~pyqtgraph.mkBrush`
        """
        brush = mkBrush(*args, **kargs)
        if self._bar_item:
            self._bar_item.setBrush(brush)
