import tushare as ts
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
#1.獲取股票代碼爲002400的股票數據
#df = ts.get_k_data('002400', start='1988')
#df.to_csv('002400.csv')
df = pd.read_csv('002400.csv', index_col='date', parse_dates=['date'])[['open', 'close', 'high', 'low']]
print(df)
#2.計算其5日和30日均線(收盤價)
df['ma5'] = np.nan
df['ma30'] = np.nan
df['ma5'] = df['close'].rolling(5).mean()
df['ma30'] = df['close'].rolling(30).mean()
#print(df)
#3.使用matplotlib模塊實現收盤價、ma5、ma30的可視化
# plt.plot(df.index, df['close'].values, 'r', label='close')
# plt.plot(df.index, df['ma5'].values, 'b', label='ma5')
# plt.plot(df.index, df['ma30'].values, 'k', label='ma30')
# plt.show()
#4.分析輸出所有金叉和死叉日期
golden_cross = []
death_cross = []
df = df['2011':]
for i in range(1, len(df)):
if (df['ma5'][i]>df['ma30'][i] and df['ma5'][i-1]<=df['ma30'][i-1]):
golden_cross.append(df.index[i].to_pydatetime())
if (df['ma5'][i] < df['ma30'][i] and df['ma5'][i - 1] >= df['ma30'][i - 1]):
death_cross.append(df.index[i].to_pydatetime())
sr1 = df['ma5'] >= df['ma30']
sr2 = df['ma5'] < df['ma30']
sr2 = sr2.shift(1)
golden_cross = df[sr1 & sr2].index
death_cross = df[~(sr1 | sr2)].index
#5.從2011/01/01開始,初始資金爲100000元,金叉儘量買入,死叉全部賣出,那麼最終我的收益爲多少?
goldsr = pd.Series(1, index=golden_cross)
deathsr = pd.Series(0, index=death_cross)
sr = goldsr.append(deathsr).sort_index()
first_money = 100000
cash = first_money
hold = 0
for i in range(len(sr)):
p = df['open'][sr.index[i]]
print('p:'+str(p))
if sr.iloc[i] == 1:#金叉,買入
buy = (cash // (100*p))
hold += buy
cash -= buy*100*p
else:
cash += hold*100*p
print('cash:'+str(cash))
hold = 0
now_money = hold*100*df['open'][-1] + cash
print(now_money)