"""Color palette generation and visualization utilities."""

import typing

import sys
from typing import Any, Iterable, Literal, overload

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.colors import Colormap, to_hex, to_rgb
from matplotlib.patches import Rectangle

from .mpl_utils import style_ax, use_color_cycle


def _require_coloraide() -> type[Any]:
    try:
        from coloraide import Color  # pyright: ignore[reportMissingImports]
    except ImportError as exc:
        raise ImportError(
            'coloraide is required for zyplib.vis.color_palette utilities'
        ) from exc
    return typing.cast(type[Any], Color)


class PredefinedColorSchemes:
    """
    为学术出版物挑选的预定义颜色方案。

    这些方案考虑了色盲友好性、灰度可区分性和视觉吸引力，
    适用于机器学习和数据科学领域的图表。

    ```py
    from zyplib.vis.color_palette import PredefinedColorSchemes, preview_palette

    for palette_name, palette in PredefinedColorSchemes.__dict__.items():
        if isinstance(palette, list):
            _ = preview_palette(palette, title=palette_name)
        ```
    """

    # --- 1. 定性调色板 (Qualitative Palettes) ---
    # 用于表示离散的、无序的类别。例如：比较不同模型的性能。
    # 强烈推荐 Tableau 10 和 Paul Tol's Muted。

    TABLEAU_10 = [
        '#4E79A7',
        '#F28E2B',
        '#E15759',
        '#76B7B2',
        '#59A14F',
        '#EDC948',
        '#B07AA1',
        '#FF9DA7',
        '#9C755F',
        '#BAB0AC',
    ]
    """Qualitative | Tableau 10: Tableau 的默认配色，色盲友好，区分度高，非常优秀。"""

    SEABORN_DEEP_10 = [
        '#4C72B0',
        '#DD8452',
        '#55A868',
        '#C44E52',
        '#8172B3',
        '#937860',
        '#DA8BC3',
        '#8C8C8C',
        '#CCB974',
        '#64B5CD',
    ]
    """Qualitative | Seaborn Deep: Seaborn 库的 "deep" 配色，经典且视觉效果好。"""

    # --- 2. 顺序调色板 (Sequential Palettes) ---
    # 用于表示连续的、有序的数据。例如：热力图、损失函数下降过程。
    # 亮度应单调变化。Viridis 和 Cividis 是黄金标准。

    VIRIDIS_10 = [
        '#440154',
        '#482878',
        '#3E4A89',
        '#31688E',
        '#26828E',
        '#1F9E89',
        '#35B779',
        '#6DCD59',
        '#B4DE2C',
        '#FDE725',
    ]
    """Sequential | Viridis: 感知均匀，色盲友好，灰度友好。Matplotlib 的默认顺序色板。"""

    CIVIDIS_10 = [
        '#00204D',
        '#00336F',
        '#00488F',
        '#005EAB',
        '#0075C4',
        '#238DDA',
        '#4BA5E8',
        '#76BEF3',
        '#A4D6F9',
        '#D5EEFB',
    ]
    """Sequential | Mono | Cividis: 专为色盲优化设计的顺序色板，以蓝色为基色的单色板，从深到浅。"""

    BLUES_10 = [
        '#F7FBFF',
        '#DEEBF7',
        '#C6DBEF',
        '#9ECAE1',
        '#6BAED6',
        '#4292C6',
        '#2171B5',
        '#08519C',
        '#08306B',
    ]
    """Sequential | Mono | ColorBrewer Blues: 经典的单色系顺序调色板，简洁明了，从浅到深。"""

    # --- 3. 发散调色板 (Diverging Palettes) ---
    # 用于表示有明确中心点（通常是 0）的数据。例如：相关系数矩阵、差异图。

    COOLWARM_10 = [
        '#67001f',
        '#b2182b',
        '#d6604d',
        '#f4a582',
        '#fddbc7',
        '#f7f7f7',
        '#d1e5f0',
        '#92c5de',
        '#4393c3',
        '#2166ac',
        '#053061',
    ]
    """Diverging | Matplotlib Coolwarm: 经典的蓝-白-红发散色板，适用于表示冷暖对比。"""


