"""
图表工具模块

提供图表生成和图片处理功能，支持：
- 图片与Base64互转
- 柱状图生成
- 折线图生成（支持双Y轴）
- 同比折线图

Example:
    >>> ChartUtil.get_bar_chart("output.png", {"A": 1, "B": 2})
    >>> ChartUtil.get_line_chart("line.png", ["Q1", "Q2"], {"data": [10, 20]})
"""
import base64
import logging
from typing import Dict, List, Optional

from matplotlib import pyplot as plt

logger = logging.getLogger(__name__)


class ChartUtil:
    """
    图表工具类

    提供静态方法进行图表生成和图片Base64编码转换。

    Example:
        >>> data = {"A": 100, "B": 200, "C": 150}
        >>> ChartUtil.get_bar_chart("chart.png", data)
    """

    # 默认颜色列表
    COLORS: List[str] = [
        "Blue", "Red", "Black", "Yellow",
        "Green", "Cyan", "Magenta", "Orange"
    ]

    @staticmethod
    def img_to_base64(img_file_path: str) -> bytes:
        """
        将图片文件转换为Base64编码

        Args:
            img_file_path: 图片文件路径

        Returns:
            Base64编码的字节数据
        """
        with open(img_file_path, "rb") as file_reader:
            return base64.b64encode(file_reader.read())

    @staticmethod
    def base64_to_img(str_base64: bytes, img_file_path: str) -> None:
        """
        将Base64编码转换为图片文件

        Args:
            str_base64: Base64编码的字节数据
            img_file_path: 输出图片文件路径
        """
        img = base64.b64decode(str_base64)
        with open(img_file_path, "wb") as file_writer:
            file_writer.write(img)

    @staticmethod
    def get_bar_chart(
        img_path: str,
        data_dict: Dict[str, float],
        width: int = 20,
        height: int = 5
    ) -> None:
        """
        生成柱状图

        Args:
            img_path: 图片保存路径
            data_dict: 数据字典，键为标签，值为数值
            width: 图表宽度
            height: 图表高度

        Example:
            >>> data = {"A": 100, "B": 200, "C": 150}
            >>> ChartUtil.get_bar_chart("bar_chart.png", data)
        """
        keys = sorted(list(data_dict.keys()))
        data_list = [data_dict[key] for key in keys]

        plt.figure(figsize=(width, height))
        plt.bar(range(len(data_list)), data_list)
        plt.xticks(range(len(keys)), keys)
        plt.savefig(img_path)
        plt.close()

    @staticmethod
    def get_line_chart(
        img_path: str,
        x_list: List[str],
        master_dict: Dict[str, List[float]],
        slave_dict: Optional[Dict[str, List[float]]] = None,
        x_label: str = "",
        y_label: str = "",
        title: str = "",
        width: int = 12,
        height: int = 3,
        bottom: float = 0.25
    ) -> None:
        """
        生成折线图

        支持双Y轴，主轴数据在左侧，副轴数据在右侧。

        Args:
            img_path: 图片保存路径
            x_list: X轴坐标标签列表
            master_dict: 主Y轴数据字典，格式为 {"系列名": [数值列表]}
            slave_dict: 副Y轴数据字典（可选）
            x_label: X轴标签
            y_label: Y轴标签
            title: 图表标题
            width: 图表宽度
            height: 图表高度
            bottom: 底部边距比例

        Example:
            >>> x = ["Q1", "Q2", "Q3", "Q4"]
            >>> master = {"Sales": [100, 150, 120, 180]}
            >>> slave = {"Growth": [0.1, 0.5, -0.2, 0.5]}
            >>> ChartUtil.get_line_chart("line.png", x, master, slave)
        """
        if slave_dict is None:
            slave_dict = {}

        plt.figure(figsize=(width, height))
        plt.subplots_adjust(bottom=bottom)

        # X轴坐标
        xx = range(len(x_list))
        plt.xlabel(x_label)
        plt.xticks(rotation=90, fontsize=8)
        plt.xticks(list(xx), x_list)

        # Y轴网格
        plt.grid(axis="y")

        # 主Y轴数据
        idx_color = 0
        for key in master_dict:
            color_idx = idx_color % len(ChartUtil.COLORS)
            yy = master_dict[key]
            plt.plot(xx, yy, label=key, color=ChartUtil.COLORS[color_idx])
            idx_color += 1

        # 副Y轴数据
        for key in slave_dict:
            color_idx = idx_color % len(ChartUtil.COLORS)
            yy = slave_dict[key]
            plt.twinx()
            plt.plot(xx, yy, label=key, color=ChartUtil.COLORS[color_idx])
            idx_color += 1

        plt.ylabel(y_label)
        plt.title(title)
        plt.savefig(img_path, bbox_inches="tight", pad_inches=0.1)
        plt.close()

    @staticmethod
    def get_line_chart_compare_base(
        date_list: List[str],
        data_dict: Dict[str, List[float]],
        save_file_path: str
    ) -> None:
        """
        生成同比折线图

        以每个系列的首个值为基数，计算相对比例。

        Args:
            date_list: 日期/X轴标签列表
            data_dict: 数据字典，格式为 {"系列名": [数值列表]}
            save_file_path: 图片保存路径

        Example:
            >>> dates = ["2024-01", "2024-02", "2024-03"]
            >>> data = {"A": [100, 120, 150], "B": [200, 180, 220]}
            >>> ChartUtil.get_line_chart_compare_base(dates, data, "compare.png")
        """
        normalized_dict: Dict[str, List[float]] = {}

        for key in data_dict:
            data_list = data_dict[key]
            base = data_list[0] if data_list else 0

            tmp_list: List[float] = []
            for value in data_list:
                if base == 0 and value == 0:
                    tmp_list.append(0.0)
                elif base == 0 and value >= 0:
                    tmp_list.append(1.0)
                    base = value
                else:
                    tmp_list.append(value / base)
            normalized_dict[key] = tmp_list

        ChartUtil.get_line_chart(save_file_path, date_list, normalized_dict)


if __name__ == "__main__":
    # 测试柱状图
    test_dict = {"A": 1, "B": 2, "C": 3}
    test_path = "../../data/test.png"
    ChartUtil.get_bar_chart(test_path, test_dict)
