import pandas as pd
from pyecharts.charts import Kline, Line, Bar, Grid
from pyecharts import options as opts
from pyecharts.commons.utils import JsCode

class StockKlineChart:
    def __init__(self, df: pd.DataFrame, stock_code: str):
        """
        初始化 K 線圖表對(duì)象

        :param df: 股票數(shù)據(jù) DataFrame（包含“日期”，“開盤”，“收盤”，“最低”，“最高”，“成交量”）
        :param stock_code: 股票代碼
        """
        self.df = df.sort_values("日期")  # 按日期排序，確保數(shù)據(jù)按時(shí)間順序排列
        self.stock_code = stock_code  # 股票代碼
        self.color_up = "#ef232a"  # 陽線（上漲）顏色
        self.color_down = "#14b143"  # 陰線（下跌）顏色
        self.ma_periods = [5, 10, 20]  # 移動(dòng)平均線周期
        self.ma_colors = {5: "#FF0000", 10: "#0000FF", 20: "#00FF00"}  # 均線顏色

    def _prepare_data(self):
        """ 處理數(shù)據(jù)：提取 K 線數(shù)據(jù)、計(jì)算移動(dòng)均線、設(shè)置成交量顏色 """
        self.dates = self.df['日期'].tolist()  # 提取日期列表
        self.kline_data = self.df[["開盤", "收盤", "最低", "最高"]].values.tolist()  # 提取 K 線數(shù)據(jù)

        # 計(jì)算移動(dòng)平均線
        for period in self.ma_periods:
            self.df[f'MA{period}'] = (
                self.df['收盤']
                .rolling(window=period)
                .mean()
                .bfill()  # 處理 NaN 值（前向填充）
                .round(2)  # 保留兩位小數(shù)
            )

        # 計(jì)算成交量顏色標(biāo)記（1: 上漲, -1: 下跌）
        self.df['color'] = self.df.apply(
            lambda x: 1 if x['收盤'] >= x['開盤'] else -1,
            axis=1
        )
        self.df['index_vol'] = range(len(self.df))  # 給成交量數(shù)據(jù)添加索引

    def create_chart(self):
        """ 生成 K 線圖表 """
        self._prepare_data()  # 處理數(shù)據(jù)

        # ================== K 線圖配置 ==================
        kline = (
            Kline()
            .add_xaxis(self.dates)  # 設(shè)置 X 軸日期
            .add_yaxis(
                series_name="K線",  # K 線名稱
                y_axis=self.kline_data,  # K 線數(shù)據(jù)（開盤、收盤、最低、最高）
                itemstyle_opts=opts.ItemStyleOpts(
                    color=self.color_up,  # 陽線顏色
                    color0=self.color_down  # 陰線顏色
                )
            )
            .set_global_opts(
                title_opts=opts.TitleOpts(
                    title="股票K線走勢圖",  # 圖表標(biāo)題
                    subtitle=f"股票代碼：{self.stock_code}",  # 副標(biāo)題
                    pos_left="left"  # 標(biāo)題位置
                ),
                legend_opts=opts.LegendOpts(
                    is_show=True,  # 是否顯示圖例
                    pos_top=10,  # 圖例位置（頂部）
                    pos_left="center"  # 居中對(duì)齊
                ),
                xaxis_opts=opts.AxisOpts(
                    type_="category",  # X 軸類型（類別）
                    axislabel_opts=opts.LabelOpts(rotate=0),  # X 軸標(biāo)簽角度
                    splitline_opts=opts.SplitLineOpts(is_show=True)  # 是否顯示網(wǎng)格線
                ),
                yaxis_opts=opts.AxisOpts(
                    is_scale=True,  # Y 軸是否自適應(yīng)縮放
                    splitarea_opts=opts.SplitAreaOpts(
                        is_show=True,  # 是否顯示網(wǎng)格背景
                        areastyle_opts=opts.AreaStyleOpts(opacity=1)  # 設(shè)置透明度
                    )
                ),
                tooltip_opts=opts.TooltipOpts(
                    trigger="axis",  # 觸發(fā)方式：鼠標(biāo)懸浮時(shí)顯示
                    axis_pointer_type="cross",  # 坐標(biāo)軸指示器類型（十字指示）
                ),
                datazoom_opts=[
                    opts.DataZoomOpts(
                        is_show=False,  # 是否顯示數(shù)據(jù)縮放控件
                        type_="inside",  # 縮放類型：內(nèi)部滑動(dòng)
                        xaxis_index=[0, 1],  # 作用于 X 軸
                        range_start=80,  # 初始顯示范圍
                        range_end=100
                    ),
                    opts.DataZoomOpts(
                        is_show=True,  # 顯示滑動(dòng)條縮放
                        xaxis_index=[0, 1],
                        type_="slider",
                        pos_top="100%",  # 位置：底部
                        range_start=80,
                        range_end=100
                    )
                ],
                # 坐標(biāo)軸指示器
                axispointer_opts=opts.AxisPointerOpts(
                    is_show=True,
                    link=[{"xAxisIndex": "all"}],
                    label=opts.LabelOpts(background_color="#777")
                )
            )
        )

        # ================== 移動(dòng)平均線配置 ==================
        line = Line().add_xaxis(self.dates)
        for period in self.ma_periods:
            line.add_yaxis(
                series_name=f"MA{period}",
                y_axis=self.df[f'MA{period}'].tolist(),
                is_smooth=True,  # 平滑曲線
                symbol="none",  # 取消數(shù)據(jù)點(diǎn)標(biāo)記
                linestyle_opts=opts.LineStyleOpts(
                    color=self.ma_colors[period],  # 顏色
                    width=2  # 線寬
                ),
                label_opts=opts.LabelOpts(is_show=False)  # 隱藏?cái)?shù)據(jù)標(biāo)簽
            )

        # 疊加 K 線和均線
        overlap_kline = kline.overlap(line)

        # ================== 成交量柱狀圖 ==================
        vol_bar = (
            Bar()
            .add_xaxis(self.dates)
            .add_yaxis(
                series_name="成交量",
                y_axis=self.df[['index_vol', '成交量', 'color']].values.tolist(),
                label_opts=opts.LabelOpts(is_show=False),  # 隱藏標(biāo)簽
                
            )
            .set_global_opts(
                xaxis_opts=opts.AxisOpts(
                    type_="category",
                    grid_index=1,
                    axislabel_opts=opts.LabelOpts(is_show=False),
                    axistick_opts=opts.AxisTickOpts(is_show=False),
                    splitline_opts=opts.SplitLineOpts(is_show=False)
                ),
                yaxis_opts=opts.AxisOpts(is_show=False),
                legend_opts=opts.LegendOpts(is_show=False),
                visualmap_opts=opts.VisualMapOpts(
                    is_show=False,
                    dimension=2,  # 顏色映射使用的維度（color）
                    series_index=4,  # 作用于第 5 個(gè)系列（成交量）
                    is_piecewise=True,  # 分段顯示
                    pieces=[
                        {"value": 1, "color": self.color_up},  # 上漲顏色
                        {"value": -1, "color": self.color_down}  # 下跌顏色
                    ]
                )
            )
        )

        # ================== 組合圖表 ==================
        grid = (
            Grid(init_opts=opts.InitOpts(
                width="98vw",
                height="95vh",
                animation_opts=opts.AnimationOpts(animation=False)  # 關(guān)閉動(dòng)畫
            ))
            .add(overlap_kline, grid_opts=opts.GridOpts(pos_top="10%", height="60%" ,pos_left="30px",pos_right="10px"))
            .add(vol_bar, grid_opts=opts.GridOpts(pos_top="73%", height="20%" ,pos_left="30px",pos_right="10px"))
        )

        return grid

    def render(self, file_path: str = "stock_kline.html"):
        """ 渲染并保存 K 線圖 """
        chart = self.create_chart()
        chart.render(file_path)
        print(f"K 線圖已保存為 {file_path}")

import akshare as ak
import StockKlineChart as skc

# 獲取股票數(shù)據(jù)
df = ak.stock_zh_a_hist(
    symbol="000001",       # 股票代碼
    period="daily",        # 日線數(shù)據(jù)
    start_date="20230101", # 起始日期
    end_date="20251201",   # 結(jié)束日期
    adjust="qfq"           # 前復(fù)權(quán)處理
)

# 創(chuàng)建 K 線圖實(shí)例
stock_chart = skc.StockKlineChart(df=df, stock_code="000001")
stock_chart.render()