import numpy as np
import pyqtgraph as pg
from pyqtgraph import mkBrush, mkPen
from qgis.PyQt import QtCore, QtGui
from qgis.PyQt.QtCore import QRectF

from openlog.datamodel.assay.generic_assay import AssayDomainType, GenericAssay
from openlog.datamodel.assay.uncertainty import UncertaintyType
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.item_tags import InspectorTag, LegendTag
from openlog.gui.assay_visualization.plot_item_utils import (
    _convert_x_values,
    _convert_y_values,
    _get_uncertainty_box_values,
    _get_uncertainty_whisker_values,
    convert_discrete_to_extended,
)
from openlog.gui.pyqtgraph.ItemSwitcher import (
    DiscreteItemSwitcher,
    ExtendedItemSwitcher,
)


# DISCRETE
class AssayPlotDataItem(pg.PlotDataItem, LegendTag, InspectorTag):
    def __init__(
        self,
        domain: AssayDomainType,
        config: AssayColumnVisualizationConfig,
        coordinates=None,
        planned_coordinates=None,
        planned_eoh=None,
        *args,
        **kargs,
    ):
        super().__init__(*args, **kargs)
        self.domain = domain
        self.config = config
        self.switcher = DiscreteItemSwitcher(
            self, coordinates, planned_coordinates, planned_eoh
        )

    def copy_item(self):
        """
        For minimap displaying.
        """

        log_mode = self.opts["logMode"]
        # get real x,y without transformations
        self.setLogMode(False, False)
        x, y = self.getData()
        config = self.config

        new_item = AssayPlotDataItem(domain=self.domain, config=config, x=x, y=y)
        new_item.setPen(pg.mkPen(color="black"))

        # apply eventual transformations
        self.setLogMode(*log_mode)
        new_item.setLogMode(*log_mode)

        return new_item

    def get_intersection_points(self, inspector):
        """
        Method to returns IntersectedPoint with inspector line.
        Args:
            - inspector : InspectorLine instance
        """

        def _get_x_y_values(inspector) -> ([float], [float]):
            """
            Get available x/y values from AssayPlotDataItem

            Args:
                c: AssayPlotDataItem

            Returns: ([float], [float]): x,y values

            """
            if inspector.domain == AssayDomainType.DEPTH:
                x = self.yData
                y = self.xData
            else:
                x = self.xData
                y = self.yData
            return x, y

        # For now define name with column name and hole display name
        name = f"{self.config.column.name} ({self.config.hole_display_name})"

        # Line value
        inspector_x = inspector.value()

        # Curve value
        x, y = _get_x_y_values(inspector)

        # Return nan if no data
        if x is None or y is None:
            return [inspector._create_nan_intersected_point(inspector_x, name)]

        # find the index of the closest point of this curve
        pix_size = inspector._get_view_pixel_size()
        adiff = np.abs(x - inspector_x)
        idx = np.argmin(adiff)

        # only add a label with defined value if the line touches the symbol
        tolerance = 0.5 * max(1, self.opts["symbolSize"]) * pix_size
        if adiff[idx] < tolerance:
            point = inspector.IntersectedPoint(
                item=self, x_value=self.xData[idx], y_value=self.yData[idx], name=name
            )
        else:
            # Check if we can interpolate values
            point = inspector._create_interpolated_intersection(
                idx, inspector_x, name, x, y
            )
            point.item = self

        return [point]


