import numpy as np
import pyqtgraph as pg


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, *args, **kargs):
        super().__init__(*args, **kargs)
        # keep trace of state
        self.x_log_transformed = self.y_log_transformed = False

    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
