import numpy
import numpy as np
import pyqtgraph as pg
from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtGui import QPen

from openlog.core import pint_utilities
from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDataExtent,
    AssayDomainType,
    AssaySeriesType,
    GenericAssay,
)
from openlog.datamodel.assay.uncertainty import UncertaintyType
from openlog.gui.assay_visualization.assay_bar_graph_item import AssayBarGraphItem
from openlog.gui.assay_visualization.assay_plot_data_item import AssayPlotDataItem
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.discrete.categorical_scatterplot_item import (
    CategoricalScatterPlotItem,
)
from openlog.gui.assay_visualization.extended.categorical_bar_graph_item import (
    CategoricalBarGraphItem,
)
from openlog.gui.assay_visualization.extended.imagery_graph_items import (
    ImageryGraphItems,
)
from openlog.gui.assay_visualization.extended.text_bar_graph_items import (
    TextBarGraphItems,
)
from openlog.gui.pyqtgraph.CustomErrorBarItem import CustomErrorBarItem
from openlog.gui.pyqtgraph.ErrorBoxItem import ErrorBoxItem


class AssayPlotItemFactory:
    """Factory class for assay column plot item creation"""

    def create_plot_items(
        self,
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> None:
        """
        Create an pg.PlotDataItem creation from generic assay

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

        Returns: (pg.PlotDataItem)

        """
        if assay.assay_definition.data_extent == AssayDataExtent.DISCRETE:
            # Create plot as extended or discrete
            if config.as_extended:
                config.propagate_hole_and_assay_properties()

                self.create_extended_plot_item_from_discrete_numerical(
                    plot, assay, column, config.extended_config
                )
            elif (
                assay.assay_definition.columns[column].series_type
                == AssaySeriesType.CATEGORICAL
            ):
                self.create_discrete_categorical_plot_items(plot, assay, column, config)

            else:
                self.create_discrete_plot_items(plot, assay, column, config)

        else:
            self.create_extended_plot_items(plot, assay, column, config)
        config.set_assay(assay)

    @staticmethod
    def create_discrete_categorical_plot_items(
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> CategoricalScatterPlotItem:
        """
        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

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

        item = CategoricalScatterPlotItem(
            assay,
            column,
            symbology=config.point_symbology,
            x=np.full(x_val.shape, 0.0),
            y=x_val,
        )

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

    @staticmethod
    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 = AssayPlotItemFactory._convert_x_values(assay, x)
        y_val = AssayPlotItemFactory._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, 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 = AssayPlotItemFactory._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
                    )
                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 = AssayPlotItemFactory._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
                    )
                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

    @staticmethod
    def _unit_conversion(
        assay_column: AssayColumn,
        config: AssayColumnVisualizationConfig,
        val: np.ndarray,
    ) -> np.ndarray:
        if (
            pint_utilities.is_pint_unit(assay_column.unit)
            and pint_utilities.can_convert(
                assay_column.unit, config.unit_parameter.value()
            )
            and assay_column.unit != config.unit_parameter.value()
        ):
            return pint_utilities.unit_conversion(
                assay_column.unit, config.unit_parameter.value(), val
            )
        else:
            return val

    @staticmethod
    def _get_uncertainty_whisker_values(
        assay: GenericAssay,
        assay_column: AssayColumn,
        config: AssayColumnVisualizationConfig,
        y_val: np.ndarray,
    ) -> (np.ndarray, np.ndarray):
        uncertainty = assay_column.uncertainty
        if uncertainty.get_uncertainty_type() == UncertaintyType.ONE_COLUMN:
            (_, unique_uncertainty) = assay.get_all_values(
                uncertainty.upper_whisker_column, related_column=assay_column.name
            )
            unique_uncertainty = AssayPlotItemFactory._unit_conversion(
                assay_column=assay_column,
                config=config,
                val=unique_uncertainty.astype(float) / 2.0,
            )
            top = unique_uncertainty
            bottom = unique_uncertainty
        elif (
            uncertainty.get_uncertainty_type() == UncertaintyType.TWO_COLUMN
            or uncertainty.get_uncertainty_type() == UncertaintyType.FOUR_COLUMN
        ):
            (_, upper_whisker) = assay.get_all_values(
                uncertainty.upper_whisker_column, related_column=assay_column.name
            )

            upper_whisker = AssayPlotItemFactory._unit_conversion(
                assay_column=assay_column,
                config=config,
                val=upper_whisker.astype(float),
            )
            (_, lower_whisker) = assay.get_all_values(
                uncertainty.lower_whisker_column, related_column=assay_column.name
            )
            lower_whisker = AssayPlotItemFactory._unit_conversion(
                assay_column=assay_column,
                config=config,
                val=lower_whisker.astype(float),
            )
            top = upper_whisker - y_val
            bottom = y_val - lower_whisker
        else:
            top = np.ndarray((0, 0))
            bottom = np.ndarray((0, 0))

        # If only NaN values are defined, don't use data
        # pyqtgraph will define invalid item bounds for errorbox
        if numpy.isnan(top).all() or numpy.isnan(bottom).all():
            top = []
            bottom = []
        return top, bottom

    @staticmethod
    def _get_uncertainty_box_values(
        assay: GenericAssay,
        assay_column: AssayColumn,
        config: AssayColumnVisualizationConfig,
        y_val: np.ndarray,
    ) -> (np.ndarray, np.ndarray):
        uncertainty = assay_column.uncertainty
        if (
            uncertainty.get_uncertainty_type() == UncertaintyType.TWO_COLUMN
            or uncertainty.get_uncertainty_type() == UncertaintyType.FOUR_COLUMN
        ):
            (_, upper_box) = assay.get_all_values(
                uncertainty.upper_box_column, related_column=assay_column.name
            )

            upper_box = AssayPlotItemFactory._unit_conversion(
                assay_column=assay_column, config=config, val=upper_box.astype(float)
            )
            (_, lower_box) = assay.get_all_values(
                uncertainty.lower_box_column, related_column=assay_column.name
            )
            lower_box = AssayPlotItemFactory._unit_conversion(
                assay_column=assay_column, config=config, val=lower_box.astype(float)
            )

            top = upper_box - y_val
            bottom = y_val - lower_box
        else:
            top = np.ndarray((0, 0))
            bottom = np.ndarray((0, 0))

        # If only NaN values are defined, don't use data
        # pyqtgraph will define invalid item bounds for errorbox
        if numpy.isnan(top).all() or numpy.isnan(bottom).all():
            top = []
            bottom = []
        return top, bottom

    @staticmethod
    def _convert_y_values(assay_column, config, y):
        y_val = []
        if assay_column.series_type == AssaySeriesType.NUMERICAL:
            y_val = numpy.array(
                [float(val) if val is not None else numpy.nan for val in y]
            )
        elif assay_column.series_type == AssaySeriesType.DATETIME:
            y_val = numpy.array(
                [val.timestamp() if val is not None else numpy.nan for val in y]
            )
        # Convert if needed
        y_val = AssayPlotItemFactory._unit_conversion(
            assay_column=assay_column, config=config, val=y_val
        )
        return y_val

    @staticmethod
    def _convert_x_values(assay, x):
        if assay.assay_definition.domain == AssayDomainType.DEPTH:
            x_val = numpy.array(
                [float(val) if val is not None else numpy.nan for val in x]
            )
        else:
            x_val = numpy.array([val.timestamp() for val in x])
        return x_val

    def create_extended_plot_items(
        self,
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> None:
        """
        Create an pg.PlotDataItem creation from extended assay

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

        Returns: (pg.PlotDataItem)

        """
        if (
            assay.assay_definition.columns[column].series_type
            == AssaySeriesType.NUMERICAL
            or assay.assay_definition.columns[column].series_type
            == AssaySeriesType.DATETIME
        ):
            # Create plot as extended or discrete
            if config.as_discrete:

                config.propagate_hole_and_assay_properties()
                self.create_discrete_plot_item_from_extended_numerical(
                    plot, assay, column, config.discrete_configuration
                )
            else:
                self.create_extended_numerical_plot_items(plot, assay, column, config)

        elif (
            assay.assay_definition.columns[column].series_type
            == AssaySeriesType.CATEGORICAL
        ):
            self.create_extended_categorical_plot_items(plot, assay, column, config)
        elif (
            assay.assay_definition.columns[column].series_type
            == AssaySeriesType.NOMINAL
        ):
            self.create_extended_nominal_plot_items(plot, assay, column, config)
        elif (
            assay.assay_definition.columns[column].series_type
            == AssaySeriesType.IMAGERY
        ):
            self.create_extended_imagery_plot_items(plot, assay, column, config)

    @staticmethod
    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 = AssayPlotItemFactory._convert_y_values(assay_column, config, y)

        if assay.get_dimension(x.shape) == 2:
            x_start = AssayPlotItemFactory._convert_x_values(assay, x[:, 0])
            x_end = AssayPlotItemFactory._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)
            )
            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_categorical_plot_items(
        self,
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> None:
        # Create item for categorical bar
        item = CategoricalBarGraphItem(assay, column, symbology=config.bar_symbology)
        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,
        )
        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)
        self.disable_tick_and_value_display(plot, "top")
        self.disable_tick_and_value_display(plot, "bottom")

    def create_extended_nominal_plot_items(
        self,
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> None:
        # For now width is hard coded to 10.0
        item = TextBarGraphItems(assay, column, 0, 10.0)
        item.add_to_plot(plot)
        config.set_plot_item(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, 10.0)
        self.disable_tick_and_value_display(plot, "top")
        self.disable_tick_and_value_display(plot, "bottom")

    def create_extended_imagery_plot_items(
        self,
        plot: pg.PlotWidget,
        assay: GenericAssay,
        column: str,
        config: AssayColumnVisualizationConfig,
    ) -> None:
        item = ImageryGraphItems(assay, column)
        item.add_to_plot(plot)
        config.set_plot_item(item)

    @staticmethod
    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 = AssayPlotItemFactory._convert_y_values(assay_column, config, y)

        if assay.get_dimension(x.shape) == 2:
            x_start = AssayPlotItemFactory._convert_x_values(assay, x[:, 0])
            x_end = AssayPlotItemFactory._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:
            # Invert X and Y for Depth visualization
            item = AssayPlotDataItem(
                domain=assay.assay_definition.domain,
                config=config,
                x=y_val,
                y=x_discrete,
            )
        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

    @staticmethod
    def create_extended_plot_item_from_discrete_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)
        x_val = AssayPlotItemFactory._convert_x_values(assay, x)
        y_val = AssayPlotItemFactory._convert_y_values(assay_column, config, y)

        x_start = []
        x_end = []

        for i, val in enumerate(x_val):
            if i == 0:
                if len(x_val) > 0:
                    half_width = (x_val[i + 1] - val) / 2.0
                else:
                    half_width = 1
                x_start.append(val)
                x_end.append(val + half_width)
            elif i == len(x) - 1:
                half_width = (val - x_val[i - 1]) / 2.0
                x_start.append(val - half_width)
                x_end.append(val)
            else:
                left_width = (val - x_val[i - 1]) / 2.0
                right_width = (x_val[i + 1] - val) / 2.0
                x_start.append(val - left_width)
                x_end.append(val + right_width)

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

        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

    @staticmethod
    def disable_tick_and_value_display(plot: pg.PlotWidget, axis: str) -> None:
        plot.getAxis(axis).setStyle(tickLength=0)
        pen = QPen()
        pen.setStyle(Qt.NoPen)
        plot.getAxis(axis).setTextPen(pen)