class ErrorBoxItem(pg.GraphicsObject):
    def __init__(self, coordinates=None, planned_coordinates=None, **opts):
        """
        All keyword arguments are passed to setData().
        """
        pg.GraphicsObject.__init__(self)
        self.picture = QtGui.QPicture()
        self.opts = dict(
            x=None,
            y=None,
            height=None,
            width=None,
            top=None,
            bottom=None,
            left=None,
            right=None,
            beam=None,
            pen=None,
        )
        self.setVisible(False)
        self.setData(**opts)
        self.switcher = DiscreteItemSwitcher(
            self, coordinates=coordinates, planned_coordinates=planned_coordinates
        )

    def setData(self, **opts):
        """
        Update the data in the item. All arguments are optional.

        Valid keyword options are:
        x, y, height, width, top, bottom, left, right, beam, pen

          * x and y must be numpy arrays specifying the coordinates of data points.
          * height, width, top, bottom, left, right, and beam may be numpy arrays,
            single values, or None to disable. All values should be positive.
          * top, bottom, left, and right specify the lengths of bars extending
            in each direction.
          * If height is specified, it overrides top and bottom.
          * If width is specified, it overrides left and right.
          * beam specifies the width of the beam at the end of each bar.
          * pen may be any single argument accepted by pg.mkPen().

        This method was added in version 0.9.9. For prior versions, use setOpts.
        """
        self.opts.update(opts)
        self.generatePicture()
        self.setVisible(all(self.opts[ax] is not None for ax in ["x", "y"]))
        self.update()
        self.prepareGeometryChange()
        self.informViewBoundsChanged()

    def setOpts(self, **opts):
        # for backward compatibility
        self.setData(**opts)
        self.generatePicture()

    def generatePicture(self):
        # pre-computing a QPicture object allows paint() to run much more quickly,
        # rather than re-drawing the shapes every time.
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)

        pen = self.opts["pen"]
        if pen is None:
            pen = pg.getConfigOption("foreground")
        p.setPen(pg.mkPen(pen))

        x, y = self.opts["x"], self.opts["y"]
        if x is None or y is None:
            return

        beam = self.opts["beam"]

        height, top, bottom = self.opts["height"], self.opts["top"], self.opts["bottom"]
        if height is not None or top is not None or bottom is not None:
            # draw vertical error box
            if height is not None:
                y1 = y - height / 2.0
                y2 = y + height / 2.0
            else:
                if bottom is None:
                    y1 = y
                else:
                    y1 = y - bottom
                if top is None:
                    y2 = y
                else:
                    y2 = y + top

            if beam is not None and beam > 0:
                x1 = x - beam / 2.0
                x2 = x + beam / 2.0

                for i in range(0, len(x1)):
                    x1_val = x1[i]
                    x2_val = x2[i]
                    y1_val = y1[i]
                    y2_val = y2[i]
                    rect = QRectF(x1_val, y1_val, x2_val - x1_val, y2_val - y1_val)
                    p.drawRect(rect)

        width, right, left = self.opts["width"], self.opts["right"], self.opts["left"]
        if width is not None or right is not None or left is not None:
            # draw vertical error box
            if width is not None:
                x1 = x - width / 2.0
                x2 = x + width / 2.0
            else:
                if left is None:
                    x1 = x
                else:
                    x1 = x - left
                if right is None:
                    x2 = x
                else:
                    x2 = x + right

            if beam is not None and beam > 0:
                y1 = y - beam / 2.0
                y2 = y + beam / 2.0

                for i in range(0, len(x1)):
                    x1_val = x1[i]
                    x2_val = x2[i]
                    y1_val = y1[i]
                    y2_val = y2[i]
                    rect = QRectF(x1_val, y1_val, x2_val - x1_val, y2_val - y1_val)
                    p.drawRect(rect)
        p.end()

    def paint(self, p, *args):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        # boundingRect _must_ indicate the entire area that will be drawn on
        # or else we will get artifacts and possibly crashing.
        # (in this case, QPicture does all the work of computing the bounding rect for us)
        return QtCore.QRectF(self.picture.boundingRect())


