Colors

class ColorEncoder:
    """
    color-encoding a categoric vector

    Example:
        ```
        >>> from ridgeplot.colors import ColorEncoder, ColorPalette
        >>> categorical_vector = ['group a','group b','group c','group a']
        >>> colors = ColorPalette["okabeito"]
        >>> ce = ColorEncoder()
        >>> ce.fit(categorical_vector, colors)
        >>> ce.encoder
        OrderedDict([('group a', '#E69F00'),
             ('group b', '#56B4E9'),
             ('group c', '#009E73')])
        >>> ce.transform(["group b", "group c", "group a"])
        ['#56B4E9', '#009E73', '#E69F00']
        ```

    or:
        ```
        >>> ce = ColorEncoder()
        >>> ce.fit_transform(categorical_vector, colors)
        ['#E69F00', '#56B4E9', '#009E73', '#E69F00']
        ```

    access color encoder:
        ```
        >>> ce.encoder
        OrderedDict([('group a', '#E69F00'),
             ('group b', '#56B4E9'),
             ('group c', '#009E73')])
        ```
    """

    def __init__(self):
        self.x: List[str] = list()
        self.distinct_categories: List[str] = []
        self.encoder: OrderedDict[str, str] = OrderedDict()

    def fit(self, categories: List[str], colors: List[str] = ColorPalette["invitae"]) -> None:
        """
        mapping colors to the unique categories in the input list
        basically fill the encoder dictionary

        Example:
            ```
            >>> categorical_vector = ['group a','group b','group c','group a']
            >>> colors = ColorPalette["okabeito"]
            >>> ce = ColorEncoder()
            >>> ce.fit(categroical_vector, colors)
            ```

        Args:
            categories: list of input values (i.e. labels of the samples), can be duplicated
            colors: list of colors, intentionally not checked for duplication
        Returns:
            NoneType
        """
        self.distinct_categories = check_color_vector_size(categories, colors)
        self.encoder = OrderedDict({category: col for category, col in zip(self.distinct_categories, colors)})

    def transform(self, categories: List[str]) -> List[str]:
        """
        mapping color to the a list of category in the input list

        Example:
            ```
            >>> categorical_vector = ['group a','group b','group c','group a']
            >>> colors = ColorPalette["okabeito"]
            >>> ce = color_encoder()
            >>> ce.fit(categroical_vector, colors)
            >>> new_categorical_vector = ["group b", "group c"]
            >>> ce.transform(new_categorical_vector)
            ['#56B4E9', '#009E73']
            ```

        Args:
            categories: list of input values (i.e. labels of the samples), can be duplicated
        Returns:
            list of colors for the input list according to the fitted color encoder
        """
        if not self.encoder:
            raise ValueError("Call color_encoder.fit() first!!")

        union_set = set(self.distinct_categories).union(set(categories))
        if len(union_set) != len(self.distinct_categories):
            unseen = union_set - set(self.distinct_categories)
            unseen_str = ", ".join(sorted(list(unseen)))
            raise ValueError(f"Input [categories] contain unseen data!!: {unseen_str}")

        return [self.encoder[category] for category in categories]

    def fit_transform(self, categories: List[str], colors: List[str] = ColorPalette["invitae"]) -> List[str]:
        """
        first map the color to the categories, and then return the
        corresponding color for each category in the input list

        Example:
            ```
            >>> categorical_vector = ["group1", "group2", "group1"]
            >>> colors = ["salmon","gold"]
            >>> ce = ColorEncoder()
            >>> ce.fit_transform(categorical_vector, colors)
            ['salmon', 'gold', 'salmon']
            ```

        Args:
            categories: list of input values (i.e. labels of the samples), can be duplicated
            colors: list of colors to be assigned to the categories
        Returns:
            list of colors corresponding to the input
        """
        self.fit(categories, colors=colors)
        return self.transform(categories)

    def show_legend(
        self,
        ax: matplotlib.axes._axes.Axes,
        sort: bool = False,
        **kwargs: Dict[str, Any],
    ) -> legend.Legend:
        """
        Adding matplotlib legend describing the color encoder to a matplotlib ax object

        Args:
            ax: matplotlib ax object
            sort: sort the legend by the category
            **kwargs: keyword arguments for matplotlib.pyplot.legend

        Returns:
            the matplotlib legend object
        """

        if sort:
            self.encoder = OrderedDict(sorted(self.encoder.items(), key=lambda item: item[0]))
        pat = [mpatches.Patch(color=col, label=lab) for lab, col in self.encoder.items()]
        lgd = ax.legend(handles=pat, **kwargs)
        return lgd