import numpy as np
from typing import Any, List, Optional, Tuple, Union

from matplotlib.axes import Axes
from matplotlib.cm import ScalarMappable
from matplotlib.lines import Line2D

from maidr.core.enum import MaidrKey


class ContainerExtractorMixin:
    @staticmethod
    def extract_container(
        ax: Axes,
        container_type: type,
        include_all: bool = False,
    ) -> Any:
        """Retrieve containers of a specified type from an Axes object."""
        if ax is None or ax.containers is None:
            return None

        # If include_all is True, return a list of all containers of the specified type.
        if include_all:
            return [
                container
                for container in ax.containers
                if isinstance(container, container_type)
            ]

        # Otherwise, return the first container of the specified type.
        return next(
            container
            for container in ax.containers
            if isinstance(container, container_type)
        )


class LevelExtractorMixin:
    @staticmethod
    def extract_level(ax: Axes, key: MaidrKey = MaidrKey.X) -> Optional[List[str]]:
        """Retrieve label texts from Axes based on the specified Maidr key."""
        if ax is None:
            return None

        level = None
        if MaidrKey.X == key:
            ticks = ax.get_xticks()
            labels = [label.get_text() for label in ax.get_xticklabels()]

            if hasattr(ax, "dataLim") and ax.dataLim.width != 0:
                # Use the actual data limits rather than padded view limits
                data_x_min, data_x_max = ax.dataLim.x0, ax.dataLim.x0 + ax.dataLim.width
                # Filter tick labels to only those within the actual data range
                valid_indices = [
                    i for i, pos in enumerate(ticks) if data_x_min <= pos <= data_x_max
                ]
                labels = [labels[i] for i in valid_indices if i < len(labels)]

            level = labels
        elif MaidrKey.Y == key:
            ticks = ax.get_yticks()
            labels = [label.get_text() for label in ax.get_yticklabels()]

            if hasattr(ax, "dataLim") and ax.dataLim.height != 0:
                # Use the actual data limits rather than padded view limits
                data_y_min, data_y_max = (
                    ax.dataLim.y0,
                    ax.dataLim.y0 + ax.dataLim.height,
                )
                # Filter tick labels to only those within the actual data range
                valid_indices = [
                    i for i, pos in enumerate(ticks) if data_y_min <= pos <= data_y_max
                ]
                labels = [labels[i] for i in valid_indices if i < len(labels)]

            level = labels
        elif MaidrKey.FILL == key:
            level = [container.get_label() for container in ax.containers]

        if len(level) == 0:  # type: ignore
            level = LevelExtractorMixin.extract_shared_xtick_labels(ax)

        return level

    @staticmethod
    def extract_shared_xtick_labels(ax):
        siblings = ax.get_shared_x_axes().get_siblings(ax)
        for shared_ax in siblings:
            labels = [label.get_text() for label in shared_ax.get_xticklabels()]
            if any(labels):
                return labels
        return []


class LineExtractorMixin:
    @staticmethod
    def extract_line(ax: Axes) -> Optional[Line2D]:
        """Retrieve the last line object from Axes, if available."""
        if ax is None or ax.get_lines() is None:
            return None

        # Since the upstream MaidrJS library currently supports only the last plot line,
        # `maidr` package supports the same.
        return ax.get_lines()[-1]

    @staticmethod
    def extract_lines(ax: Axes) -> Optional[List[Line2D]]:
        """Retrieve all line objects from Axes, if available."""
        if ax is None or ax.get_lines() is None:
            return None

        # Since the upstream MaidrJS library currently supports only the last plot line,
        # `maidr` package supports the same.
        return ax.get_lines()

    @staticmethod
    def extract_line_data_with_categorical_labels(
        ax: Axes, line: Line2D
    ) -> Optional[List[Tuple[Union[str, float], float]]]:
        """
        Extract line data with proper handling of categorical x-axis labels.

        Parameters
        ----------
        ax : Axes
            The matplotlib axes object
        line : Line2D
            The line object to extract data from

        Returns
        -------
        Optional[List[Tuple[Union[str, float], float]]]
            List of (x, y) tuples where x values are categorical labels if available,
            or numeric values if no categorical labels are found
        """
        if ax is None or line is None:
            return None

        xydata = line.get_xydata()
        if xydata is None:
            return None

        # Convert to numpy array for easier handling
        xy_array = np.asarray(xydata)
        if xy_array.size == 0:
            return None

        # Extract x-axis labels for categorical data
        x_labels = LevelExtractorMixin.extract_level(ax, MaidrKey.X)

        # If we have categorical labels, map numeric coordinates to labels
        if x_labels:
            result = []
            for x, y in xy_array:
                # Map numeric x-coordinate to categorical label if available
                if isinstance(x, (int, float)) and 0 <= int(x) < len(x_labels):
                    x_value = x_labels[int(x)]
                else:
                    x_value = float(x)
                result.append((x_value, float(y)))
            return result
        else:
            # No categorical labels, return numeric values
            return [(float(x), float(y)) for x, y in xy_array]


class CollectionExtractorMixin:
    @staticmethod
    def extract_collection(ax: Axes, collection_type: type) -> Any:
        """Retrieve the first collection of a specified type from an Axes object."""
        if ax is None or ax.collections is None:
            return None

        # We assume only one collection of each type is present to avoid plot clutter,
        # even though multiples are technically possible.
        return next(
            collection
            for collection in ax.collections
            if isinstance(collection, collection_type)
        )


class ScalarMappableExtractorMixin:
    @staticmethod
    def extract_scalar_mappable(ax: Axes) -> Optional[ScalarMappable]:
        """Retrieve the first collection ScalarMappable from an Axes object."""
        if ax is None or ax.get_children() is None:
            return None

        # We assume only one ScalarMappable is present to avoid plot clutter,
        # even though multiples are technically possible.
        return next(sm for sm in ax.get_children() if isinstance(sm, ScalarMappable))