def rgb_to_hex(rgba: Any) -> str:
    return to_hex(rgba)


def hex_to_rgb(
    hex_color: str,
) -> tuple[float, float, float] | tuple[float, float, float, float]:
    return to_rgb(hex_color)


def generate_another_color(
    base_color: str,
    scheme: Literal[
        'complementary',
        'triadic',
        'split_complementary',
        'analogous',
        'tetradic',
        'rectangle',
        'custom',
    ] = 'complementary',
    angle: float | typing.Sequence[float] | None = None,
    preserve_lightness: bool = True,
    preserve_chroma: bool = True,
    space: str = 'oklch',
    include_base_in_return: bool = True,
) -> list[str]:
    """
    基于色轮模型生成配色方案。

    根据指定的配色方案或自定义角度，在色相环上生成和谐的配色。
    使用感知均匀的色彩空间确保视觉效果的一致性。

    Parameters
    ----------
    - `base_color` : `str`
        - 基色，支持 hex 字符串、CSS 颜色名等 (e.g., '#3498db', 'blue')
    - `scheme` : `Literal['complementary', 'triadic', 'split_complementary', 'analogous', 'tetradic', 'rectangle', 'custom']`, optional
        - 配色方案类型，by default 'complementary'
        - 'complementary': 对比色 (180°)
        - 'triadic': 三分色 (120°, 240°)
        - 'split_complementary': 分裂对比色 (150°, 210°)
        - 'analogous': 类似色 (30°, -30°)
        - 'tetradic': 方形色 (90°, 180°, 270°)
        - 'rectangle': 矩形色 (60°, 180°, 240°)
        - 'custom': 使用自定义 angle 参数
    - `angle` : `float | list[float] | None`, optional
        - 色相偏移角度，by default None
        - 当 scheme='custom' 时必须提供
        - 可以是单个角度或角度列表
    - `preserve_lightness` : `bool`, optional
        - 是否保持原始颜色的亮度，by default True
    - `preserve_chroma` : `bool`, optional
        - 是否保持原始颜色的饱和度，by default True
    - `space` : `str`, optional
        - 工作色彩空间，推荐 'oklch'、'hsl'，by default 'oklch'

    Returns
    -------
    - `str | list[str]`
        - 单个颜色的 hex 字符串，或多个颜色的列表
        - 当方案生成多个颜色时返回列表，单个颜色时返回字符串

    Examples
    --------
    >>> # 生成对比色
    >>> generate_another_color('#3498db', 'complementary')
    '#db7334'

    >>> # 生成三分色
    >>> generate_another_color('#3498db', 'triadic')
    ['#9834db', '#34db98']

    >>> # 自定义角度
    >>> generate_another_color('#3498db', 'custom', angle=[45, 90])
    ['#3473db', '#3434db']
    """
    Color = _require_coloraide()

    # 转换基色到指定色彩空间
    base = Color(base_color).convert(space)

    SCHEME_ANGLES = {
        'complementary': [180],
        'triadic': [120, 240],
        'split_complementary': [150, 210],
        'analogous': [30, -30],
        'tetradic': [90, 180, 270],
        'rectangle': [60, 180, 240],
    }

    if scheme == 'custom':
        if angle is None:
            raise ValueError("Custom scheme requires 'angle' parameter")
        angles = [angle] if isinstance(angle, (int, float)) else angle
    elif scheme in SCHEME_ANGLES:
        angles = SCHEME_ANGLES[scheme]
    else:
        raise ValueError(f'Unknown scheme: {scheme}')

    # 生成新颜色
    new_colors = []
    if include_base_in_return:
        new_colors.append(base.convert('srgb').to_string(hex=True))

    for offset_angle in angles:
        new_color = base.clone()

        if space == 'oklch':
            # 在 OKLCH 空间中调整色相
            current_hue = base['h'] if base['h'] is not None else 0
            new_hue = (current_hue + offset_angle) % 360
            new_color['h'] = new_hue

            # 根据参数决定是否保持亮度和色度
            if not preserve_lightness:
                # 可以在这里添加亮度调整逻辑
                pass
            if not preserve_chroma:
                # 可以在这里添加饱和度调整逻辑
                pass

        elif space == 'hsl':
            # 在 HSL 空间中调整色相
            current_hue = base['h'] if base['h'] is not None else 0
            new_hue = (current_hue + offset_angle) % 360
            new_color['h'] = new_hue

        else:
            # 对于其他空间，转换到 OKLCH 进行色相操作
            temp_color = base.convert('oklch')
            current_hue = temp_color['h'] if temp_color['h'] is not None else 0
            new_hue = (current_hue + offset_angle) % 360
            temp_color['h'] = new_hue
            new_color = temp_color.convert(space)

        # 转换回 sRGB hex 格式
        hex_color = new_color.convert('srgb').to_string(hex=True)
        new_colors.append(hex_color)

    return new_colors


