
import pandas as pd
import numpy as np
import pickle
import math
import os
import tushare as ts
token='2876ea85cb005fb5fa17c809a98174f2d5aae8b1f830110a5ead6211'
pro=ts.pro_api(token)

# 复权函数
def adjust_price(df: pd.DataFrame, label = 'front') -> pd.DataFrame:
    df_adj = df.copy()
    adj_factor = df['close'] / df['pre_close']
    adj_factor = adj_factor.cumprod()
    
    if label == 'front':
        df_adj['close'] = adj_factor / adj_factor.iloc[-1] * df['close'].iloc[-1]
    elif label == 'back':
        df_adj['close'] = adj_factor / adj_factor.iloc[0] * df['close'].iloc[0]
        
    df_adj['open'] = df['open'] / df['close'] * df_adj['close']
    df_adj['high'] = df['high'] / df['close'] * df_adj['close']
    df_adj['low'] = df['low'] / df['close'] * df_adj['close']
    return df_adj


#%% 起止日期，这里的起止日期可以根据需要进行修改
start_date, end_date = '20180101', '20260516'

# 趋势回归计算时回溯的天数
reg_num = 25  
# 上证180ETF: 510180.SH； 创业板ETF: 159915.SZ； 纳指ETF: 513100.SH； 黄金ETF: 518880.SH
etf_libs = ['510180.SH', '159915.SZ', '513100.SH', '518880.SH']

# 更新ETF的行情数据，数据源是Tushare
df_all = []
os.makedirs("data_ts", exist_ok=True)
for code in etf_libs:
    df = pro.fund_daily(ts_code=code, start_date=start_date, end_date=end_date)
    df = df.sort_values(by='trade_date').ffill()
    df = adjust_price(df) # 复权
    df_all.append(df)
    df.to_csv('data_ts/'+ code + '.csv')
df_all = pd.concat(df_all)   

# 交易日列表
trade_day = pro.trade_cal(exchange='', start_date=start_date, end_date=end_date)
trade_day = trade_day[trade_day['is_open'] == 1]
trade_day = trade_day.sort_values(by='cal_date', ascending=True)    
trade_day = pd.to_datetime(trade_day['cal_date']).tolist()


# 将数据字段改成backtrader需要的格式
df_all.rename(columns={'vol': 'volume'}, inplace=True) # 改名以满足backtrader要求
df_all['trade_date'] = pd.to_datetime(df_all['trade_date'])
df_all = df_all.set_index(['ts_code','trade_date'])
df_all = df_all.sort_index()

df_bt = {}
for code in etf_libs:
    df_temp = df_all.loc[code]
    df_temp = df_temp.sort_index()
    df_temp = df_temp.reindex(index=trade_day)
    df_temp = df_temp.ffill()
    df_bt[code] = df_temp
with open('data_ts/daily_data_ETF_bt.pkl', "wb") as file:
    pickle.dump(df_bt, file) 
print('更新数据完成')
   

   
#%% 计算每天需要买入的ETF
close = df_all['close'].swaplevel('ts_code', 'trade_date').unstack().sort_index()
close = close.reindex(index=trade_day).ffill()

# 计算每一天四只ETF的得分，持仓得分最高的ETF
etf_tohold = pd.Series(index = trade_day, dtype = str)
os.makedirs('data_ts/每日结果', exist_ok=True)

for idx, day in enumerate(close.index):
    if idx < reg_num: continue
    scores = pd.Series(index = close.columns, dtype = 'float64')
    for stk in close.columns:
        temp_close = close.iloc[idx-reg_num: idx][stk]
        y =  np.log(temp_close)
        x =  np.arange(y.size)
        slope, intercept = np.polyfit(x, y, 1)
        annualized_returns = math.pow(math.exp(slope), 250) - 1
        r_squared = 1 - (sum((y - (slope * x + intercept))**2) / ((len(y) - 1) * np.var(y, ddof=1)))
        scores[stk] = annualized_returns * r_squared 
    # print(day, scores)
    print(day, scores.idxmax())
    etf_tohold[day] = scores.idxmax()
    
    # 每天的结果单独保存成一个txt文件
    txt_name = day.strftime('%Y%m%d') + '.txt'
    with open('data_ts/每日结果/' + txt_name, 'w', encoding='utf-8') as file:
        file.write(etf_tohold[day])
    
# 保存每天需要持仓的ETF
etf_tohold.to_pickle('data_ts/etf_tohold.pkl')
