import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['font.serif'] = ['SimHei']
import seaborn as sns
df = pd.read_csv('fut_zz500.csv')
df['date']=pd.to_datetime(df['TrdDt'])
df.set_index('date',inplace=True)
df['RSV'] = (df['close'] - df['low'].rolling(7).min()) / (df['high'].rolling(7).max() - df['low'].rolling(7).min()) * 100
RSV = df['RSV'][7:]
K = pd.Series(0.0,index=RSV.index)
K[0] = 50
D = pd.Series(0.0,index=RSV.index)
D[0] = 50
for i in range(1,len(RSV)):
    K[i] = (RSV[i] + K[i-1]*2)/3
    D[i] = (K[i] + D[i-1]*2)/3
J = 3*K - 2*D
plt.figure(figsize=(40, 8), dpi=80)
plt.plot(K, label='K线快线')
plt.plot(D, label='D线慢线')
plt.plot(J, label='J线')
plt.legend()
plt.show()
KSignal = K.apply(lambda x: -1 if x>85 else 1 if x<20 else 0)
DSignal = D.apply(lambda x: -1 if x>80 else 1 if x<20 else 0)
JSignal = J.apply(lambda x: -1 if x>100 else 1 if x<0 else 0)
KDJSignal = KSignal + DSignal + JSignal
KDJSignal = KDJSignal.apply(lambda x: 1 if x>=2 else -1 if x<=-2 else 0)
df['fut_ret'] = ((df['close']-df['close'].shift(1))/df['close'].shift(1))[1:]
df['fut_ret'][0] = 0.
KDJSignal = KDJSignal.shift(1)[1:]
tradeRet = df['fut_ret'] * KDJSignal
tradeRet = tradeRet[1:]
#考虑开仓平仓交易成本
#根据中金所，中证500股指期货开仓成本为万分之零点二三，平仓成本为万分之三点四五
for i in range(1,len(tradeRet)-1):
    if KDJSignal[i] == 0 and KDJSignal[i+1] == 1:
        tradeRet[i+1] = tradeRet[i+1] - 0.23/10000
    elif KDJSignal[i] == 1 and KDJSignal[i+1] == 0:
        tradeRet[i+1] = tradeRet[i+1] - 3.45/10000
    elif KDJSignal[i] == 0 and KDJSignal[i+1] == -1:
        tradeRet[i+1] = tradeRet[i+1] - 0.23/10000
    elif KDJSignal[i] == -1 and KDJSignal[i+1] == 0:
        tradeRet[i+1] = tradeRet[i+1] - 3.45/10000
df['Strategy_Ret'] = np.zeros(len(df['close']))
df['Strategy_Ret'][1:] = tradeRet
df['Strategy_NetValue'] = np.ones(len(df['close']))
for i in range(len(df['close'])-1):
    df['Strategy_NetValue'][i+1] = df['Strategy_NetValue'][i]*(1+df['Strategy_Ret'][i+1])

df['zz500_NetValue'] = np.ones(len(df['close']))
for i in range(len(df['close'])-1):
    df['zz500_NetValue'][i+1] = df['zz500_NetValue'][i]*(1+df['IdxDRet'][i+1])