def generate_monochromatic_palette(
    base_color: str,
    n: int = 5,
    direction: Literal['lighter', 'darker', 'vibrant', 'muted', 'both'] = 'both',
    spread: float = 0.3,
    space: str = 'oklch',
) -> list[str]:
    """
    基于一个基色生成单色系调色板。

    在指定色彩空间中，沿某个方向生成视觉上和谐的同色系颜色。

    Parameters
    ----------
    - `base_color` : `str`
        - 基色，支持 hex 字符串、CSS 颜色名等 (e.g., '#3498db', 'blue')
    - `n` : `int`, optional
        - 要生成的颜色数量，by default 5
    - `direction` : `Literal['lighter', 'darker', 'vibrant', 'muted', 'both']`, optional
        - 在色彩空间中的变化方向，by default 'both'
        - 'lighter': 向更亮方向；'darker': 向更暗方向
        - 'vibrant': 向更饱和方向；'muted': 向更柔和方向
        - 'both': 围绕基色双向展开
    - `spread` : `float`, optional
        - 颜色间的扩散程度 (0.1-1.0)，by default 0.3
    - `space` : `str`, optional
        - 工作色彩空间，推荐 'oklch'、'oklab'、'hsl'，by default 'oklch'

    Returns
    -------
    - `list[str]`
        - 包含 n 个 hex 颜色字符串的列表
    """
    Color = _require_coloraide()

    # --- 1. 定义参数和安全边界 ---
    # 使用常量提高可读性，并确保颜色值在有效范围内
    MIN_LUMINANCE = 0.05
    MAX_LUMINANCE = 0.98
    MIN_CHROMA = 0.005
    # OKLCH 色度在 sRGB 色域内通常低于 0.4，这是一个安全上限
    MAX_CHROMA_SRGB = 0.4

    # --- 2. 初始化颜色对象 ---
    base = Color(base_color)
    # 统一在 oklch 空间计算 L (Lightness) 和 C (Chroma) 的目标值，因为它感知均匀
    base_oklch = base.convert('oklch')
    l_base = base_oklch['l']

    # --- FIX: 使用直接的键访问来代替错误的 .get() 调用 ---
    # 在 oklch 空间中，'c' 通道总是存在的，对于灰度色其值为 0。
    c_base = base_oklch['c']

    # --- 3. 根据 direction 定义插值的起点和终点 ---
    start_color, end_color = base.clone(), base.clone()

    if direction == 'both':
        # 在亮度上向两边扩展，总范围为 spread
        half_spread = spread / 2
        start_color.set('oklch.l', max(MIN_LUMINANCE, l_base - half_spread))
        end_color.set('oklch.l', min(MAX_LUMINANCE, l_base + half_spread))

    elif direction == 'lighter':
        # 从基色变为更亮的颜色
        start_color = base
        end_color.set('oklch.l', min(MAX_LUMINANCE, l_base + spread))

    elif direction == 'darker':
        # 从更暗的颜色变为基色
        start_color.set('oklch.l', max(MIN_LUMINANCE, l_base - spread))
        end_color = base

    elif direction == 'vibrant':
        # 从基色变为更鲜艳的颜色
        start_color = base
        end_color.set('oklch.c', min(MAX_CHROMA_SRGB, c_base + spread))

    elif direction == 'muted':
        # 从更柔和的颜色变为基色
        start_color.set('oklch.c', max(MIN_CHROMA, c_base - spread))
        end_color = base

    # --- 4. 执行插值并生成调色板 ---
    # 在用户指定的空间(space)中进行插值
    interpolator = Color.interpolate([start_color, end_color], space=space)

    # 通过在 [0, 1] 区间上均匀采样来生成 n 个颜色
    palette = (
        [interpolator(i / (n - 1)) for i in range(n)] if n > 1 else [interpolator(0.5)]
    )

    # --- 5. 格式化输出 ---
    return [color.convert('srgb').to_string(hex=True) for color in palette]


