import pandas as pd
import os
from paste.chart.line import gen_lines, gen_splines


class LineChartExample:
    """
    折线图与平滑曲线测试管理器。
    封装 gen_lines 和 gen_splines 的调用，统一输出管理。
    """

    def __init__(self, output_directory="./charts"):
        """
        初始化测试器，定义所有测试数据。
        使用 'ME' 替代 'M' 以兼容 pandas >=2.0。
        """
        self.output_directory = output_directory
        os.makedirs(self.output_directory, exist_ok=True)

        # =========================
        # gen_lines 测试数据（使用 'ME' 替代废弃的 'M'）
        # =========================
        # 时间序列数据：3个年份，每月最后一个交易日（12个月）
        dates_2022 = pd.date_range('2022-01-01', periods=12, freq='ME')
        dates_2023 = pd.date_range('2023-01-01', periods=12, freq='ME')
        dates_2024 = pd.date_range('2024-01-01', periods=12, freq='ME')

        self.data_dict_lines = {
            '2022': pd.Series([100 + i * 2 for i in range(12)], index=dates_2022),
            '2023': pd.Series([110 + i * 1.5 for i in range(12)], index=dates_2023),
            '2024': pd.Series([120 + i * 1 for i in range(12)], index=dates_2024),
        }

        self.color_palette_lines = 'BuPu'
        self.dpi_lines = 128

        # =========================
        # gen_splines 测试数据（数值索引，无频率问题）
        # =========================
        self.data_dict_splines = {
            'A组': pd.Series([10, 15, 12, 18, 16], index=[0, 1, 2, 3, 4]),
            'B组': pd.Series([12, 14, 13, 17, 15], index=[0, 1, 2, 3, 4]),
            'C组': pd.Series([8, 20, 10, 22, 11], index=[0, 1, 2, 3, 4]),
        }

        self.total_splines = pd.Series({
            'A组': 71,
            'B组': 71,
            'C组': 71
        })

        self.color_palette_splines = 'viridis'
        self.dpi_splines = 128
        self.smooth_points = 100
        self.spline_k = 3
        self.markevery = 30

    def generate_lines(self) -> str:
        """调用 gen_lines，参数顺序完全匹配原始函数"""
        return gen_lines(
            self.data_dict_lines,
            self.color_palette_lines,
            self.dpi_lines
        )

    def generate_splines(self) -> str:
        """调用 gen_splines，参数顺序完全匹配原始函数"""
        return gen_splines(
            self.data_dict_splines,
            self.total_splines,
            None,
            self.color_palette_splines,
            self.dpi_splines
        )

    def save_svg(self, svg_data: str, filename: str) -> None:
        """
        将 SVG 的 base64 Data URL 写入文件（保留原始 SVG 格式）。
        注意：svg_data 是 "data:image/svg+xml;base64,..."，需提取真实 SVG 内容。
        """
        if not svg_data or not isinstance(svg_data, str):
            print(f"生成的 SVG 数据无效（为空或非字符串）: {filename}")
            return

        # 提取 base64 编码部分（去除 data URL 前缀）
        if svg_data.startswith("data:image/svg+xml;base64,"):
            base64_content = svg_data[len("data:image/svg+xml;base64,"):]
            try:
                # 解码 base64 得到原始 SVG 字符串
                import base64
                svg_content = base64.b64decode(base64_content).decode('utf-8')
            except Exception as e:
                print(f"解码 base64 失败: {e}")
                svg_content = svg_data  # 退化为直接写入
        else:
            # 如果不是标准格式，直接写入（兼容调试）
            svg_content = svg_data

        filepath = os.path.join(self.output_directory, filename)
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(svg_content)
        print(f"已保存: {filepath}")

    def run(self) -> None:
        """按顺序生成并保存所有图表"""
        print("生成折线图...")
        svg1 = self.generate_lines()
        self.save_svg(svg1, "lines.svg")

        print("生成平滑曲线图...")
        svg2 = self.generate_splines()
        self.save_svg(svg2, "splines.svg")

        print("\n所有图表已生成。")
        print(f"输出目录: {self.output_directory}")
        print("文件列表:")
        print("  - lines.svg")
        print("  - splines.svg")


# 程序入口
if __name__ == "__main__":
    tester = LineChartExample()
    tester.run()