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

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

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

from .mpl_utils import style_ax, use_color_cycle


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 | list[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']
    """
    from coloraide import Color

    # 转换基色到指定色彩空间
    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 颜色字符串的列表
    """
    from coloraide import Color

    # --- 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: list[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 颜色字符串的列表
    """
    from coloraide import Color

    # 使用 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: list[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'
        """
        from coloraide import Color

        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)

    def __call__(self, value: float | Iterable[float]) -> 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
    ) -> plt.cm:
        """
        将 coloraide 的颜色映射转换为 Matplotlib 可用的 colormap。

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

        Returns
        -------
        - `plt.cm`
            - Matplotlib colormap 对象
        """

        # 生成颜色列表
        vals = np.linspace(self.value_range[0], self.value_range[1], bins).tolist()
        colors = self.__call__(vals)
        return to_mpl_cmap(colors, type=type)


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

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

    Returns
    -------
    - `plt.cm`
        - 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: list[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)
    """
    from coloraide import Color

    # 转换调色板为标准格式（确保为 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 内创建整个布局和绘制示例图表
    with use_color_cycle(standard_palette, 'context'):
        # 使用 subplot_mosaic 创建复杂布局
        layout = [['T', 'T'], ['L', 'R']]
        fig, axes = plt.subplot_mosaic(
            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(plt.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()