def generate_interpolated_palette(
    colors: typing.Sequence[str],
    n: int = 5,
    space: str = 'oklch',
    hue: Literal['shorter', 'longer'] = 'shorter',
) -> list[str]:
    """
    在多个颜色之间进行感知均匀的插值，生成渐变调色板。

    使用 coloraide 的高质量插值算法，支持多颜色插值和色相路径控制。

    Parameters
    ----------
    - `colors` : `list[str]`
        - 颜色列表，支持 hex 字符串、CSS 颜色名等
    - `n` : `int`, optional
        - 要生成的颜色数量，by default 5
    - `space` : `str`, optional
        - 进行插值的色彩空间，推荐 'oklch'、'oklab'、'hsl'，by default 'oklch'
    - `hue` : `Literal['shorter', 'longer']`, optional
        - 色相插值路径选择，by default 'shorter'
        - 'shorter': 在色相环上选择短路径
        - 'longer': 在色相环上选择长路径

    Returns
    -------
    - `list[str]`
        - 包含 n 个 hex 颜色字符串的列表
    """
    Color = _require_coloraide()

    # 使用 coloraide 的原生插值方法
    interpolator = Color.interpolate(colors, space=space, hue=hue)

    # 生成指定数量的颜色
    palette = [
        interpolator(i / (n - 1)).convert('srgb').to_string(hex=True) for i in range(n)
    ]

    return palette


