# Better Colour Class for matplotlib

__author__ = "Cal Wing"
__version__ = "0.0.2"

from cycler import cycler, Cycler
from typing import Sequence, ItemsView, KeysView
from functools import lru_cache
import matplotlib.colors as colors

import numpy as np

type ValidColour = ColourValue | \
    str | \
    tuple[float, float, float, float] | \
    tuple[float, float, float]

class ColourValue(str):
    def __new__(
            cls, 
            value: ValidColour,
            /, 
            name: str | None = None
        ) -> "ColourValue":
        
        hex_val = colors.to_hex(value, True)
        instance = super().__new__(cls, hex_val)
        
        instance._name = name

        return instance
           
    def __repr__(self) -> str:
        if self.name:
            name = f"\"{self.name}\" "
        else:
            name = ""
        
        return f"<{self.__class__.__name__}" \
            f" {name}{self} " \
            f"({self.rgba[0]:.2f}, {self.rgba[1]:.2f}, " \
            f"{self.rgba[2]:.2f}, {self.rgba[3]:.2f})>"
    
    def alpha_adj(self, alpha) -> tuple[float, float, float, float]:
        r, g, b, _ = self.rgba
        return (r, g, b, alpha)

    @property
    @lru_cache
    def rgba(self) -> tuple[float, float, float, float]:
        return colors.to_rgba(self)

    @property
    @lru_cache
    def rgb(self) -> tuple[float, float, float]:
        return colors.to_rgb(self)
    
    @property
    @lru_cache
    def hsv(self) -> np.ndarray:
        return colors.rgb_to_hsv(self.rgb)
    
    @property
    @lru_cache
    def name(self) -> str:
        if hasattr(self, "_name") and self._name is not None:
            name = self._name
        else:
            name = ""

        assert isinstance(name, str), \
            TypeError(f"`name` should be an instance of `str` not `{type(name)}`")

        return name
    
    @property
    def hex(self) -> str:
        return self


type ValidColourDicts = dict[str, str | ColourValue] | \
    dict[str, tuple[float, float, float, float]] | \
    dict[str, tuple[float, float, float]]

class ColourContainer:
    def __init__(
            self, 
            colours: ValidColourDicts | Sequence[ColourValue],
            *,
            lookup_replacement: dict[str, str] = {" ": "_"}
        ) -> None:
        
        self._colours: dict[str, ColourValue] = {} 
        self._lookup_replacement = lookup_replacement

        if isinstance(colours, dict):
            for key, value in colours.items():
                assert isinstance(key, str), \
                    TypeError(f"`key` must be an instance of `str` not type `{type(key)}`")

                self._colours[key] = ColourValue(value, name=key)  # ty:ignore[invalid-argument-type]
                
                setattr(self, key, self._colours[key])
        
        else:
            for colour in colours:
                self._colours[colour.name] = colour

                setattr(self, colour.name, self._colours[colour.name])

    def __repr__(self) -> str:
        return f"Colour List of {len(self._colours)} " \
            f"colour{'s' if len(self._colours) > 0 else ''}: " + \
            str(list(self._colours.keys()))

    def __getitem__(self, key: str) -> ColourValue:
        for alt_char, original_char in self._lookup_replacement.items():
            if key.replace(alt_char, original_char) in self.keys():
                key = key.replace(alt_char, original_char)
                break

        return self._colours[key]
    
    def __iter__(self):
        return iter(self._colours)

    def items(self) -> ItemsView[str, ColourValue]:
        return self._colours.items()
    
    def keys(self) -> KeysView[str]:
        return self._colours.keys()
    
    def make_cycler(
        self, 
        order: Sequence[str] | None = None, 
        /, 
        hex: bool = True
    ) -> Cycler[str, str]:
        
        cycle_order = []

        if order is None:
            for _, colour in self.items():
                c = colour.hex if hex else colour 
                cycle_order.append(c)
        else:
            for key in order:
                c = self[key].hex if hex else self[key] 
                cycle_order.append(c)

        return cycler(color=cycle_order)
    
    # Generate a matplotlib colour mapping dictionary
    def mpl_colour_mapping(
            self, 
            prefix: str,
            /,
            hex: bool = True,
            *,
            add_separator: bool = True
        ) -> dict[str, str]:
        
        mapping: dict[str, str] = {}

        if add_separator:
            prefix += ":"

        for name, colour in self.items():
            c = colour.hex if hex else colour            
            mapping[prefix + name] = c

            # Allow all alt names
            for alt_char, original_char in self._lookup_replacement.items():
                mapping[prefix + name.replace(alt_char, original_char)] = c
        
        return mapping