import tushare as ts
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np


def getDailyData(code, startDate, endDate):
    # ---------------------- 1. 初始化 Tushare ----------------------
    # 把这里替换成你自己的 TOKEN
    ts.set_token("你的key")
    pro = ts.pro_api()

    # ---------------------- 2. 获取 A 股日线数据 ----------------------
    # 示例：江波龙 301308.SZ
    # df = pro.daily(ts_code="301308.SZ", start_date="20260401", end_date="20260511")
    df = pro.daily(ts_code=code, start_date=startDate, end_date=endDate)

    # 按交易日期排序
    df = df.sort_values("trade_date")

    # ---------------------- 3. 查看数据 ----------------------
    print(df.head())
    return df


def drawPlot(df, *args):
    # 画图：以时间为X轴
    plt.figure(figsize=(10, 5))

    # 绘制股价收盘价，并添加图例标签：收盘价
    plt.plot(df['trade_date'], df['close'], linewidth=2, color='red', label='收盘价')

    # 循环绘制均线（args 是你传入的均线列名，如 ma5, ma10, svg_1 等）
    for arg in args:
        if arg in df.columns:
            # 绘制均线，并使用列名作为图例标签
            plt.plot(df['trade_date'], df[arg], '--', linewidth=2, label=arg)

    # ========== 添加图表标注信息 ==========
    plt.title('股价走势与均线图', fontsize=14, fontproperties='SimHei')  # 图表标题
    plt.xlabel('交易日期', fontsize=12, fontproperties='SimHei')  # X轴标签
    plt.ylabel('价格', fontsize=12, fontproperties='SimHei')  # Y轴标签

    # 显示图例（必须有 label 才会显示）
    plt.legend(prop={'family': 'SimHei'})  # 支持中文图例

    # 网格 + 日期旋转
    plt.grid(True)
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

def strategy1(df):
    #双移动平均策略
    #复制一份用作策略计算
    strategy=pd.DataFrame(index=df.index)
    strategy['date']=df['trade_date']
    #添加一个signal字段，用来存储交易信号
    strategy['signal']=0
    #将5日均线存到avg_5
    strategy['avg_5']=df['close'].rolling(5).mean()
    #将10日均线存到avg_10
    strategy['avg_10']=df['close'].rolling(10).mean()
    #计算买入信号，当5日均值大于10日均值时设为1 反之为0
    strategy['signal']=np.where(strategy['avg_5']>strategy['avg_10'],1,0)
    #根据交易信号变化下单，当交易信号从0变为1时买入，当从1变为0时卖出
    strategy['order']=strategy['signal'].diff()
    print("--------------双移动平均策略-----------")
    print(strategy)
    return strategy
    
def strategytest(cash, df):
    # 回测
    # 执行双移动平均策略
    strategy = strategy1(df)

    # 初始启动资金
    initial_cash = cash

    # 新建一个 strategy 的拷贝 仓位表
    positions = pd.DataFrame(index=strategy.index)
    positions['stock'] = 0  # 初始持仓为 0

    # 设置最小买入 100 股
    positions['stock'] = strategy['signal'] * 100

    # 创建投资组合数据库
    portfolio = pd.DataFrame(index=strategy.index)

    # 持仓市值 = 持股数量 × 收盘价
    portfolio['stock value'] = positions['stock'] * df['close']

    # 仓位变化 = 下单的数量
    order = positions.diff()

    # 剩余资金 = 初始资金 - 累计下单金额
    portfolio['cash'] = initial_cash - (order['stock'] * df['close']).cumsum()

    # 总资产 = 现金 + 持仓市值
    portfolio['total'] = portfolio['cash'] + portfolio['stock value']

    # 打印结果
    print("--------------回测结果-----------")
    print(portfolio)

    return portfolio  # 可选：返回结果方便后续画图


if __name__ == '__main__':
    # 获取股票数据
    df = getDailyData('600276.SH', '20260401', '20260501')
    # 增加10日均线列
    df['avg_10'] = df['close'].rolling(10).mean()
    df['avg_5'] = df['close'].rolling(5).mean()
    # 画出来
    drawPlot(df, 'avg_10', 'avg_5')
    #执行双移动平均策略回测
    strategytest(20000,df)