class ColorMap:
    """
    基于 coloraide 的连续颜色映射器。

    将数值范围映射到调色板上，支持多种超出范围的处理策略。
    """

    def __init__(
        self,
        palette: typing.Sequence[str],
        value_range: tuple[float, float] = (0, 1),
        out_range_policy: Literal['circular', 'clip', 'raise', 'none'] = 'circular',
        space: str = 'oklch',
        hue: Literal['shorter', 'longer'] = 'shorter',
    ):
        """
        初始化颜色映射器。

        Parameters
        ----------
        - `value_range` : `tuple[float, float]`
            - 数值映射范围 (min_value, max_value)
        - `palette` : `list[str]`
            - 调色板，hex 颜色字符串列表
        - `out_range_policy` : `Literal['circular', 'clip', 'raise', 'none']`, optional
            - 超出范围时的处理策略，by default 'circular'
            - 'circular': 循环映射；'clip': 裁剪到边界
            - 'raise': 抛出异常；'none': 返回 None
        - `space` : `str`, optional
            - 颜色插值空间，by default 'oklch'
        """
        Color = _require_coloraide()

        self.value_range = value_range
        self.palette = palette
        self.out_range_policy = out_range_policy
        self.space = space

        # 创建 coloraide 插值器
        self._interpolator = Color.interpolate(palette, space=space, hue=hue)

        # 缓存范围计算
        self._range_size = value_range[1] - value_range[0]
        if self._range_size == 0:
            raise ValueError('value_range cannot have zero size')

    def _normalize_value(self, value: float) -> float | None:
        """
        将数值标准化到 [0, 1] 范围，并处理超出范围的情况。

        Parameters
        ----------
        - `value` : `float`
            - 输入数值

        Returns
        -------
        - `float | None`
            - 标准化后的值，或根据策略返回 None
        """
        # 计算标准化值
        normalized = (value - self.value_range[0]) / self._range_size

        # 处理在范围内的情况
        if 0 <= normalized <= 1:
            return normalized

        # 处理超出范围的情况
        if self.out_range_policy == 'clip':
            return max(0, min(1, normalized))
        elif self.out_range_policy == 'circular':
            return normalized % 1.0
        elif self.out_range_policy == 'raise':
            raise ValueError(
                f'Value {value} is outside range {self.value_range}. '
                f'Normalized: {normalized}'
            )
        elif self.out_range_policy == 'none':
            return None
        else:
            raise ValueError(f'Unknown out_range_policy: {self.out_range_policy}')

    def cmap(self, value: float) -> str | None:
        """
        进行颜色映射，将数值映射到调色板上的颜色。

        Parameters
        ----------
        - `value` : `float`
            - 要映射的数值

        Returns
        -------
        - `str | None`
            - 映射后的 hex 颜色字符串，超出范围且策略为 'none' 时返回 None
        """
        normalized = self._normalize_value(value)
        if normalized is None:
            return None

        # 使用 coloraide 插值器获取颜色
        color = self._interpolator(normalized)
        return color.convert('srgb').to_string(hex=True)

    @overload
    def __call__(self, value: float) -> str | None: ...

    @overload
    def __call__(self, value: Iterable[float]) -> list[str | None]: ...

    def __call__(self, value: float | Iterable[float]) -> str | None | list[str | None]:
        """允许直接调用实例进行颜色映射。"""
        if isinstance(value, Iterable):
            return [self.cmap(v) for v in value]
        return self.cmap(value)

    def to_mpl_cmap(
        self, type: Literal['discrete', 'continuous'] = 'continuous', bins: int = 10
    ) -> Colormap:
        """
        将 coloraide 的颜色映射转换为 Matplotlib 可用的 colormap。

        Parameters
        ----------
        - `bins` : `int`, optional
            - 分段数，by default 256

        Returns
        -------
        - `matplotlib.colors.Colormap`
            - Matplotlib colormap 对象
        """

        # 生成颜色列表
        vals = np.linspace(self.value_range[0], self.value_range[1], bins).tolist()
        colors = self.__call__(vals)
        if not isinstance(colors, list):
            raise AssertionError('Expected a list of colors for matplotlib colormap')
        mpl_colors = [color for color in colors if color is not None]
        return to_mpl_cmap(mpl_colors, type=type)


def to_mpl_cmap(
    color: typing.Sequence[str], type: Literal['discrete', 'continuous'] = 'discrete'
) -> Colormap:
    """
    将颜色列表转换为 Matplotlib 可用的 colormap。

    Parameters
    ----------
    - `color` : `list[str]`
        - 颜色列表，支持 hex 字符串或 RGB/RGBA 元组
    - `type` : `Literal['discrete', 'continuous']`, optional
        - 指定 colormap 类型，by default 'discrete'

    Returns
    -------
    - `matplotlib.colors.Colormap`
        - Matplotlib colormap 对象
    """
    if type == 'discrete':
        from matplotlib.colors import ListedColormap

        return ListedColormap(color)
    elif type == 'continuous':
        from matplotlib.colors import LinearSegmentedColormap

        return LinearSegmentedColormap.from_list('custom_cmap', color)
    else:
        raise ValueError(
            f'Unknown colormap type: {type}. Use "discrete" or "continuous".'
        )


