# 导入库

import talib

import numpy as np

import pandas as pd

# 策略参数

short_window = 10  # 短期均线周期

long_window = 30   # 长期均线周期

stop_loss = 0.02   # 止损比例

take_profit = 0.05 # 止盈比例

# 数据准备

data = pd.read_csv('rb_futures.csv')  # 读取螺纹钢期货数据

data['short_ma'] = talib.SMA(data['close'], short_window)  # 计算短期均线

data['long_ma'] = talib.SMA(data['close'], long_window)    # 计算长期均线

# 信号生成

data['signal'] = 0

data.loc[data['short_ma'] > data['long_ma'], 'signal'] = 1  # 短期均线上穿长期均线，买入信号

data.loc[data['short_ma'] < data['long_ma'], 'signal'] = -1 # 短期均线下穿长期均线，卖出信号

# 回测

initial_capital = 100000  # 初始资金

position = 0              # 持仓数量

portfolio_value = []      # 组合价值

for i in range(1, len(data)):

   # 计算当前持仓价值

   current_value = position * data['close'][i]

   # 生成交易信号

   signal = data['signal'][i]

   # 执行交易

   if signal == 1 and position == 0:  # 买入

       position = initial_capital / data['close'][i]

       stop_loss_price = data['close'][i] * (1 - stop_loss)

       take_profit_price = data['close'][i] * (1 + take_profit)

   elif signal == -1 and position > 0:  # 卖出

       initial_capital = position * data['close'][i]

       position = 0

   elif position > 0:  # 持仓

       if data['close'][i] <= stop_loss_price or data['close'][i] >= take_profit_price:  # 止损或止盈

           initial_capital = position * data['close'][i]

           position = 0

   # 记录组合价值

   portfolio_value.append(initial_capital + current_value)

# 回测结果分析

print('最终组合价值:', portfolio_value[-1])

print('收益率:', (portfolio_value[-1] - 100000) / 100000 * 100, '%')