import numpy as np

import matplotlib as plt

import pandas as pd

import tushare as ts

#获取股票号为601318的股票信息

df=ts.get_k_data("")

#将数据保存到本地,方便处理

df.to_csv("601318.csv",index=False)

#获取数据内的有用列,并将date列作为index

df = pd.read_csv("601318.csv",index_col="date")[["open","close","high","low","volume"]]

#获取5日均线,10日均线

df["ma5"]=np.nan

df["ma30"]=np.nan

# 方式一

# for i in range(4, len(df)):

#     df.loc[df.index[i], 'ma5'] = df['close'][i-4:i+1].mean()

# for i in range(29, len(df)):

#     df.loc[df.index[i], 'ma30'] = df['close'][i-29:i+1].mean()

# 方式二

df["ma5"]=df["close"].rolling(5).mean()

df["ma30"]=df["close"].rolling(30).mean()

df[["close","ma5","ma30"]].plot()

# 获取金叉点,死叉点

golden_cross = []

death_cross = []

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])

    if df['ma5'][i] < df['ma30'][i] and df['ma5'][i-1] >= df['ma30'][i-1]:

        death_cross.append(df.index[i])

first_money = 100000

money = first_money

hold = 0 # 持有的股票数

sr1 = pd.Series(1, index=golden_cross)

sr2 = pd.Series(0, index=death_cross)

sr = sr2.append(sr1).sort_index()

for i in range(0, len(sr)):

    p = df['close'][sr.index[i]]

    if sr.iloc[i] == 1:

        #金叉

        buy = (money // p*100)

        hold += buy*100

        money -= buy*100*p

    else:

        money += hold*p

        hold = 0

p = df['close'][-1]

now_money = hold*p + money

now_money - first_money