from ssquant.SQDATA import TakeData
from pyecharts import options as opts
from pyecharts.charts import Kline, Bar, Grid
import pandas as pd


def plotK(data):
    # 示例数据（您需要替换为您的实际数据）
    kline_data = data[['open', 'close', 'low', 'high']].values.tolist()
    dates = data.index.strftime('%Y-%m-%d %H:%M:%S').tolist()
    symbol_data = data['symbol'].values.tolist()
    # 标记 symbol 变化的位置
    markline_data = []
    for i in range(1, len(symbol_data)):
        if symbol_data[i] != symbol_data[i-1]:
            # 当前 symbol 与前一个不同时，添加红色竖线
            markline_data.append(opts.MarkLineItem(x=dates[i], name=f'前一个合约{symbol_data[i-1]},当前合约{symbol_data[i]}'))


    # 数据缩放组件配置
    datazoom_slider = opts.DataZoomOpts(type_="slider", xaxis_index=[0, 1, 2, 3,4], range_start=50, range_end=100)
    datazoom_inside = opts.DataZoomOpts(type_="inside", xaxis_index=[0, 1, 2, 3,4])

    # 创建 K 线图
    kline = (
        Kline(init_opts=opts.InitOpts(width="100%", height="900px"))
        .add_xaxis(dates)
        .add_yaxis('K线图表', kline_data,markline_opts=opts.MarkLineOpts(data=markline_data, symbol='none', linestyle_opts=opts.LineStyleOpts(color="red")))#"ssss",
        .set_global_opts(
            datazoom_opts=[datazoom_slider, datazoom_inside],
            toolbox_opts=opts.ToolboxOpts(is_show=True, pos_top="0%", pos_right="80%"),
            legend_opts=opts.LegendOpts(pos_left='40%'),  # 调整图例位置到底部
        )
    )
    kline.render('K线图.html')



'''
获取数据-
品种:symbol,
起始时间:start_date,
结束时间:end_date(包含当天),
周期kline_period:1M..5M..NM(分钟),1D(天),1W(周),1Y(月)
复权adjust_type:0(不复权)1(后复权)
'''

# 请在下方输入你的俱乐部账号密码，username='俱乐部账号' password='密码'

client = TakeData(username='1234@qq.com', password='123')
data = client.get_data(
    symbol='rb888',
    start_date='2023-12-28',
    end_date='2024-01-17',
    kline_period='60M',
    adjust_type=1
)
data.set_index("datetime", inplace=True)
data.index = pd.to_datetime(data.index)
print(data)
#生产K线图表到脚本同目录下
plotK(data)