def preview_palette(
    palette: typing.Sequence[str | tuple],
    title: str = 'Palette Preview',
    figsize: tuple[float, float] = (8, 5),
):
    """
    预览调色板，展示颜色条、Hex 值以及在实际图表中的应用示例。

    生成包含三部分的图：顶部颜色条显示每种颜色及其代码，
    左下方多系列折线图示例，右下方多系列柱状图示例。

    Parameters
    ----------
    - `palette` : `list[str | tuple]`
        - 颜色列表，支持 hex 字符串或 RGB/RGBA 元组
    - `title` : `str`, optional
        - 预览图的总标题，by default 'Palette Preview'
    - `figsize` : `tuple[float, float]`, optional
        - 整个 Figure 的大小，by default (10, 6)
    """
    Color = _require_coloraide()

    # 转换调色板为标准格式（确保为 hex 字符串）
    standard_palette = [
        color if isinstance(color, str) else rgb_to_hex(color) for color in palette
    ]
    n_colors = len(standard_palette)

    # 在 color_cycle_context 内创建整个布局和绘制示例图表
    color_cycle_ctx = use_color_cycle(standard_palette, 'context')
    if color_cycle_ctx is None:
        raise AssertionError('Expected a matplotlib color cycle context manager')

    with color_cycle_ctx:
        # 使用 subplot_mosaic 创建复杂布局
        layout = [['T', 'T'], ['L', 'R']]
        fig, axes = plt.subplot_mosaic(
            typing.cast(typing.Any, layout),
            figsize=figsize,
            height_ratios=[1, 4],
            gridspec_kw={'hspace': 0.4},
        )

        fig.suptitle(title, fontsize=16, y=0.98)

        # --- 1. 绘制顶部的颜色条 ---
        ax_top = axes['T']
        ax_top.set_xlim(0, n_colors)
        ax_top.set_ylim(0, 1)

        for i, color in enumerate(standard_palette):
            # 绘制颜色矩形
            ax_top.add_patch(Rectangle((i, 0), 1, 1, color=color))

            # 处理不同的颜色格式并添加 Hex 值
            display_text = color.upper()
            text_color = 'white' if Color(color).luminance() < 0.4 else 'black'

            ax_top.text(
                i + 0.5,
                0.5,
                display_text,
                ha='center',
                va='center',
                color=text_color,
                fontsize=10,
                family='monospace',
            )

        ax_top.set_xticks([])
        ax_top.set_yticks([])
        style_ax(ax_top, hide_spines='all')

        # --- 2. 绘制左下方的多系列折线图 ---
        ax_left = axes['L']
        x = np.linspace(0, 10, 50)
        for i in range(n_colors):
            np.random.seed(42 + i)  # 固定随机种子以确保一致性
            y = np.sin(x + i) + np.random.randn(50) * 0.2
            ax_left.plot(x, y, label=f'Series {i+1}')
        style_ax(
            ax_left, title='Multi-line Plot Example', xlabel='X-axis', ylabel='Y-axis'
        )
        ax_left.legend(loc='best', fontsize='small')

        # --- 3. 绘制右下方的多系列柱状图 ---
        ax_right = axes['R']
        bar_width = 0.8 / n_colors
        group_positions = np.arange(4)
        for i in range(n_colors):
            offsets = (i - (n_colors - 1) / 2) * bar_width
            bar_positions = group_positions + offsets
            np.random.seed(42 + i)  # 固定随机种子
            heights = np.random.rand(4) * 0.8 + 0.2
            ax_right.bar(bar_positions, heights, width=bar_width, label=f'Series {i+1}')
        style_ax(ax_right, title='Multi-bar Plot Example', xlabel='Group')
        ax_right.set_xticks(group_positions, [f'G{j+1}' for j in group_positions])
        ax_right.legend(loc='best', fontsize='small')

        plt.tight_layout()
        plt.show()