class CustomErrorBarItem(pg.ErrorBarItem):
    """
    pg.ErrorBarItem override to support log scale.
    When log scale is checked, plotitem.UpdateLogMode() iterate through each item and call setLogMode method.
    """

    def __init__(
        self,
        coordinates: np.ndarray = None,
        planned_coordinates: np.ndarray = None,
        *args,
        **kargs,
    ):
        super().__init__(*args, **kargs)
        # keep trace of state
        self.x_log_transformed = self.y_log_transformed = False
        self.switcher = DiscreteItemSwitcher(self, coordinates, planned_coordinates)

    def setLogMode(self, x: bool = None, y: bool = None) -> None:
        """
        .setData() automatically set visibility to True. So we have to keep trace of it.
        """
        visible = self.isVisible()
        x_val = self.opts["x"]
        y_val = self.opts["y"]
        top = self.opts["top"]
        bottom = self.opts["bottom"]
        left = self.opts["left"]
        right = self.opts["right"]

        if not y and self.y_log_transformed:
            y_val, bottom, top = self._disable_log_values(y_val, bottom, top)
            self.y_log_transformed = not self.y_log_transformed
        if y and not self.y_log_transformed:
            y_val, bottom, top = self._enable_log_values(y_val, bottom, top)
            self.y_log_transformed = not self.y_log_transformed
        if not x and self.x_log_transformed:
            x_val, left, right = self._disable_log_values(x_val, left, right)
            self.x_log_transformed = not self.x_log_transformed
        if x and not self.x_log_transformed:
            x_val, left, right = self._enable_log_values(x_val, left, right)
            self.x_log_transformed = not self.x_log_transformed

        self.setData(x=x_val, y=y_val, top=top, bottom=bottom, left=left, right=right)
        self.setVisible(visible)

    def _disable_log_values(
        self, values: np.ndarray, lower: np.ndarray, upper: np.ndarray
    ) -> tuple:
        values, upper, lower = (
            10**values,
            10 ** (values + upper),
            10 ** (values - lower),
        )
        upper = upper - values
        lower = values - lower
        return values, lower, upper

    def _enable_log_values(
        self, values: np.ndarray, lower: np.ndarray, upper: np.ndarray
    ) -> tuple:
        values, upper, lower = (
            np.log10(values),
            np.log10(values + upper),
            np.log10(values - lower),
        )
        upper = upper - values
        lower = values - lower
        return values, lower, upper


# EXTENDED


class CustomBarGraphItem(pg.BarGraphItem):
    def __init__(
        self,
        coordinates=None,
        planned_coordinates=None,
        altitude=None,
        planned_altitude=None,
        planned_eoh=None,
        **opts,
    ):
        """
        Override pg.BarGraphItem to add setPen and setBrush methods

        Args:
            **opts:
        """
        pg.BarGraphItem.__init__(self, **opts)
        self.switcher = ExtendedItemSwitcher(
            self,
            coordinates,
            planned_coordinates,
            altitude,
            planned_altitude,
            planned_eoh,
        )

    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>`.
        """
        pen = mkPen(*args, **kargs)
        self.setOpts(pen=pen)
        self.picture = None

    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)
        self.setOpts(brush=brush)
        self.picture = None


class AssayBarGraphItem(CustomBarGraphItem, LegendTag, InspectorTag):
    def __init__(self, config: AssayColumnVisualizationConfig, **opts):
        super().__init__(**opts)
        self.config = config

    def copy_item(self):
        """
        For minimap displaying.
        """
        config = self.config
        y1 = self.opts["y1"]
        # time domain
        if y1 is None:
            x0 = self.opts["x0"]
            x1 = self.opts["x1"]
            height = self.opts["height"]
            y0 = self.opts["y0"]
            new_item = AssayBarGraphItem(
                config=config, x0=x0, x1=x1, height=height, y0=y0
            )
        else:
            y0 = self.opts["y0"]

            width = self.opts["width"]
            x0 = self.opts["x0"]
            new_item = AssayBarGraphItem(
                config=config, y0=y0, y1=y1, width=width, x0=x0
            )
        new_item.setBrush(pg.mkBrush("black"))
        new_item.setPen(pg.mkPen(color="black"))

        return new_item

    def get_intersection_points(self, inspector):
        """
        Get intersection points for a AssayBarGraphItem

        Args:
            b: (AssayBarGraphItem)

        Returns: [IntersectedPoint]

        """
        points = []
        # For now define nam with column name and hole display name
        name = f"{self.config.column.name} ({self.config.hole_display_name})"

        if inspector.domain == AssayDomainType.DEPTH:
            y1, y0 = self.opts.get("y1"), self.opts.get("y0")
            x = zip(y1, y0) if inspector.switcher.is_altitude else zip(y0, y1)
            y = self.opts.get("width")
        else:
            x = zip(self.opts.get("x0"), self.opts.get("x1"))
            y = self.opts.get("height")

        inspector_x = inspector.value()
        mask = [x_min <= inspector_x <= x_max for x_min, x_max in x]
        for i, intersect in enumerate(mask):
            if intersect:
                if inspector.domain == AssayDomainType.DEPTH:
                    x_val = y[i]
                    y_val = inspector_x
                else:
                    x_val = inspector_x
                    y_val = y[i]

                points.append(
                    inspector.IntersectedPoint(
                        item=self, x_value=x_val, y_value=y_val, name=name
                    )
                )
        if not len(points):
            points.append(inspector._create_nan_intersected_point(inspector_x, name))
        return points


# FUNCTIONS


def create_discrete_plot_items(
    plot: pg.PlotWidget,
    assay: GenericAssay,
    column: str,
    config: AssayColumnVisualizationConfig,
) -> AssayPlotDataItem:
    """
    Create an AssayPlotDataItem from discrete assay

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

    Returns: (AssayPlotDataItem)

    """
    # For now display all values available
    assay_column = assay.assay_definition.columns[column]
    (x, y) = assay.get_all_values(column)
    x_val = _convert_x_values(assay, x)
    y_val = _convert_y_values(assay_column, config, y)

    if assay.assay_definition.domain == AssayDomainType.DEPTH:
        # Invert X and Y for Depth visualization
        item = AssayPlotDataItem(
            domain=assay.assay_definition.domain,
            config=config,
            coordinates=assay.get_coordinates(),
            planned_coordinates=assay.get_coordinates(planned=True),
            planned_eoh=assay.get_planned_eoh(),
            x=y_val,
            y=x_val,
        )
    else:
        item = AssayPlotDataItem(
            domain=assay.assay_definition.domain,
            config=config,
            x=x_val,
            y=y_val,
        )

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

    # Add uncertainty plot item
    if (
        assay_column.uncertainty.get_uncertainty_type() != UncertaintyType.UNDEFINED
        and len(x_val)
    ):
        uncertainty_plot_items = []

        # Add whisker
        top, bottom = _get_uncertainty_whisker_values(
            assay, assay_column, config, y_val
        )
        if len(top) and len(bottom):
            if assay.assay_definition.domain == AssayDomainType.DEPTH:
                # Invert X and Y for Depth visualization
                err = CustomErrorBarItem(
                    x=y_val,
                    y=x_val,
                    left=bottom,
                    right=top,
                    beam=0.5,
                    coordinates=assay.get_coordinates(),
                    planned_coordinates=assay.get_coordinates(planned=True),
                )
            else:
                err = CustomErrorBarItem(
                    x=x_val, y=y_val, top=top, bottom=bottom, beam=100
                )

            plot.addItem(err)
            uncertainty_plot_items.append(err)

        # Add box
        top, bottom = _get_uncertainty_box_values(assay, assay_column, config, y_val)
        if len(top) and len(bottom):
            if assay.assay_definition.domain == AssayDomainType.DEPTH:
                # Invert X and Y for Depth visualization
                err_box = ErrorBoxItem(
                    x=y_val,
                    y=x_val,
                    left=bottom,
                    right=top,
                    beam=0.5,
                    coordinates=assay.get_coordinates(),
                    planned_coordinates=assay.get_coordinates(planned=True),
                )
            else:
                err_box = ErrorBoxItem(
                    x=x_val,
                    y=y_val,
                    top=top,
                    bottom=bottom,
                    beam=100,
                )
            plot.addItem(err_box)
            uncertainty_plot_items.append(err_box)

        config.set_uncertainty_plot_items(uncertainty_plot_items)

    return item


def create_extended_numerical_plot_items(
    plot: pg.PlotWidget,
    assay: GenericAssay,
    column: str,
    config: AssayColumnVisualizationConfig,
) -> AssayBarGraphItem:
    # For now display all values available
    (x, y) = assay.get_all_values(column)
    assay_column = assay.assay_definition.columns[column]
    y_val = _convert_y_values(assay_column, config, y)

    if assay.get_dimension(x.shape) == 2:
        x_start = _convert_x_values(assay, x[:, 0])
        x_end = _convert_x_values(assay, x[:, 1])
    else:
        x_start = []
        x_end = []

    if assay.assay_definition.domain == AssayDomainType.DEPTH:
        # Invert X and Y  / with and height for Depth visualization
        item = AssayBarGraphItem(
            config=config,
            y0=x_start,
            y1=x_end,
            width=y_val,
            x0=np.full(len(x), 0),
            coordinates=assay.get_coordinates(),
            planned_coordinates=assay.get_coordinates(planned=True),
            planned_eoh=assay.get_planned_eoh(),
        )
        plot.addItem(item)
    else:
        item = AssayBarGraphItem(
            config=config,
            x0=x_start,
            x1=x_end,
            height=y_val,
            y0=np.full(len(x), 0),
        )
        plot.addItem(item)
    config.set_plot_item(item)

    return item


def create_extended_plot_item_from_discrete_numerical(
    plot: pg.PlotWidget,
    assay: GenericAssay,
    column: str,
    config: AssayColumnVisualizationConfig,
) -> AssayBarGraphItem:
    """
    Create an AssayBarGraphItem from discrete assay_name as extended assay_name

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

    Returns: (AssayPlotDataItem)

    """
    # For now display all values available
    assay_column = assay.assay_definition.columns[column]
    (x, y) = assay.get_all_values(column)
    coordinates = assay.get_coordinates()
    planned_coordinates = assay.get_coordinates(planned=True)

    x_val = _convert_x_values(assay, x)
    y_val = _convert_y_values(assay_column, config, y)

    x_start, x_end, y_values = convert_discrete_to_extended(x_val, y_val)

    if coordinates is not None:
        ext_coordinates = [
            convert_discrete_to_extended(coordinates[:, i], y_val) for i in range(3)
        ]
        ext_coordinates = np.concatenate(ext_coordinates)
        ext_coordinates = np.vstack([ext_coordinates[i] for i in [0, 3, 6, 1, 4, 7]]).T
    else:
        ext_coordinates = None

    if planned_coordinates is not None:
        ext_planned_coordinates = [
            convert_discrete_to_extended(planned_coordinates[:, i], y_val)
            for i in range(3)
        ]
        ext_planned_coordinates = np.concatenate(ext_planned_coordinates)
        ext_planned_coordinates = np.vstack(
            [ext_planned_coordinates[i] for i in [0, 3, 6, 1, 4, 7]]
        ).T
    else:
        ext_planned_coordinates = None

    if assay.assay_definition.domain == AssayDomainType.DEPTH:
        # Invert X and Y  / with and height for Depth visualization
        item = AssayBarGraphItem(
            config=config,
            y0=x_start,
            y1=x_end,
            width=y_values,
            x0=np.full(len(x), 0),
            coordinates=ext_coordinates,
            planned_coordinates=ext_planned_coordinates,
            planned_eoh=assay.get_planned_eoh(),
        )

    else:
        item = AssayBarGraphItem(
            config=config,
            x0=x_start,
            x1=x_end,
            height=y_values,
            y0=np.full(len(x), 0),
        )

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

    return item


def create_discrete_plot_item_from_extended_numerical(
    plot: pg.PlotWidget,
    assay: GenericAssay,
    column: str,
    config: AssayColumnVisualizationConfig,
) -> AssayPlotDataItem:
    """
    Create an AssayPlotDataItem from extended assay_name as discrete assay_name

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

    Returns: (AssayPlotDataItem)

    """
    # For now display all values available
    assay_column = assay.assay_definition.columns[column]
    (x, y) = assay.get_all_values(column)
    y_val = _convert_y_values(assay_column, config, y)

    if assay.get_dimension(x.shape) == 2:
        x_start = _convert_x_values(assay, x[:, 0])
        x_end = _convert_x_values(assay, x[:, 1])
    else:
        x_start = []
        x_end = []

    # TODO : define how we choose between start or stop or middle
    x_discrete = x_start

    if assay.assay_definition.domain == AssayDomainType.DEPTH:
        coordinates = assay.get_coordinates()
        if coordinates is not None:
            coordinates[:, :3]

        planned_coordinates = assay.get_coordinates(planned=True)
        if planned_coordinates is not None:
            planned_coordinates[:, :3]

        # Invert X and Y for Depth visualization
        item = AssayPlotDataItem(
            domain=assay.assay_definition.domain,
            config=config,
            x=y_val,
            y=x_discrete,
            coordinates=coordinates,  # from
            planned_coordinates=planned_coordinates,
            planned_eoh=assay.get_planned_eoh(),
        )
    else:
        item = AssayPlotDataItem(
            domain=assay.assay_definition.domain,
            config=config,
            x=x_discrete,
            y=y_val,
        )

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

    return item
